module Language.Atom.Elaboration
  (
  -- * Atom monad and container.
    Atom
  , AtomDB     (..)
  , Global     (..)
  , Rule       (..)
  , buildAtom
  -- * Type Aliases and Utilities
  , UID
  , Name
  , Path
  , elaborate
  , var
  , var'
  , array
  , addName
  , get
  , put
  ) where

import Control.Monad.Trans
import Data.Function (on)
import Data.List
import Data.Maybe
import Data.Word
import Language.Atom.Expressions
import System.IO

type UID = Int

-- | A name.
type Name = String

-- | A hierarchical name.
type Path = [Name]

data Global = Global
  { gId      :: Int
  , gProbes  :: [(String, Type, E Word64)]
  , gInit8   :: [Const]
  , gInit16  :: [Const]
  , gInit32  :: [Const]
  , gInit64  :: [Const]
  , gPeriod  :: Int
  }

data AtomDB = AtomDB
  { atomId          :: Int
  , atomName        :: Name
  , atomNames       :: [Name]      -- Names used at this level.
  , atomEnable      :: UE          -- Enabling condition.
  , atomSubs        :: [AtomDB]    -- Sub atoms.
  , atomPeriod      :: Int
  , atomAssigns     :: [(UV, UE)]
  , atomActions     :: [([String] -> String, [UE])]
  }

data Rule = Rule
  { ruleId        :: Int
  , ruleName      :: Name
  , ruleEnable    :: UE
  , ruleAssigns   :: [(UV, UE)]
  , ruleActions   :: [([String] -> String, [UE])]
  , rulePeriod    :: Int
  }

instance Show AtomDB where show = atomName
instance Eq   AtomDB where (==) = (==) `on` atomId
instance Ord  AtomDB where compare a b = compare (atomId a) (atomId b)
instance Show Rule   where show = ruleName
instance Eq   Rule   where (==) = (==) `on` ruleId
instance Ord  Rule   where compare a b = compare (ruleId a) (ruleId b)

elaborateRules:: UE -> AtomDB -> [Rule]
elaborateRules parentEnable atom = if isRule then rule : rules else rules
  where
  isRule = not (null $ atomAssigns atom) || not (null $ atomActions atom)
  enable = uand parentEnable $ atomEnable atom
  rule = Rule
    { ruleId        = atomId   atom
    , ruleName      = atomName atom
    , ruleEnable    = enable
    , ruleAssigns   = map enableAssign $ atomAssigns atom
    , ruleActions   = atomActions atom
    , rulePeriod    = atomPeriod atom
    }
  rules = concatMap (elaborateRules enable) (atomSubs atom)
  enableAssign :: (UV, UE) -> (UV, UE)
  enableAssign (uv, ue) = (uv, umux enable ue $ UVRef uv)

reIdRules :: [Rule] -> [Rule]
reIdRules rules = map (\ r -> r { ruleId = fromJust $ lookup (ruleId r) ids } ) rules
  where
  ids = zip (map ruleId rules) [0..]

buildAtom :: Global -> Name -> Atom a -> IO (a, (Global, AtomDB))
buildAtom g name (Atom f) = f (g { gId = gId g + 1 }, AtomDB
  { atomId        = gId g
  , atomName      = name
  , atomNames     = []
  , atomEnable    = ubool True
  , atomSubs      = []
  , atomPeriod    = gPeriod g
  , atomAssigns   = []
  , atomActions   = []
  })

-- | The Atom monad holds variable and rule declarations.
data Atom a = Atom ((Global, AtomDB) -> IO (a, (Global, AtomDB)))

instance Monad Atom where
  return a = Atom (\ s -> return (a, s))
  (Atom f1) >>= f2 = Atom f3
    where
    f3 s = do
      (a, s) <- f1 s
      let Atom f4 = f2 a
      f4 s

instance MonadIO Atom where
  liftIO io = Atom f
    where
    f s = do
      a <- io
      return (a, s)

get :: Atom (Global, AtomDB)
get = Atom (\ s -> return (s, s))

put :: (Global, AtomDB) -> Atom ()
put s = Atom (\ _ -> return ((), s))

-- | A Relation is used for relative performance constraints between 'Action's.
-- data Relation = Higher UID | Lower UID deriving (Show, Eq)

