{-# LANGUAGE LambdaCase #-}

module Funcons.Patterns where

import Funcons.MSOS
import Funcons.Types
import Funcons.Substitution
import Funcons.Exceptions

import Control.Monad (foldM)
import Data.Function (on)
import Data.List (sortBy)
import Data.Monoid
import Data.Text (unpack)
import qualified Data.BitVector as BV

-- pattern matching
type Matcher a = [a] -> Int -> Env -> Rewrite [(Int, Env)]
type SeqVarInfo = (MetaVar, SeqSortOp, Maybe FTerm)

singleMatcher :: (a -> b -> Env -> Rewrite Env) -> b -> Matcher a
singleMatcher p pat str k env = case drop k str of
    []  -> return []
    f:_ -> eval_catch (p f pat env) >>= \case
            Left ie | failsRule ie  -> return []
                    | otherwise     -> rewrite_rethrow ie
            Right env'              -> return [(k+1,env')]

seqMatcher :: (a -> Maybe FTerm -> Env -> Rewrite Bool) -> ([a] -> Levelled) 
                -> SeqVarInfo -> Matcher a
seqMatcher p level (var, op, mty) str k env = case op of
    QuestionMarkOp -> makeResults ((<=1) . length)
    PlusOp         -> case str of  
                        [] -> return []
                        _  -> makeResults ((>=1) . length)
    StarOp         -> makeResults (const True)
    where makeResults filter_op = do
            furthest <- takeWhileM (\a -> p a mty env) (drop k str)
            return (map ins (filter filter_op $ ordered_subsequences furthest))
            where
              ins fs  = (k+length fs, envInsert var (level fs) env)

              takeWhileM :: (a -> Rewrite Bool) -> [a] -> Rewrite [a]
              takeWhileM _ [] = return []
              takeWhileM p (x:xs) = eval_catch (p x) >>= \case
                Right True          -> (x:) <$> takeWhileM p xs
                Right False         -> return []
                Left ie | failsRule ie  -> return []
                        | otherwise     -> rewrite_rethrow ie

