{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeApplications #-}
module Compile where

import Control.Monad.Trans.State.Strict
import Data.Foldable (toList)
import Data.Functor.Const
import qualified Data.Functor.Product as Product
import Data.List (intersperse, intercalate)
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)

import AST
import AST.Pretty (ppTy)
import Data


-- In shape and index arrays, the innermost dimension is on the right (last index).


data StructDecl = StructDecl
    String  -- ^ name
    String  -- ^ contents
    String  -- ^ comment
  deriving (Show)

data Stmt
  = SVarDecl Bool String String CExpr  -- ^ const, type, variable name, right-hand side
  | SVarDeclUninit String String  -- ^ type, variable name (no initialiser)
  | SAsg String CExpr  -- ^ variable name, right-hand side
  | SBlock [Stmt]
  | SIf CExpr [Stmt] [Stmt]
  | SVerbatim String
  deriving (Show)

data CExpr
  = CELit String
  | CEStruct String [(String, CExpr)]
  | CEProj CExpr String
  | CECall String [CExpr]
  | CEBinop CExpr String CExpr
  | CEIf CExpr CExpr CExpr
  deriving (Show)

printStructDecl :: StructDecl -> ShowS
printStructDecl (StructDecl name contents comment) =
  showString "typedef struct { " . showString contents . showString " } " . showString name
  . showString (";  // " ++ comment)

printStmt :: Int -> Stmt -> ShowS
printStmt indent = \case
  SVarDecl cnst typ name rhs -> showString ((if cnst then "const " else "") ++ typ ++ " " ++ name ++ " = ") . printCExpr rhs . showString ";"
  SVarDeclUninit typ name -> showString (typ ++ " " ++ name ++ ";")
  SAsg name rhs -> showString (name ++ " = ") . printCExpr rhs . showString ";"
  SBlock stmts ->
    showString "{"
    . compose [showString ("\n" ++ replicate (2*indent+2) ' ') . printStmt (indent+1) stmt | stmt <- stmts]
    . showString ("\n" ++ replicate (2*indent) ' ' ++ "}")
  SIf cond b1 b2 ->
    showString "if (" . printCExpr cond . showString ") "
    . printStmt indent (SBlock b1) . showString " else " . printStmt indent (SBlock b2)
  SVerbatim s -> showString s

printCExpr :: CExpr -> ShowS
printCExpr = \case
  CELit s -> showString s
  CEStruct name pairs ->
    showString ("(" ++ name ++ "){")
    . compose (intersperse (showString ", ") [showString ("." ++ n ++ " = ") . printCExpr e
                                             | (n, e) <- pairs])
    . showString "}"
  CEProj e name -> showString "(" . printCExpr e . showString (")." ++ name)
  CECall n es ->
    showString (n ++ "(") . compose (intersperse (showString ", ") (map printCExpr es)) . showString ")"
  CEBinop e1 n e2 ->
    showString "(" . printCExpr e1 . showString (") " ++ n ++ " (") . printCExpr e2 . showString ")"
  CEIf e1 e2 e3 ->
    printCExpr e1 . showString " ? " . printCExpr e2 . showString " : " . printCExpr e3

repTy :: STy t -> String
repTy (STScal st) = case st of
  STI32 -> "int32_t"
  STI64 -> "int64_t"
  STF32 -> "float"
  STF64 -> "double"
  STBool -> "bool"
repTy t = genStructName t

genStructName :: STy t -> String
genStructName = \t -> "ty_" ++ gen t where
  -- all tags start with a letter, so the array mangling is unambiguous.
  gen :: STy t -> String
  gen STNil = "n"
  gen (STPair a b) = 'P' : gen a ++ gen b
  gen (STEither a b) = 'E' : gen a ++ gen b
  gen (STMaybe t) = 'M' : gen t
  gen (STArr n t) = "A" ++ show (fromSNat n) ++ gen t
  gen (STScal st) = case st of
    STI32 -> "i"
    STI64 -> "j"
    STF32 -> "f"
    STF64 -> "d"
    STBool -> "b"
  gen (STAccum t) = 'C' : gen t

