module Feldspar.Core.Types where



import Data.Bits
import Data.Complex
import Data.Int
import Data.List
import Data.Tagged
import Data.Proxy
import Data.Typeable (Typeable)
import Data.Word

import Feldspar.Set
import Feldspar.Range



--------------------------------------------------------------------------------
-- * Heterogenous lists
--------------------------------------------------------------------------------

-- | Heterogeneous list
data a :> b = a :> b
  deriving (Eq, Ord, Show)

infixr 5 :>

instance (Set a, Set b) => Set (a :> b)
  where
    empty     = empty :> empty
    universal = universal :> universal
    (a1:>a2) \/ (b1:>b2) = (a1 \/ b1) :> (a2 \/ b2)
    (a1:>a2) /\ (b1:>b2) = (a1 /\ b1) :> (a2 /\ b2)



--------------------------------------------------------------------------------
-- * Integers
--------------------------------------------------------------------------------

-- | Platform-independent unsigned integers
newtype DefaultWord = DefaultWord Word32
  deriving (Eq, Ord, Num, Enum, Real, Integral, Bits, Bounded, Typeable)
  -- TODO Find better name

-- | Platform-independent signed integers
newtype DefaultInt = DefaultInt Int32
  deriving (Eq, Ord, Num, Enum, Real, Integral, Bits, Bounded, Typeable)
  -- TODO Find better name

-- TODO Should really be defined as:
--
--     data DefaultWord
--         = DefWord32 Word32
--         | DefWord16 Word16
--
--     data DefaultInt
--         = DefInt32 Int32
--         | DefInt16 Int16

type Length = DefaultWord
type Index  = DefaultWord

instance Show DefaultWord
  where
    show (DefaultWord a) = show a

instance Show DefaultInt
  where
    show (DefaultInt a) = show a

-- | The set of signed integer types
class Signed a

instance Signed Int8
instance Signed Int16
instance Signed Int32
instance Signed DefaultInt



--------------------------------------------------------------------------------
-- * Type/data representation
--------------------------------------------------------------------------------

-- | Representation of types
data TypeRep
  = BoolType
  | forall a . (BoundedInt a, Typeable a) => IntType { intRange :: Range a }
  | FloatType
  | UserType String
  | ComplexType TypeRep
  | ArrayType (Range Length) TypeRep
  | StructType [TypeRep]

-- | Representation of data
data DataRep
  = BoolData Bool
  | IntData Integer
  | FloatData Float
  | ComplexData DataRep DataRep
  | ArrayData [DataRep]
  | StructData [DataRep]
  deriving (Eq, Show)

class (Eq a, Show a, Typeable a, Eq (Size a), Show (Size a), Set (Size a)) => Type a
  where
    type Size a

    dataRep :: a -> DataRep

    -- | Gives the type representation of a storable value.
    typeRep :: Tagged a (Size a) -> TypeRep

    -- | Gives the size of a storable value.
    sizeOf :: a -> Size a

instance Type ()
  where
    type Size () = ()
    dataRep _    = BoolData False
    typeRep _    = BoolType
    sizeOf _     = ()

instance Type Bool
  where
    type Size Bool = ()
    dataRep        = BoolData
    typeRep _      = BoolType
    sizeOf _       = ()

instance Type Word8
  where
    type Size Word8 = Range Word8
    dataRep         = IntData . toInteger
    typeRep         = IntType . untag
    sizeOf a        = singletonRange a

instance Type Int8
  where
    type Size Int8 = Range Int8
    dataRep        = IntData . toInteger
    typeRep        = IntType . untag
    sizeOf a       = singletonRange a

instance Type Word16
  where
    type Size Word16 = Range Word16
    dataRep          = IntData . toInteger
    typeRep          = IntType . untag
    sizeOf a         = singletonRange a

instance Type Int16
  where
    type Size Int16 = Range Int16
    dataRep         = IntData . toInteger
    typeRep         = IntType . untag
    sizeOf a        = singletonRange a

instance Type Word32
  where
    type Size Word32 = Range Word32
    dataRep          = IntData . toInteger
    typeRep          = IntType . untag
    sizeOf a         = singletonRange a

instance Type Int32
  where
    type Size Int32 = Range Int32
    dataRep         = IntData . toInteger
    typeRep         = IntType . untag
    sizeOf a        = singletonRange a

instance Type DefaultWord
  where
    type Size DefaultWord = Range DefaultWord
    dataRep               = IntData . toInteger
    typeRep               = IntType . untag
    sizeOf a              = singletonRange a

instance Type DefaultInt
  where
    type Size DefaultInt = Range DefaultInt
    dataRep              = IntData . toInteger
    typeRep              = IntType . untag
    sizeOf a             = singletonRange a

instance Type Float
  where
    type Size Float = ()
    dataRep         = FloatData
    typeRep _       = FloatType
    sizeOf _        = ()

