module UHC.Light.Compiler.Pred.ToCHR
( CHRClassDecl, CHRScopedInstanceDecl
, initCHRStore
, mkScopedCHR2
, SimplifyResult, SimplifyResult'' (..), emptySimplifyResult
, simplifyResultResetForAdditionalWork
, partitionUnresolved2AssumableAndOthers
, patchUnresolvedWithAssumption
, chrSimplifySolveToRedGraph
, chrSimplifyRedGraphToEvidence )
where
import UHC.Light.Compiler.Opts
import UHC.Light.Compiler.Base.Common
import UHC.Light.Compiler.Base.TermLike
import UHC.Light.Compiler.Ty
import UHC.Light.Compiler.Ty.Ftv
import UHC.Light.Compiler.Error
import UHC.Light.Compiler.VarMp
import UHC.Light.Compiler.Substitutable
import Data.Maybe
import qualified Data.Set as Set
import qualified Data.Map as Map
import UHC.Util.CHR
import UHC.Light.Compiler.CHR.Constraint
import UHC.Light.Compiler.CHR.Guard
import UHC.Light.Compiler.CHR.Solve
import UHC.Light.Compiler.CHR.Key
import UHC.Light.Compiler.Pred.Heuristics
import UHC.Light.Compiler.Pred.Evidence
import UHC.Light.Compiler.Pred.RedGraph
import UHC.Light.Compiler.Ty.FitsInCommon2
import UHC.Light.Compiler.Ty.Trf.Canonic
import UHC.Util.Pretty
import UHC.Util.Utils
import UHC.Light.Compiler.Base.HsName.Builtin


