{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Array where

import qualified Data.Array.RankedU as U
import Data.Kind
import Data.Proxy
import Data.Type.Equality
import qualified Data.Vector.Unboxed as VU
import Unsafe.Coerce (unsafeCoerce)

import Nats


type family l1 ++ l2 where
  '[] ++ l2 = l2
  (x : xs) ++ l2 = x : xs ++ l2

lemAppNil :: l ++ '[] :~: l
lemAppNil = unsafeCoerce Refl

lemAppAssoc :: Proxy a -> Proxy b -> Proxy c -> (a ++ b) ++ c :~: a ++ (b ++ c)
lemAppAssoc _ _ _ = unsafeCoerce Refl


type IxX :: [Maybe Nat] -> Type
data IxX sh where
  IZX :: IxX '[]
  (::@) :: Int -> IxX sh -> IxX (Just n : sh)
  (::?) :: Int -> IxX sh -> IxX (Nothing : sh)
deriving instance Show (IxX sh)

type StaticShapeX :: [Maybe Nat] -> Type
data StaticShapeX sh where
  SZX :: StaticShapeX '[]
  (:$@) :: SNat n -> StaticShapeX sh -> StaticShapeX (Just n : sh)
  (:$?) :: () -> StaticShapeX sh -> StaticShapeX (Nothing : sh)

type KnownShapeX :: [Maybe Nat] -> Constraint
class KnownShapeX sh where
  knownShapeX :: StaticShapeX sh
instance KnownShapeX '[] where
  knownShapeX = SZX
instance (KnownNat n, KnownShapeX sh) => KnownShapeX (Just n : sh) where
  knownShapeX = knownNat :$@ knownShapeX
instance KnownShapeX sh => KnownShapeX (Nothing : sh) where
  knownShapeX = () :$? knownShapeX

type family Rank sh where
  Rank '[] = Z
  Rank (_ : sh) = S (Rank sh)

type XArray :: [Maybe Nat] -> Type -> Type
data XArray sh a = XArray (U.Array (GNat (Rank sh)) a)

zeroIdx :: StaticShapeX sh -> IxX sh
zeroIdx SZX = IZX
zeroIdx (_ :$@ ssh) = 0 ::@ zeroIdx ssh
zeroIdx (_ :$? ssh) = 0 ::? zeroIdx ssh

zeroIdx' :: IxX sh -> IxX sh
zeroIdx' IZX = IZX
zeroIdx' (_ ::@ sh) = 0 ::@ zeroIdx' sh
zeroIdx' (_ ::? sh) = 0 ::? zeroIdx' sh

ixAppend :: IxX sh -> IxX sh' -> IxX (sh ++ sh')
ixAppend IZX idx' = idx'
ixAppend (i ::@ idx) idx' = i ::@ ixAppend idx idx'
ixAppend (i ::? idx) idx' = i ::? ixAppend idx idx'

ixDrop :: IxX (sh ++ sh') -> IxX sh -> IxX sh'
ixDrop sh IZX = sh
ixDrop (_ ::@ sh) (_ ::@ idx) = ixDrop sh idx
ixDrop (_ ::? sh) (_ ::? idx) = ixDrop sh idx

ssxAppend :: StaticShapeX sh -> StaticShapeX sh' -> StaticShapeX (sh ++ sh')
ssxAppend SZX idx' = idx'
ssxAppend (n :$@ idx) idx' = n :$@ ssxAppend idx idx'
ssxAppend (() :$? idx) idx' = () :$? ssxAppend idx idx'

shapeSize :: IxX sh -> Int
shapeSize IZX = 1
shapeSize (n ::@ sh) = n * shapeSize sh
shapeSize (n ::? sh) = n * shapeSize sh