genStruct :: STy t -> Map String StructDecl
genStruct topty = case topty of
  STNil ->
    Map.singleton (genStructName STNil) (StructDecl (genStructName STNil) "" com)
  STPair a b ->
    let name = genStructName (STPair a b)
    in Map.singleton name (StructDecl name (repTy a ++ " a; " ++ repTy b ++ " b;") com)
  STEither a b ->
    let name = genStructName (STEither a b)  -- 0 -> a, 1 -> b
    in Map.singleton name (StructDecl name ("uint8_t tag; union { " ++ repTy a ++ " a; " ++ repTy b ++ " b; };") com)
  STMaybe t ->
    let name = genStructName (STMaybe t)  -- 0 -> nothing, 1 -> just
    in Map.singleton name (StructDecl name ("uint8_t tag; " ++ repTy t ++ " a;") com)
  STArr n t ->
    let name = genStructName (STArr n t)
    in Map.singleton name (StructDecl name ("size_t sh[" ++ show (fromSNat n) ++ "]; " ++ repTy t ++ " *a;") com)
  STScal _ -> mempty
  STAccum t ->
    let name = genStructName (STAccum t)
    in Map.singleton name (StructDecl name (repTy t ++ " a;") com)
         <> genStruct t
  where
    com = ppTy 0 topty

data CompState = CompState
  { csStructs :: Map String StructDecl
  , csStmts :: Bag Stmt
  , csNextId :: Int }
  deriving (Show)

type CompM a = State CompState a

genId :: CompM Int
genId = state $ \s -> (csNextId s, s { csNextId = csNextId s + 1 })

genName :: CompM String
genName = ('x' :) . show <$> genId

emit :: Stmt -> CompM ()
emit stmt = modify $ \s -> s { csStmts = csStmts s <> pure stmt }

scope :: CompM a -> CompM (a, [Stmt])
scope m = do
  stmts <- state $ \s -> (csStmts s, s { csStmts = mempty })
  res <- m
  innerStmts <- state $ \s -> (csStmts s, s { csStmts = stmts })
  return (res, toList innerStmts)

emitStruct :: STy t -> CompM String
emitStruct ty = do
  modify $ \s -> s { csStructs = genStruct ty <> csStructs s }
  return (genStructName ty)

nameEnv :: SList f env -> SList (Const String) env
nameEnv = flip evalState (0::Int) . slistMapA (\_ -> state $ \i -> (Const ("arg" ++ show i), i + 1))

