{-# LANGUAGE Rank2Types #-}

module LOAG.AOAG where

import LOAG.Common
import LOAG.Graphs
import LOAG.Rep

import AbstractSyntax
import CommonTypes
import Control.Arrow ((&&&), (***))
import Control.Monad (forM, forM_, MonadPlus(..), when, unless)
import Control.Monad.ST
import Control.Monad.State (MonadState(..))
import Data.Maybe (fromMaybe, catMaybes, fromJust, isNothing)
import Data.List (elemIndex, foldl', delete, (\\), insert, nub)
import Data.STRef
import Data.Tuple (swap)
import qualified Data.Set as S
import qualified Data.IntSet as IS
import qualified Data.IntMap as IM
import qualified Data.Map as M
import           Data.Array.MArray
import qualified Data.Array as A
import Data.Array.ST
import ErrorMessages as Err
import Pretty
import UU.Scanner.Position

data Settings = Settings  
                { -- current minimum ads size
                  prune_val  :: Int 
                  -- whether to minimize the number of fake dependencies
                  -- could be very costly
                , min_ads :: Bool
                }
default_settings = Settings 999 False

type AOAG s a = ST s a

-- | Catch a type 3 cycle-error made by a given constructor
-- |  two alternatives are given to proceed
type ADS = [Edge]
type AOAGRes =  Either Error LOAGRes
-- | Calculate a total order if the semantics given 
--    originate from a linearly-ordered AG

type2error,limiterror,aoagerror :: Error
type2error = Err.CustomError False noPos $ text "Type 2 cycle"
limiterror = Err.CustomError False noPos $ text "Limit reached"
aoagerror  = Err.CustomError False noPos $ text "Not an LOAG/AOAG"

schedule :: LOAGRep -> Grammar -> Ag -> [Edge] -> AOAGRes
schedule sem gram@(Grammar _ _ _ _ dats _ _ _ _ _ _ _ _ _) 
                ag@(Ag bounds_s bounds_p de nts) ads 
    = runST $ aoag default_settings ads
 where
    -- get the maps from semantics and translate them to functions    
    nmp  = (nmp_LOAGRep_LOAGRep  sem)     
    ofld = (ofld_LOAGRep_LOAGRep sem)       
    genA = gen_LOAGRep_LOAGRep sem             
    inss = inss_LOAGRep_LOAGRep sem        
 
    -- select candidates, using the edge that caused the cycle
    -- from the list of intra-thread dependencies 
    -- (intra-visit dependencies without edges in ids)
    candidates :: Edge -> Cycle -> [Edge] -> [Edge]
    candidates _ c = foldr (\(f,t) acc -> 
                                if f `IS.member` c &&t `IS.member` c
                                    then (t,f):acc else acc) []
    -- | Move occurrence to its corresponding attribute 
    gen :: Vertex -> Vertex
    gen v = genA A.! v

    genEdge :: Edge -> Edge
    genEdge (f,t) = (gen f, gen t)

    -- | Decide for a given production edge whether the vertices 
    --      belong to the same field
    siblings :: Edge -> Bool
    siblings (f, t) = ofld A.! f == ofld A.! t

    -- | Given an nonterminal-edge, instantiate it
    --   assumes that the occurrences of fields are added in the same order
    instEdge :: Edge -> [Edge]
    instEdge (f, t) = zip (inss A.! f) (inss A.! t)
 
    aoag :: Settings -> [Edge] -> AOAG s AOAGRes
    aoag cfg init_ads = run
        where
            run :: AOAG s AOAGRes
            run = induced ads >>= detect

            detect (Left err) = return $ Left err
            detect (Right (dp,idp,ids@(idsf,idst))) = do
                -- Attribute -> TimeSlot
                schedA <- mapArray (const Nothing) idsf
                -- map TimeSlot -> [Attribute]
                schedS <- newSTRef $ 
                    foldr (\(Nonterminal nt _ _ _ _) -> M.insert (getName nt) 
                                (IM.singleton 1 [])) M.empty dats
                fr_ids <- freeze_graph ids
                threads <- completing fr_ids (schedA, schedS) nts
                let (ivd, comp) = fetchEdges fr_ids threads nts
                eRoC <- m_edp dp init_ads ivd comp (schedA, schedS)
                case eRoC of
                    Left res -> return $ Right res
                    Right (e,c,T3 cs) -> find_ads dp idp ids (schedA, schedS) e c cs

            find_ads :: Graph s -> Graph s -> Graph s -> SchedRef s -> 
                         Edge -> Cycle -> [Edge] -> AOAG s AOAGRes
            find_ads dp idp ids sched e cycle comp = do
                pruner <- newSTRef 999
                explore dp idp ids sched init_ads pruner e cycle comp
 
            explore :: Graph s -> Graph s -> Graph s -> SchedRef s -> 
                        [Edge] -> STRef s Int -> Edge -> Cycle -> [Edge] ->
                            AOAG s AOAGRes
            explore dp idp ids sched@(schedA, schedS) ads pruner e c comp =
                explore' dp idp ids sched ads (candidates e c comp) pruner
             where
              explore' :: Graph s -> Graph s -> Graph s -> SchedRef s -> 
                          [Edge] -> [Edge] -> STRef s Int -> 
                            AOAG s AOAGRes
              explore' _  _   _   _  _   [] _ = return $ Left aoagerror
              explore' dp idp ids sched@(schedA,schedS) ads (fd:cs) pruner = do
                  p_val <- readSTRef pruner
                  if length ads >= p_val -1
                   then return $ Left limiterror 
                   else do 
                    idpf_clone <- mapArray id (fst idp)
                    idpt_clone <- mapArray id (snd idp)
                    let idp_c  =  (idpf_clone, idpt_clone)
                    idsf_clone <- mapArray id (fst ids)
                    idst_clone <- mapArray id (snd ids)
                    let ids_c  =  (idsf_clone, idst_clone)
                    schedA_c   <- mapArray id schedA
                    schedS_v   <- readSTRef schedS
                    schedS_c   <- newSTRef schedS_v
                    let sched_c = (schedA_c, schedS_c)
    
                    let runM = reschedule dp idp ids sched 
                                                (fd:ads) fd pruner
                    let backtrack = explore' dp idp_c ids_c sched_c ads cs 
                                        pruner
                    maoag <- runM
                    case maoag of 
                      Left _                 -> backtrack
                      Right (tdp1,inf1,ads1) -> 
                            if LOAG.AOAG.min_ads cfg 
                             then do
                              writeSTRef pruner (length ads1)
                              maoag' <- backtrack
                              case maoag' of
                                Right (tdp2,inf2,ads2)
                                          -> return $ Right (tdp2,inf2,ads2)
                                otherwise -> return $ Right (tdp1,inf1,ads1)
                             else return $ Right (tdp1,inf1,ads1)

            -- step 1, 2 and 3
            induced :: [Edge] -> AOAG s (Either Error (Graph s, Graph s, Graph s))
            induced ads = do
                dpf  <- newArray bounds_p IS.empty
                dpt  <- newArray bounds_p IS.empty
                idpf <- newArray bounds_p IS.empty
                idpt <- newArray bounds_p IS.empty
                idsf <- newArray bounds_s IS.empty
                idst <- newArray bounds_s IS.empty
                let ids = (idsf,idst)
                let idp = (idpf,idpt)
                let dp  = (dpf ,dpt)
                inducing dp idp ids (de ++ ads) 

            inducing :: Graph s -> Graph s -> Graph s -> [Edge] 
                            -> AOAG s (Either Error (Graph s, Graph s, Graph s))
            inducing dp idp ids es = do
                res <- adds (addD dp idp ids) [] es
                case res of 
                    Left _ -> return $ Left $ type2error
                    Right _  -> return $ Right (dp, idp, ids)
            addD :: Graph s -> Graph s -> Graph s -> Edge -> AOAG s (Either Error [Edge])
            addD dp' idp' ids' e = do
                resd <- e `insErt` dp'
                resdp <- e `inserT` idp'
                case resdp of 
                  Right es  -> adds (addN idp' ids') [] (e:es)
                  Left c    -> return $ Left $ type2error

            addI :: Graph s -> Graph s -> Edge -> AOAG s (Either Error [Edge])
            addI idp' ids' e = do
                exists <- member e idp'
                if not exists then do
                    res <- e `inserT` idp'
                    case res of
                        Right es -> adds (addN idp' ids') [] es
                        Left c   -> return $ Left $ type2error
                 else return $ Right []

            adds f acc [] = return $ Right acc
            adds f acc (e:es) = do
                mes <- f e
                case mes of 
                    Left err -> return $ Left err
                    Right news -> adds f (acc++news) es

            addN :: Graph s -> Graph s -> Edge -> AOAG s (Either Error [Edge])
            addN idp' ids' e = do
                    if (siblings e) then do
                        let s_edge = genEdge e
                        exists <- member s_edge ids'
                        if not exists then do
                            _ <- inserT s_edge ids'
                            let es = instEdge s_edge
                            addedEx <- adds (addI idp' ids') [] es
                            case addedEx of
                                Right news -> return $ Right (s_edge : news)
                                Left err   -> return $ Left err
                         else return $ Right []
                     else return $ Right []            

            
            -- step 6, 7
            m_edp :: Graph s -> [Edge] -> [Edge] -> [Edge] -> SchedRef s ->
                        AOAG s (Either LOAGRes (Edge,Cycle,CType))
            m_edp (dpf, dpt) ads ivd comp sched = do
                edpf <- mapArray id dpf
                edpt <- mapArray id dpt
                mc   <- addEDs (edpf,edpt) (concatMap instEdge ivd) 
                case mc of
                  Just (e, c) -> return $ Right (e,c,T3 $ concatMap instEdge comp)
                  Nothing     -> do 
                        tdp  <- freeze edpt
                        infs <- readSTRef (snd sched)
                        return $ Left (Just tdp,infs,ads)

            reschedule :: Graph s -> Graph s -> Graph s -> SchedRef s -> 
                           [Edge] -> Edge -> STRef s Int 
                            -> AOAG s AOAGRes
            reschedule dp idp ids sched@(_,threadRef) ads e pruner = do
                extra <- addN idp ids e
                case extra of 
                    Left err -> return $ Left err
                    Right extra -> do
                        forM_ extra $ swap_ivd ids sched
                        fr_ids <- freeze_graph ids
                        threads <- readSTRef threadRef
                        let (ivd, comp) = fetchEdges fr_ids threads nts 
                        eRoC <- m_edp dp ads ivd comp sched
                        case eRoC of
                            Left res -> return $ Right res 
                            Right (e,c,(T3 cs)) -> explore dp idp ids sched ads pruner e c cs
             where 
              swap_ivd :: Graph s -> SchedRef s -> Edge -> AOAG s ()
              swap_ivd ids@(idsf, idst) sr@(schedA, schedS) (f,t) = do
                --the edge should point from higher to lower timeslot
                assigned <- freeze schedA
                let oldf = maybe (error "unassigned f") id $ assigned A.! f
                    oldt = maybe (error "unassigned t") id $ assigned A.! t
                    dirf = snd $ alab $ nmp M.! f
                    dirt = snd $ alab $ nmp M.! t
                    newf | oldf < oldt = oldt + (if dirf /= dirt then 1 else 0)
                         | otherwise   = oldf
                    nt   = show $ typeOf $ nmp M.! f

                -- the edge was pointing in wrong direction so we moved 
                -- the attribute to a new interaction, now some of its
                -- predecessors/ancestors might need to be moved too
                unless (oldf == newf) $ do
                    writeArray schedA f (Just newf)
                    modifySTRef schedS 
                            (M.adjust (IM.update (Just . delete f) oldf) nt)
                    modifySTRef schedS 
                        (M.adjust(IM.alter(Just. maybe [f] (insert f))newf)nt)
                    predsf <- readArray idst f
                    succsf <- readArray idsf f
                    let rest = (map (flip (,) f) $ IS.toList predsf) ++ 
                               (map ((,) f)      $ IS.toList succsf)
                     in mapM_ (swap_ivd ids sr) rest