instance (Type a, RealFloat a) => Type (Complex a)
  -- 'RealFloat' comes from the constraint on the 'Complex' data type. It
  -- implies 'Floating'
  where
    type Size (Complex a) = ()
    dataRep (r :+ i) = ComplexData (dataRep r) (dataRep i)
    typeRep sz = ComplexType $ typeRep (Tagged universal :: Tagged a (Size a))
    sizeOf  _  = ()

instance Type a => Type [a]
  where
    type Size [a]            = Range Length :> Size a
    dataRep as               = ArrayData (map dataRep as)
    typeRep (Tagged (l:>sz)) = ArrayType l (typeRep sz')
      where
        sz' = Tagged sz :: Tagged a (Size a)
    sizeOf as = singletonRange (genericLength as) :> unions (map sizeOf as)

instance (Type a, Type b) => Type (a,b)
  where
    type Size (a,b)            = (Size a, Size b)
    dataRep (a,b)              = StructData [dataRep a, dataRep b]
    typeRep (Tagged (sza,szb)) = StructType [typeRep sza', typeRep szb']
      where
        sza' = Tagged sza :: Tagged a (Size a)
        szb' = Tagged szb :: Tagged b (Size b)
    sizeOf (a,b) = (sizeOf a, sizeOf b)

instance (Type a, Type b, Type c) => Type (a,b,c)
  where
    type Size (a,b,c) = (Size a, Size b, Size c)
    dataRep (a,b,c)   = StructData [dataRep a, dataRep b, dataRep c]
    typeRep (Tagged (sza,szb,szc)) = StructType [typeRep sza', typeRep szb', typeRep szc']
      where
        sza' = Tagged sza :: Tagged a (Size a)
        szb' = Tagged szb :: Tagged b (Size b)
        szc' = Tagged szc :: Tagged c (Size c)
    sizeOf (a,b,c) = (sizeOf a, sizeOf b, sizeOf c)

instance (Type a, Type b, Type c, Type d) => Type (a,b,c,d)
  where
    type Size (a,b,c,d) = (Size a, Size b, Size c, Size d)
    dataRep (a,b,c,d)   = StructData [dataRep a, dataRep b, dataRep c, dataRep d]
    typeRep (Tagged (sza,szb,szc,szd)) = StructType [typeRep sza', typeRep szb', typeRep szc', typeRep szd']
      where
        sza' = Tagged sza :: Tagged a (Size a)
        szb' = Tagged szb :: Tagged b (Size b)
        szc' = Tagged szc :: Tagged c (Size c)
        szd' = Tagged szd :: Tagged d (Size d)
    sizeOf (a,b,c,d) = (sizeOf a, sizeOf b, sizeOf c, sizeOf d)

-- TODO Document
class MetaType role a
  where
    listTypes :: [Int] -> Proxy role -> Proxy a -> [([Int], TypeRep)]

instance Type a => MetaType () a
  where
    listTypes path _ _ =
        [(path, typeRep (Tagged universal :: Tagged a (Size a)))]

instance (MetaType ra a, MetaType rb b) => MetaType (ra,rb) (a,b)
  where
    listTypes path _ _
        =  listTypes (1:path) (Proxy :: Proxy ra) (Proxy :: Proxy a)
        ++ listTypes (2:path) (Proxy :: Proxy rb) (Proxy :: Proxy b)

instance (MetaType ra a, MetaType rb b, MetaType rc c) =>
    MetaType (ra,rb,rc) (a,b,c)
  where
    listTypes path _ _
        =  listTypes (1:path) (Proxy :: Proxy ra) (Proxy :: Proxy a)
        ++ listTypes (2:path) (Proxy :: Proxy rb) (Proxy :: Proxy b)
        ++ listTypes (3:path) (Proxy :: Proxy rc) (Proxy :: Proxy c)

instance (MetaType ra a, MetaType rb b, MetaType rc c, MetaType rd d) =>
    MetaType (ra,rb,rc,rd) (a,b,c,d)
  where
    listTypes path _ _
        =  listTypes (1:path) (Proxy :: Proxy ra) (Proxy :: Proxy a)
        ++ listTypes (2:path) (Proxy :: Proxy rb) (Proxy :: Proxy b)
        ++ listTypes (3:path) (Proxy :: Proxy rc) (Proxy :: Proxy c)
        ++ listTypes (4:path) (Proxy :: Proxy rd) (Proxy :: Proxy d)



-- | A version of 'typeRep' that gets the 'Size' implicitly from the argument.
typeRep' :: forall a . Type a => a -> TypeRep
typeRep' a = typeRep (Tagged (sizeOf a) :: Tagged a (Size a))

isNil :: Type a => a -> Bool
isNil a = case dataRep a of
    ArrayData [] -> True
    _            -> False



--------------------------------------------------------------------------------
-- * Size propagation
--------------------------------------------------------------------------------

class FullProp a
  where
    -- | Size propagation function that maps any number of arguments to
    -- 'universal'.
    fullProp :: a

instance FullProp ()
  where
    fullProp = universal

instance BoundedInt a => FullProp (Range a)
  where
    fullProp = universal

instance FullProp b => FullProp (a -> b)
  where
    fullProp = const fullProp