{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs               #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeInType                 #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances       #-}

module TensorOps.Backend.NTensor
  ( NTensor
  , NTensorL
  , NTensorV
  )
  where

import           Control.DeepSeq
import           Control.Monad.Primitive
import           Data.Distributive
import           Data.Kind
import           Data.List.Util
import           Data.Nested
import           Data.Proxy
import           Data.Singletons
import           Data.Singletons.Prelude.List hiding (Length, Reverse)
import           Data.Type.Combinator
import           Data.Type.Combinator.Util
import           Data.Type.Length                    as TCL
import           Data.Type.Length.Util               as TCL
import           Data.Type.Product                   as TCP
import           Data.Type.Product.Util              as TCP
import           Data.Type.Sing
import           Data.Type.Uniform
import           GHC.Generics
import           Statistics.Distribution
import           System.Random.MWC
import           TensorOps.NatKind
import           TensorOps.Types
import           Type.Class.Higher
import           Type.Class.Higher.Util
import           Type.Class.Witness
import           Type.Family.List
import qualified Data.Type.Vector                    as TCV
import qualified Data.Type.Vector.Util               as TCV
import qualified Data.Vector.Sized                   as VS
import qualified Type.Family.Nat                     as TCN

newtype NTensor :: (k -> Type -> Type) -> Type -> [k] -> Type where
    NTensor
        :: { getNVec :: Nested v js a }
        -> NTensor v a js

deriving instance (Num a, SingI ns, Nesting1 Proxy Functor v, Nesting1 Sing Applicative v)
        => Num (NTensor v a ns)
deriving instance Generic (NTensor v a ns)
deriving instance (NFData a, Nesting Proxy NFData v) => NFData (NTensor v a ns)
instance (NFData a, Nesting Proxy NFData v) => Nesting1 w NFData (NTensor v a) where
    nesting1 _ = Wit

instance (NFData a, Nesting Proxy NFData v) => NFData1 (NTensor v a)

genNTensor
    :: forall k (v :: k -> Type -> Type) (ns :: [k]) (a :: Type). Vec v
    => Sing ns
    -> (Prod (IndexN k) ns -> a)
    -> NTensor v a ns
genNTensor s f = NTensor $ genNested s f
{-# INLINE genNTensor #-}

genNTensorA
    :: forall k (v :: k -> Type -> Type) (ns :: [k]) (a :: Type) f. (Applicative f, Vec v)
    => Sing ns
    -> (Prod (IndexN k) ns -> f a)
    -> f (NTensor v a ns)
genNTensorA s f = NTensor <$> genNestedA s f
{-# INLINE genNTensorA #-}

indexNTensor
    :: forall k (v :: k -> Type -> Type) (ns :: [k]) (a :: Type). Vec v
    => Prod (IndexN k) ns
    -> NTensor v a ns
    -> a
indexNTensor i = indexNested i . getNVec
{-# INLINE indexNTensor #-}

overNVec2
    :: (Nested v ns a -> Nested w ms b -> Nested u os c)
    -> NTensor v a ns
    -> NTensor w b ms
    -> NTensor u c os
overNVec2 f x y = NTensor $ f (getNVec x) (getNVec y)
{-# INLINE overNVec2 #-}

ntNVec
    :: Functor f
    => (Nested v ns a -> f (Nested w ms b))
    -> NTensor v a ns
    -> f (NTensor w b ms)
ntNVec f = fmap NTensor . f . getNVec
{-# INLINE ntNVec #-}

nvecNT
    :: Functor f
    => (NTensor v a ns -> f (NTensor w b ms))
    -> Nested v ns a
    -> f (Nested w ms b)
nvecNT f = fmap getNVec . f . NTensor
{-# INLINE nvecNT #-}

overNVec
    :: (Nested v ns a -> Nested v ms a)
    -> NTensor v a ns
    -> NTensor v a ms
overNVec f = getI . ntNVec (I . f)
{-# INLINE overNVec #-}

underNVec
    :: (NTensor v a ns -> NTensor v a ms)
    -> Nested v ns a
    -> Nested v ms a
underNVec f = getNVec . f . NTensor


instance
      ( Vec (v :: k -> Type -> Type)
      , RealFloat a
      , Nesting1 Proxy Functor      v
      , Nesting1 Sing  Applicative  v
      , Nesting1 Proxy Foldable     v
      , Nesting1 Proxy Traversable  v
      , Nesting1 Sing  Distributive v
      , Eq1 (IndexN k)
      ) => Tensor (NTensor v a) where
    type ElemT  (NTensor v a) = a

    liftT
        :: forall (n :: TCN.N) (o :: [k]). SingI o
        => (TCV.Vec n a -> a)
        -> TCV.Vec n (NTensor v a o)
        -> NTensor v a o
    liftT f = NTensor . liftNested f . fmap getNVec
    {-# INLINE liftT #-}

    sumT = sum'
    {-# INLINE sumT #-}

    scaleT α = overNVec (fmap (α *))
    {-# INLINE scaleT #-}

    transp
        :: forall ns. (SingI ns, SingI (Reverse ns))
        => NTensor v a ns
        -> NTensor v a (Reverse ns)
    -- transp = overNVec (transpose sing)
    transp = overNVec (transpose' (singLength sing) sing)
    {-# INLINE transp #-}

    gmul
        :: forall ms os ns. SingI (Reverse os ++ ns)
        => Length ms
        -> Length os
        -> Length ns
        -> NTensor v a (ms         ++ os)
        -> NTensor v a (Reverse os ++ ns)
        -> NTensor v a (ms         ++ ns)
    gmul lM lO lN = overNVec2 (gmul' lM lO lN)
                      \\ sO'
                      \\ sN
      where
        sO' :: Sing (Reverse os)
        sN  :: Sing ns
        (sO', sN) = splitSing (TCL.reverse' lO)
                              (sing :: Sing (Reverse os ++ ns))
    {-# INLINE gmul #-}

    -- TODO: this is a dumb implementation
    diag
        :: forall n ns. SingI (n ': ns)
        => Uniform n ns
        -> NTensor v a '[n]
        -> NTensor v a (n ': ns)
    diag u d
        = genNTensor sing (\i -> case TCV.uniformVec (prodToVec I (US u) i) of
                                   Nothing     -> 0
                                   Just (I i') -> indexNTensor (i' :< Ø) d
                          )
            \\ (produceEq1 :: Eq1 (IndexN k) :- Eq (IndexN k n))
    {-# INLINE diag #-}

    -- TODO: this can be just done using genNTensor
    getDiag
        :: forall n ns. SingI '[n]
        => Uniform n ns
        -> NTensor v a (n ': n ': ns)
        -> NTensor v a '[n]
    getDiag u = overNVec (diagNV sing u)
                  \\ sHead (sing :: Sing '[n])
    {-# INLINE getDiag #-}

    genRand
        :: forall m d (ns :: [k]). (ContGen d, PrimMonad m, SingI ns)
        => d
        -> Gen (PrimState m)
        -> m (NTensor v a ns)
    genRand d g = generateA (\_ -> realToFrac <$> genContVar d g)
    {-# INLINE genRand #-}

    generateA
        :: forall f ns. (Applicative f, SingI ns)
        => (Prod (IndexN k) ns -> f a)
        -> f (NTensor v a ns)
    generateA = genNTensorA sing
    {-# INLINE generateA #-}

    (!) = flip indexNTensor
    {-# INLINE (!) #-}

    ixRows
        :: Applicative f
        => Length ms
        -> Length os
        -> (Prod (IndexN k) ms -> NTensor v a ns -> f (NTensor v a os))
        -> NTensor v a (ms ++ ns)
        -> f (NTensor v a (ms ++ os))
    ixRows l _ f = ntNVec $ fmap joinNested . nIxRows l (\i -> nvecNT (f i))
    {-# INLINE ixRows #-}

    sumRows
        :: forall n ns. SingI ns
        => NTensor v a (n ': ns)
        -> NTensor v a ns
    sumRows = overNVec sumRowsNested
                \\ (nesting1 Proxy :: Wit (Foldable (v n)))
    {-# INLINE sumRows #-}

    mapRows
        :: Length ns
        -> (NTensor v a ms -> NTensor v a ms)
        -> NTensor v a (ns ++ ms)
        -> NTensor v a (ns ++ ms)
    mapRows l f = overNVec $ joinNested . mapNVecSlices (underNVec f) l
    {-# INLINE mapRows #-}

type NTensorL = NTensor (Flip2 TCV.VecT   I) Double
type NTensorV = NTensor (Flip2 VS.VectorT I) Double