-- | Given a top level name and design, elaborates design and returns a design database.
elaborate :: Name -> Atom () -> IO (Maybe ([Rule], ([Const], [Const], [Const], [Const])))
elaborate name atom = do
  putStrLn "Starting atom elaboration..."
  hFlush stdout
  (_, (g, atomDB)) <- buildAtom Global { gId = 0, gInit8 = [], gInit16 = [], gInit32 = [], gInit64 = [], gProbes = [], gPeriod = 1 } name atom
  let rules = reIdRules $ elaborateRules (ubool True) atomDB
  mapM_ checkEnable rules
  ok <- checkAssignConflicts atomDB
  if not ok
    then return Nothing
    else return $ Just (rules, (gInit8 g, gInit16 g, gInit32 g, gInit64 g))

checkEnable :: Rule -> IO ()
checkEnable rule | ruleEnable rule == ubool False = putStrLn $ "WARNING: Rule will never execute: " ++ show rule
                 | otherwise                      = return ()

checkAssignConflicts :: AtomDB -> IO Bool
checkAssignConflicts atom =
  if length vars /= length vars'
    then do
      putStrLn $ "ERROR: Atom " ++ show atom ++ " contains multiple assignments to the same variable(s)."
      return False
    else do
      subs <- mapM checkAssignConflicts $ atomSubs atom
      return $ and subs
  where
  vars = fst $ unzip $ atomAssigns atom
  vars' = nub vars

-- | Generic local variable declaration.
var :: Expr a => Name -> a -> Atom (V a)
var name init = do
  A a <- array name [init]
  return $ V $ UV $ Array a $ UConst $ CWord8 0

-- | Generic external variable declaration.
var' :: Name -> Type -> Atom (V a)
var' name t = return $ V $ UV $ External name t

-- | Generic array declaration.
array :: Expr a => Name -> [a] -> Atom (A a)
array name [] = error $ "ERROR: arrays can not be empty: " ++ name
array name init = do
  name <- addName name
  (g, atom) <- get
  let constants = map constant init
      (addr, g') = case width $ head constants of
        1  -> (length $ gInit8  g, g { gInit8  = gInit8  g ++ constants })
        8  -> (length $ gInit8  g, g { gInit8  = gInit8  g ++ constants })
        16 -> (length $ gInit16 g, g { gInit16 = gInit16 g ++ constants })
        32 -> (length $ gInit32 g, g { gInit32 = gInit32 g ++ constants })
        64 -> (length $ gInit64 g, g { gInit64 = gInit64 g ++ constants })
        _  -> error "Elaboration.array: unknown width"
      ua = UA addr name constants
  put (g', atom)
  return $ A ua

addName :: Name -> Atom Name
addName name = do
  (g, atom) <- get
  if elem name (atomNames atom)
    then error $ "ERROR: Name \"" ++ name ++ "\" not unique in " ++ show atom ++ "."
    else do
      put (g, atom { atomNames = name : atomNames atom })
      return $ atomName atom ++ "." ++ name

{-
ruleGraph :: Name -> [Rule] -> [UV] -> IO ()
ruleGraph name rules uvs = do
  putStrLn $ "Writing rule graph (" ++ name ++ ".dot)..."
  writeFile (name ++ ".dot") g
  --system $ "dot -o " ++ name ++ ".png -Tpng " ++ name ++ ".dot"
  return ()
  where
  adminUVs =
    [ UV (-1) "__clock" (External Word64)
    , UV (-2) "__coverage_index" (External Word32)
    , UV (-3) "__coverage[__coverage_index]" (External Word32)
    ]

  g = unlines
    [ "digraph " ++ name ++ "{"
    , concat [ "  r" ++ show (ruleId r) ++ " [label = \"" ++ show r ++ "\" shape = ellipse];\n" | r <- rules ]
    , concat [ "  v" ++ show i ++ " [label = \"" ++ n ++ "\" shape = box];\n" | (UV i n _) <- adminUVs ++ uvs ]
    , concat [ "  r" ++ show (ruleId r) ++ " -> v" ++ show i ++ "\n" | r <- rules,  (UV i _ _, _) <- ruleAssigns r ]
    , concat [ "  v" ++ show i ++ " -> r" ++ show (ruleId r) ++ "\n" | r <- rules,  (UV i _ _) <- ruleUVRefs r ]
    , "}"
    ]

  ruleUVRefs r = nub $ concatMap uvSet ues
    where
    ues = ruleEnable r : snd (unzip (ruleAssigns r)) ++ concat (snd (unzip (ruleActions r)))
-}