{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Interpreter.Rep where

import Data.List (intersperse, intercalate)
import Data.Foldable (toList)
import Data.IORef
import GHC.TypeError

import Array
import AST
import AST.Pretty
import Data


type family Rep t where
  Rep TNil = ()
  Rep (TPair a b) = (Rep a, Rep b)
  Rep (TEither a b) = Either (Rep a) (Rep b)
  Rep (TMaybe t) = Maybe (Rep t)
  Rep (TArr n t) = Array n (Rep t)
  Rep (TScal sty) = ScalRep sty
  Rep (TAccum t) = RepAcSparse t

-- Mutable, and has a zero. The zero may not be O(1), but RepAcSparse (D2 t) will have an O(1) zero.
type family RepAcSparse t where
  RepAcSparse TNil = ()
  RepAcSparse (TPair a b) = IORef (RepAcSparse a, RepAcSparse b)
  RepAcSparse (TEither a b) = TypeError (Text "Non-sparse coproduct is not a monoid")
  RepAcSparse (TMaybe t) = IORef (Maybe (RepAcDense t))  -- allow the value to be dense, because the Maybe's zero can be used for the contents
  -- TODO: an empty array is invalid for a zero-dimensional array, so zero-dimensional arrays don't actually have an O(1) zero.
  RepAcSparse (TArr n t) = IORef (Array n (RepAcSparse t))  -- empty array is zero
  RepAcSparse (TScal sty) = IORef (ScalRep sty)
  RepAcSparse (TAccum t) = TypeError (Text "RepAcSparse: Nested accumulators")

-- Immutable, and does not necessarily have a zero.
type family RepAcDense t where
  RepAcDense TNil = ()
  -- RepAcDense (TPair a b) = (RepAcSparse a, RepAcSparse b)
  RepAcDense (TEither a b) = Either (RepAcSparse a) (RepAcSparse b)
  -- RepAcDense (TMaybe t) = RepAcSparse (TMaybe t)  -- ^ This can be optimised to TMaybe (RepAcSparse t), but that makes accumAddDense very hard to write. And in any case, we don't need it because D2 will not produce Maybe of Maybe.
  -- RepAcDense (TArr n t) = Array n (RepAcSparse t)
  -- RepAcDense (TScal sty) = ScalRep sty
  -- RepAcDense (TAccum t) = TypeError (Text "RepAcDense: Nested accumulators")

newtype Value t = Value { unValue :: Rep t }

liftV :: (Rep a -> Rep b) -> Value a -> Value b
liftV f (Value x) = Value (f x)

liftV2 :: (Rep a -> Rep b -> Rep c) -> Value a -> Value b -> Value c
liftV2 f (Value x) (Value y) = Value (f x y)

vPair :: Value a -> Value b -> Value (TPair a b)
vPair = liftV2 (,)

vUnpair :: Value (TPair a b) -> (Value a, Value b)
vUnpair (Value (x, y)) = (Value x, Value y)

showValue :: Int -> STy t -> Rep t -> ShowS
showValue _ STNil () = showString "()"
showValue _ (STPair a b) (x, y) = showString "(" . showValue 0 a x . showString "," . showValue 0 b y . showString ")"
showValue d (STEither a _) (Left x) = showParen (d > 10) $ showString "Left " . showValue 11 a x
showValue d (STEither _ b) (Right y) = showParen (d > 10) $ showString "Right " . showValue 11 b y
showValue _ (STMaybe _) Nothing = showString "Nothing"
showValue d (STMaybe t) (Just x) = showParen (d > 10) $ showString "Just " . showValue 11 t x
showValue d (STArr _ t) arr = showParen (d > 10) $
  showString "arrayFromList " . showsPrec 11 (arrayShape arr)
  . showString " ["
  . foldr (.) id (intersperse (showString ",") $ map (showValue 0 t) (toList arr))
  . showString "]"
showValue _ (STScal sty) x = case sty of
  STF32 -> shows x
  STF64 -> shows x
  STI32 -> shows x
  STI64 -> shows x
  STBool -> shows x
showValue _ (STAccum t) _ = showString $ "<accumulator for " ++ ppTy 0 t ++ ">"

showEnv :: SList STy env -> SList Value env -> String
showEnv = \env vals -> "[" ++ intercalate ", " (showEntries env vals) ++ "]"
  where
    showEntries :: SList STy env -> SList Value env -> [String]
    showEntries SNil SNil = []
    showEntries (t `SCons` env) (Value x `SCons` xs) = showValue 0 t x "" : showEntries env xs