-- | Atom code generation.
module Language.Atom.Code
  ( writeC
  , ruleComplexity
  ) where

import Data.Char
import Data.List
import Data.Maybe
import System.IO

import Language.Atom.Elaboration
import Language.Atom.Expressions

declareMemory :: UV -> String
declareMemory (UV id name init) = cType (constType init) ++ " " ++ e id ++ " = " ++ c ++ ";  /* " ++ name ++ " */\n"
  where
  c = case init of
    CBool   c -> if c then "1" else "0" 
    CInt8   c -> show c
    CInt16  c -> show c
    CInt32  c -> show c
    CInt64  c -> show c
    CWord8  c -> show c
    CWord16 c -> show c
    CWord32 c -> show c
    CWord64 c -> show c
    CFloat  c -> show c
    CDouble c -> show c

declareUE :: String -> (UE, Int) -> String
declareUE d (ue, i) = case ue of
  UVRef _              -> ""
  UConst (CBool True ) -> d ++ "const " ++ cType (ueType ue) ++ " " ++ e i ++ " = 1;\n"
  UConst (CBool False) -> d ++ "const " ++ cType (ueType ue) ++ " " ++ e i ++ " = 0;\n"
  UConst c             -> d ++ "const " ++ cType (ueType ue) ++ " " ++ e i ++ " = " ++ show c ++ ";\n"
  _                    -> d ++ cType (ueType ue) ++ " " ++ e i ++ ";\n"

cType :: Type -> String
cType t = case t of
  Bool   -> "unsigned char"
  Int8   -> "signed char"
  Int16  -> "signed short int"
  Int32  -> "signed long int"
  Int64  -> "signed long long int"
  Word8  -> "unsigned char"
  Word16 -> "unsigned short int"
  Word32 -> "unsigned long int"
  Word64 -> "unsigned long long int"
  Float  -> "float"
  Double -> "double"

codeUE :: [(UE, Int)] -> String -> (UE, Int) -> String
codeUE ues d (ue, i) = case ue of
  UConst _       -> ""
  UVRef _        -> ""
  _              -> d ++ e i ++ " = " ++ basic operands ++ ";\n"
  where
  operands = map (fromJust . flip lookup ues) $ ueUpstream ue
  basic :: [Int] -> String
  basic operands = case ue of
    UVRef _              -> error "Code.ueStmt: should not get here."
    UCust _ c            -> c
    UCast _ _            -> "(" ++ cType (ueType ue) ++ ") " ++ a
    UConst _             -> error "Code.ueStmt: should not get here."
    UAdd _ _             ->  a ++ " + " ++ b
    USub _ _             ->  a ++ " - " ++ b
    UMul _ _             ->  a ++ " * " ++ b
    UDiv _ _             ->  a ++ " / " ++ b
    UMod _ _             ->  a ++ " % " ++ b
    UNot _               ->  "! " ++ a
    UAnd _               ->  drop 4 $ concat [ " && " ++ e a | a <- operands ]
    UBWNot _             ->  "~ " ++ a
    UBWAnd _ _           ->  a ++ " & " ++ b
    UBWOr  _ _           ->  a ++ " | " ++ b
    UShift _ n           -> (if n >= 0 then a ++ " << " ++ show n else a ++ " >> " ++ show (0 - n))
    UEq  _ _             -> a ++ " == " ++ b
    ULt  _ _             -> a ++ " < " ++ b
    UMux _ _ _           -> a ++ " ? " ++ b ++ " : " ++ c
    UF2B _               -> "*((unsigned long int *) &" ++ a ++ ")"
    UD2B _               -> "*((unsigned long long int *) &" ++ a ++ ")"
    UB2F _               -> "*((float *) &" ++ a ++ ")"
    UB2D _               -> "*((double *) &" ++ a ++ ")"
    where
    a = e $ operands !! 0
    b = e $ operands !! 1
    c = e $ operands !! 2

writeC :: Name -> String -> String -> String -> [[[Rule]]] -> [UV] -> IO ()
writeC name include preCode postCode periods uvs = do
  writeFile (name ++ ".c") c
  where
  c = unlines
    [ include
    , cType Word64 ++ " globalClock = 0;"
    , concatMap declareMemory uvs
    , concatMap (codeRule topo') $ concat $ concat periods
    , "void " ++ name ++ "(void) {"
    , preCode
    , concatMap codePeriod $ zip [1..] periods
    , postCode
    , "  globalClock = globalClock + 1;"
    , "}"
    ]

  topo' = topo $ maximum (map (\ (UV i _ _) -> i) uvs) + 1

codeRule :: ([UE] -> [(UE, Int)]) -> Rule -> String
codeRule topo rule = 
  "/* " ++ show rule ++ " */\n" ++
  "void r" ++ show (ruleId rule) ++ "(void) {\n" ++
  concatMap (declareUE  "  ") ues ++
  concatMap (codeUE ues "  ") ues ++
  concatMap codeAction (ruleActions rule) ++
  concatMap (\ (UV i _ _, ue) -> "  " ++ e i ++ " = " ++ e (id ue) ++ ";\n") (ruleAssigns rule) ++
  "}\n\n"
  where
  ues = topo $ ruleEnable rule : snd (unzip (ruleAssigns rule)) ++ concat (snd (unzip (ruleActions rule)))
  id ue = fromJust $ lookup ue ues
  codeAction :: (([String] -> String), [UE]) -> String
  codeAction (f, args) = "  if (" ++ e (id (ruleEnable rule)) ++ ") " ++ f (map (e . id) args) ++ ";\n"

codePeriod :: (Int, [[Rule]]) -> String
codePeriod (period, cycles) = concatMap (codeCycle period) $ zip [0..] cycles

codeCycle :: Int -> (Int, [Rule]) -> String
codeCycle period (cycle, _) | cycle >= period = error "Code.codeCycle"
codeCycle _ (_, rules)      | null rules = ""
codeCycle period (cycle, rules) =
  "  if (globalClock % " ++ show period ++ " == " ++ show cycle ++ ") {\n" ++
  concatMap (\ r -> "    r" ++ show (ruleId r) ++ "();  /* " ++ show r ++ " */\n") rules ++
  "  }\n"

e :: Int -> String
e i = "e" ++ show i


-- | Topologically sorts a list of expressions and subexpressions.
topo :: Int -> [UE] -> [(UE, Int)]
topo start ues = reverse ues'
  where
  (_, ues') = foldl collect (start, []) ues
  collect :: (Int, [(UE, Int)]) -> UE -> (Int, [(UE, Int)])
  collect (n, ues) ue | any ((== ue) . fst) ues = (n, ues)
  collect (n, ues) ue = case ue of
    UVRef (UV i _ _) -> (n,      (ue, i)  : ues )
    _              -> (n' + 1, (ue, n') : ues')
      where
      (n', ues') = foldl collect (n, ues) $ ueUpstream ue

-- | Number of UE's computed in rule.
ruleComplexity :: Rule -> Int
ruleComplexity rule = length $ topo 0 $ ruleEnable rule : snd (unzip (ruleAssigns rule)) ++ concat (snd (unzip (ruleActions rule)))