{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

module Linear.Vector (
  module Data.Vector.Fixed.Boxed,
  Vec1,
  (!?),
  delete,
  element_sum )
where

import qualified Algebra.Additive as Additive ( C )
import qualified Algebra.Ring as Ring ( C )
import Data.Vector.Fixed (
  Dim,
  N1,
  S,
  Vector(..),
  fromList,
  toList )
import Data.Vector.Fixed (
  (!),
  foldl,
  length )
import Data.Vector.Fixed.Boxed (
  Vec,
  Vec2,
  Vec3,
  Vec4,
  Vec5 )
import NumericPrelude hiding ( abs, length, foldl )


type Vec1 = Vec N1



-- | Safe indexing.
--
--   Examples:
--
--   >>> import Data.Vector.Fixed (mk3)
--   >>> let v1 = mk3 1 2 3 :: Vec3 Int
--   >>> v1 !? 2
--   Just 3
--   >>> v1 !? 3
--   Nothing
--
(!?) :: (Vector v a) => v a -> Int -> Maybe a
(!?) v1 idx
  | idx < 0 || idx >= length v1 = Nothing
  | otherwise                   = Just $ v1 ! idx


-- | Remove an element of the given vector.
--
--   Examples:
--
--   >>> import Data.Vector.Fixed (mk3)
--   >>> let b = mk3 1 2 3 :: Vec3 Int
--   >>> delete b 1 :: Vec2 Int
--   fromList [1,3]
--
delete :: (Vector v a,
           Vector w a,
           Dim v ~ S (Dim w))
       => v a
       -> Int
       -> w a
delete v1 idx =
  fromList (lhalf ++ rhalf')
  where
    (lhalf, rhalf) = splitAt idx (toList v1)
    rhalf' = tail rhalf


-- | We provide our own sum because sum relies on a Num instance
--   from the Prelude that we don't have.
--
--   Examples:
--
--   >>> import Data.Vector.Fixed (mk3)
--   >>> let b = mk3 1 2 3 :: Vec3 Int
--   >>> element_sum b
--   6
--
element_sum :: (Additive.C a, Ring.C a, Vector v a) => v a -> a
element_sum = foldl (+) (fromInteger 0)