{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DefaultSignatures #-}

{-# OPTIONS_GHC -Wall #-}

module CodeGen.C
  where

import           Deep.Expr

import           Data.Typeable

import           Data.Semigroup
import           Control.Monad.Writer
import           Control.Monad.State
import           Data.List

import           Data.Map (Map)
import qualified Data.Map as Map

import qualified Data.Set as Set

import           Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NE

import           Data.Maybe (isNothing)

import           GHC.Stack (HasCallStack)

import           Data.Ord

import           Data.Complex

import           CodeGen.C.Types
import           CodeGen.C.Helper

import Debug.Trace

nub' :: Ord a => [a] -> [a]
nub' = Set.toList . Set.fromList

deleteAll :: Eq a => a -> [a] -> [a]
deleteAll _ [] = []
deleteAll x (y:ys)
  | y == x    =     deleteAll x ys
  | otherwise = y : deleteAll x ys

data SomeName = forall a. Typeable a => SomeName (Name a)

-- | NOTE: This instance will only take the unique value into account, not
-- the type tag of the 'Name'.
instance Eq SomeName where
  (SomeName (Name a)) == (SomeName (Name b)) = a == b

instance Ord SomeName where
  compare (SomeName (Name a)) (SomeName (Name b)) = compare a b

instance Show SomeName where
  show (SomeName n) = show n

newtype LocalEnv = LocalEnv (Map Int CName) deriving Show

-- Takes a unique from a 'Name a' and gives back the current CName that
-- references it.
le_lookup :: HasCallStack => Int -> LocalEnv -> CName
le_lookup uniq (LocalEnv env) =
  case Map.lookup uniq env of
    Nothing -> error $ "le_lookup: Cannot find " ++ show uniq
    Just x  -> x

le_modifyName :: LocalEnv -> Int -> CName -> LocalEnv
le_modifyName (LocalEnv env) uniq cname = LocalEnv $ Map.insert uniq cname env

le_modifyNames :: LocalEnv -> [(Int, CName)] -> LocalEnv
le_modifyNames le [] = le
le_modifyNames (LocalEnv env) ((uniq, cname):changes) = le_modifyNames (LocalEnv (Map.insert uniq cname env)) changes

le_fromList :: [(Int, CName)] -> LocalEnv
le_fromList = LocalEnv . Map.fromList

le_initial :: Int -> LocalEnv
le_initial varCount = LocalEnv $ Map.fromList $ map (\n -> (n, "vars[" <> show n <> "]")) [0..varCount-1]

data CodeGenState =
  CodeGenState
    { cg_le :: LocalEnv
    , cg_currUniq :: Int
    , cg_lams :: [SomeLambda]
    }

newtype CodeGen a =
  CodeGen (State CodeGenState a)
  deriving (Functor, Applicative, Monad)

deriving instance MonadState CodeGenState CodeGen

cg_modifyName :: Int -> CName -> CodeGen ()
cg_modifyName uniq cname =
  modify (\cg -> cg { cg_le = le_modifyName (cg_le cg) uniq cname })

cg_modifyNames :: [(Int, CName)] -> CodeGen ()
cg_modifyNames changes =
  modify (\cg -> cg { cg_le = le_modifyNames (cg_le cg) changes })

cg_lookup :: HasCallStack => Int -> CodeGen CName
cg_lookup uniq =
  fmap (le_lookup uniq . cg_le) get

withLocalEnv :: LocalEnv -> CodeGen a -> CodeGen a
withLocalEnv newEnv m = do
  oldEnv <- fmap cg_le get
  modify (\cg -> cg { cg_le = newEnv })
  r <- m
  modify (\cg -> cg { cg_le = oldEnv })
  return r

freshCName :: CodeGen CName
freshCName = do
  uniq <- fmap cg_currUniq get

  modify (\cg -> cg { cg_currUniq = uniq+1 })

  return ("x" ++ show uniq)

lookupLambda :: Int -> CodeGen SomeLambda
lookupLambda uniq = do
  lams <- fmap cg_lams get

  case find (\(SomeLambda lam) -> getNameUniq (lambda_name lam) == uniq) lams of
    Nothing -> error "lookupLambda"
    Just x  -> return x

runTopLevelCodeGen :: Int -> [SomeLambda] -> CodeGen a -> a
runTopLevelCodeGen varCount lams (CodeGen m) =
  evalState m (CodeGenState { cg_le = le_initial varCount, cg_currUniq = 0, cg_lams = lams })

data Lambda a b
  = Lambda
    { lambda_fvs :: [SomeName] -- | Free variables. These must remain sorted by unique, in ascending order
    , lambda_name :: Name a -- | Lambdas are named by their argument name
      -- TODO: Should we have a [(SomeName, CName)] (possibly replacing
      -- lambda_fvs?)
    , lambda_body :: GPUExp b
    , lambda_isTailRec :: Bool
    }

data SomeLambda = forall a b. {- Typeable a => -} SomeLambda (Lambda a b)

getFVs :: GPUExp a -> [SomeName]
getFVs = (`execState` []) . go
  where
    go :: forall s. GPUExp s -> State [SomeName] ()
    go (Var n) = modify (SomeName n:)

    go (CaseExp x y) = go x *> go y
    go (ProdMatchExp x) = go x
    go (SumMatchExp x y) = go x *> go y
    go (NullaryMatch x) = go x
    go (OneSumMatch x) = go x
    go (OneProdMatch x) = go x
    go EmptyMatch = pure ()
    go FalseExp = pure ()
    go TrueExp = pure ()
    go (Repped {}) = pure ()
    go (Lam n x) = do
      go x
      modify (deleteAll (SomeName n))
    go (App f x) = go f *> go x
    go (Lit {}) = pure ()
    go (Add x y) = go x *> go y
    go (Sub x y) = go x *> go y
    go (Mul x y) = go x *> go y
    go (FromEnum x) = go x
    go (FromIntegral x) = go x
    go (Sqrt x) = go x
    go (Equal x y) = go x *> go y
    go (Lte x y) = go x *> go y
    go (Gt x y) = go x *> go y
    go (Not x) = go x
    go (LeftExp x) = go x
    go (RightExp x) = go x
    go (PairExp x y) = go x *> go y
    go (FstExp x) = go x
    go (SndExp x) = go x
    go (StepExp x) = go x
    go (DoneExp x) = go x
    go (IfThenElse x y z) = go x *> go y *> go z
    go (TailRec x) = go x
    go (Construct {}) = pure ()
    go (ConstructAp x y) = go x *> go y
    go (Ord x) = go x
    go (CharLit {}) = pure ()
    go UnitExp = pure ()
    go (ConstructRep x) = go x
    go (CastExp x) = go x

-- | Precondition: The maximum unique of all the variables must be
-- equal to the number of variables.
--
-- TODO: Eliminate this precondition.
--
getVarCount :: GPUExp a -> Int
getVarCount = getMax . execWriter . go
  where
    -- TODO: Abstract this traversal pattern
    go :: forall s. GPUExp s -> Writer (Max Int) ()
    go (Var (Name n)) = tell (Max n)

    go (CaseExp x y) = go x *> go y
    go (ProdMatchExp x) = go x
    go (SumMatchExp x y) = go x *> go y
    go (NullaryMatch x) = go x
    go (OneSumMatch x) = go x
    go (OneProdMatch x) = go x
    go EmptyMatch = pure ()
    go FalseExp = pure ()
    go TrueExp = pure ()
    go (Repped {}) = pure ()
    go (Lam n x) = go x
    go (App f x) = go f *> go x
    go (Lit {}) = pure ()
    go (Add x y) = go x *> go y
    go (Sub x y) = go x *> go y
    go (Mul x y) = go x *> go y
    go (FromEnum x) = go x
    go (FromIntegral x) = go x
    go (Sqrt x) = go x
    go (Equal x y) = go x *> go y
    go (Lte x y) = go x *> go y
    go (Gt x y) = go x *> go y
    go (Not x) = go x
    go (LeftExp x) = go x
    go (RightExp x) = go x
    go (PairExp x y) = go x *> go y
    go (FstExp x) = go x
    go (SndExp x) = go x
    go (StepExp x) = go x
    go (DoneExp x) = go x
    go (IfThenElse x y z) = go x *> go y *> go z
    go (TailRec x) = go x
    go (Construct {}) = pure ()
    go (ConstructAp x y) = go x *> go y
    go (Ord x) = go x
    go (CharLit {}) = pure ()
    go UnitExp = pure ()
    go (ConstructRep x) = go x
    go (CastExp x) = go x

collectLambdas :: GPUExp a -> [SomeLambda]
collectLambdas = execWriter . go
  where
    go :: forall s. GPUExp s -> Writer [SomeLambda] ()
    go (Var {}) = pure ()
    go (CaseExp x y) = go x *> go y
    go (ProdMatchExp x) = go x
    go (SumMatchExp x y) = go x *> go y
    go (NullaryMatch x) = go x
    go (OneSumMatch x) = go x
    go (OneProdMatch x) = go x
    go EmptyMatch = pure ()
    go FalseExp = pure ()
    go TrueExp = pure ()
    go (Repped {}) = pure ()
    go (TailRec lamExp@(Lam n body)) = tell [mkLambda True (SomeName n) body] *> go body
    go (Lam n x) = tell [mkLambda False (SomeName n) x] *> go x
    go (App f x) = go f *> go x
    go (Lit {}) = pure ()
    go (Add x y) = go x *> go y
    go (Sub x y) = go x *> go y
    go (Mul x y) = go x *> go y
    go (FromEnum x) = go x
    go (FromIntegral x) = go x
    go (Sqrt x) = go x
    go (Equal x y) = go x *> go y
    go (Lte x y) = go x *> go y
    go (Gt x y) = go x *> go y
    go (Not x) = go x
    go (LeftExp x) = go x
    go (RightExp x) = go x
    go (PairExp x y) = go x *> go y
    go (FstExp x) = go x
    go (SndExp x) = go x
    go (StepExp x) = go x
    go (DoneExp x) = go x
    go (IfThenElse x y z) = go x *> go y *> go z
    go (TailRec x) = go x
    go (Construct {}) = pure ()
    go (ConstructAp x y) = go x *> go y
    go (Ord x) = go x
    go (CharLit {}) = pure ()
    go UnitExp = pure ()
    go (ConstructRep x) = go x
    go (CastExp x) = go x

mkLambda :: Bool -> SomeName -> GPUExp a -> SomeLambda
mkLambda isTailRec (SomeName argName) body =
  SomeLambda $
    Lambda
      { lambda_fvs = trace "lambda_fvs:" $ traceShowId $
          nub' $
          sortBy (comparing (\(SomeName a) -> getNameUniq a)) $
          deleteAll
            (SomeName argName)
            (getFVs body)
      , lambda_name = argName
      , lambda_body = body
      , lambda_isTailRec = isTailRec
      }

genProgram :: HasCallStack => GPUExp a -> CCode
genProgram e =
  let varCount = getVarCount e
      lambdas = collectLambdas e
  in
  runTopLevelCodeGen varCount lambdas $ do

    resultName <- freshCName

    lambdasCode <- genLambdas e
    exprCode <- genExp e resultName
    return $ unlines
      [ prelude varCount
      , unlines
          (map (\(SomeLambda lam) -> "var_t " <> lambdaCName_byInt (getNameUniq (lambda_name lam)) <> "(var_t, struct closure_t*);")
               lambdas)
      , ""
      , lambdasCode
      , "var_t top_level() {"
      , "  var_t " <> resultName <> ";"
      , exprCode
      , "  return " <> resultName <> ";"
      , "}"
      ]

prelude :: Int -> CCode
prelude varCount =
  unlines
    [ "#include <stdio.h>"
    , "#include <stdlib.h>"
    , "#include <stdbool.h>"
    , "#include <string.h>"
    , "#include <assert.h>"
    , "#include <math.h>"
    , ""
    , "typedef enum var_type_tag {"
    , "  EXPR_INT"
    , ", EXPR_FLOAT"
    , ", EXPR_DOUBLE"
    , ", EXPR_CHAR"
    , ", EXPR_BOOL"
    , ", EXPR_CLOSURE"
    , ", EXPR_EITHER_LEFT"
    , ", EXPR_EITHER_RIGHT"
    , ", EXPR_PAIR"
    , ", EXPR_UNIT"
    , ", EXPR_STEP"
    , ", EXPR_DONE"
    -- , ", EXPR_UNBOXED"
    , ", EXPR_NIL"
    , "} var_type_tag;"
    , ""
    , "typedef enum semantic_type_tag {"
    , "  NO_SEMANTIC_TAG"
    , ", EXPR_COMPLEX     // Complex numbers"
    , "} semantic_type_tag;"
    , ""
    , "struct closure_t;"
    , ""
    , "typedef struct var_t {"
    , "  var_type_tag tag;"
    , "  semantic_type_tag semantic_tag;  // For optional additional information about the \"un-deconstructed\" type (for instance, it could be the case that tag=EXPR_PAIR, but specific_tag=EXPR_COMPLEX)"
    , "  void* value;"
    , "} var_t;"
    , ""
    , "typedef struct closure_t {"
    , "  int env_size;"
    , "  var_t* fv_env;"
    , "  var_t (*fn)(var_t, struct closure_t*);"
    , "} closure_t;"
    , ""
    , "void copyClosure(closure_t*, closure_t*);"
    , ""
    , "void copyVar(var_t* dest, var_t* src) {"
    , "  dest->tag = src->tag;"
    , "  dest->semantic_tag = src->semantic_tag;"
    , "  if (src->tag == EXPR_PAIR) {"
    , "    dest->value = malloc(2*sizeof(var_t));"
    , "    copyVar(&((var_t*)dest->value)[0], &((var_t*)src->value)[0]);"
    , "    copyVar(&((var_t*)dest->value)[1], &((var_t*)src->value)[1]);"
    , "  } else {"
    , "    switch (src->tag) {"
    , "      case EXPR_INT:"
    , "       dest->value = malloc(sizeof(int));"
    , "       *(int*)(dest->value) = *(int*)(src->value);"
    , "       break;"
    , "      case EXPR_FLOAT:"
    , "       dest->value = malloc(sizeof(float));"
    , "       *(float*)(dest->value) = *(float*)(src->value);"
    , "       break;"
    , "      case EXPR_DOUBLE:"
    , "       dest->value = malloc(sizeof(double));"
    , "       *(double*)(dest->value) = *(double*)(src->value);"
    , "       break;"
    , "      case EXPR_CHAR:"
    , "       dest->value = malloc(sizeof(char));"
    , "       *(char*)(dest->value) = *(char*)(src->value);"
    , "       break;"
    , "      case EXPR_CLOSURE:"
    , "       dest->value = malloc(sizeof(closure_t));"
    -- , "       *(closure_t*)(dest->value) = *(closure_t*)(src->value);"
    , "       copyClosure((closure_t*)dest->value, (closure_t*)src->value);"
    , "       break;"
    , "      case EXPR_EITHER_LEFT:"
    , "      case EXPR_EITHER_RIGHT:"
    , "      case EXPR_STEP:"
    , "      case EXPR_DONE:"
    , "      case EXPR_UNIT:"
    , "       dest->value = malloc(sizeof(var_t));"
    , "       copyVar((var_t*)(dest->value), (var_t*)(src->value));"
    -- , "       *(var_t*)(dest->value) = *(var_t*)(src->value);"
    , "       break;"
    , "      case EXPR_NIL:"
    , "       dest->value = NULL;"
    , "       break;"
    , "      default:"
    , "        assert(0 && \"invalid tag\");"
    , "    }"
    , "  }"
    , "}"
    , ""
    , "void copyClosure(closure_t* dest, closure_t* src) {"
    , "  dest->env_size = src->env_size;"
    , "  dest->fn = src->fn;"
    , "  dest->fv_env = malloc(src->env_size * sizeof(var_t));"
    , "  for (int i = 0; i < src->env_size; ++i) {"
    , "    copyVar(&(dest->fv_env[i]), &(src->fv_env[i]));"
    , "  }"
    , "}"
    , ""
    , "var_t car(var_t v) {"
    , "  return ((var_t*)v.value)[0];"
    , "}"
    , "var_t cdr(var_t v) {"
    , "  return ((var_t*)v.value)[1];"
    , "}"
    , ""
    , "var_t vars[" <> show varCount <> "];"
    , ""
    , "#define GET_MATH_OP_ADD +"
    , "#define GET_MATH_OP_MUL *"
    , "#define GET_MATH_OP_SUB -"
    , "#define GET_MATH_OP_DIV /"
    , ""
    , "#define GET_COMPARE_OP_LTE <="
    , "#define GET_COMPARE_OP_LT <"
    , "#define GET_COMPARE_OP_GTE >="
    , "#define GET_COMPARE_OP_GT >"
    , "#define GET_COMPARE_OP_EQUAL =="
    , ""
    , "#define assert_msg(cond, msg) do { if (!(cond)) { printf msg; assert(0 && #cond); } } while(0)"
    , ""
    , "#define MATH_OP(op, result, a, b)\\"
    , "  do {\\"
    , "    assert_msg((a).tag == (b).tag, (\"(a,b).tag = (%i,%i)\\n\", (int)(a).tag, (int)(b).tag));\\"
    , "    (result).semantic_tag = (a).semantic_tag;\\"
    , "    (result).tag = (a).tag;\\"
    , "    switch ((a).tag) {\\"
    , "      case EXPR_INT:\\"
    , "        (result).value = malloc(sizeof(int));\\"
    , "        *((int*)(result).value) = *(int*)((a).value) GET_MATH_OP_##op *(int*)((b).value);\\"
    , "        break;\\"
    , "      case EXPR_FLOAT:\\"
    , "        (result).value = malloc(sizeof(float));\\"
    , "        *((float*)(result).value) = *(float*)((a).value) GET_MATH_OP_##op *(float*)((b).value);\\"
    , "        break;\\"
    , "      case EXPR_DOUBLE:\\"
    , "        (result).value = malloc(sizeof(double));\\"
    , "        *((double*)(result).value) = *(double*)((a).value) GET_MATH_OP_##op *(double*)((b).value);\\"
    , "        break;\\"
    , "      default:\\"
    , "        fprintf(stderr, \"%s type tag = %d\\n\", #a, (a).tag);\\"
    , "        assert(0 && \"Attempting to perform arithmetic on non-numeric types\");\\"
    , "    }\\"
    , "  } while (0);"
    , ""
    , "#define COMPARE(op, result, a, b)\\"
    , "  do {\\"
    , "   assert((a).tag == (b).tag);\\"
    , "    switch ((a).tag) {\\"
    , "      case EXPR_INT:\\"
    , "        *((bool*)(result).value) = *(int*)((a).value) GET_COMPARE_OP_##op *(int*)((b).value);\\"
    , "        break;\\"
    , "      case EXPR_FLOAT:\\"
    , "        *((bool*)(result).value) = *(float*)((a).value) GET_COMPARE_OP_##op *(float*)((b).value);\\"
    , "        break;\\"
    , "      case EXPR_DOUBLE:\\"
    , "        *((bool*)(result).value) = *(double*)((a).value) GET_COMPARE_OP_##op *(double*)((b).value);\\"
    , "        break;\\"
    , "      case EXPR_BOOL:\\"
    , "        *((bool*)(result).value) = *(bool*)((a).value) GET_COMPARE_OP_##op *(bool*)((b).value);\\"
    , "        break;\\"
    , "      case EXPR_CHAR:\\"
    , "        *((bool*)(result).value) = *(char*)((a).value) GET_COMPARE_OP_##op *(char*)((b).value);\\"
    , "        break;\\"
    , "      default:\\"
    , "        fprintf(stderr, \"%s type tag = %d\\n\", #a, (a).tag);\\"
    , "        assert(0 && \"Cannot compare given type for equality\");\\"
    , "    }\\"
    , "  } while (0);"
    , ""
    , "#define CAST_TO(result, r_type, a)\\"
    , "  do {\\"
    , "    switch ((a).tag) {\\"
    , "      case EXPR_INT:\\"
    , "        *((r_type*)(result).value) = (r_type) *(int*)((a).value);\\"
    , "        break;\\"
    , "      case EXPR_FLOAT:\\"
    , "        *((r_type*)(result).value) = (r_type) *(float*)((a).value);\\"
    , "        break;\\"
    , "      case EXPR_DOUBLE:\\"
    , "        *((r_type*)(result).value) = (r_type) *(double*)((a).value);\\"
    , "        break;\\"
    , "      default:\\"
    , "        fprintf(stderr, \"%s type tag = %d\\n\", #a, (a).tag);\\"
    , "        assert(0 && \"CAST_FROM: Invalid type tag\");\\"
    , "    }\\"
    , "  } while (0);"
    , ""
    , "#define MATH_SQRT(result, a)\\"
    , "  do {\\"
    , "      case EXPR_FLOAT:\\"
    , "        *((float*)(result).value) = sqrt(*(float*)((a).value));\\"
    , "        break;\\"
    , "      case EXPR_DOUBLE:\\"
    , "        *((double*)(result).value) = sqrt(*(double*)((a).value));\\"
    , "        break;\\"
    , "  } while(0);"
    , ""
    , "#define INIT_COMPLEX_PAIR(result)\\"
    , "  do {\\"
    , "    (result).semantic_tag = EXPR_COMPLEX;\\"
    , "    (result).tag = EXPR_PAIR;\\"
    , "    (result).value = malloc(2*sizeof(var_t));\\"
    , "    ((var_t*)((result).value))[1].value = malloc(2*sizeof(var_t));\\"
    , "  } while (0);"
    , ""
    , "#define INIT_COMPLEX(a, type, eTag)\\"
    , "  do {\\"
    , "    (a).semantic_tag = EXPR_COMPLEX;\\"
    , "    (a).tag = EXPR_PAIR;\\"
    , "    (a).value = malloc(2*sizeof(var_t));\\"
    , "    ((var_t*)((a).value))[0].tag = eTag;\\"
    , "    ((var_t*)((a).value))[0].value = malloc(sizeof(var_t));\\"

    , "    ((var_t*)((a).value))[1].tag = EXPR_PAIR;\\"
    , "    ((var_t*)((a).value))[1].semantic_tag = NO_SEMANTIC_TAG;\\"
    , "    ((var_t*)((a).value))[1].value = malloc(2*sizeof(var_t));\\"

    , "    ((var_t*)((var_t*)((a).value))[1].value)[0].tag = eTag;\\"
    , "    ((var_t*)((var_t*)((a).value))[1].value)[0].semantic_tag = NO_SEMANTIC_TAG;\\"
    , "    ((var_t*)((var_t*)((a).value))[1].value)[0].value = malloc(sizeof(var_t));\\"

    , "    ((var_t*)((var_t*)((a).value))[1].value)[1].tag = EXPR_NIL;\\"
    , "    ((var_t*)((var_t*)((a).value))[1].value)[1].semantic_tag = NO_SEMANTIC_TAG;\\"
    , "    ((var_t*)((var_t*)((a).value))[1].value)[1].value = NULL;\\"

    -- , "    ((var_t*)((a).value))[1].tag = eTag;\\"
    -- , "    ((var_t*)((a).value))[1].value = malloc(sizeof(var_t));\\"
    , "  } while (0);"
    , ""
    , "#define PAIR_FST(result, p)\\"
    , "  do {\\"
    , "    assert((p).tag == EXPR_PAIR);\\"
    , "    " <> stmt (cCall "copyVar" [addrOf "result", addrOf (cCast "var_t*" ("(p)" # "value") ! 0)]) <> "\\"
    -- , "    (result) = ((var_t*)((p).value))[0];\\"
    , "  } while(0);"
    , ""
    , "#define PAIR_SND(result, p)\\"
    , "  do {\\"
    , "    assert((p).tag == EXPR_PAIR);\\"
    , "    " <> stmt (cCall "copyVar" [addrOf "result", addrOf (cCast "var_t*" ((cCast "var_t*" ("(p)" # "value") ! 1) # "value") ! 0)]) <> "\\"
    -- , "    (result) = ((var_t*)((var_t*)((result).value))[1].value)[0];\\"
    , "  } while(0);"
    , ""
    , "#define PAIR_ASSIGN_FST(result, x)\\"
    , "  do {\\"
    , "    " <> stmt (cCall "copyVar" [addrOf (cCast "var_t*" ("(result)" # "value") ! 0), addrOf "x"]) <> "\\"
    -- , "    ((var_t*)((result).value))[0] = (x);\\"
    , "  } while(0);"
    , ""
    , "#define PAIR_ASSIGN_SND(result, x)\\"
    , "  do {\\"
    , "    " <> stmt (cCall "copyVar" [addrOf (cCast "var_t*" ((cCast "var_t*" ("(result)" # "value") ! 1) # "value") ! 0), addrOf "x"]) <> "\\"
    -- , "    ((var_t*)((var_t*)((result).value))[1].value)[0] = (x); \\"
    , "  } while(0);"
    , ""
    , "#define GET_COMPLEX_REAL(c) ((var_t*)((c).value))[0]"
    , "#define GET_COMPLEX_IMAG(c) ((var_t*)((var_t*)((c).value))[1].value)[0]"
    , ""
    , "#define COMPLEX_REAL(result, c)\\"
    , "  do {\\"
    , "    assert((p).semantic_tag == EXPR_COMPLEX);\\"
    , "    PAIR_FST(result, c);\\"
    , "  } while (0);"
    , ""
    , "#define COMPLEX_IMAG(result, c)\\"
    , "  do {\\"
    , "    assert((c).semantic_tag == EXPR_COMPLEX);\\"
    , "    PAIR_SND(result, c);\\"
    , "  } while (0);"
    , ""
    , "#define COMPLEX_ASSIGN_REAL(result, x)\\"
    , "  do {\\"
    , "    PAIR_ASSIGN_FST(result, x);\\"
    , "  } while(0);"
    , ""
    , "#define COMPLEX_ASSIGN_IMAG(result, x)\\"
    , "  do {\\"
    , "    PAIR_ASSIGN_SND(result, x);\\"
    , "  } while(0);"
    , ""
    , "bool isIterTag(var_type_tag tag) {"
    , "  return tag == EXPR_STEP || tag == EXPR_DONE;"
    , "}"
    ]

data ShowablePair = forall a. (Typeable a, Show a) => ShowablePair a a

mainCode :: [CCode] -> CCode
mainCode body =
  unlines
    [ "int main(int argc, char* argv[]) {"
    , unlines (map (++"  ") body)
    , "}"
    ]

genExp :: HasCallStack => GPUExp a -> CName -> CodeGen CCode

genExp (ConstructRep x :: GPUExp t) resultName =
  case getComplexType_maybe (Proxy :: Proxy t) of
    Nothing -> do
      r <- genExp x resultName
      return $ unlines
        [ "// ConstructRep (Non-Complex)"
        , r
        ]

    Just (ty, tag) -> do
      pairName <- freshCName
      pairCode <- genExp x pairName

      return $ unlines
        [ "// Complex ConstructRep"
        , "var_t " <> pairName <> ";"
        , pairCode

        -- , (resultName#"value") =: cCall "malloc" ["2*sizeof(var_t)"]
        , cCall "INIT_COMPLEX" [resultName, ty, tag]
        , resultName <> ".semantic_tag = EXPR_COMPLEX;"
        , resultName <> ".tag = EXPR_PAIR;"
        , cCall "COMPLEX_ASSIGN_REAL" [resultName, cCall "GET_COMPLEX_REAL" [pairName]]
        , cCall "COMPLEX_ASSIGN_IMAG" [resultName, cCall "GET_COMPLEX_IMAG" [pairName]]

        -- , cCall "COMPLEX_ASSIGN_REAL" [resultName, (cCast "var_t*" (pairName # "value")) ! 0]
        -- , cCall "COMPLEX_ASSIGN_IMAG" [resultName, (cCast "var_t*" (pairName # "value")) ! 1]

        -- , "COMPLEX_ASSIGN_REAL(" <> resultName <> ", " <> pairName <> "
        -- , resultName <> ".value = malloc(sizeof(var_t));"

        -- , stmt $ cCall "copyVar" [cCast "var_t*" (resultName # "value"), addrOf pairName]

        -- , "*((var_t*)(" <> resultName <> ".value)) = " <> pairName <> ";"

        ]

genExp (Var name@(Name n)) resultName = do
  nCName <- cg_lookup n

  return $ unlines
    [ "  // Var with Name id " <> show n <> " and Haskell type " <> show (typeRep name)
    , "if (isIterTag(" <> nCName <> ".tag)) {"
    , stmt $ cCall "copyVar" [addrOf resultName, cCast "var_t*" (nCName # "value")]
    , "} else {"
    , stmt $ cCall "copyVar" [addrOf resultName, addrOf nCName]
    , "}"

    -- , "if (isIterTag(" <> nCName <> ".tag)) {"
    -- , resultName <> " = " <> "*(var_t*)(" <> nCName <> ".value);"
    -- , "} else {"
    -- , resultName <> " = " <> nCName <> ";"
    -- , "}"
    ]

genExp (CaseExp s body) resultName = do
  sResultName <- freshCName
  computeS <- genExp s sResultName

  caseCode <- genCaseExp sResultName body resultName

  return $ unlines
    [ "var_t " <> sResultName <> ";"
    , computeS
    , caseCode
    ]

genExp (Lam (Name n) body) resultName = do
  closurePtrName <- freshCName

  sLam@(SomeLambda lam) <- lookupLambda n
  closureCode <- buildClosure sLam ("(*" <> closurePtrName <> ")")

  -- let fvs = lambda_fvs lam

  return $ unlines
    [ "closure_t* " <> closurePtrName <> " = malloc(sizeof(closure_t));"
    , closureCode
    , ""
    , resultName <> ".tag = EXPR_CLOSURE;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = (void*)" <> closurePtrName <> ";"
    ]

genExp (ProdMatchExp f) resultName = genExp f resultName

genExp (OneProdMatch f) resultName = genExp f resultName

-- TODO: Generalize
genExp (App lamExp@(Lam (Name n) _) x) resultName = do
  xName <- freshCName
  xCode <- genExp x xName

  le <- fmap cg_le get

  traceM $ "/////////// le = " ++ show le
  traceM $ "/////////// le' = " ++ show (le_modifyName le n xName)

  closureName <- freshCName
  closureVarName <- freshCName
  closureCode <- withLocalEnv (le_modifyName le n xName) $ genExp lamExp closureVarName

  lam <- lookupLambda n

  callCode <- callClosure lam closureName xName resultName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , "closure_t " <> closureName <> ";"
    , "var_t " <> closureVarName <> ";"
    , ""
    , xCode
    , closureCode
    , "copyClosure(&" <> closureName <> ", (closure_t*)(" <> closureVarName <> ".value));"
    , callCode
    ]

genExp (App (TailRec lam) x) resultName = genExp (App lam x) resultName

genExp (App _ _) _ = error "genExp: App"

-- genExp (App f x) resultName = do
--   rName <- freshCName
--   xName <- freshCName

--   rCode <- genExp f fName
--   xCode <- genExp x xName

--   callCode <- callClosure

--   return $ unlines
--     [ "var_t " <> fName <> ";"
--     , "var_t " <> xName <> ";"
--     ,

genExp (Lit x :: GPUExp t) resultName =
    case getComplex_maybe x of
      Just (ShowablePair a b, (ty, tag)) -> do
        aName <- freshCName
        bName <- freshCName

        aCode <- oneDimNumCode a aName ty tag
        bCode <- oneDimNumCode b bName ty tag

        return $ unlines
          [ "var_t " <> aName <> ";"
          , "var_t " <> bName <> ";"
          , aCode
          , bCode
          , "// Complex Lit"
          , "INIT_COMPLEX(" <> resultName <> ", " <> ty <> ", " <> tag <> ");"
          , "COMPLEX_ASSIGN_REAL(" <> resultName <> ", " <> aName <> ");"
          , "COMPLEX_ASSIGN_IMAG(" <> resultName <> ", " <> bName <> ");"
          -- , resultName <> ".value = malloc(sizeof(2*sizeof(var_t)));"
          -- , resultName <> ".tag = " <> tag <> ";"
          -- , resultName <> ".value[0] = " <> "*(" <> ty <> ")(" <> aName <> ".value);"
          -- , resultName <> ".value[1] = " <> "*(" <> ty <> ")(" <> bName <> ".value);"
          ]
      Nothing ->
        uncurry (oneDimNumCode x resultName) (numTypeInfo x)
  where

    oneDimNumCode x' resultName' ty tag =
      return $ unlines
        [ "// oneDimNumCode"
        , resultName' <> ".value = malloc(sizeof(" <> ty <> "));"
        , resultName' <> ".tag = " <> tag <> ";"
        , resultName' <> ".semantic_tag = NO_SEMANTIC_TAG;"
        , "*(" <> ty <>"*)(" <> resultName' <> ".value) = " <> show x' <> ";"
        ]

genExp TrueExp resultName =
  return $ unlines
    [ resultName <> ".tag = EXPR_BOOL;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(bool));"
    , "*(bool*)(" <> resultName <> ".value) = true;"
    ]

genExp FalseExp resultName =
  return $ unlines
    [ resultName <> ".tag = EXPR_BOOL;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(bool));"
    , "*(bool*)(" <> resultName <> ".value) = false;"
    ]

genExp (DoneExp x) resultName = do
  xName <- freshCName

  xCode <- genExp x xName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , xCode
    , resultName <> ".tag = EXPR_DONE;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(var_t));"
    , stmt $ cCall "copyVar" [cCast "var_t*" (resultName # "value"), addrOf xName]
    -- , "*(var_t*)(" <> resultName <> ".value) = " <> xName <> ";"
    ]

genExp (StepExp x) resultName = do
  xName <- freshCName

  xCode <- genExp x xName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , xCode
    , resultName <> ".tag = EXPR_STEP;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(var_t));"
    , stmt $ cCall "copyVar" [cCast "var_t*" (resultName # "value"), addrOf xName]
    -- , "*(var_t*)(" <> resultName <> ".value) = " <> xName <> ";"
    ]

genExp (Sub x y) resultName = do -- TODO: Fix complex number support (see Add)
  xName <- freshCName
  yName <- freshCName

  xCode <- genExp x xName
  yCode <- genExp y yName

  return $ unlines
    [ "// Sub"
    , "var_t " <> xName <> ";"
    , "var_t " <> yName <> ";"
    , xCode
    , yCode
    , "assert(" <> xName <> ".tag == " <> yName <> ".tag);"
    , ""
    , "if (" <> xName <> ".semantic_tag == EXPR_COMPLEX) {"
    , "  INIT_COMPLEX_PAIR(" <> resultName <> ");"
    , "  MATH_OP(SUB, GET_COMPLEX_REAL(" <> resultName <> "), GET_COMPLEX_REAL(" <> xName <> "), GET_COMPLEX_REAL(" <> yName <> "));"
    , "  MATH_OP(SUB, GET_COMPLEX_IMAG(" <> resultName <> "), GET_COMPLEX_IMAG(" <> xName <> "), GET_COMPLEX_IMAG(" <> yName <> "));"
    -- , "  MATH_OP(SUB, GET_COMPLEX_IMAG(" <> resultName <> "), ((var_t*)(" <> xName <> ".value))[1], ((var_t*)(" <> yName <> ".value)[1]); "
    , "} else {"
    , "  MATHOP(SUB, " <> resultName <> ", " <> xName <> ", " <> yName <> ");"
    , "}"
    ]

genExp (Mul x y) resultName = do
  xName <- freshCName
  yName <- freshCName

  x0y0Name <- freshCName
  x0y1Name <- freshCName
  x1y0Name <- freshCName
  x1y1Name <- freshCName

  realName <- freshCName
  imagName <- freshCName

  xCode <- genExp x xName
  yCode <- genExp y yName

  return $ unlines
    [ "// Mul"
    , "var_t " <> xName <> ";"
    , "var_t " <> yName <> ";"
    , xCode
    , yCode
    -- , "assert(" <> xName <> ".tag == " <> yName <> ".tag);"
    , ""
    , "if (" <> xName <> ".semantic_tag == EXPR_COMPLEX) {"
    , "  var_t " <> x0y0Name <> ";"
    , "  var_t " <> x0y1Name <> ";"
    , "  var_t " <> x1y0Name <> ";"
    , "  var_t " <> x1y1Name <> ";"
    , "  var_t " <> realName <> ";"
    , "  var_t " <> imagName <> ";"
    , ""
    , "  MATH_OP(MUL, " <> x0y0Name <> ", GET_COMPLEX_REAL(" <> xName <> "), GET_COMPLEX_REAL(" <> yName <> "));"
    , "  MATH_OP(MUL, " <> x0y1Name <> ", GET_COMPLEX_REAL(" <> xName <> "), GET_COMPLEX_IMAG(" <> yName <> "));"
    , "  MATH_OP(MUL, " <> x1y0Name <> ", GET_COMPLEX_IMAG(" <> xName <> "), GET_COMPLEX_REAL(" <> yName <> "));"
    , "  MATH_OP(MUL, " <> x1y1Name <> ", GET_COMPLEX_IMAG(" <> xName <> "), GET_COMPLEX_IMAG(" <> yName <> "));"
    -- , "  MATH_OP(MUL, " <> x0y0Name <> ", ((var_t*)(" <> xName <> ".value))[0], ((var_t*)(" <> yName <> ".value))[0]);"
    -- , "  MATH_OP(MUL, " <> x0y1Name <> ", ((var_t*)(" <> xName <> ".value))[0], ((var_t*)(" <> yName <> ".value))[1]);"
    -- , "  MATH_OP(MUL, " <> x1y0Name <> ", ((var_t*)(" <> xName <> ".value))[1], ((var_t*)(" <> yName <> ".value))[0]);"
    -- , "  MATH_OP(MUL, " <> x1y1Name <> ", ((var_t*)(" <> xName <> ".value))[1], ((var_t*)(" <> yName <> ".value))[1]);"
    , ""
    , "  MATH_OP(SUB, " <> realName <> ", " <> x0y0Name <> ", " <> x1y1Name <> ");"
    , "  MATH_OP(ADD, " <> imagName <> ", " <> x0y1Name <> ", " <> x1y0Name <> ");"
    -- , "  assert(*(var_t*)(" <> realName <> ".value).tag == *(var_t*)(" <> imagName <> ".value).tag);"
    , "  INIT_COMPLEX_PAIR(" <> resultName <> ");"
    , "  COMPLEX_ASSIGN_REAL(" <> resultName <> ", " <> realName <> ");"
    , "  COMPLEX_ASSIGN_IMAG(" <> resultName <> ", " <> imagName <> ");"
    -- , "  ((var_t*)(" <> resultName <> ".value))[0] = *((var_t*)(" <> realName <> ".value));"
    -- , "  ((var_t*)(" <> resultName <> ".value))[1] = *((var_t*)(" <> imagName <> ".value));"
    , "} else {"
    , "  MATH_OP(MUL, " <> resultName <> ", " <> xName <> ", " <> yName <> ");"
    , "}"
    ]

genExp (Add x y) resultName = do
  xName <- freshCName
  yName <- freshCName

  complexPairNameX <- freshCName
  complexPairNameY <- freshCName

  realResultName <- freshCName
  imagResultName <- freshCName

  xCode <- genExp x xName
  yCode <- genExp y yName

  return $ unlines
    [ "// Add"
    , "var_t " <> xName <> ";"
    , "var_t " <> yName <> ";"
    , "var_t " <> complexPairNameX <> ";"
    , "var_t " <> complexPairNameY <> ";"
    , "var_t " <> realResultName <> ";"
    , "var_t " <> imagResultName <> ";"
    , xCode
    , yCode
    , "assert(" <> xName <> ".tag == " <> yName <> ".tag);"
    , ""
    , "if (" <> xName <> ".semantic_tag == EXPR_COMPLEX) {"
    , "  INIT_COMPLEX_PAIR(" <> resultName <> ");"
    , ""
    -- , "  MATH_OP(ADD, " <> realResultName <> ", ((var_t*)(" <> xName <> ".value))[0], ((var_t*)(" <> yName <> ".value))[0]);"
    -- , "  MATH_OP(ADD, " <> imagResultName <> ", ((var_t*)(" <> xName <> ".value))[1], ((var_t*)(" <> yName <> ".value))[1]);"
    , "  MATH_OP(SUB, " <> realResultName <> ", GET_COMPLEX_REAL(" <> xName <> "), GET_COMPLEX_REAL(" <> yName <> "));"
    , "  MATH_OP(SUB, " <> imagResultName <> ", GET_COMPLEX_IMAG(" <> xName <> "), GET_COMPLEX_IMAG(" <> yName <> "));"
    , ""
    , "  COMPLEX_ASSIGN_REAL(" <> resultName <> ", " <> realResultName <> ");"
    , "  COMPLEX_ASSIGN_IMAG(" <> resultName <> ", " <> imagResultName <> ");"
    , "} else {"
    , "  MATH_OP(ADD, " <> resultName <> ", " <> xName <> ", " <> yName <> ");"
    , "}"
    ]

genExp (FromIntegral (x :: GPUExp a) :: GPUExp b) resultName = do
  xName <- freshCName

  xCode <- genExp x xName

  case getComplexType_maybe (Proxy :: Proxy b) of
    Just (ty, tag) -> do

      let (tyX, tagX) = numTypeInfo (undefined :: a)

      zeroName <- freshCName

      castedXName <- freshCName

      return $ unlines
        [ "// Complex FromIntegral with ty = " <> ty
        , "var_t " <> xName <> ";"
        , "var_t " <> castedXName <> ";"
        , "var_t " <> zeroName <> ";"
        , zeroName <> ".tag = " <> tag <> ";"
        , zeroName <> ".semantic_tag = NO_SEMANTIC_TAG;"
        , zeroName <> ".value = malloc(sizeof(" <> ty <> "));"
        , "*((" <> ty <> "*)(" <> zeroName <> ".value)) = 0;"
        , xCode
        , castedXName <> ".tag = " <> tag <> ";"
        , castedXName <> ".semantic_tag = NO_SEMANTIC_TAG;"
        , castedXName <> ".value = malloc(sizeof(" <> ty <> "));"
        , "CAST_TO(" <> castedXName <> ", " <> ty <> ", " <> xName <> ");"
        -- , "*(" <> ty <> "*)(" <> castedXName <> ".value) = *(" <> ty <> "*)(" <> xName <> ".value);"
        , "INIT_COMPLEX(" <> resultName <> ", " <> ty <> ", " <> tag <> ");"
        , "COMPLEX_ASSIGN_REAL(" <> resultName <> ", " <> castedXName <> ");"
        , "COMPLEX_ASSIGN_IMAG(" <> resultName <> ", " <> zeroName <> ");"
        -- , resultName <> ".value = malloc(2*sizeof(var_t));"
        -- , "((var_t*)(" <> resultName <> ".value))[0] = " <> xName <> ";"
        -- , "*(" <> ty <> "*)(((var_t*)(" <> resultName <> ".value))[1].value) = 0;"
        ]
    Nothing -> do
      let (ty, tag) = numTypeInfo (undefined :: b)

      return $ unlines
        [ "var_t " <> xName <> ";"
        , xCode
        , ""
        , "CAST_TO(" <> resultName <> ", " <> ty <> ", " <> xName <> ");"
        ]

genExp (Sqrt x) resultName = do
  xName <- freshCName

  xCode <- genExp x xName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , xCode
    , ""
    , "MATH_SQRT(" <> resultName <> ", " <> xName <>");"
    ]

genExp (Lte x y) resultName = do
  xName <- freshCName
  yName <- freshCName

  xCode <- genExp x xName
  yCode <- genExp y yName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , "var_t " <> yName <> ";"
    , xCode
    , yCode
    , "assert(" <> xName <> ".tag == " <> yName <> ".tag);"
    , resultName <> ".tag = EXPR_BOOL;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(bool));"
    , "COMPARE(LTE, " <> resultName <> ", " <> xName <> ", " <> yName <> ");"
    ]

genExp (Gt x y) resultName = do
  xName <- freshCName
  yName <- freshCName

  xCode <- genExp x xName
  yCode <- genExp y yName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , "var_t " <> yName <> ";"
    , xCode
    , yCode
    , "assert(" <> xName <> ".tag == " <> yName <> ".tag);"
    , resultName <> ".tag = EXPR_BOOL;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(bool));"
    , "COMPARE(GT, " <> resultName <> ", " <> xName <> ", " <> yName <> ");"
    ]

genExp (Equal x y) resultName = do
  xName <- freshCName
  yName <- freshCName

  xCode <- genExp x xName
  yCode <- genExp y yName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , "var_t " <> yName <> ";"
    , xCode
    , yCode
    , "assert(" <> xName <> ".tag == " <> yName <> ".tag);"
    , ""
    , resultName <> ".tag = EXPR_BOOL;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(bool));"
    , "if (" <> resultName <> ".tag == EXPR_COMPLEX) {"
    , "  COMPARE(EQUAL, " <> resultName <> ", ((var_t*)(" <> xName <> ".value))[0], ((var_t*)(" <> yName <> ".value))[0]);"
    , "  COMPARE(EQUAL, " <> resultName <> ", ((var_t*)(" <> xName <> ".value))[1], ((var_t*)(" <> yName <> ".value))[1]);"
    , "} else {"
    , "  COMPARE(EQUAL, " <> resultName <> ", " <> xName <> ", " <> yName <> ");"
    , "}"
    ]

genExp (IfThenElse b t f) resultName = do
  bName <- freshCName
  tName <- freshCName
  fName <- freshCName

  bCode <- genExp b bName
  tCode <- genExp t resultName
  fCode <- genExp f resultName

  return $ unlines
    [ "var_t " <> bName <> ";"
    , "var_t " <> tName <> ";"
    , "var_t " <> fName <> ";"
    , ""
    , bCode
    , ""
    , "if (*(bool*)(" <> bName <> ".value)) {"
    , tCode
    , "} else {"
    , fCode
    , "}"
    ]

genExp (TailRec lamExp@(Lam (Name n) body)) resultName = do
  genExp lamExp resultName

genExp (PairExp x y) resultName = do
  xName <- freshCName
  yName <- freshCName

  xCode <- genExp x xName
  yCode <- genCdr y yName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , "var_t " <> yName <> ";"
    , xCode
    , yCode
    , resultName <> ".tag = EXPR_PAIR;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(var_t)*2);"
    , stmt $ cCall "copyVar" [addrOf (cCast "var_t*" (resultName # "value") ! 0), addrOf xName]
    , stmt $ cCall "copyVar" [addrOf (cCast "var_t*" (resultName # "value") ! 1), addrOf yName]
    -- , "((var_t*)(" <> resultName <> ".value))[0] = " <> xName <> ";"
    -- , "((var_t*)(" <> resultName <> ".value))[1] = " <> yName <> ";"
    ]

genExp (Repped x) resultName = error "genExp: Repped" --genExp (rep x) resultName

genExp (LeftExp x) resultName = do
  xName <- freshCName

  xCode <- genExp x xName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , xCode
    , resultName <> ".tag = EXPR_EITHER_LEFT;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(var_t));"
    , "*(var_t*)(" <> resultName <> ".value) = " <> xName <> ";"
    ]

genExp (RightExp x) resultName = do
  xName <- freshCName

  xCode <- genExp x xName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , xCode
    , resultName <> ".tag = EXPR_EITHER_RIGHT;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(var_t));"
    , "*(var_t*)(" <> resultName <> ".value) = " <> xName <> ";"
    ]

genExp UnitExp resultName =
  return $ unlines
    [ resultName <> ".tag = EXPR_UNIT;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = 0;"
    ]

genExp (CharLit c) resultName =
  return $ unlines
    [ resultName <> ".tag = EXPR_CHAR;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , "*(char*)(" <> resultName <> ".value) = " <> show c <> ";"
    ]

genExp (Ord x) resultName = do
  xName <- freshCName

  xCode <- genExp x xName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , xCode
    , ""
    , resultName <> ".tag = EXPR_INT;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , "*(int*)(" <> resultName <> ".value) = *(int*)(" <> xName <> ".value);"
    ]
genExp (CastExp x) resultName = genExp x resultName

genExp (FromEnum x) resultName = error "genExp: TODO: implement FromEnum"

genExp (SumMatchExp _ _) _ = error "genExp: SumMatchExp"
genExp (NullaryMatch _) _ = error "genExp: NullaryMatch"
genExp (OneSumMatch _) _ = error "genExp: OneSumMatch"
genExp EmptyMatch _ = error "genExp: EmptyMatch"

genCdr :: forall a. GPUExp a -> CName -> CodeGen CCode
genCdr x@(PairExp {}) resultName
  | Nothing <- eqT :: Maybe (a :~: Complex Double)
      = genExp x resultName
genCdr x resultName = do
  xName <- freshCName

  xCode <- genExp x xName

  return $ unlines
    [ "var_t " <> xName <> ";"
    , xCode
    , resultName <> ".tag = EXPR_PAIR;"
    , resultName <> ".semantic_tag = NO_SEMANTIC_TAG;"
    , resultName <> ".value = malloc(sizeof(var_t)*2);"
    -- , "((var_t*)(" <> resultName <> ".value))[0] = " <> xName <> ";"
    , stmt $ cCall "copyVar" [addrOf (cCast "var_t*" (resultName # "value") ! 0), addrOf xName]
    , "((var_t*)(" <> resultName <> ".value))[1].tag = EXPR_NIL;"
    , "((var_t*)(" <> resultName <> ".value))[1].semantic_tag = NO_SEMANTIC_TAG;"
    , "((var_t*)(" <> resultName <> ".value))[1].value = NULL;"
    ]

genCaseExp :: CName -> GPUExp (SumMatch a r) -> CName -> CodeGen CCode
genCaseExp s (CastExp x) resultName = genCaseExp s x resultName
genCaseExp s (OneSumMatch p) resultName = genProdMatch s p resultName
genCaseExp s (SumMatchExp x y) resultName = do
  s' <- freshCName

  thenPart <- genProdMatch s' x resultName
  elsePart <- genCaseExp s' y resultName

  return $ unlines
    [ "assert(" <> s <> ".tag == EXPR_EITHER_LEFT || " <> s <> ".tag == EXPR_EITHER_RIGHT);"
    -- , "var_t " <> s' <> " = *(var_t*)(" <> s <> ".value);"
    , "var_t " <> s' <> ";"
    , stmt $ cCall "copyVar" [addrOf s', cCast "var_t*" (s # "value")]

    , "if (" <> s <> ".tag == EXPR_EITHER_LEFT) {"
    , thenPart
    , "} else if (" <> s <> ".tag == EXPR_EITHER_RIGHT) {"
    , elsePart
    , "}"
    ]

-- TODO: Implement with multiple actual function calls?
genProdMatch :: CName -> GPUExp (ProdMatch a r) -> CName -> CodeGen CCode
genProdMatch s p resultName = do
  r <- getProdMatchLams p
  case r of
    Right x ->
      -- Nullary
      genExp x resultName

    Left (argUniqs0, innerLam) -> do
      -- Not nullary
      let lastArg = last argUniqs0
          argUniqs = init argUniqs0

      let itemNames = map (\i -> pairCar i s <> "/* item #" <> show i <> " */") [0..length argUniqs0-1]

      -- TODO: Should the enivronment be saved and restored?
      -- cg_modifyNames (zip argUniqs itemNames)

      le <- fmap cg_le get
      withLocalEnv (le_modifyNames le (zip argUniqs itemNames))
        $ buildAndCall innerLam (pairCar (length argUniqs0-1) s <> "/* item #" <> show (length argUniqs0-1) <> " */") resultName

pairCar :: Int -> CName -> CCode
pairCar 0 p = "((var_t*)" <> p <> ".value)[0]"
pairCar 1 p = "((var_t*)" <> pairCdr 0 p <> ".value)[0]"
pairCar depth p =
  "((var_t*)(" <> (pairCdr (depth-1) p) <> ").value)[0]"

pairCdr :: Int -> CName -> CCode
pairCdr 0 p = "((var_t*)" <> p <> ".value)[1]"
pairCdr depth p =
  "((var_t*)" <> (pairCdr (depth-1) p) <> ".value)[1]"

-- | Get argument uniques for each lambda and also the innermost lambda
getProdMatchLams :: GPUExp (ProdMatch a r) -> CodeGen (Either ([Int], SomeLambda) (GPUExp r))
getProdMatchLams (CastExp x) = getProdMatchLams x
getProdMatchLams (OneProdMatch (Lam (Name n) _)) = do
  lam <- lookupLambda n
  return $ Left ([n], lam)

getProdMatchLams (ProdMatchExp (Lam (Name n) x)) = do
  r <- getProdMatchLams x
  case r of
    Left (restNs, lam) -> return $ Left (n:restNs, lam)
    Right _ -> error "getProdMatchLams"

getProdMatchLams (NullaryMatch x) = return $ Right x

lambdaCName_byInt :: Int -> CName
lambdaCName_byInt i = "lam_" <> show i

lambdaCName :: SomeLambda -> CName
lambdaCName (SomeLambda cl) = lambdaCName_byInt (getNameIdent (lambda_name cl))

genLambda :: SomeLambda -> CodeGen CCode
genLambda sc@(SomeLambda c) = do
  rName <- freshCName

  currEnv <- fmap cg_le get

  let modifyFvEnv =
            le_modifyName
                (le_modifyNames
                  currEnv
                  (zipWith
                    (\(SomeName fv) i -> (getNameUniq fv, fvEnvItem i))
                    (lambda_fvs c)
                    [0..]))

                    -- (zip (map (\(SomeName n) -> getNameUniq n) (lambda_fvs c)) -- TODO: Does this put things in the correct order?
                    --    fvIxes))

  let localEnv
        | lambda_isTailRec c =
            modifyFvEnv
                (getNameUniq (lambda_name c))
                rName
        | otherwise =
            modifyFvEnv
                (getNameUniq (lambda_name c))
                "arg"

  body <- withLocalEnv localEnv (genExp (lambda_body c) rName)

  tempName <- freshCName

  let body'
        | lambda_isTailRec c =
            unlines
              [ rName <> ".tag = EXPR_STEP;"
              , rName <> ".semantic_tag = NO_SEMANTIC_TAG;"
              , rName <> ".value = malloc(sizeof(var_t));"
              -- , "*(var_t*)(" <> rName <> ".value) = arg;"
              , stmt $ cCall "copyVar" [cCast "var_t*" (rName # "value"), addrOf "arg"]
              , "while (" <> rName <> ".tag != EXPR_DONE) {"
              , body
              , "}"
              , ""
              , "var_t " <> tempName <> ";"
              , stmt $ cCall "copyVar"
                          [addrOf tempName, cCast "var_t*" (rName # "value")]
              -- , "var_t " <> tempName <> " = *(var_t*)(" <> rName <> ".value);"
              , rName <> ".tag = " <> tempName <> ".tag;"
              , rName <> ".semantic_tag = " <> tempName <> ".semantic_tag;"
              , stmt $ cCall "copyVar" [cCast "var_t*" (rName # "value"), cCast "var_t*" (tempName # "value")]
              -- , rName <> ".value = " <> tempName <> ".value;"
              ]
        | otherwise = body

  return $
    unlines
      [ "var_t " <> lambdaCName sc <> "(var_t arg, struct closure_t* self) {"
      , "  var_t " <> rName <> ";"
      , body'
      , "  return " <> rName <> ";"
      , "}"
      ]
  where
    fvEnvItem i = "self->fv_env[" <> show i <> "]"

    Name n = lambda_name c

genLambdas :: GPUExp a -> CodeGen CCode
genLambdas e = do
  let lambdas = collectLambdas e

  unlines <$> mapM genLambda lambdas

buildAndCall :: SomeLambda -> CName -> CName -> CodeGen CCode
buildAndCall sLam argName resultName = do
  closureVarName <- freshCName

  closureCode <- buildClosure sLam closureVarName
  callCode <- callClosure sLam closureVarName argName resultName

  return $ unlines
    [ "closure_t " <> closureVarName <> ";"
    , closureCode
    , callCode
    ]

buildClosure :: HasCallStack => SomeLambda -> CName -> CodeGen CCode
buildClosure sc@(SomeLambda c) closureVarName = do
  currEnv <- fmap cg_le get

  let fvs = lambda_fvs c
  traceM $ "////////// fvs = " ++ show fvs
  init_fvEnv <- forM (zip [0..] fvs) (\(i, n@(SomeName fv)) -> do
                            fvName <- cg_lookup (getNameUniq fv)
                            return (cCall "copyVar" [addrOf ((closureVarName # "fv_env") ! i), addrOf fvName] <> ";  // For FV with id " <> (show (getNameUniq fv) <> " with Haskell type " <> showNameType n)))
                            -- return (closureVarName <> ".fv_env[" <> show i <> "] = " <> fvName <> ";  // For FV with id " <> (show (getNameUniq fv) <> " with Haskell type " <> showNameType n)))

  return $
    unlines
      [ closureVarName <> ".fv_env = malloc(sizeof(var_t)*" <> show fvCount <> ");"
      , closureVarName <> ".env_size = " <> show fvCount <> ";"
      , closureVarName <> ".fn = &" <> lambdaCName sc <> ";"
      , unlines init_fvEnv
      ]
  where
    fvCount = length (lambda_fvs c)

callClosure :: SomeLambda -> CName -> CName -> CName -> CodeGen CCode
callClosure (SomeLambda (Lambda { lambda_name })) closureName argName resultName = do
  r <- freshCName
  -- copiedArg <- freshCName
  return $
    unlines
      [
       -- "var_t " <> copiedArg <> ";"
      -- , stmt $ cCall "copyVar" [addrOf copiedArg, addrOf argName]
        "var_t " <> r <> " = " <> closureName <> ".fn(" <> argName <> ", &" <> closureName <> ");"
      , stmt $ cCall "copyVar" [addrOf resultName, addrOf r]
      ]

showNameType :: SomeName -> String
showNameType (SomeName n) = show (typeRep n)

-- TODO: Determine if the warning this gives is a sign of a problem (it
-- is likely to be, at least, a minor bug in GHC):
--
-- Update: Despite the warning, it seems ok.
numTypeInfo :: forall n a. (HasCallStack, Typeable n) => n -> (CName, CName)
numTypeInfo arg =
 case eqT :: Maybe (n :~: Int) of
   Just _ -> ("int", "EXPR_INT")
   Nothing ->
     case eqT :: Maybe (n :~: Float) of
       Just _ -> ("float", "EXPR_FLOAT")
       Nothing ->
         case eqT :: Maybe (n :~: Double) of
           Just _ -> ("double", "EXPR_DOUBLE")
           Nothing ->
             case eqT :: Maybe (n :~: Integer) of
               Just _ -> trace "Warning: using int to represent Integer" ("int", "EXPR_INT")
               Nothing -> error $ "numTypeInfo: invalid lit type: " ++ show (typeOf arg)

getComplexType_maybe :: forall a. Typeable a => Proxy a -> Maybe (CName, CName)
getComplexType_maybe Proxy =
  case eqT :: Maybe (a :~: Complex Float) of
    Just Refl -> Just ("float", "EXPR_FLOAT")
    Nothing ->
      case eqT :: Maybe (a :~: Complex Double) of
        Just Refl -> Just ("double", "EXPR_DOUBLE")
        Nothing -> Nothing

getComplex_maybe :: Typeable a => a -> Maybe (ShowablePair, (CName, CName))
getComplex_maybe x =
  case cast x :: Maybe (Complex Float) of
    Just (a :+ b) -> Just (ShowablePair a b, ("float", "EXPR_FLOAT"))
    Nothing ->
      case cast x :: Maybe (Complex Double) of
        Just (a :+ b) -> Just (ShowablePair a b, ("double", "EXPR_DOUBLE"))
        Nothing -> Nothing

-- Tests and examples --
cTest :: GPUExp Int
cTest =
  CaseExp (Var (Name 0) :: GPUExp (Either Int Float))
    (SumMatchExp
      (OneProdMatch (Lam (Name 1) (Var (Name 1))))
      (OneSumMatch (OneProdMatch (Lam (Name 2) (Var (Name 0))))))

cTest2 :: GPUExp Int
cTest2 =
  CaseExp (Var (Name 0) :: GPUExp (Int, Float))
    (OneSumMatch
      (ProdMatchExp
        (Lam (Name 1)
          (OneProdMatch
            (Lam (Name 2)
              (Var (Name 1)))))))

cTest3 :: GPUExp Int
cTest3 =
  CaseExp (Var (Name 0) :: GPUExp (Int, Float))
    (OneSumMatch
      (ProdMatchExp
        (Lam (Name 1)
          (OneProdMatch
            (Lam (Name 2)
              (Mul (Var (Name 1)) (Var (Name 1))))))))

cTest4 :: GPUExp Int
cTest4 =
  CaseExp (Var (Name 0) :: GPUExp (Int, Float))
    (OneSumMatch
      (ProdMatchExp
        (Lam (Name 1)
          (OneProdMatch
            (Lam (Name 2)
              (Mul (Var (Name 1)) (Lit 11)))))))

(&) :: GPURep a => GPUExp a -> GPUExp (a -> b) -> GPUExp b
(&) = flip App

-- | isEven 6
cTest5 :: GPUExp Bool
cTest5 = Lit 6 &
  TailRec (Lam (Name 0 :: Name Int)
    (IfThenElse (Equal (Var (Name 0 :: Name Int)) (Lit 0))
      (DoneExp TrueExp)
      (IfThenElse (Equal (Var (Name 0 :: Name Int)) (Lit 1))
        (DoneExp FalseExp)
        (StepExp (Sub (Var (Name 0 :: Name Int)) (Lit 2))))))

-- | factorial 5
cTest6 :: GPUExp Int
cTest6 = --(PairExp (Lit 5) (Lit 1)) &
  (construct ((5,1) :: (Int, Int))) &
    TailRec (Lam (Name 0 :: Name (Int, Int))
      (CaseExp (Var (Name 0 :: Name (Int, Int)))
        (OneSumMatch
          (ProdMatchExp
            (Lam (Name 1 :: Name Int)
              (OneProdMatch
                (Lam (Name 2 :: Name Int)
                  (IfThenElse (Equal (Var (Name 1)) (Lit 0 :: GPUExp Int))
                    (DoneExp (Var (Name 2)))
                    (StepExp (PairExp (Sub (Var (Name 1)) (Lit 1)) (Mul (Var (Name 1)) (Var (Name 2)))))
                    ))))))))


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE FlexibleContexts #-}

module CPVO.IO.Fortran (
  loadFile
  ) where

--import Numeric.LinearAlgebra
--import Numeric.LinearAlgebra.Data hiding (find)
--import Data.List (findIndex,groupBy)
--import Data.Maybe (fromJust)
--import qualified Language.C.Inline as C

loadFile :: IO ()
loadFile = do
  putStrLn "=======done : loadFile@CPVO.IO.Fortran"


{-# LANGUAGE NoImplicitPrelude #-}
module Test where
import Data.Complex.Compat


-----------------------------------------------------------------------------
-- |
-- Module      :  Graphics.Rendering.Plot.Render.Text
-- Copyright   :  (c) A. V. H. McPhail 2010
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Rendering 'Figure's
--
-----------------------------------------------------------------------------

module Graphics.Rendering.Plot.Render.Text (
                                       -- * Rendering
                                       renderText
                                       , renderTextVertical
                                       -- * Internal
                                       , textSize
                                       , textSizeVertical
                                       , showText
                                       , formatText
                                       ) where

-----------------------------------------------------------------------------

--import Data.Either

--import Data.Packed.Vector
--import Numeric.LinearAlgebra.Linear

--import Data.Word

--import Data.Maybe

--import Data.Colour.SRGB
--import Data.Colour.Names

---import qualified Data.Array.IArray as A

import qualified Graphics.Rendering.Cairo as C
import qualified Graphics.Rendering.Pango as P

import Control.Monad.Reader
--import Control.Monad.State
--import Control.Monad.Trans

import Graphics.Rendering.Plot.Types
--import Graphics.Rendering.Plot.Defaults

import Graphics.Rendering.Plot.Figure.Text

import Graphics.Rendering.Plot.Render.Types

--import qualified Text.Printf as Printf

--import Prelude hiding(min,max)
--import qualified Prelude(max)

-----------------------------------------------------------------------------

textSize :: P.PangoLayout -> TextXAlign -> TextYAlign -> Double -> Double -> C.Render ((Double,Double),(Double,Double))
textSize l xa ya x y = do
                       (_,P.PangoRectangle _ _ w h) <- pango $ P.layoutGetExtents l
                       return ((xStart xa x w h,yStart ya y w h),(w,h))
    where xStart TLeft   x' w' _  = x' - w'
          xStart Centre  x' w' _  = x' - (w'/2)
          xStart TRight  x' _  _  = x'
          yStart TBottom y' _  h' = y' - h'
          yStart Middle  y' _  h' = y' - (h'/2)
          yStart TTop    y' _  _  = y'

textSizeVertical :: P.PangoLayout -> TextXAlign -> TextYAlign -> Double -> Double -> C.Render ((Double,Double),(Double,Double))
textSizeVertical l xa ya x y = do
                       (_,P.PangoRectangle _ _ w h) <- pango $ P.layoutGetExtents l
                       return ((xStart xa x w h,yStart ya y w h),(w,h))
    where xStart TLeft   x' _  w' = x' - w'
          xStart Centre  x' _  w' = x' - (w'/2)
          xStart TRight  x' _  _  = x'
          yStart TBottom y' _  _  = y'
          yStart Middle  y' h' _  = y' + (h'/2)
          yStart TTop    y' h' _  = y' + (h')

showText :: P.PangoLayout -> Double -> Double -> C.Render ()
showText pl x y = do
                  C.moveTo x y
                  P.showLayout pl

-----------------------------------------------------------------------------

formatText :: TextEntry -> Render TextEntry
formatText te@NoText          = return te
formatText (BareText s)       = do
                                to <- asks (_textoptions . _renderoptions)
                                return (FontText to s)
formatText (SizeText fz c s)  = do
                                to <- asks (_textoptions . _renderoptions)
                                return $ (FontText (changeFontSize fz $ changeFontColour c to) s)
formatText te@(FontText _ _)  = return te

{-
getTextSize :: Text -> Render (Double,Double)
getTextSize (Text Nothing s) = do
                               to <- asks _text
                               getTextSize (Text to s)
getTextSize (Text (Just (TextOptions (FontOptions ff fs fw) fz _)) s) = cairo $ do
                               C.selectFontFace ff fs fw
                               C.setFontSize fz
                               te <- C.textExtents s
                               return (C.textExtentsWidth te,C.textExtentsHeight te)
-}
renderText :: TextEntry -> TextXAlign -> TextYAlign -> Double -> Double -> Render (Double,Double)
renderText NoText              _  _  _ _ = return (0,0)
renderText te@(BareText _)     xa ya x y = do
                                           te' <- formatText te
                                           renderText te' xa ya x y
renderText te@(SizeText _ _ _) xa ya x y = do
                                           te' <- formatText te
                                           renderText te' xa ya x y
renderText (FontText to s)     xa ya x y = do
                                           pc <- asks _pangocontext
                                           cairo $ do
                                                   lo <- pango $ P.layoutText pc s
                                                   setTextOptions to lo
                                                   ((x',y'),twh) <- textSize lo xa ya x y
                                                   showText lo x' y'
                                                   return twh

renderTextVertical :: TextEntry -> TextXAlign -> TextYAlign -> Double -> Double -> Render (Double,Double)
renderTextVertical NoText              _  _  _ _ = return (0,0)
renderTextVertical te@(BareText _)     xa ya x y = do
                                                 te' <- formatText te
                                                 renderTextVertical te' xa ya x y
renderTextVertical te@(SizeText _ _ _) xa ya x y = do
                                                   te' <- formatText te
                                                   renderTextVertical te' xa ya x y
renderTextVertical (FontText to s)     xa ya x y = do
                                                   pc <- asks _pangocontext
                                                   cairo $ do
                                                           lo <- pango $ P.layoutText pc s
                                                           setTextOptions to lo
                                                           C.save
                                                           C.rotate (-pi/2)
                                                           P.updateLayout lo
                                                           ((x',y'),twh) <- textSizeVertical lo xa ya x y
                                                           showText lo (-y') (x')
                                                           C.restore
                                                           return twh

-----------------------------------------------------------------------------


{-# LANGUAGE AllowAmbiguousTypes, CPP, ConstraintKinds, DataKinds, FlexibleContexts, FlexibleInstances, GADTs, MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels, OverloadedStrings, PartialTypeSignatures, PolyKinds, QuasiQuotes, RankNTypes, ScopedTypeVariables     #-}
{-# LANGUAGE TypeApplications, TypeFamilyDependencies, TypeInType, TypeOperators, UndecidableInstances                               #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures -fconstraint-solver-iterations=10 #-}

{-| Talk to Tensorboard. You will need to run it manually with
'tensorboard --logdir <yourdir>'

It will occasionally update on its own.  We have fairly comprehensive
Tensorboard support and integration with showing images, plots, and graphs.
-}

module Torch.Tensorboard where
import           Control.Monad
import           Control.Monad.Logger
import           Data.Binary.Put
import           Data.Bits
import qualified Data.ByteString              as B
import qualified Data.ByteString              as BS
import qualified Data.ByteString.Internal     as BS
import qualified Data.ByteString.Lazy         as BL
import           Data.Coerce
import           Data.IORef
import qualified Data.Map                     as M
import           Data.ProtoLens
import           Data.ProtoLens.Labels        ()
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Text                    (Text)
import qualified Data.Text                    as T
import qualified Data.Text.Encoding           as T
import           Data.Time.Clock.POSIX
import           Data.Time.Format
import           Data.Time.LocalTime
import           Data.Vector.Storable         (Vector)
import qualified Data.Vector.Storable         as V
import           Foreign.C.Types
import           Foreign.ForeignPtr
import           Foreign.Ptr
import           GHC.Float
import           GHC.Int
import           GHC.Word
import qualified Graphics.Matplotlib          as M
import           Lens.Micro
import           Network.HostName
import qualified Statistics.Quantile          as S
import           System.Directory
import           System.FilePath
import           System.IO
import qualified Tensorboard.Proto.Attributes as PA
import qualified Tensorboard.Proto.Event      as PE
import qualified Tensorboard.Proto.Graph      as PG
import qualified Tensorboard.Proto.Summary    as PS
import qualified Tensorboard.Proto.Tensor     as PT
import qualified Torch.C.Tensor               as C
import qualified Torch.C.Types                as C
import qualified Torch.C.Variable             as C
import           Torch.Images
import           Torch.Internal.CRC32C
import           Torch.Tensor                 as Tensor
import           Torch.Types                  as Tensor
import           Torch.Visualization

-- | `SummaryWriter` keeps track of some basic information so we can talk to
-- Tensorboard. You'll be passing this around everywhere.
data SummaryWriter = SummaryWriter { runName      :: Text
                                   , runDirectory :: Text
                                   , runFilename  :: Text
                                   , lastStep     :: IORef Int
                                   }

-- | Make a new summary writer for a given directory and run name. We will also
-- add some information to the run time like the current host and local timestamp.
summaryWriter :: Text -> Text -> IO SummaryWriter
summaryWriter directory runname = do
  now <- getCurrentTime
  tz <- getCurrentTimeZone
  let time = formatTime defaultTimeLocale "%Y-%m-%d_%H:%M:%S" (utcToZonedTime tz now)
  let timeUnix = formatTime defaultTimeLocale "%s" (utcToZonedTime tz now)
  host <- T.pack <$> getHostName
  let name = runname <> "_" <> host <> "_" <> T.pack time
  let location = T.unpack directory </> T.unpack name
  createDirectoryIfMissing True location
  step <- newIORef 0
  pure $ SummaryWriter { runName = name
                       , runDirectory = T.pack location
                       , runFilename = "events.haskelltorch." <> T.pack timeUnix <> ".tfevents." <> host
                       , lastStep = step
                       }

-- | Advance the writer to the next step. This helps you keep track of the
-- iteration number without needing a manual counter.
nextStep :: SummaryWriter -> IO SummaryWriter
nextStep r = do
  modifyIORef' (lastStep r) (+ 1)
  pure r

-- | Explicitly set the step counter.
setStep :: SummaryWriter -> Int -> IO SummaryWriter
setStep sw s = do
  writeIORef (lastStep sw) s
  pure sw

-- | Add a scalar to Tensorboard with a given name and value.
addScalar :: Real a => SummaryWriter -> Text -> a -> IO ()
addScalar summary name value =
  writeEvent summary (EventSummary name $ SummaryValue $ fromRational $ toRational value)

-- | Add a histogram to Tensorboard, we do some internal work to make sure the
-- number of bins and their sizes make sense.
addHistogram ::
  (TensorConstraints ty 'KCpu sz) =>
  SummaryWriter -> Text -> Tensor ty 'KCpu sz -> IO ()
addHistogram summary name tensor = do
  bins  <- histogramBinsAuto <$> V.map toDoubleC <$> toVector tensor
  withSomeSing (fromIntegral bins) (\ n@(SNat :: SNat m) -> do
    hist_ <- toVector =<< histc @m tensor 0 0
    let hist = V.map toDoubleC hist_
    min   <- toDouble <$> (fromScalar =<< Tensor.min tensor)
    max   <- toDouble <$> (fromScalar =<< Tensor.max tensor)
    let step = (max-min)/2
    let binLimits = V.generate bins (\i -> step * fromIntegral i)
    sum     <- fromScalar =<< Tensor.sum tensor
    sumSq <- fromScalar =<< Tensor.sum =<< Tensor.pow tensor =<< toScalar 2
    writeEvent summary
      (EventSummary name
       (SummaryHistogram
         { shMin = min
         , shMax = max
         , shNum = fromIntegral $ numel tensor
         , shSum = toDouble sum
         , shSumSq = toDouble sumSq
         , shBucketLimit = binLimits
         , shBucket = hist })))

-- | Add an RGB image tensor to the Tensorboard.
addImageRGB :: (KnownNat szw, KnownNat szh, SingI ty)
            => SummaryWriter -> Text -> Tensor ty 'KCpu '[3, szh, szw] -> IO ()
addImageRGB summary name image = do
    (bytes, len) <- V.unsafeToForeignPtr0 <$> rgbTensorToVector image
    (h,w,c,b) <- pure $ case V.toList $ size image of
       [3,h,w] -> (h,w,3,bytes)
       -- TODO Move this to compile time
       _       -> error $ "Don't know how to display a tensor with these dimensions as an RGB image " ++ show (size image)
    writeEvent summary
      (EventSummary name
       (SummaryImage
        { siHeight     = fromIntegral h
        , siWidth      = fromIntegral w
        , siColorspace = c
        , siBytes      = BS.PS b 0 len}))

-- | Add a greyscale image tensor to the Tensorboard.
addImageGrey :: (KnownNat szw, KnownNat szh, SingI ty)
             => SummaryWriter -> Text -> Tensor ty 'KCpu '[1, szh, szw] -> IO ()
addImageGrey summary name image = do
    (bytes, len) <- V.unsafeToForeignPtr0 <$> greyTensorToVector image
    (h,w,c,b) <- pure $ case V.toList $ size image of
       [1,h,w] -> (h,w,1,bytes)
       -- TODO Move this to compile time
       _       -> error $ "Don't know how to display a tensor with these dimensions as an Greyscale image " ++ show (size image)
    writeEvent summary
      (EventSummary name
       (SummaryImage
        { siHeight     = fromIntegral h
        , siWidth      = fromIntegral w
        , siColorspace = c
        , siBytes      = BS.PS b 0 len}))

-- | Add a tensor that contains a grid of RGB images, like a batch of RGB images
-- with the leading dimension being the number of images. Grids are padded with
-- the given Size, constructed with the @Size@ datatype. They are padded with a
-- given constant, usually 0.
addImageRGBGrid ::
  forall (imagesPerRow :: Nat) (padding :: Nat) (rows :: Nat) (szh :: Nat) (szw :: Nat) (nr :: Nat) (ty :: TensorType).
    _
  => SummaryWriter
  -> Text
  -> Size imagesPerRow
  -> Padding padding
  -> TensorTyToHs ty
  -> Tensor ty 'KCpu '[nr, 3, szh, szw]
  -> IO ()
addImageRGBGrid summary name sz padding padValue images = do
  grid <- makeRGBGrid sz padding padValue images
  addImageRGB summary name grid

-- | Add a matplotlib plot to Tensorboard. This uses the Haskell matplotlib
-- bindings from the `matplotlib` package,
-- `https://github.com/abarbu/matplotlib-haskell`.
addPlot :: SummaryWriter -> Text -> M.Matplotlib -> IO ()
addPlot summary name mplot = do
  (fname, handle) <- openTempFile "/tmp/" "plot.png"
  hClose handle
  M.file fname mplot
  (h,w,_) <- fileImageProperties (T.pack fname)
  withSomeSing (fromIntegral h)
    (\ szh@(SNat :: SNat szh) ->
        withSomeSing (fromIntegral w)
          (\ szw@(SNat :: SNat szw) -> do
              tensor <- typed @TByte <$> readRGBTensorFromFile @szh @szw (T.pack fname)
              addImageRGB summary name tensor))
  fileExists <- doesFileExist fname
  when fileExists (removeFile fname)

-- | Add the trace of a computation to Tensorboard as a graph.
-- TODO Support naming nodes, more type information, more device information
addGraph :: SummaryWriter -> Text -> ForeignPtr C.CTracingState -> IO ()
addGraph summary name trace = do
  (inputs, nodes, outputs, ret, block) <-
    withForeignPtr trace C.tracing_state_graph
  inputs' <- mapM (\i -> do
                    b <- cbool <$> C.check_value_tensor i
                    na <- T.pack <$> C.value_name i
                    if b then do
                      sz <- C.value_sizes i
                      ty <- C.value_scalar_type i
                      pure (na, Just sz, Just ty)
                      else
                      pure (na, Nothing, Nothing))
            $ V.toList inputs
  nodes' <- mapM (\(n,nr) -> do
                   ins <- C.node_inputs n
                   outs <- C.node_outputs n
                   kind <- C.node_kind n
                   b <- C.node_has_attribute n "value"
                   attr <- if cbool b then do
                            ak <- C.node_attribute_kind n "value"
                            case ak of
                              C.AttributeKindFloat -> do
                                f <- C.node_get_attribute_float n "value"
                                pure $ Just $ defMessage & #f .~ coerce f
                              C.AttributeKindInt -> do
                                i <- C.node_get_attribute_int n "value"
                                pure $ Just $ defMessage & #i .~ fromIntegral i
                              C.AttributeKindString -> do
                                s <- C.node_get_attribute_string n "value"
                                pure $ Just $ defMessage & #s .~ T.encodeUtf8 (T.pack s)
                              C.AttributeKindTensor -> do
                                tensor <- C.node_get_attribute_tensor n "value"
                                ty <- C.getType tensor
                                tensor' <- C.contiguous_mm tensor (fromIntegral $ fromEnum C.MemoryFormatContiguous)
                                ptr <- newForeignPtr_ =<< castPtr <$> C.data_ptr tensor'
                                len <- fromIntegral <$> C.numel tensor'
                                let get :: V.Storable a => [a]
                                    get = V.toList $ V.unsafeFromForeignPtr0 (castForeignPtr ptr) len
                                pure $ Just $ defMessage & #tensor .~
                                  case ty of
                                    C.ScalarTypeByte ->
                                      defMessage & #dtype .~ PS.DT_UINT8
                                                 & #intVal .~ map (fromIntegral :: Word8 -> Int32) get
                                    C.ScalarTypeChar ->
                                      defMessage & #dtype .~ PS.DT_INT8
                                                 & #intVal .~ map (fromIntegral :: Int8 -> Int32) get
                                    C.ScalarTypeShort ->
                                      defMessage & #dtype .~ PS.DT_INT16
                                                 & #intVal .~ map (fromIntegral :: Int16 -> Int32) get
                                    C.ScalarTypeInt ->
                                      defMessage & #dtype .~ PS.DT_INT32
                                                 & #intVal .~ get
                                    C.ScalarTypeLong ->
                                      defMessage & #dtype .~ PS.DT_INT64
                                                 & #int64Val .~ get
                                    C.ScalarTypeHalf ->
                                      defMessage & #dtype .~ PS.DT_HALF
                                                 & #halfVal .~ get
                                    C.ScalarTypeFloat ->
                                      defMessage & #dtype .~ PS.DT_FLOAT
                                                 & #floatVal .~ get
                                    C.ScalarTypeDouble ->
                                      defMessage & #dtype .~ PS.DT_DOUBLE
                                                 & #doubleVal .~ get
                                    C.ScalarTypeUndefined ->
                                      defMessage & #dtype .~ PS.DT_INVALID
                              _ -> pure Nothing
                     else pure Nothing
                   ins' <- mapM (\v -> do
                                n <- C.value_name v
                                t <- C.value_type v
                                ct <- convertCType t
                                pure $ TempGraphValue v n t ct)
                          $ V.toList ins
                   outs' <- mapM (\v -> do
                                n <- C.value_name v
                                t <- C.value_type v
                                ct <- convertCType t
                                pure $ TempGraphValue v n t ct)
                           $ V.toList outs
                   pure (n, TempGraphNode nr kind ins' outs' attr)) $ zip (V.toList nodes) [0..]
  let nodesMap = M.fromList nodes'
  let insToNodes  = M.fromListWith (++)
                  $ concatMap (\(n, TempGraphNode nr kinds ins outs attr) ->
                             map (\(TempGraphValue v _ _ _) -> (v,[n])) ins) nodes'
  let outsToNodes = M.fromListWith (++)
                  $ concatMap (\(n, TempGraphNode nr kinds ins outs attr) ->
                             map (\(TempGraphValue v _ _ _) -> (v,[n])) outs) nodes'
  let nodeName nr kind = T.replace "prim::" "" $ T.replace "aten::" "" $ T.pack kind <> "_"<> T.pack (show nr)
  writeEvent summary
    (EventGraph name
      (map (\(name, size, ty) ->
               case (size, ty) of
                 (Just sz, Just t) ->
                   defMessage & #name .~ "input_" <> name
                              & #op .~ "Input"
                              & #attr .~ M.fromList [("_output_shapes",
                                                      defMessage & #list .~
                                                       (defMessage & #shape .~
                                                         [defMessage & #dim .~
                                                          map (\x -> defMessage & #size .~ x) (V.toList sz)]))]
                 (_, _) ->
                   defMessage & #name .~ "input_" <> name
                              & #op .~ "Input")
         inputs'
       ++
       map (\(node,(TempGraphNode nr kind ins outs attr)) ->
               defMessage & #name .~ nodeName nr kind
                          & #op .~ T.pack kind
                          & #device .~ T.pack ""
                          & #input .~ map (\(TempGraphValue v n t ct) ->
                                                case M.lookup v outsToNodes of
                                                  Nothing -> T.pack ("input_" ++ n)
                                                  Just [outNode] -> case M.lookup outNode nodesMap of
                                                                   Just (TempGraphNode nr kind _ _ _) -> nodeName nr kind
                                                                   Nothing                            -> error "Can't find node") ins
                          & #attr .~ M.fromList([("_output_shapes",
                                                   defMessage
                                                   & #list .~ (defMessage
                                                               & #shape .~
                                                                  (map (\(TempGraphValue v n t ct) ->
                                                                           case ct of
                                                                             TraceTypeSimple _ _ ->
                                                                               defMessage & #dim .~ [defMessage & #size .~ 1]
                                                                             TraceTypeTensor _ (Just sz) _ _ ->
                                                                               defMessage & #dim .~ (map (\x -> defMessage & #size .~ x)
                                                                                (V.toList sz))
                                                                             _ -> defMessage)
                                                                    outs)))]
                                                ++ (case attr of
                                                     Nothing -> []
                                                     Just a  -> [("value",a)])))
           nodes'))

-- | A convenience function, add the summary statistics and histograms for a
-- tensor.
addTensorSummary :: (Num (TensorTyToHs ty), Num (TensorTyToHsC ty),
                     V.Storable (TensorTyToHs ty), V.Storable (TensorTyToHsC ty),
                     SingI sz, SingI ty) =>
                   SummaryWriter -> Text -> Tensor ty 'KCpu sz -> IO ()
addTensorSummary summary name tensor = do
  mean   <- fromScalar =<< Tensor.mean tensor
  median <- fromScalar =<< Tensor.median tensor
  var    <- fromScalar =<< Tensor.var tensor True
  min    <- fromScalar =<< Tensor.min tensor
  max    <- fromScalar =<< Tensor.max tensor
  addScalar    summary ("summaries/" <> name <> "/mean")   $ toDouble mean
  addScalar    summary ("summaries/" <> name <> "/median") $ toDouble median
  addScalar    summary ("summaries/" <> name <> "/stddev") $ Prelude.sqrt $ toDouble var
  addScalar    summary ("summaries/" <> name <> "/min")    $ toDouble min
  addScalar    summary ("summaries/" <> name <> "/max")    $ toDouble max
  addHistogram summary ("summaries/" <> name <> "/histogram") tensor

-- * Basic heuristics

-- | The numpy auto histogram mode
histogramBinsAuto :: Vector Double -> Int
histogramBinsAuto x | V.length x < 1000 = histogramBinsSturges x
                    | otherwise = histogramBinsFD x

-- | The Freedman-Diaconis histogram bin estimator.
-- bins = 2 \frac{IQR}{n^{1/3}}
histogramBinsFD :: Vector Double -> Int
histogramBinsFD v =
  ceiling $ (V.maximum v - V.minimum v) /
               (2 * (S.weightedAvg 3 4 v - S.weightedAvg 1 4 v)
                  / (fromIntegral (V.length v) ** (1/3)))

-- | Sturges's histogram bin estimator.
histogramBinsSturges :: Vector Double -> Int
histogramBinsSturges v = Prelude.round $ logBase 2 (fromIntegral $ V.length v) + 1

-- * Internal

-- | This is internal, you never need to use it manually. It's what we fill out
-- with the various functions in this module and then serialize for Tensorboard.
data Summary = SummaryValue Double
             | SummaryHistogram { shMin         :: Double
                                , shMax         :: Double
                                , shNum         :: Double
                                , shSum         :: Double
                                , shSumSq       :: Double
                                , shBucketLimit :: Vector Double
                                , shBucket      :: Vector Double }
             | SummaryImage { siHeight     :: Int32
                            , siWidth      :: Int32
                            , siColorspace :: Int32
                            , siBytes      :: BS.ByteString }

-- | This is internal.
data Node = Node { nName   :: Text
                 , nOp     :: Text
                 , nInput  :: [Text]
                 , nDevice :: Text
                 }

-- | Internal. Used to capture the various kinds of entities that Tensorboard
-- can understand.
data Event = EventVersion
           | EventGraph Text [PG.NodeDef]
           | EventSummary Text Summary
           | EventMessage LogLevel Text
           | EventSessionLog
           | EventMetadata
           | EventMetaGraph

-- | Internal
data TempGraphValue = TempGraphValue (Ptr C.CJitValue) String (Ptr C.CType) TraceType
  deriving (Show)
-- | Internal
data TempGraphNode  = TempGraphNode Int String [TempGraphValue] [TempGraphValue] (Maybe PA.AttrValue)
  deriving (Show)

-- | Internal
eventToProto :: Int -> Event -> IO PE.Event
eventToProto step event = do
  (time :: Double) <- realToFrac <$> getPOSIXTime
  let d = case event of
            EventVersion -> defMessage & #fileVersion .~ "brain.Event:2"
            EventMessage lvl txt ->
              defMessage & #logMessage .~ (defMessage & #level .~ (case lvl of
                                                        LevelDebug         -> PE.LogMessage'DEBUG
                                                        LevelInfo          -> PE.LogMessage'INFO
                                                        LevelWarn          -> PE.LogMessage'WARN
                                                        LevelError         -> PE.LogMessage'ERROR
                                                        LevelOther "fatal" -> PE.LogMessage'FATAL
                                                        LevelOther _       -> PE.LogMessage'UNKNOWN)
                                        & #message .~ txt)
            EventSummary tag (SummaryValue val) ->
              defMessage & #summary .~
              (defMessage & #value .~
                [defMessage & #tag .~ tag
                -- TODO Metadata!
                     & #simpleValue .~ double2Float val])
            EventSummary tag sh@SummaryHistogram {} ->
              defMessage & #summary .~
              (defMessage & #value .~
                [defMessage & #tag .~ tag
                -- TODO Metadata!
                     & #histo .~
                     (defMessage
                          & #min .~ shMin sh
                          & #max .~ shMax sh
                          & #num .~ shNum sh
                          & #sum   .~ shSum sh
                          & #sumSquares .~ shSumSq sh
                          & #bucketLimit .~ V.toList (shBucketLimit sh)
                          & #bucket .~ V.toList (shBucket sh))])
            EventSummary tag sh@SummaryImage{} ->
              defMessage & #summary .~
              (defMessage & #value .~
                [defMessage & #tag .~ tag
                -- TODO Metadata!
                     & #image .~
                     (defMessage
                          & #height             .~ siHeight sh
                          & #width              .~ siWidth sh
                          & #colorspace         .~ siColorspace sh
                          & #encodedImageString .~ siBytes sh)])
            EventGraph tag nodes ->
              defMessage & #graphDef .~
                encodeMessage ((defMessage :: PG.GraphDef)
                                & #versions .~ (defMessage & #producer .~ 1
                                                           & #minConsumer .~ 1
                                                           & #badConsumers .~ [])
                                & #version .~ 3
                                & #node .~ nodes)
  pure $ d & #wallTime .~ time
           & #step .~ fromIntegral step

-- | Internal
encodeEvent :: PE.Event -> B.ByteString
encodeEvent ev = BL.toStrict $ runPut $ do
  putByteString len
  putWord32le (maskedCRC32 len)
  putByteString msg
  putWord32le (maskedCRC32 msg)
  where msg = encodeMessage ev
        len = BL.toStrict $ runPut $ putWord64le (fromIntegral $ B.length msg)

-- | Internal
writeEvent :: SummaryWriter -> Event -> IO ()
writeEvent r e = do
  step <- readIORef $ lastStep r
  let fname = T.unpack (runDirectory r) </> (T.unpack $ runFilename r)
  b <- doesFileExist fname
  unless b $ do
    bs <- encodeEvent <$> eventToProto step EventVersion
    B.writeFile fname bs
  bs <- encodeEvent <$> eventToProto step e
  B.appendFile fname bs

-- | Internal
maskedCRC32 :: BS.ByteString -> Word32
maskedCRC32 bs = (crc32c bs `rotateR` 15) + 0xa282ead8