matching :: [a] -> [Matcher a] -> Env -> Rewrite Env
matching str ps env = do 
    matches <- (seqms ps) str 0 env 
    let rule_fail = PatternMismatch ("Pattern match failed: " ++ show (map fst matches))
    case matches of
        [] -> rewrite_throw rule_fail
        [(_,env')] -> return env'
        _  -> internal ("ambiguity not resolved") 
    where   m = length str
        
            seqms :: [Matcher a] -> Matcher a
            seqms = foldr seqlongest lastMatcher

            -- sequencing of matchers specifically to disambiguate safely
            lastMatcher :: Matcher a
            lastMatcher _ k env | k == m       = return [(k,env)]
                                | otherwise    = return []

            seqlongest :: Matcher a -> Matcher a -> Matcher a
            seqlongest p q str k env = do
                matches <- p str k env
                -- implement `longest match' such that it always returns at least one
                -- pattern match (if at least one exists).
                -- (in combination with (`seqm` lastMatcher) it will always
                -- produce exactly one match)
                -- Strategy: try all `pivots' from largest to smallest and `use'
                -- the first that does not yield an empty result list
                foldM tryLargest [] (sortBy (((flip compare) `on` fst)) matches)
             where  tryLargest acc (r, env)
                        | null acc  = q str r env
                        | otherwise = return acc

ordered_subsequences :: [a] -> [[a]]
ordered_subsequences xs = ordered_subsequences' xs []
    where   ordered_subsequences' [] acc = [acc]
            ordered_subsequences' (x:xs) acc = acc : ordered_subsequences' xs (acc++[x])


-- | Patterns for matching funcon terms ('FTerm').
data FPattern   = PValue VPattern
                | PMetaVar MetaVar 
                | PSeqVar MetaVar SeqSortOp
                | PAnnotated FPattern FTerm 
                | PWildCard

f2vPattern :: FPattern -> VPattern
f2vPattern (PValue v) = v
f2vPattern (PMetaVar var) = VPMetaVar var
f2vPattern (PSeqVar var op) = VPSeqVar var op
f2vPattern (PAnnotated fp t) = VPAnnotated (f2vPattern fp) t
f2vPattern PWildCard = VPWildCard

-- | Patterns for matching values ('Values').
data VPattern   = PADT Name [VPattern]
                | VPWildCard
                | PEmptySet
                | PTuple [VPattern]
                | PList [VPattern]
                | VPMetaVar MetaVar 
                | VPAnnotated VPattern FTerm 
                | VPSeqVar MetaVar SeqSortOp
                | VPLit Values 

-- | Variant of 'vsMatch' that is lifted into the 'MSOS' monad.
lifted_vsMatch str pats env = liftRewrite $ vsMatch str pats env
-- | Matching values with value patterns patterns.
-- If the list of patterns is a singleton list, then 'vsMatch' attempts
-- to match the values as a tuple against the pattern as well.
vsMatch :: [Values] -> [VPattern] -> Env -> Rewrite Env 
vsMatch str pats env = case pats of
    [pat]   -> do
        e_ie_env <- eval_catch (strict_vsMatch str [pat] env)
        case e_ie_env of
            Right env' -> return env'
            Left ie | failsRule ie -> vMatch (safe_tuple_val str) pat env
                    | otherwise -> rewrite_rethrow ie
    _       -> strict_vsMatch str pats env

-- | Match stricly values with patterns.
strict_vsMatch :: [Values] -> [VPattern] -> Env -> Rewrite Env
strict_vsMatch str pats env = matching str matchers env
        where   matchers = map (toMatcher vMatch vpSeqVarInfo) pats
                toMatcher prop minfo pat = case minfo pat of
                    Just info   -> seqMatcher isInMaybeTermType ValuesTerm info
                    Nothing     -> singleMatcher prop pat 

-- | Variant of 'premiseStep' that applies substitute and pattern-matching.
premise :: FTerm -> FPattern -> Env -> MSOS Env
premise x pat env = do
    f <- liftRewrite (substitute x env)
    case isVal f of
        True -> msos_throw (SideCondFail "attempting to step a value")
        False -> do f' <- premiseStep f 
                    liftRewrite $ (fMatch f' pat env)

-- | Variant of 'fsMatch' that is lifted into the 'MSOS' monad.
-- If all given terms are values, then 'vsMatch' is used instead.
lifted_fsMatch str pats env = liftRewrite $ fsMatch str pats env
-- | Match a sequence of terms to a sequence of patterns.
fsMatch = fsMatchStrictness False
strict_fsMatch = fsMatchStrictness True
fsMatchStrictness :: Bool -> [Funcons] -> [FPattern] -> Env -> Rewrite Env
fsMatchStrictness strict str pats env 
    -- if all the given funcons are values, then perform value matching instead.
    | not strict && all isVal str = vsMatch (map downcastValue str) (map f2vPattern pats) env
    | otherwise     = matching str matchers env
    where   matchers = map (toMatcher fMatch fpSeqVarInfo) pats
            toMatcher prop minfo pat = case minfo pat of
                Just info   -> seqMatcher (\_ _ _ -> return True) FunconsTerm info
                Nothing     -> singleMatcher prop pat

fMatch :: Funcons -> FPattern -> Env -> Rewrite Env 
fMatch _ PWildCard env = return env
fMatch f (PMetaVar var) env = return (envInsert var (FunconTerm f) env)
fMatch f (PAnnotated pat term) env = do
    ty <- subsAndRewrite term env
    let fail = rewrite_throw (PatternMismatch ("pattern annotation check failed: " ++ show ty))
    rewriteFuncons f >>= \case  
        ValTerm v   -> do   b <- isIn v ty
                            if b then vMatch v (f2vPattern pat) env
                                 else fail   
        otherwise   -> fail 
-- * a sequence variable can match the singleton sequence 
fMatch f pat@(PSeqVar _ _) env = fsMatch [f] [pat] env
-- if the pattern is a value attempt evaluation by rewrite
fMatch f (PValue pat) env = rewriteFuncons f >>= 
    \case   ValTerm v -> vMatch v pat env
            CompTerm _ _ -> rewrite_throw --important, should remain last 
                (PatternMismatch ("could not rewrite to value: " ++ showFuncons f))

lifted_vMaybeMatch mv mp env = liftRewrite $ vMaybeMatch mv mp env
vMaybeMatch :: Maybe Values -> Maybe VPattern -> Env -> Rewrite Env
vMaybeMatch Nothing Nothing env = return env
vMaybeMatch (Just v) (Just p) env = vMatch v p env
vMaybeMatch _ _ env = rewrite_throw (PatternMismatch ("vMaybeMatch")) 

lifted_vMatch v p env = liftRewrite $ vMatch v p env
vMatch :: Values -> VPattern -> Env -> Rewrite Env
vMatch _ (VPWildCard) env = return env
vMatch v (VPMetaVar var) env = return (envInsert var (ValueTerm v) env)
vMatch (Set s) PEmptySet env | null s = return env
vMatch EmptyTuple (PTuple pats) env = vsMatch [] pats env
vMatch (NonEmptyTuple v1 v2 vs) (PTuple pats) env = vsMatch (v1:v2:vs) pats env
vMatch (ADTVal str vs) (PADT con pats) env = adtMatch str con vs pats env
-- strict because we do not want to match the sequence "inside" the list
vMatch (List vs) (PList ps) env = strict_vsMatch vs ps env 
vMatch v (VPAnnotated pat term) env = do
    ty <- subsAndRewrite term env
    isIn v ty >>= \case
        True  -> vMatch v pat env
        False -> rewrite_throw (PatternMismatch ("pattern annotation check failed: " ++ show ty))
vMatch v (VPLit v2) env | v == v2 = return env
-- special treatment for sequence variables:
-- * a (single) sequence variable can match a tuple
vMatch EmptyTuple pat@(VPSeqVar _ _) env = vsMatch [] [pat] env
vMatch (NonEmptyTuple v1 v2 vs) pat@(VPSeqVar _ _) env = vsMatch (v1:v2:vs) [pat] env
-- * a sequence variable can match the singleton sequence
vMatch v pat@(VPSeqVar _ _) env = vsMatch [v] [pat] env
-- * a single value can match a tuple of patterns if it contains sequences
vMatch v (PTuple pats) env = vsMatch [v] pats env 
vMatch v _ _ = rewrite_throw (PatternMismatch ("failed to match"))


adtMatch :: Name -> Name -> [Values] -> [VPattern] -> Env -> Rewrite Env
adtMatch con pat_con vs pats env 
    | con /= pat_con = rewrite_throw (PatternMismatch ("failed to match constructors: (" ++ show (con,pat_con) ++ ")"))
    | otherwise = vsMatch vs pats env


fpSeqVarInfo :: FPattern -> Maybe SeqVarInfo
fpSeqVarInfo (PSeqVar var op) = Just (var, op, Nothing)
fpSeqVarInfo (PAnnotated (PSeqVar var op) _) = Just (var, op, Nothing)
fpSeqVarInfo _ = Nothing

vpSeqVarInfo :: VPattern -> Maybe SeqVarInfo 
vpSeqVarInfo (VPSeqVar var op) = Just (var, op, Nothing)
vpSeqVarInfo (VPAnnotated (VPSeqVar var op) term) = Just (var, op, Just term)
vpSeqVarInfo _ = Nothing


-- | CSB supports five kinds of side conditions.
-- Each of the side conditions are explained below.
-- When a side condition is not accepted an exception is thrown that 
-- is caught by the backtrackign procedure 'evalRules'.
-- A value is a /ground value/ if it is not a thunk (and not composed out of
--  thunks).
data SideCondition  
    -- | /T1 == T2/. Accepted only when /T1/ and /T2/ rewrite to /equal/ ground values.
    = SCEquality FTerm FTerm 
    -- | /T1 =\/= T2/. Accepted only when /T1/ and /T2/ rewrite to /unequal/ ground values.
    | SCInequality FTerm FTerm 
    -- | /T1 : T2/. Accepted only when /T2/ rewrites to a type and /T1/ rewrites to a value of that type.
    | SCIsInSort FTerm FTerm
    -- | /~(T1 : T2)/. Accepted only when /T2/ rewrites to a type and /T1/ rewrites to a value /not/ of that type.
    | SCNotInSort FTerm FTerm
    -- | /T = P/. Accepted only when /T/ rewrites to a value that matches the pattern /P/. (May produce new bindings in 'Env').
    | SCPatternMatch FTerm VPattern

-- | Variant of 'sideCondition' that is lifted into the 'MSOS' monad.
lifted_sideCondition sc env = liftRewrite $ sideCondition sc env

-- | Executes a side condition, given an 'Env' environment, throwing possible exceptions, and 
-- possibly extending the environment.
sideCondition :: SideCondition -> Env -> Rewrite Env
sideCondition cs env = case cs of
    SCEquality term1 term2 -> 
        prop "equality condition" (\a b -> return (a === b)) term1 term2 env
    SCInequality term1 term2 ->     
        prop "inequality condition" (\a b -> return (a =/= b))term1 term2 env
    SCIsInSort term1 term2 -> prop "type-check condition" isIn term1 term2 env
    SCNotInSort term1 term2 -> 
        prop "type-check condition" (\a b -> isIn a b >>= return . not) term1 term2 env
    SCPatternMatch term vpat -> do
        -- special treatment of pattern-matching condition 
        f <- substitute term env
        eval_catch (rewriteFuncons f) >>= \case
            Right (ValTerm v)       -> vMatch v vpat env
            Right (CompTerm lf _)   -> fMatch lf pat env
            Left (_,_,PartialOp _)  -> fMatch f pat env
            Left ie                 -> rewrite_rethrow ie
      where pat = case vpat of
                      VPMetaVar var   -> PMetaVar var 
                      value_pat       -> PValue value_pat
  where prop msg op term1 term2 env = do
            v1 <- subsAndRewrite term1 env
            v2 <- subsAndRewrite term2 env
            b <- op v1 v2
            if b then return env
                 else rewrite_throw (SideCondFail (msg ++ " fails"))

-- piggy back on 
matchTypeParams :: [Types] -> [TypeParam] -> Rewrite Env
matchTypeParams tys tparams = 
    let param_pats = map mkPattern tparams
         where mkPattern (Nothing, kind)  = VPAnnotated VPWildCard kind
               mkPattern (Just var, kind) = VPAnnotated (VPMetaVar var) kind
    in vsMatch (map typeVal tys) param_pats emptyEnv 


-- type checking
isInMaybeTermType :: Values -> (Maybe FTerm) -> Env -> Rewrite Bool
isInMaybeTermType v Nothing _ = return True
isInMaybeTermType v (Just term) env = 
    subsAndRewrite term env >>= isIn v

isIn :: Values -> Values -> Rewrite Bool
isIn v mty = case castType mty of
    Nothing -> sortErr (FValue mty) "rhs of annotation is not a type"
    Just ty -> isInType v ty

isInType :: Values -> Types -> Rewrite Bool
isInType v (ADT nm tys) = do
    DataTypeMembers tparams alts <- typeEnvLookup nm
    env <- matchTypeParams tys tparams 
    or <$> mapM (isInAlt env) alts 
 where isInAlt env (DataTypeInclusion ty_term) = do 
            subsAndRewrite ty_term env >>= isIn v 
       isInAlt env (DataTypeConstructor cons ty_term) = case v of
            ADTVal cons' arg | cons' == cons ->
                subsAndRewrite ty_term env >>= isIn (safe_tuple_val arg)
            _ -> return False
isInType (ADTVal _ _) ADTs = return True
isInType (Atom _) Atoms = return True
isInType (Ascii _) AsciiCharacters = return True
isInType (Bit bv) (Bits n) = return (BV.size bv == n)
isInType v (BoundedIntegers m n) 
    | Int i <- upcastIntegers v = return (i >= m && i <= n)
isInType (ComputationType (ComputesFromType _ _)) ComputationTypes = return True
isInType (ComputationType (ComputesType _)) ComputationTypes = return True
isInType _ EmptyType = return False
isInType (IEEE_Float_32 _) (IEEEFloats Binary32) = return True
isInType (IEEE_Float_64 _) (IEEEFloats Binary64) = return True
isInType v Integers | Int _ <- upcastIntegers v = return True
isInType (List _) (Lists _) = return True 
isInType (Map _) (Maps _ _) = return True
isInType (Multiset _) (Multisets _) = return True 
isInType v Naturals | Nat _ <- upcastNaturals v = return True
isInType v Rationals | Rational _ <- upcastRationals v = return True 
isInType (Set _) (Sets _) = return True
isInType (String _) Strings = return True
isInType (Thunk _) (Thunks _) = return True
isInType v (Tuples ttparams) = case v of
    EmptyTuple              -> isInTupleType [] ttparams
    NonEmptyTuple v1 v2 vs  -> isInTupleType (v1:v2:vs) ttparams
    _                       -> isInTupleType [v] ttparams
isInType (ComputationType (Type _)) Types = return True
isInType v UnicodeCharacters | Char _ <- upcastUnicode v = return True
isInType v (Union ty1 ty2) = (||) <$> isInType v ty1 <*> isInType v ty2
isInType _ Values = return True
isInType (Vector _) (Vectors _) = return True
isInType _ _ = return False

isInTupleType :: [Values] -> [TTParam] -> Rewrite Bool
isInTupleType vs ttparams = 
    eval_catch (vsMatch vs (map mkPattern ttparams) emptyEnv) >>= \case
        Right env' -> return True
        Left (_,_,PatternMismatch _) -> return False
        Left ie -> rewrite_rethrow ie 
    where mkPattern (ty, mop) = VPAnnotated ty_pat (TFuncon ty_funcon)
            where ty_pat = case mop of 
                                Nothing -> VPMetaVar "Dummy"
                                Just op -> VPSeqVar "Dummy" op
                  ty_funcon = type_ ty

typeEnvLookup :: Name -> Rewrite DataTypeMembers 
typeEnvLookup con = Rewrite $ \ctxt st -> 
    case typeLookup con (ty_env ctxt) of
      Nothing -> (Left (evalctxt2exception(Internal "type lookup failed") ctxt)
                        , st, mempty)
      Just members -> (Right members, st, mempty)

-- | 
-- Parameterisable evaluation function function for types.
rewriteType :: Name -> [Values] -> Rewrite Rewritten
rewriteType nm vs 
    | all isType_ vs = 
        rewritten (ComputationType(Type(ADT nm (map downcastValueType vs))))
    | otherwise = sortErr (applyFuncon nm (map FValue vs)) 
                    ("argument of type " <> unpack nm <> " is not a type")