{-# LANGUAGE TypeFamilies          #-}

{-# LANGUAGE MultiParamTypeClasses #-}

{-# LANGUAGE TypeSynonymInstances  #-}           

{-# LANGUAGE DataKinds             #-}

{-# LANGUAGE KindSignatures        #-}

{-# LANGUAGE TypeOperators         #-}

{-# LANGUAGE FlexibleInstances     #-}

{-# LANGUAGE FlexibleContexts      #-}           

{-# LANGUAGE ScopedTypeVariables   #-}

{-# LANGUAGE DeriveGeneric         #-}

{-# LANGUAGE DeriveAnyClass        #-}

{-# LANGUAGE TemplateHaskell       #-}

{-# LANGUAGE StandaloneDeriving    #-}           

{-# LANGUAGE UndecidableInstances  #-}

{-# LANGUAGE TypeApplications      #-}

{-# LANGUAGE TypeInType            #-}

{-# LANGUAGE AllowAmbiguousTypes   #-}

{-# LANGUAGE ConstraintKinds       #-}



-----------------------------------------------------------------------------

-- |

-- Module      :  Data.Vector.Static

-- Copyright   :  (C) 2017 Alexey Vagarenko

-- License     :  BSD-style (see LICENSE)

-- Maintainer  :  Alexey Vagarenko (vagarenko@gmail.com)

-- Stability   :  experimental

-- Portability :  non-portable

--

----------------------------------------------------------------------------



module Data.Vector.Static (

    -- * Vector types

      Vector

    , VectorConstructor

    , IsVector

    -- ** Construction

    , vector

    -- ** Vector Operations

    , normalize

    , NormalizedVector

    , Normalize

    , unNormalizedVector

    , vectorLenSquare

    , VectorLenSquare

    , vectorLen

    , VectorLen

    , dot

    , Dot

    , cross

    -- ** Generating vector instances

    , genVectorInstance

) where

  

import Data.Containers       (MonoZip(..))

import Data.MonoTraversable  (omap, osum)

import Data.Tensor.Static    ( IsTensor(..), scale, Scale, TensorConstructor, withTensor

                             , MonoFunctorCtx, MonoFoldableCtx, MonoZipCtx)

import Data.Tensor.Static.TH (genTensorInstance)

import GHC.Generics          (Generic)

import GHC.TypeLits          (Nat)

import Language.Haskell.TH   (Q, Name, Dec)

import qualified Data.List.NonEmpty as N



---------------------------------------------------------------------------------------------------

-- | N-dimensional vector.

type Vector n e = Tensor '[n] e



-- | Type of vector data constructor.

type VectorConstructor n e = TensorConstructor '[n] e



-- | Vector constraint.

type IsVector n e = IsTensor '[n] e



-- | Normalized vector.

newtype NormalizedVector n e =

    NormalizedVector

        { unNormalizedVector :: Vector n e

            -- ^ unwrap 'NormalizedVector'. Note: this does not give you original vector back.

            --

            -- @unNormalizedVector . normalize /= id@

            --

        }

    deriving (Generic)

    

deriving instance (Eq   (Vector n e)) => Eq   (NormalizedVector n e)

deriving instance (Show (Vector n e)) => Show (NormalizedVector n e)



---------------------------------------------------------------------------------------------------

-- | Alias for a conrete vector data constructor.

vector :: forall n e. (IsVector n e) => VectorConstructor n e

vector = tensor @'[n] @e

{-# INLINE vector #-}



-- | Get square of length of a vector.

vectorLenSquare :: (VectorLenSquare n e) => Vector n e -> e

vectorLenSquare = osum . omap (\x -> x * x)

{-# INLINE vectorLenSquare #-}



-- | Constraints for 'vectorLenSquare' function.

type VectorLenSquare (n :: Nat) e =

    ( Num e

    , IsVector n e

    , MonoFunctorCtx '[n] e

    , MonoFoldableCtx '[n] e

    )



-- | Get length of a vector.    

vectorLen :: (VectorLen n e) => Vector n e -> e

vectorLen = sqrt . vectorLenSquare

{-# INLINE vectorLen #-}



-- | Constraints for 'vectorLen' function.

type VectorLen (n :: Nat) e =

    ( Floating e

    , VectorLenSquare n e

    )



-- | Normalize vector.

normalize :: (Normalize n e) => Vector n e -> NormalizedVector n e

normalize v = NormalizedVector $ scale v (1 / vectorLen v)

{-# INLINE normalize #-}



-- | Constraints for 'normalize' function.

type Normalize (n :: Nat) e =

    ( VectorLen n e

    , Scale '[n] e

    )



-- | Dot product of two vectors.

dot :: (Dot n e) => Vector n e -> Vector n e -> e

dot v1 v2 = osum $ ozipWith (*) v1 v2

{-# INLINE dot #-}



-- | Constraints for 'dot' function.

type Dot (n :: Nat) e =

    ( Num e

    , IsVector n e

    , MonoFunctorCtx '[n] e

    , MonoFoldableCtx '[n] e

    , MonoZipCtx '[n] e

    )



---------------------------------------------------------------------------------------------------

-- | Cross product is only defined for 3-dimensional vectors.

cross :: (Num e, IsVector 3 e) => Vector 3 e -> Vector 3 e -> Vector 3 e

cross t0 t1 =

    withTensor t0 $ \x0 y0 z0 ->

        withTensor t1 $ \x1 y1 z1 ->

            vector @3 (y0 * z1 - z0 * y1) (z0 * x1 - x0 * z1) (x0 * y1 - y0 * x1)

{-# INLINE cross #-}



---------------------------------------------------------------------------------------------------

-- | Generate instance of a vector.

genVectorInstance :: Int       -- ^ Size of the vector.

                  -> Name      -- ^ Type of elements.

                  -> Q [Dec]

genVectorInstance size elemTypeName = genTensorInstance (N.fromList [size]) elemTypeName