{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
module SC.Defs where

import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Type
import Data.Either (lefts, rights)

import qualified Language.C as C
import Language.C (Name(..))
import SC.Monad


-- ENVIRONMENTS
-- ------------

data AVarEnv env where
    AVENil :: AVarEnv ()
    AVEPush :: ShNames sh -> ANames t -> AVarEnv env -> AVarEnv (env, Array sh t)

aveprj :: AVarEnv env -> Idx env (Array sh t) -> (ShNames sh, ANames t)
aveprj (AVEPush shn n _) ZeroIdx = (shn, n)
aveprj (AVEPush _ _ aenv) (SuccIdx idx) = aveprj aenv idx

data VarEnv env where
    VENil :: VarEnv ()
    VEPush :: Name -> VarEnv env -> VarEnv (env, t)

veprj :: VarEnv env -> Idx env t -> Name
veprj (VEPush n _) ZeroIdx = n
veprj (VEPush _ env) (SuccIdx idx) = veprj env idx


-- IGNORE TUPLES
-- -------------

data ITup s t where
    ITupPair :: ITup s a -> ITup s b -> ITup s (a, b)
    ITupSingle :: s -> ITup s a
    ITupIgnore :: ITup s a

itupfold :: (forall a. f a) -> (forall a. s -> f a) -> (forall a b. f a -> f b -> f (a, b))
         -> ITup s t -> f t
itupfold z _ _ ITupIgnore = z
itupfold _ f _ (ITupSingle x) = f x
itupfold z f g (ITupPair a b) = g (itupfold z f g a) (itupfold z f g b)

itupmap :: (s1 -> s2) -> ITup s1 t -> ITup s2 t
itupmap f = itupfold ITupIgnore (ITupSingle . f) ITupPair

itupList :: ITup s t -> [s]
itupList (ITupPair t1 t2) = itupList t1 ++ itupList t2
itupList (ITupSingle x) = [x]
itupList ITupIgnore = []

data TypedName = TypedName C.Type Name
type Names = ITup TypedName
type ANames = ITup TypedAName

type Exprs = ITup C.Expr

itupEvars :: ITup TypedName t -> Exprs t
itupEvars = itupmap (\(TypedName _ n) -> C.EVar n)

-- Type is the pointer type of the type that this name is supposed to be according to the type index.
data TypedAName = TypedAName C.Type Name

data TupANames t where
    ANPair :: TupANames a -> TupANames b -> TupANames (a, b)
    ANArray :: ShNames sh -> ANames t -> TupANames (Array sh t)
    ANIgnore :: TupANames a

-- Shape names and data array names
tupanamesList :: TupANames t -> ([TypedName], [TypedAName])
tupanamesList an = let l = tupanamesList' an in (lefts l, rights l)

tupanamesList' :: TupANames t -> [Either TypedName TypedAName]
tupanamesList' (ANPair a b) = tupanamesList' a ++ tupanamesList' b
tupanamesList' (ANArray shn ns) =
    map Left (shnamesList shn) ++ map Right (itupList ns)
tupanamesList' ANIgnore = []

data ShNames sh where
    ShZ :: ShNames ()
    ShS :: ShNames sh -> Name -> ShNames (sh, Int)
deriving instance Show (ShNames sh)

shnamesList :: ShNames sh -> [TypedName]
shnamesList = reverse . go
  where
    go :: ShNames sh -> [TypedName]
    go ShZ = []
    go (ShS shns n) = TypedName (C.TInt C.B64) n : go shns

makeShNames :: ShapeR sh -> ITup TypedName sh -> ShNames sh
makeShNames ShapeRz ITupIgnore = ShZ
makeShNames (ShapeRsnoc sht) (ITupPair ns (ITupSingle (TypedName _ n))) =
    ShS (makeShNames sht ns) n
makeShNames _ _ = error "wat"

fromShNames :: ShNames sh -> ITup TypedName sh
fromShNames ShZ = ITupIgnore
fromShNames (ShS ns n) = ITupPair (fromShNames ns) (ITupSingle (TypedName (C.TInt C.B64) n))

shNamesShape :: ShNames sh -> ShapeR sh
shNamesShape ShZ = ShapeRz
shNamesShape (ShS ns _) = ShapeRsnoc (shNamesShape ns)


-- GENERATING VARIABLE NAMES
-- -------------------------

genName :: String -> SC Name
genName prefix = Name . (prefix ++) . show <$> genId


-- TYPE CONVERSION
-- ---------------

cvtType :: ScalarType t -> SC C.Type
cvtType (SingleScalarType (NumSingleType (IntegralNumType it))) = return (cvtIT it)
  where cvtIT :: IntegralType t -> C.Type
        cvtIT TypeInt = C.TInt C.B64
        cvtIT TypeInt8 = C.TInt C.B8
        cvtIT TypeInt16 = C.TInt C.B16
        cvtIT TypeInt32 = C.TInt C.B32
        cvtIT TypeInt64 = C.TInt C.B64
        cvtIT TypeWord = C.TUInt C.B64
        cvtIT TypeWord8 = C.TUInt C.B8
        cvtIT TypeWord16 = C.TUInt C.B16
        cvtIT TypeWord32 = C.TUInt C.B32
        cvtIT TypeWord64 = C.TUInt C.B64
cvtType (SingleScalarType (NumSingleType (FloatingNumType ft))) = cvtFT ft
  where cvtFT :: FloatingType t -> SC C.Type
        cvtFT TypeHalf = throw "Half floats not supported"
        cvtFT TypeFloat = return C.TFloat
        cvtFT TypeDouble = return C.TDouble
cvtType VectorScalarType{} = throw "Vector types not supported"