compile :: SList STy env -> Ex env t -> String
compile env expr =
  let args = nameEnv env
      (res, s) = runState (compile' args expr) (CompState mempty mempty 1)
  in ($ "") $ compose
       [compose $ map (\sd -> printStructDecl sd . showString "\n") (Map.elems (csStructs s))
       ,showString "\n"
       ,showString $
          repTy (typeOf expr) ++ " kernel(" ++
            intercalate ", " (reverse (unSList (\(Product.Pair t n) -> repTy t ++ " " ++ getConst n) (slistZip env args))) ++
            ") {\n"
       ,compose $ map (\st -> showString "  " . printStmt 1 st . showString "\n") (toList (csStmts s))
       ,showString ("  return ") . printCExpr res . showString ";\n}\n"]

compile' :: SList (Const String) env -> Ex env t -> CompM CExpr
compile' env = \case
  EVar _ _ i -> return $ CELit (getConst (slistIdx env i))

  ELet _ rhs body -> do
    e <- compile' env rhs
    var <- genName
    emit $ SVarDecl True (repTy (typeOf rhs)) var e
    compile' (Const var `SCons` env) body

  EPair _ a b -> do
    name <- emitStruct (STPair (typeOf a) (typeOf b))
    e1 <- compile' env a
    e2 <- compile' env b
    return $ CEStruct name [("a", e1), ("b", e2)]

  EFst _ e -> CEProj <$> compile' env e <*> pure "a"

  ESnd _ e -> CEProj <$> compile' env e <*> pure "b"

  ENil _ -> do
    name <- emitStruct STNil
    return $ CEStruct name []

  EInl _ t e -> do
    name <- emitStruct (STEither (typeOf e) t)
    e1 <- compile' env e
    return $ CEStruct name [("tag", CELit "0"), ("a", e1)]

  EInr _ t e -> do
    name <- emitStruct (STEither t (typeOf e))
    e2 <- compile' env e
    return $ CEStruct name [("tag", CELit "1"), ("b", e2)]

  ECase _ (EOp _ OIf e) a b -> do
    e1 <- compile' env e
    (e2, stmts2) <- scope $ compile' (Const undefined `SCons` env) a  -- don't access that nil, stupid you
    (e3, stmts3) <- scope $ compile' (Const undefined `SCons` env) b
    retvar <- genName
    emit $ SVarDeclUninit (repTy (typeOf a)) retvar
    emit $ SIf e1
             (stmts2 <> pure (SAsg retvar e2))
             (stmts3 <> pure (SAsg retvar e3))
    return (CELit retvar)

  ECase _ e a b -> do
    let STEither t1 t2 = typeOf e
    e1 <- compile' env e
    var <- genName
    fieldvar <- genName
    (e2, stmts2) <- scope $ compile' (Const fieldvar `SCons` env) a
    (e3, stmts3) <- scope $ compile' (Const fieldvar `SCons` env) b
    retvar <- genName
    emit $ SVarDeclUninit (repTy (typeOf a)) retvar
    emit $ SBlock (pure (SVarDecl True (repTy (typeOf e)) var e1)
                <> pure (SIf (CEBinop (CEProj (CELit var) "tag") "==" (CELit "0"))
                           (pure (SVarDecl True (repTy t1) fieldvar
                                           (CEProj (CELit var) "a"))
                            <> stmts2
                            <> pure (SAsg retvar e2))
                           (pure (SVarDecl True (repTy t2) fieldvar
                                           (CEProj (CELit var) "b"))
                            <> stmts3
                            <> pure (SAsg retvar e3))))
    return (CELit retvar)

  ENothing _ t -> do
    name <- emitStruct (STMaybe t)
    return $ CEStruct name [("tag", CELit "0")]

  EJust _ e -> do
    name <- emitStruct (STMaybe (typeOf e))
    e1 <- compile' env e
    return $ CEStruct name [("tag", CELit "1"), ("a", e1)]

  EMaybe _ a b e -> do
    e1 <- compile' env e
    var <- genName
    fieldvar <- genName
    (e2, stmts2) <- scope $ compile' env a
    (e3, stmts3) <- scope $ compile' (Const fieldvar `SCons` env) b
    retvar <- genName
    emit $ SVarDeclUninit (repTy (typeOf a)) retvar
    emit $ SBlock (pure (SVarDecl True (repTy (typeOf e)) var e1)
                <> pure (SIf (CEBinop (CEProj (CELit var) "tag") "==" (CELit "0"))
                           (stmts2
                            <> pure (SAsg retvar e2))
                           (pure (SVarDecl True (repTy (typeOf b)) fieldvar
                                           (CEProj (CELit var) "a"))
                            <> stmts3
                            <> pure (SAsg retvar e3))))
    return (CELit retvar)

  EConstArr _ n t arr -> do
    name <- emitStruct (STArr n (STScal t))
    error "TODO"

  EBuild _ n a b -> error "TODO" -- genStruct (STArr n (typeOf b)) <> EBuild ext n (compile' a) (compile' b)

  EFold1Inner _ a b c -> error "TODO" -- EFold1Inner ext (compile' a) (compile' b) (compile' c)

  ESum1Inner _ e -> error "TODO" -- ESum1Inner ext (compile' e)

  EUnit _ e -> error "TODO" -- EUnit ext (compile' e)

  EReplicate1Inner _ a b -> error "TODO" -- EReplicate1Inner ext (compile' a) (compile' b)

  EMaximum1Inner _ e -> error "TODO" -- EMaximum1Inner ext (compile' e)

  EMinimum1Inner _ e -> error "TODO" -- EMinimum1Inner ext (compile' e)

  EConst _ t x -> case t of
    STI32 -> return $ CELit $ "(int32_t)" ++ show x
    STI64 -> return $ CELit $ "(int64_t)" ++ show x
    STF32 -> return $ CELit $ show x ++ "f"
    STF64 -> return $ CELit $ show x
    STBool -> return $ CELit $ if x then "true" else "false"

  EIdx0 _ e -> error "TODO" -- EIdx0 ext (compile' e)

  EIdx1 _ a b -> error "TODO" -- EIdx1 ext (compile' a) (compile' b)

  EIdx _ a b -> error "TODO" -- EIdx ext (compile' a) (compile' b)

  EShape _ e -> error "TODO" -- EShape ext (compile' e)

  EOp _ op (EPair _ e1 e2) -> do
    e1' <- compile' env e1
    e2' <- compile' env e2
    compileOpPair op e1' e2'

  EOp _ op e -> do
    e' <- compile' env e
    compileOpGeneral op e'

  ECustom _ t1 t2 t3 a b c e1 e2 -> error "TODO" -- ECustom ext t1 t2 t3 (compile' a) (compile' b) (compile' c) (compile' e1) (compile' e2)

  EWith a b -> error "TODO" -- EWith (compile' a) (compile' b)

  EAccum n a b e -> error "TODO" -- EAccum n (compile' a) (compile' b) (compile' e)

  EError t s -> do
    name <- emitStruct t
    -- using 'show' here is wrong, but it's good enough for me.
    emit $ SVerbatim $ "fprintf(stderr, \"ERROR: %s\\n\", " ++ show s ++ "); exit(1);"
    return $ CEStruct name []

  EZero{} -> error "Compile: monoid operations should have been eliminated"
  EPlus{} -> error "Compile: monoid operations should have been eliminated"
  EOneHot{} -> error "Compile: monoid operations should have been eliminated"

compileOpGeneral :: SOp a b -> CExpr -> CompM CExpr
compileOpGeneral op e1 = do
  let unary cop = return @(State CompState) $ CECall cop [e1]
  let binary cop = do
        name <- genName
        emit $ SVarDecl True (repTy (opt1 op)) name e1
        return $ CEBinop (CEProj (CELit name) "a") cop (CEProj (CELit name) "b")
  case op of
    OAdd _ -> binary "+"
    OMul _ -> binary "*"
    ONeg _ -> unary "-"
    OLt _ -> binary "<"
    OLe _ -> binary "<="
    OEq _ -> binary "=="
    ONot -> unary "!"
    OAnd -> binary "&&"
    OOr -> binary "||"
    OIf -> do
      name <- emitStruct (STEither STNil STNil)
      _ <- emitStruct STNil
      return $ CEIf e1 (CEStruct name [("tag", CELit "0")])
                       (CEStruct name [("tag", CELit "1")])
    ORound64 -> unary "(int64_t)round"  -- ew
    OToFl64 -> unary "(double)"
    ORecip _ -> return $ CEBinop (CELit "1.0") "/" e1
    OExp STF32 -> unary "expf"
    OExp STF64 -> unary "exp"
    OLog STF32 -> unary "logf"
    OLog STF64 -> unary "log"
    OIDiv _ -> binary "/"

compileOpPair :: SOp a b -> CExpr -> CExpr -> CompM CExpr
compileOpPair op e1 e2 = do
  let binary cop = return @(State CompState) $ CEBinop e1 cop e2
  case op of
    OAdd _ -> binary "+"
    OMul _ -> binary "*"
    OLt _ -> binary "<"
    OLe _ -> binary "<="
    OEq _ -> binary "=="
    OAnd -> binary "&&"
    OOr -> binary "||"
    OIDiv _ -> binary "/"
    _ -> error "compileOpPair: got unary operator"

compose :: Foldable t => t (a -> a) -> a -> a
compose = foldr (.) id