content: {- | Implementation of mathematical vector.
  
  'Vec' is only a newtype wrapper around 'Data.Vector', which
  implements several mathematical operations on itself.
  
  'Vec' offers meaningful abstraction and easy interface
  (you can unwrap it to perform more complex tasks).
  -}
  
  
  -- 'TypeFamilies' are needed to instantiate 'Indexable', 'ElementwiseScalarOps', 'SingletonOps', 'VecOps' typeclasses.
  {-# LANGUAGE TypeFamilies #-}
  
  
  module Synapse.Tensors.Vec
      ( -- * 'Vec' datatype and simple getters.
  
        Vec (Vec, unVec)
  
      , size
  
        -- * Constructors
  
      , empty
      , singleton
      , fromList
      , generate
      , replicate
  
        -- * Concatenation and splitting
  
      , cons
      , snoc
      , (++)
      , concat
  
      , splitAt
  
        -- * Combining
  
      , map
      , imap
      , for
      , zipWith
      , zip
  
      -- * Mathematics
  
      , zeroes
      , ones
  
      , squaredMagnitude
      , magnitude
      , clampMagnitude
      , normalized
  
      , linearCombination
      , dot
  
      , angleBetween
      , lerp
      ) where
  
  
  import Synapse.Tensors (DType, Indexable(..), ElementwiseScalarOps(..), SingletonOps(..), VecOps(..))
  
  import Prelude hiding ((++), concat, splitAt, map, replicate, zip, zipWith)
  import Data.Foldable (Foldable(..))
  import Data.Ord (clamp)
  
  import qualified Data.Vector as V
  
  
  -- | Mathematical vector (collection of elements).
  newtype Vec a = Vec
      { unVec :: V.Vector a  -- ^ Internal representation.
      } deriving (Eq, Read)
  
  -- | Size of a vector - number of elements.
  size :: Vec a -> Int
  size = V.length . unVec
  
  
  -- Typeclasses
  
  instance Show a => Show (Vec a) where
      show (Vec x) = show x
  
  
  type instance DType (Vec a) = a
  
  
  instance Indexable (Vec a) where
      type Index (Vec a) = Int
  
      unsafeIndex (Vec x) = V.unsafeIndex x
      (!) (Vec x) = (V.!) x
      (!?) (Vec x) = (V.!?) x
  
  
  instance Num a => Num (Vec a) where
      (+) = zipWith (+)
      (-) = zipWith (-)
      negate = fmap (0 -)
      (*) = zipWith (*)
      abs = fmap abs
      signum = fmap signum
      fromInteger = singleton . fromInteger
  
  instance Fractional a => Fractional (Vec a) where
      (/) = zipWith (/)
      recip = fmap (1/)
      fromRational = singleton . fromRational
  
  instance Floating a => Floating (Vec a) where
      pi = singleton pi
      (**) = zipWith (**)
      sqrt = fmap sqrt
      exp = fmap exp
      log = fmap log
      sin = fmap sin
      cos = fmap cos
      asin = fmap asin
      acos = fmap acos
      atan = fmap atan
      sinh = fmap sinh
      cosh = fmap cosh
      asinh = fmap asinh
      acosh = fmap acosh
      atanh = fmap atanh
  
  instance ElementwiseScalarOps (Vec a) where
      (+.) x n = fmap (+ n) x
      (-.) x n = fmap (subtract n) x
      (*.) x n = fmap (* n) x
      (/.) x n = fmap (/ n) x
      (**.) x n = fmap (** n) x
  
      elementsMin x n = fmap (min n) x
      elementsMax x n = fmap (max n) x
  
  instance SingletonOps (Vec a) where
      singleton = pure
      isSingleton vec = size vec == 1
      unSingleton vec
          | not $ isSingleton vec = error "Vector is not a singleton"
          | otherwise             = unsafeIndex vec 0
  
      extendSingleton vec reference = replicate (size reference) (unSingleton vec)
  
      elementsSum = singleton . V.sum . unVec
      elementsProduct = singleton . V.product . unVec
      mean x = elementsSum x /. fromIntegral (size x)
      norm x = sqrt $ elementsSum $ x * x
  
  instance Functor Vec where
      fmap f = Vec . V.map f . unVec
      (<$) = fmap . const
  
  instance Applicative Vec where
      pure = Vec . V.singleton
      (<*>) = zipWith (\f x -> f x)
  
  instance Foldable Vec where
      foldr f x = V.foldr f x . unVec
      foldl f x = V.foldl f x . unVec
  
      foldr' f x = V.foldr' f x . unVec
      foldl' f x = V.foldl' f x . unVec
  
      foldr1 f = V.foldr1 f . unVec
      foldl1 f = V.foldl1 f . unVec
  
      toList = V.toList . unVec
  
      null x = size x == 0
  
      length = size
  
  instance Traversable Vec where
      traverse f (Vec x) = Vec <$> traverse f x
  
  
  -- Constructors
  
  -- | Creates empty 'Vec'.
  empty :: Vec a
  empty = Vec V.empty
  
  -- | Creates 'Vec' from list.
  fromList :: [a] -> Vec a
  fromList = Vec . V.fromList
  
  -- | Creates 'Vec' of given length using generating function.
  generate :: Int -> (Int -> a) -> Vec a
  generate n = Vec . V.generate n
  
  -- | Creates 'Vec' of given length filled with given element.
  replicate :: Int -> a -> Vec a
  replicate n = generate n . const
  
  
  -- Concatenation and splitting
  
  -- | Prepend 'Vec' with given element.
  cons :: a -> Vec a -> Vec a
  cons x = Vec . V.cons x . unVec
  
  -- | Append 'Vec' with given element.
  snoc :: Vec a -> a -> Vec a
  snoc (Vec vec) x = Vec $ V.snoc vec x
  
  -- | Concatenate two 'Vec's.
  infixr 5 ++
  (++) :: Vec a -> Vec a -> Vec a
  (++) (Vec x) (Vec y) = Vec $ x V.++ y
  
  -- | Concatenate all 'Vec's.
  concat :: [Vec a] -> Vec a
  concat = foldr1 (++)
  
  -- | Splits 'Vec' into two 'Vec's at a given index.
  splitAt :: Int -> Vec a -> (Vec a, Vec a)
  splitAt i (Vec v) = let (v1, v2) = V.splitAt i v
                      in (Vec v1, Vec v2)
  
  
  -- Combining
  
  -- | Map a function over a 'Vec'.
  map :: (a -> b) -> Vec a -> Vec b
  map = fmap
  
  -- | Apply a function to every element of a 'Vec' and its index.
  imap :: (Int -> a -> b) -> Vec a -> Vec b
  imap f = Vec . V.imap f . unVec
  
  -- | 'map' with its arguments flipped.
  for :: Vec a -> (a -> b) -> Vec b
  for = flip fmap
  
  -- | Zips two 'Vec's with the given function.
  zipWith :: (a -> b -> c) -> Vec a -> Vec b -> Vec c
  zipWith f (Vec a) (Vec b) = Vec $ V.zipWith f a b
  
  -- | Zips two 'Vec's.
  zip :: Vec a -> Vec b -> Vec (a, b)
  zip = zipWith (,)
  
  
  -- Functions that work on mathematical vector (type constraint refers to a number)
  
  -- | Creates 'Vec' of given length filled with zeroes.
  zeroes :: Num a => Int -> Vec a
  zeroes = flip generate (const 0)
  
  -- | Creates 'Vec' of given length filled with ones.
  ones :: Num a => Int -> Vec a
  ones = flip generate (const 1)
  
  
  -- | Squared magnitude of a 'Vec'.
  squaredMagnitude :: Num a => Vec a -> a
  squaredMagnitude x = sum (fmap (^ (2 :: Int)) x)
  
  -- | Magnitude of a 'Vec'.
  magnitude :: Floating a => Vec a -> a
  magnitude = sqrt . squaredMagnitude
  
  -- | Clamps 'Vec' magnitude.
  clampMagnitude :: (Floating a, Ord a) => a -> Vec a -> Vec a
  clampMagnitude m x = x *. (min (magnitude x) m / magnitude x)
  
  -- | Normalizes 'Vec' by dividing each component by 'Vec' magnitude.
  normalized :: Floating a => Vec a -> Vec a
  normalized x = x /. magnitude x
  
  
  -- | Computes linear combination of 'Vec's. Returns empty 'Vec' if empty list was passed to this function.
  linearCombination :: Num a => [(a, Vec a)] -> Vec a
  linearCombination [] = empty
  linearCombination (x:xs) = foldl' (\acc (a, vec) -> acc + vec *. a) (snd x *. fst x) xs
  
  
  instance Num a => VecOps (Vec a) where
      dot a b = elementsSum $ a * b
  
  
  -- | Calculates an angle between two 'Vec's.
  angleBetween :: Floating a => Vec a -> Vec a -> a
  angleBetween a b = acos $ unSingleton (a `dot` b) / (magnitude a * magnitude b)
  
  -- | Linearly interpolates between two 'Vec's. Given parameter will be clamped between [0.0, 1.0].
  lerp :: (Floating a, Ord a) => a -> Vec a -> Vec a -> Vec a
  lerp k a b = b - (b - a) *. clamp (0.0, 1.0) k
  