fromLinearIdx :: IxX sh -> Int -> IxX sh
fromLinearIdx = \sh i -> case go sh i of
  (idx, 0) -> idx
  _ -> error $ "fromLinearIdx: out of range (" ++ show i ++
               " in array of shape " ++ show sh ++ ")"
  where
    -- returns (index in subarray, remaining index in enclosing array)
    go :: IxX sh -> Int -> (IxX sh, Int)
    go IZX i = (IZX, i)
    go (n ::@ sh) i =
      let (idx, i') = go sh i
          (upi, locali) = i' `quotRem` n
      in (locali ::@ idx, upi)
    go (n ::? sh) i =
      let (idx, i') = go sh i
          (upi, locali) = i' `quotRem` n
      in (locali ::? idx, upi)

toLinearIdx :: IxX sh -> IxX sh -> Int
toLinearIdx = \sh i -> fst (go sh i)
  where
    -- returns (index in subarray, size of subarray)
    go :: IxX sh -> IxX sh -> (Int, Int)
    go IZX IZX = (0, 1)
    go (n ::@ sh) (i ::@ ix) =
      let (lidx, sz) = go sh ix
      in (sz * i + lidx, n * sz)
    go (n ::? sh) (i ::? ix) =
      let (lidx, sz) = go sh ix
      in (sz * i + lidx, n * sz)

enumShape :: IxX sh -> [IxX sh]
enumShape = \sh -> go 0 sh id []
  where
    go :: Int -> IxX sh -> (IxX sh -> a) -> [a] -> [a]
    go _ IZX _ = id
    go i (n ::@ sh) f
      | i < n = go (i + 1) (n ::@ sh) f . go 0 sh (f . (i ::@))
      | otherwise = id
    go i (n ::? sh) f
      | i < n = go (i + 1) (n ::? sh) f . go 0 sh (f . (i ::?))
      | otherwise = id

shapeLshape :: IxX sh -> U.ShapeL
shapeLshape IZX = []
shapeLshape (n ::@ sh) = n : shapeLshape sh
shapeLshape (n ::? sh) = n : shapeLshape sh

lemKnownNatRank :: IxX sh -> Dict KnownNat (Rank sh)
lemKnownNatRank IZX = Dict
lemKnownNatRank (_ ::@ sh) | Dict <- lemKnownNatRank sh = Dict
lemKnownNatRank (_ ::? sh) | Dict <- lemKnownNatRank sh = Dict

lemKnownShapeX :: StaticShapeX sh -> Dict KnownShapeX sh
lemKnownShapeX SZX = Dict
lemKnownShapeX (n :$@ ssh) | Dict <- lemKnownShapeX ssh, Dict <- snatKnown n = Dict
lemKnownShapeX (() :$? ssh) | Dict <- lemKnownShapeX ssh = Dict

lemAppKnownShapeX :: StaticShapeX sh1 -> StaticShapeX sh2 -> Dict KnownShapeX (sh1 ++ sh2)
lemAppKnownShapeX SZX ssh' = lemKnownShapeX ssh'
lemAppKnownShapeX (n :$@ ssh) ssh'
  | Dict <- lemAppKnownShapeX ssh ssh'
  , Dict <- snatKnown n
  = Dict
lemAppKnownShapeX (() :$? ssh) ssh'
  | Dict <- lemAppKnownShapeX ssh ssh'
  = Dict

shape :: forall sh a. KnownShapeX sh => XArray sh a -> IxX sh
shape (XArray arr) = go (knownShapeX @sh) (U.shapeL arr)
  where
    go :: StaticShapeX sh' -> [Int] -> IxX sh'
    go SZX [] = IZX
    go (n :$@ ssh) (_ : l) = fromIntegral (unSNat n) ::@ go ssh l
    go (() :$? ssh) (n : l) = n ::? go ssh l
    go _ _ = error "Invalid shapeL"

fromVector :: forall sh a. U.Unbox a => IxX sh -> VU.Vector a -> XArray sh a
fromVector sh v
  | Dict <- lemKnownNatRank sh
  , Dict <- gknownNat (Proxy @(Rank sh))
  = XArray (U.fromVector (shapeLshape sh) v)

toVector :: U.Unbox a => XArray sh a -> VU.Vector a
toVector (XArray arr) = U.toVector arr

generate :: U.Unbox a => IxX sh -> (IxX sh -> a) -> XArray sh a
generate sh f = fromVector sh $ VU.generate (shapeSize sh) (f . fromLinearIdx sh)

-- generateM :: (Monad m, U.Unbox a) => IxX sh -> (IxX sh -> m a) -> m (XArray sh a)
-- generateM sh f | Dict <- lemKnownNatRank sh =
--   XArray . U.fromVector (shapeLshape sh)
--     <$> VU.generateM (shapeSize sh) (f . fromLinearIdx sh)

indexPartial :: U.Unbox a => XArray (sh ++ sh') a -> IxX sh -> XArray sh' a
indexPartial (XArray arr) IZX = XArray arr
indexPartial (XArray arr) (i ::@ idx) = indexPartial (XArray (U.index arr i)) idx
indexPartial (XArray arr) (i ::? idx) = indexPartial (XArray (U.index arr i)) idx

index :: forall sh a. U.Unbox a => XArray sh a -> IxX sh -> a
index xarr i
  | Refl <- lemAppNil @sh
  = let XArray arr' = indexPartial xarr i :: XArray '[] a
    in U.unScalar arr'