
module Igor2.Data.IgorMonad (
    
    initIgor, setupTarget,
    Igor(..), 
    
    -- Igor Monad
    runIM, IM,
    
    -- Context
    context, instances, ctxFile, 
    applyC, ctorsOf,

    setTarget, getTarget, addBgk, remBgk, getEvidence, tick, tickT, loopCount,
    loopsCount, tierCount, isDebug, inEnhanced, ifIsSet, doGreedySplt,
    introduceAccums,
    doGreedyMtch,
    usePara, atMaxLoops, atMaxTiers, getPatComparison,
    background, targets, predicates, addIO, breakupM, coverAll, evalIO,

    module Igor2.Data.IOData,
  --  module Data.HypoSpace
    module Igor2.Data.Hypotheses
    
    )where

import Prelude hiding ((<$>))

import Control.Monad.State
import Control.Monad.Error

import Igor2.Config
import Igor2.Logging
import Igor2.Data.IOData
import Igor2.Data.HypoSpace
import Igor2.Data.Hypotheses

import Syntax
import Igor2.Ppr

import Data.Maybe (fromMaybe)
import Data.List ( (\\), transpose )
import Data.Map (toList)
import Data.Function

 
-- | a record encapsulating all necessary data for synthesis
data Igor = Igor
    { igor_io  :: !IOData
    , igor_lc  :: [(Name,Int)]
    , igor_tc  :: Int
    , igor_cnf :: SCR
    , igor_ctx :: Context
    }
    deriving (Show)
   
-- | Igor in a Monad
type IM a = StateT Igor LM a 

-- | evaluate an 'Igor' 'StateT' 'IM'
runIM :: IM a -> Igor -> LM a
runIM = evalStateT

--------------------------------------------------------------------------------
-- Monadic operations on 'Igor'
--------------------------------------------------------------------------------

setTarget :: Name -> IM()
setTarget n = get >>= lift . (setupTarget n) >>= put 

getTarget :: IM Name
getTarget = gets $ fst.head.igor_lc

--------------------------------------------------------------------------------
-- Non-Monadic operations on 'Igor'
--------------------------------------------------------------------------------

-- | Set the target for the given 'Igor' to a function with name 'Name'. If no
--   evidence for such a function exists in 'Igor' an error is thrown. 
setupTarget :: Name -> Igor -> LM Igor
setupTarget n igor@(Igor { igor_lc = lc }) = do
    return $ igor { igor_lc = ((n,0):lc)
                  , igor_tc = 0
                  }

-- | Initialise an 'Igor' (data object) with '[(Name,Rules)]' as evidence for 
--   the target function and background knowledge. Each '(Name,Rules)' is the
--   the evidence of a function with 'Name' and IO examples 'Rules' 
initIgor :: [(Name,Rules)] -> SCR -> Context -> Igor
initIgor nr scr ctx = Igor
    { igor_io = (initIOData nr)
    , igor_lc = []
    , igor_tc = 0
    , igor_cnf = scr
    , igor_ctx = ctx
    }

--------------------------------------------------------------------------------
-- Monadic operations on 'Context'
--------------------------------------------------------------------------------

context :: IM Context
context = gets igor_ctx

instances :: IM [(Name,[Type])]
instances = gets $ toList . ctx_instances . igor_ctx

--------------------------------------------------------------------------------
-- Monadic operations on 'Config', inspecting the current settings
--------------------------------------------------------------------------------

-- | are we debugging?
isDebug :: IM Bool
isDebug = gets $ scr_debug.igor_cnf

-- | are we in enhanced mode?
inEnhanced :: IM Bool
inEnhanced = gets $ scr_enhanced.igor_cnf

-- | use paramoprhisms?
usePara :: IM Bool
usePara = gets $ scr_para.igor_cnf

-- | split greedily
doGreedySplt :: IM Bool
doGreedySplt = gets $ scr_greedySplt.igor_cnf

-- | introduce accumulator variables
introduceAccums :: IM Bool
introduceAccums = gets $ scr_accum.igor_cnf

-- | whether to match calls greedily
doGreedyMtch :: IM Bool
doGreedyMtch = gets $ scr_greedyMtch . igor_cnf

-- | monadic if_then_else
ifIsSet :: IM Bool -> IM a -> a ->  IM a
ifIsSet c f d = do
    b <- c
    if b then f else return d
    
-- | a threshold for a maximum number of loops, i.e. algorithm cycles
maxLoops :: IM Integer
maxLoops = gets $ scr_maxLoops.igor_cnf   

-- | have we reached the maximum loop theshold
atMaxLoops :: IM Bool
atMaxLoops = liftM2 ((==).((+1).fromInteger)) maxLoops loopCount

-- | how do we compare recursive arguments
recArgComp :: IM RecArgComp
recArgComp = gets $ scr_redOrder.igor_cnf

-- | get the comparison function according to how we compare recursive arguments
getPatComparison :: IM (LHS -> LHS -> Ordering)
getPatComparison = do
     howCompare <- recArgComp
     case howCompare of 
        Linear -> return $ compareSizeLinear
        AWise  -> return $ compareSizePairwise

-- | what predicates are we using for conditional rules
--   NOT IMPLEMENTED YET        
predicates :: IM [(Name,Type)]
predicates = gets $ scr_preds.igor_cnf

