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

import Control.Monad.State hiding (join)
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)

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


-- | The 'Atom' container holds top level IO, 'Var', and 'Rule' definitions.
type Atom = StateT (Global, AtomDB) IO

-- | 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 = elaborateRules (ubool True) atomDB
  mapM_ checkEnable rules
  ok <- checkAssignConflicts atomDB
  if not ok
    then return Nothing
    else return $ Just (rules, sort $ gUVs g, 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 variable declaration.
var :: Name -> Const -> Atom (V a)
var name const = do
  name <- addName name
  (g, atom) <- get
  let uv = UV (gId g) name const
  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