{-# LINE 36 "src/ehc/Pred/ToCHR.chs" #-}
type CHRClassDecl'          a info      = ([a], a, [info])
type CHRClassDecl                       = CHRClassDecl' Pred RedHowAnnotation
type CHRScopedInstanceDecl' a info sc   = ([a], a, info, sc)
type CHRScopedInstanceDecl              = CHRScopedInstanceDecl' Pred RedHowAnnotation PredScope

{-# LINE 47 "src/ehc/Pred/ToCHR.chs" #-}
type MkRes1 = (CHRStore, ([CHRPredOcc],CHRPredOcc) )
type MkResN = (CHRStore,[([CHRPredOcc],CHRPredOcc)])

{-# LINE 56 "src/ehc/Pred/ToCHR.chs" #-}
([sc1,sc2,sc3]
 ,[pr1,pr2,pr3]
 ,[ty1,ty2,ty3,ty4]
 ,[lab1]
 ,[off1]
 ,[pr1v]
 ,[pa1]
 )
  = ( map PredScope_Var [u1,u2,u3]
    , map Pred_Var [u7,u8,u9]
    , map mkTyMetaVar {- mkTyVar -} [u10,u11,u14,u15]
    , map Label_Var [u12]
    , map LabelOffset_Var [u13]
    , [u7]
    , map PredSeq_Var [u4]
    )
  where [u1,u2,u3,u4,u5,u6,u7,u8,u9,u10,u11,u12,u13,u14,u15] = mkNewLevUIDL 15 uidStart

{-# LINE 101 "src/ehc/Pred/ToCHR.chs" #-}
-- | The basic initial set of CHRs
initCHRStore :: CHRStore
initCHRStore
  = chrStoreFromElems $
      [ scopeProve, {- scopeAssum1, -} scopeAssum2 ]
      ++ [ labelProve1, labelProve2 ]
      ++ [ instForall, predArrow, predSeq1, predSeq2 ]
  where p1s1         = mkCHRPredOcc pr1 sc1
        p1s2         = mkCHRPredOcc pr1 sc2
        p1s3         = mkCHRPredOcc pr1 sc3
        scopeProve   = [mkProve p1s1, mkProve p1s2]
                         ==> [mkReduction p1s2 (RedHow_ByScope (ByScopeRedHow_Other $ AlwaysEq "prv")) [p1s3]]
                          |> [IsStrictParentScope sc3 sc1 sc2]
{-
        scopeAssum1  = [mkProve p1s1, mkAssume p1s2]
                         ==> [mkReduction p1s1 (RedHow_Assumption sc2) []]
                          |> [EqualScope sc1 sc2]
-}
        scopeAssum2  = [mkProve p1s1, mkAssume p1s2]
                         ==> [mkReduction p1s1 (RedHow_ByScope ByScopeRedHow_Assume) [p1s2]]
                          |> [NotEqualScope sc1 sc2,IsVisibleInScope sc2 sc1]
        l1s1         = mkCHRPredOcc (Pred_Lacks ty1 lab1) sc1
        l2s1         = mkCHRPredOcc (Pred_Lacks ty2 lab1) sc1
        l3s1         = mkCHRPredOcc (Pred_Lacks recRowEmp lab1) sc1
        labelProve1  = [mkProve l1s1]
                         ==> [mkProve l2s1, mkReduction l1s1 (RedHow_ByLabel lab1 off1 sc1) [l2s1]]
                          |> [NonEmptyRowLacksLabel ty2 off1 ty1 lab1]
        labelProve2  = [mkProve l3s1]
                         ==> [mkReduction l3s1 (RedHow_ByLabel lab1 (LabelOffset_Off 0) sc1) []]
        f1s1         = mkCHRPredOcc (tyPred $ mkTyQu tyQu_Forall [(pr1v,kiStar)] $ mkTyPr pr1) sc1	-- TBD
        f2s1         = mkCHRPredOcc pr1 sc1
        instForall   = [mkAssume f1s1]
                         ==> [mkAssume f2s1]
        a1s1         = mkCHRPredOcc (Pred_Arrow pa1 pr1) sc1
        a2s1         = mkCHRPredOcc pr1 sc1
        a3s1         = mkCHRPredOcc (Pred_Preds pa1) sc1
        predArrow    = [mkAssume a1s1, mkProve a2s1]
                         ==> [mkProve a3s1, mkReduction a2s1 (RedHow_ByInstance hsnUnknown pr1 sc1) [a1s1,a3s1]]
        s1s1         = mkCHRPredOcc (Pred_Preds (PredSeq_Cons pr1 pa1)) sc1
        s2s1         = mkCHRPredOcc pr1 sc1
        s3s1         = mkCHRPredOcc (Pred_Preds pa1) sc1
        predSeq1     = [mkProve s1s1]
                         <==> [mkProve s2s1, mkProve s3s1]
        predSeq2     = [mkProve $ mkCHRPredOcc (Pred_Preds PredSeq_Nil) sc1]
                         <==> ([] :: [Constraint])
        -- inclSc       = ehcCfgCHRInclScope $ feEHCOpts $ fiEnv env

{-# LINE 190 "src/ehc/Pred/ToCHR.chs" #-}
-- | Construct CHRs from class and instance decls
mkScopedCHR2
  :: FIIn -> [CHRClassDecl] -> [CHRScopedInstanceDecl]
       -> CHRStore -> (CHRStore,CHRStore)
mkScopedCHR2 env clsDecls insts prevStore
  = (chrStoreUnions [store2,instSimplStore], chrStoreUnions [assumeStore,instSimplStore])
  where  ucls        = mkNewLevUIDL (length clsDecls) $ fiUniq env
         ((assumeStore,assumePredOccs), (instStore,_))
                     = mkScopedChrs clsDecls canonInsts
         store2      = chrStoreUnions [assumeStore,prevStore]
         simplStores = zipWith (\u (cx,h,i) -> mkClassSimplChrs (env {fiUniq = u}) store2 (cx,h,i)) ucls clsDecls
         instSimplStore
         			 = chrStoreUnions $ instStore : simplStores
         canonInsts  = [ (map mkC cx, mkC hd, info, sc) | (cx,hd,info,sc) <- insts ]
                     where mkC = fst . predCanonic (emptyTyBetaRedEnv {tbredFI=env})

{-# LINE 208 "src/ehc/Pred/ToCHR.chs" #-}
-- | Construct simplification CHRs from class decls, building upon a given CHR store
mkClassSimplChrs :: FIIn -> CHRStore -> CHRClassDecl -> CHRStore
mkClassSimplChrs env rules (context, head, infos)
  = simps
  where simps        = chrStoreFromElems $ mapTrans (Set.fromList [head1]) [] head1 (zip infos (map (\p -> Red_Pred $ mkCHRPredOcc p sc1) context))
        (superClassesWork, superClassesDone, _ :: SolveTrace FIIn Constraint Guard VarMp)
                     = chrSolve' env rules (map (\p -> toSolverConstraint $ mkAssume $ mkCHRPredOcc p sc1) context)
        superClasses = superClassesWork ++ superClassesDone
        graph        = mkRedGraphFromReductions $ filterMb fromSolverConstraint superClasses
        head1        = mkCHRPredOcc head sc1
        head2        = mkCHRPredOcc head sc2
        head3        = mkCHRPredOcc head sc3

        mapTrans done reds subClass
          = concatMap (transClosure done reds subClass)
            . filter (\(_,x) -> not (rednodePred x `Set.member` done))

        transClosure done reds par (info, pr@(Red_Pred p@(CHRPredOcc {cpoPr = super})))
          = [superRule] ++ ({-if ehcCfgCHRScoped opts >= CHRScopedMutualSuper then -} [scopeRule1, scopeRule2] {- else [] -}) ++ rules
          where super1     = mkCHRPredOcc super sc1
                super2     = mkCHRPredOcc super sc2
                super3     = mkCHRPredOcc super sc3
                superRule  = [mkProve head1, mkProve p] ==> reds'
                scopeRule1 = [mkProve head1, mkProve super2]
                               ==> [mkProve head3, mkReduction head1 (RedHow_ByScope (ByScopeRedHow_Other $ AlwaysEq "sup1")) [head3]]
                                 |> [HasStrictCommonScope sc3 sc1 sc2]
                scopeRule2 = [mkProve head2, mkProve super1]
                               ==> [mkProve super3, mkReduction super1 (RedHow_ByScope (ByScopeRedHow_Other $ AlwaysEq "sup2")) [super3]]
                                 |> [HasStrictCommonScope sc3 sc1 sc2]
                reds'      = mkReduction p info [par] : reds
                rules      = mapTrans (Set.insert p done) reds' p (predecessors graph pr)

        opts          = feEHCOpts $ fiEnv env

mkScopedChrs :: [CHRClassDecl] -> [CHRScopedInstanceDecl] -> (MkResN,MkResN)
mkScopedChrs clsDecls insts
  = ((chrStoreUnions assumeStores,assumePredOccs), instChrs)
  where (assumeStores,assumePredOccs) = unzip $ mapMaybe mkAssumeChrs clsDecls
        instChrs   = mkInstanceChrs insts

mkAssumeChrs :: CHRClassDecl -> Maybe MkRes1
mkAssumeChrs ([]     ,  _  , _    ) = Nothing
mkAssumeChrs (context, head, infos) =
  let prThis = mkCHRPredOcc head sc1
      super prSuper info = [mkAssume prSuper, mkReduction prSuper info [prThis]]
      prSuper = map (\c -> mkCHRPredOcc c sc1) context
  in  Just ( chrStoreSingletonElem $ [mkAssume prThis] ==> concat (zipWith super prSuper infos)
           , (prSuper,prThis)
           )

mkInstanceChrs :: [CHRScopedInstanceDecl] -> MkResN
mkInstanceChrs insts
  = (chrStoreUnions instStores,instChrs)
  where (instStores,instChrs) = unzip $ map mkInstanceChr insts

mkInstanceChr :: CHRScopedInstanceDecl -> MkRes1
mkInstanceChr (context, hd, i, s)
  = ( chrStoreSingletonElem
      $ [mkProve constraint]
          ==> mkReduction constraint i body : map mkProve body
            |> [s `IsVisibleInScope` sc1]
    , (body,constraint)
    )
  where constraint = mkCHRPredOcc hd sc1
        body = map (\p -> mkCHRPredOcc p sc1) context

{-# LINE 282 "src/ehc/Pred/ToCHR.chs" #-}
data SimplifyResult'' p i g s
  = SimplifyResult
      { simpresSolveState		:: SolveState FIIn (Constraint' p i) g s
      , simpresRedGraph			:: RedGraph' p i

      -- for debugging only:
      , simpresRedAlts			:: [HeurAlts' p i]
      , simpresRedTrees			:: [[(i, Evidence' p i)]]
      , simpresRedGraphs		:: [(String,RedGraph' p i)]
      , simpresRemPredL         :: [p]							-- remaining pred occurrences, which cannot be proven, as a list
      }

type SimplifyResult' g s = SimplifyResult'' CHRPredOcc RedHowAnnotation g s

type SimplifyResult = SimplifyResult' Guard VarMp

emptySimplifyResult :: SimplifyResult
emptySimplifyResult
  = SimplifyResult
      emptySolveState emptyRedGraph
      [] [] [] []

{-# LINE 306 "src/ehc/Pred/ToCHR.chs" #-}
simplifyResultResetForAdditionalWork :: SimplifyResult -> SimplifyResult
simplifyResultResetForAdditionalWork r = r {simpresRedGraph = emptyRedGraph}

{-# LINE 315 "src/ehc/Pred/ToCHR.chs" #-}
mkEvidence
  :: Heuristic
  -> ConstraintToInfoMap
  -> RedGraph
  -> ( -- ConstraintToInfoMap						-- remaining constraints
       ConstraintToInfoTraceMp						-- remaining constraints
     , InfoToEvidenceMap							-- mapping to evidence
     , [Err]											-- errors
     )
mkEvidence heur cnstrMp redGraph
  = ( {- (cnstrMp `Map.intersection` remCnstrMp) `Map.union` -}
      cnstrTraceMpFromList remCnstrMp
    , evidMp
    , err
    )
  where (remCnstrMp,evidMp,err,dbg)
          = foldl (\(cm,em,err,dbg) c -> let (cm',em',err',dbg') = mk c in (cm' ++ cm, evidMpUnion em' em, err' ++ err, dbg' ++ dbg))
                  ([],Map.empty,[],[])
            $ Map.toList cnstrMp
        mk (Prove p, infos)
          = (remCnstrMp,evidMp,[],[(redAlts,redTrees)])
          where redAlts    = redAlternatives redGraph p
                redTrees   = heur infos redAlts
                evidMp     = foldr (uncurry evidMpInsert) Map.empty redTrees
                remCnstrMp = [ (Prove (utraceRedFrom u),(i,[u])) | (i,t) <- redTrees, u <- evidUnresolved t ]
        mk (c, infos)
          = ([ (c,(i,[])) | i <- infos],Map.empty,[],[])

{-# LINE 375 "src/ehc/Pred/ToCHR.chs" #-}
partitionUnresolved2AssumableAndOthers :: ConstraintToInfoTraceMp -> ([CHRIntermediateUntilAssume],ConstraintToInfoTraceMp)
partitionUnresolved2AssumableAndOthers unresCnstrMp
  = (unres,cannotResCnstrMp)
  where (unresCnstrMp',cannotResCnstrMp)
                      = Map.partitionWithKey canAssume unresCnstrMp
                      where -- if p only ranges over non-fixed tvars, we potentially can assume them (if not found ambiguous later)
                            canAssume (Prove p) _ = Map.null $ Map.filter (tvCatIsFixed . tvinfoCateg) $ tyFtvMp $ predTy $ cpoPr p
                            canAssume _         _ = True
        unres         = [ (p,x) | (Prove p,x) <- shareUnresolvedAssumptionsByScope (unresCnstrMp') ]

{-# LINE 387 "src/ehc/Pred/ToCHR.chs" #-}
-- | Group unresolved constraints, reducing the various scopes to the outermost scope.
--   Find assume's wich have a common scope prefix, then share these.
--   Assumption: we will never share outer scopes because we only get passed inner scopes, because these will be abstracted over in bindings of a let expression.
shareUnresolvedAssumptionsByScope :: ConstraintToInfoTraceMp -> [(Constraint,(PredScope,ConstraintToInfoTraceMp))]
shareUnresolvedAssumptionsByScope unres
  = [ ( c
      , ( -- the common prefix off all scopes, i.e. the most global scope
          foldr1 (\s1 s2 -> panicJust "shareUnresolvedAssumptionsByScope" $ pscpCommon s1 s2)
                 [ cpoScope $ cnstrPred c | (c,_) <- cs ]
          -- all original info for this predicate
        , Map.fromList cs
        )
      )
    | cs@((c,_):_) <-
        -- sort, then group by predicate, so we get all scopes for each predicate
        groupSortOn (cpoPr . cnstrPred . fst) $ Map.toList unres
    ]

{-# LINE 411 "src/ehc/Pred/ToCHR.chs" #-}
-- | Transform unresolved Prove constraints to Assume variants, used either for quantification over, or for error messages about unresolved predicates
patchUnresolvedWithAssumption :: FIIn -> [CHRIntermediateUntilAssume] -> RedGraph -> InfoToEvidenceMap -> (ConstraintToInfoTraceMp,InfoToEvidenceMap)
patchUnresolvedWithAssumption env unres redGraph evidMp
  = ( assumeCnstrs
    , evidMpSubst (\p -> Map.lookup p assumeSubstMp) evidMp
    )
  where us = mkNewLevUIDL (length unres) $ fiUniq env
        assumeCnstrs
          = cnstrMpUnions
              [ cnstrTraceMpSingleton c i (concat [ us | es <- Map.elems trMap, (_,us) <- es ])
              | ((p,(sc,trMap)),u) <- zip unres us
              , let (c,i) = rngLift emptyRange mkAssumeConstraint (cpoPr p) u sc
              ]
        assumeSubstMp
          = Map.fromList [ (p,Evid_Proof p i []) | (Assume p,((i,_):_)) <- Map.toList assumeCnstrs ]

{-# LINE 433 "src/ehc/Pred/ToCHR.chs" #-}
chrSimplifySolveToRedGraph
  ::   FIIn
    -> CHRStore
    -> ConstraintToInfoMap
    -> ConstraintToInfoMap
    -> SimplifyResult
    -> ( ConstraintToInfoMap
       , SimplifyResult
       )
chrSimplifySolveToRedGraph env chrStore cnstrInfoMpPrev cnstrInfoMp prevRes
  = ( cnstrInfoMpAll
    , emptySimplifyResult
        { simpresSolveState = solveState
        , simpresRedGraph   = redGraph
        , simpresRedGraphs  = ("chrSimplifySolveToRedGraph",redGraph) : simpresRedGraphs prevRes
        }
    )
  where (_,u1,u2) = mkNewLevUID2 $ fiUniq env
        solveState = chrSolve'' (env {fiUniq = u1}) chrStore (map toSolverConstraint $ Map.keys $ cnstrInfoMp `Map.difference` cnstrInfoMpPrev) (simpresSolveState prevRes)
        cnstrInfoMpAll = cnstrMpUnion cnstrInfoMp cnstrInfoMpPrev
        redGraph
          = addToRedGraphFromReductions (filterMb fromSolverConstraint $ chrSolveStateDoneConstraints solveState)
            $ addToRedGraphFromAssumes cnstrInfoMpAll
            $ simpresRedGraph prevRes

{-# LINE 464 "src/ehc/Pred/ToCHR.chs" #-}
chrSimplifyRedGraphToEvidence
  :: Heuristic
  -> ConstraintToInfoMap
  -> SimplifyResult
  -> ( ( ConstraintToInfoTraceMp, InfoToEvidenceMap, [Err] )
     , SimplifyResult
     )
chrSimplifyRedGraphToEvidence heur cnstrInfoMpAll simpRes
  = ( (chrSolveRemCnstrMp,chrSolveEvidMp,chrSolveErrs)
    , simpRes
    )
  where ( chrSolveRemCnstrMp,chrSolveEvidMp,chrSolveErrs
         )
          = mkEvidence heur cnstrInfoMpAll (simpresRedGraph simpRes)