-- | the given target functions
targets :: IM [Name]
targets = gets $ scr_tgts.igor_cnf

-- | the backgrou nd functions
background :: IM [Name]
background = gets $ scr_bgks.igor_cnf

-- | modify the StateT and add a name of a backgroubd function
addBgk :: [Name] -> IM ()
addBgk ns = do
    bgk <- gets $ scr_bgks.igor_cnf
    let bgk' = bgk ++ ns
    modify (\igor@(Igor { igor_cnf = conf }) -> igor { igor_cnf = conf { scr_bgks = bgk' } })
    
-- | modify the StateT and remove a name of a backgroubd function
remBgk :: [Name] -> IM ()
remBgk ns = do
    bgk <- gets $ scr_bgks.igor_cnf
    let bgk' = bgk \\ ns
    modify (\igor@(Igor { igor_cnf = conf }) -> igor { igor_cnf = conf { scr_bgks = bgk' } })
    
-- | what is our context file, i.e. the specification given
ctxFile :: IM FilePath
ctxFile = gets $ scr_ctxFile . igor_cnf

--------------------------------------------------------------------------------
-- Monadic operations on the SearchSpace
-- DEPRECATED
--------------------------------------------------------------------------------
--currentBestHypos :: IM Hypos
--currentBestHypos = liftM bestHypos $! gets igor_sp
--
--getSearchSpace :: IM HSpace
--getSearchSpace = gets igor_sp
--
--modifyHS :: (HSpace -> HSpace) -> IM()
--modifyHS f = 
--    modify $ \igor@(Igor _ sp _ _ _) ->
--      igor{igor_sp = f sp}
--    
--propagate :: Hypo -> Hypos -> IM ()  
--propagate  = (modifyHS .) . replaceHypos
      
--------------------------------------------------------------------------------
-- Monadic operations on IOData
--------------------------------------------------------------------------------

-- | get all evidence data
getEvidence :: IM IOData
getEvidence = gets igor_io

-- | apply a function to the evidence data in 'IM'
modifyIO ::(IOData -> IOData) -> IM ()
modifyIO f = 
    modify $ \igor@(Igor { igor_io = io }) ->
      igor{igor_io = f io}

-- | increase the loop count by 1
tick :: IM ()
tick = get >>= \igor@(Igor { igor_lc = ((n,lc):lcs) }) ->
          put igor { igor_lc = (n,lc+1):lcs }

-- | increase the tier count by one
tickT :: IM ()
tickT = get >>= \igor@(Igor { igor_tc = tc }) ->
            put igor { igor_tc = tc+1 }

-- | get the tier count
tierCount :: IM Int
tierCount = gets $ igor_tc 
  
-- | threshold for maximum number of tiers            
maxTiers :: IM Integer
maxTiers = gets $ scr_maxTiers.igor_cnf   

-- | have we reached the maximum tier threshold?
atMaxTiers :: IM Bool
atMaxTiers = liftM2 ((==).fromInteger) maxTiers tierCount
            
-- | get the loop count of current target function
loopCount :: IM Int
loopCount = gets $ snd.head.igor_lc  
  
-- | get the loop count of all synthesised target functions so far            
loopsCount :: IM [(Name,Int)]
loopsCount = gets $ igor_lc

--hypoCount :: IM Int
--hypoCount = gets $ countHypos.igor_sp

-- | add the rules 'rs' as IO examples to 'IM', the name under wich they are 
--   stored is returned.
addIO :: Rules -> IM Name
addIO rs = do
    igor@(Igor { igor_io = iod }) <- get
    let (n, iod') = insertRules rs iod
    put igor { igor_io = iod' }
    return n
    
-- | a monadic helper to get all covered IO examples from a 'CovrRule'
breakupM :: CovrRule -> IM [CovrRule]
breakupM rf = gets igor_io >>= return . (breakup rf)

-- | returns a 'CovrRule' covering all examples with 'Name' 'n'
coverAll :: Name -> IM CovrRule
coverAll n = getEvidence >>= lift . fuseByName n

--------------------------------------------------------------------------------
-- Monadic operations on Context
--------------------------------------------------------------------------------
-- | evaluate function 'n' with arguments 'args' using the IO examples in 
--   context
evalIO :: Name -> LHS -> IM (Maybe RHS)
evalIO n args = do
    ios <- liftM (getAll n) getEvidence    
    r   <- (applyC $ matchEvals args $ map crul ios) :: (IM (Either String RHS))
    return $ either (const Nothing) Just r
    
-- | resolve an monadic context function, by applying the context given in the 
--   Igor Monad.
applyC :: Monad m => C m a -> IM (m a)
applyC a = liftM (withC a) $ gets igor_ctx

-- | get the names of constructors by name of a type (may fail)
ctorsOf :: Name -> IM [Name]
ctorsOf n = (lift . lift . getCtors $ n) >>=
             maybe (fail $ "No ctors for " ++ (show n)) return 

--------------------------------------------------------------------------------
-- Instance declarations
--------------------------------------------------------------------------------          
instance Pretty Igor where
    pretty i = text "IGOR" <$> 
               parens ( indent 2 $ 
                        pretty (igor_io i) <$>
                        --pretty (igor_sp i) <$>
                        text "LOOPCOUNT: " <$> pretty (igor_lc i))
    
