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

import Control.Monad.Trans
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 heirarchical name.
type Path = [Name]

data Global = Global
  { gId      :: Int
  , gProbes  :: [(String, Type, E Word64)]
  , gUVs     :: [UV]
  , gPeriod  :: Int
  , gAsserts :: [(String, UE)]
  }

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 a == b = atomId a == atomId b
instance Ord  AtomDB where compare a b = compare (atomId a) (atomId b)
instance Show Rule   where show = ruleName
instance Eq   Rule   where a == b = ruleId a == ruleId b
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, elabortes design and returns a design database.
elaborate :: Name -> Atom () -> IO (Maybe ([Rule], [UV], [(String, UE)]))
elaborate name atom = do
  putStrLn "Starting atom elaboration..."
  hFlush stdout
  (_, (g, atomDB)) <- buildAtom (Global { gId = 0, gProbes = [], gUVs = [], gPeriod = 1, gAsserts = [] }) name atom
  let rules = reIdRules $ elaborateRules (ubool True) atomDB
  mapM_ checkEnable rules
  ok <- checkAssignConflicts atomDB
  if not ok
    then return Nothing
    else do
      let uvs = sort $ gUVs g
      ruleGraph name rules uvs
      return $ Just (rules, uvs, gAsserts 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 = do
  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 :: Name -> Const -> Atom (V a)
var name const = do
  name <- addName name
  (g, atom) <- get
  let uv = UV (gId g) name $ Local const
  put (g { gId = gId g  + 1, gUVs =  uv : gUVs g }, atom)
  return $ V uv

-- | Generic external variable declaration.
var' :: Name -> Type -> Atom (V a)
var' name t = do
  (g, atom) <- get
  let uv = UV (gId g) name $ External t
  put (g { gId = gId g  + 1, gUVs =  uv : gUVs g }, atom)
  return $ V uv

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)))