{-# OPTIONS_GHC -XTypeOperators -XRank2Types #-}
{-# LANGUAGE FlexibleContexts, ScopedTypeVariables, DeriveDataTypeable #-} 
module Language.Haskell.FreeTheorems.Variations.PolySeq.AlgCommon
    (Cont(..)
    , emptyCont
    , addTermVar
    , addTypVar
    , getLabTVar
    , getTypVarInCont
    , makeLabel
    , getElemType
    , getArrowComps
    , isInt
    , isBool
    , getAllComps
    , instantiateWithEpsilon
    , annotate
    , removeTrue
    , getEquals
    , substTyp
--    , parseTerm
--    , collectAll
    , collectOne
    , substLabel
    , getUsedLabels
    , getUsedExtraLabels
    ) where

import Data.Generics(Data, everywhere, everywhereM, mkT, mkQ, mkM, gmapT, gmapQ)
import Data.List(nub)
import Control.Monad(mplus)

import Language.Haskell.FreeTheorems.Variations.PolySeq.M
import Language.Haskell.FreeTheorems.Variations.PolySeq.Parser.ParseTerm

import Language.Haskell.FreeTheorems.Variations.PolySeq.Syntax

-- |representation of \Gamma
data Cont = Cont { tVars :: [(TypVar,Label)], vars :: [TypedVar] } deriving (Show,Eq)

emptyCont :: Cont
emptyCont = Cont [] []

-- * auxiliar functions

addTermVar :: Cont -> TypedVar -> Cont
addTermVar (Cont tVars vars) v = Cont tVars (v:vars)

addTypVar :: Cont -> (TypVar,Label) -> Cont
addTypVar (Cont tVars vars) tv = Cont (tv:tVars) vars

getTypVarInCont :: Cont -> TermVar -> M Typ
getTypVarInCont gamma v = 
    case filter (\x->(fst x == v)) (vars gamma) of
      [(_,tau)] -> return tau
      _        -> abort

getLabTVar :: Cont -> TypVar -> M Label
getLabTVar gamma tv =
    case filter (\x->(fst x == tv)) (tVars gamma) of
      [(_,lab)] -> return lab
      _         -> abort

makeLabel :: M Label
makeLabel = do{ i <- newLab;
		return (LVar (LabVar i))
	      }

getElemType :: Typ -> M Typ
getElemType tau =
    case tau of
      TList tau' -> return tau'
      _          -> abort

getArrowComps :: Typ -> M (Typ,Typ)
getArrowComps tau =
    case tau of
      TArrow _ tau1 tau2 -> return (tau1,tau2)
      _                  -> abort
    
getAllComps :: Typ -> M (Label,TypVar,Typ)
getAllComps tau =
    case tau of
      TAll lab tv tau' -> return (lab,tv,tau')
      _                -> abort

isInt :: Typ -> M ()
isInt tau = if tau == TInt then return () else abort

isBool :: Typ -> M ()
isBool tau = if tau == TBool then return () else abort

annotate dat = everywhereM (mkM nonToVar) dat

nonToVar :: Label -> M Label
nonToVar l = if l == Non then makeLabel else return l

substLabel :: (Data a) => Label -> Label -> a -> a
substLabel new old dat = everywhere (mkT (substLab new old)) dat

substLab :: Label -> Label -> Label -> Label
substLab new old recent =
    if old == recent then new else recent

substTyp :: Typ -> Typ -> TypVar -> Typ
substTyp tau new old = (justUnboundTypVars old) (mkT (substTV new old)) tau

justUnboundTypVars :: TypVar -> (Typ -> Typ) -> Typ -> Typ
justUnboundTypVars tv sub tau =
    case tau of
      TAll _ tv' _ -> if tv == tv' then tau else gmapT (mkT (justUnboundTypVars tv sub)) (sub tau)
      _            -> gmapT (mkT (justUnboundTypVars tv sub)) (sub tau)

substTV :: Typ -> TypVar -> Typ -> Typ
substTV new old tau = case tau of
                      TVar tv -> if tv == old then new else tau
                      _            -> tau


getEquals :: Constraint -> [(Label,Label)]
getEquals c = getEquals' (mkQ [] rmEq) c

getEquals' f c =
    case c of
      Impl c1 c2 -> concat ((f c1) : (gmapQ  (mkQ [](getEquals' f)) c1))
      _          -> concat ((f c) : (gmapQ (mkQ [] (getEquals' f)) c))

rmEq c =
    case c of
      Eq l1 l2 -> case l1 of
                    LVar (LabVar i) -> case l2 of
                                         LVar (LabVar j) -> if i<j then [(l1,l2)] else [(l2,l1)]
                                         _               -> [(l1,l2)]
                    _               -> [(l2,l1)]
      _        -> []

collectAll :: (Data a) => (forall a1. (Data a1) => a1 -> [u]) -> a -> [u]
collectAll f c = concat ((f c) : (gmapQ (collectAll f) c)) 

collectOne :: (Data a) => (forall a1. (Data a1) => a1 -> Maybe u) -> a -> Maybe u
collectOne f c = foldl mplus (f c) (gmapQ (collectOne f) c) 


removeTrue :: Constraint -> Constraint
removeTrue = everywhere (mkT remTru)

remTru :: Constraint -> Constraint
remTru c =
    case c of
      Conj Tru c2 -> c2
      Conj c1 Tru -> c1
      Impl Tru c2 -> c2
      _           -> c


instantiateWithEpsilon :: Typ -> Typ
instantiateWithEpsilon tau = everywhere (mkT instWithEpsilon) tau

instWithEpsilon :: Label -> Label
instWithEpsilon l = LVal Epsilon

getUsedLabels :: (Data a) => a -> [Int]
getUsedLabels = getUsedExtraLabels []

getUsedExtraLabels ::(Data a) => [Int] -> a -> [Int]
getUsedExtraLabels notIn x = nub (collectAll (mkQ [] (getLabVar notIn)) x)

getLabVar :: [Int] -> LabVar -> [Int]
getLabVar ls (LabVar i) = if i `elem` ls then [] else [i]