{-# OPTIONS --safe --without-K #-}
module Tactic.ClauseBuilder where
 
open import Meta.Prelude hiding ([_,_])
open import Meta.Init hiding (sort)
 
import Data.List.NonEmpty as NE
open import Data.List using (zipWith)
open import Data.Nat.Properties using (≤-totalOrder; ≤-decTotalOrder)
open import Data.Product using (map₁)
 
-- Note: using sort from Data.List.Sort directly doesn't work for
-- metaprogramming, since it's `abstract` for whatever reason
open import Data.List.Sort using (SortingAlgorithm)
open import Data.List.Sort.MergeSort using (mergeSort)
open SortingAlgorithm ≤-decTotalOrder (mergeSort ≤-decTotalOrder) public
 
open import Reflection.QuotedDefinitions
open import Reflection.Utils
open import Reflection.Utils.TCI
 
open import Class.Show
open import Class.DecEq
open import Class.Functor
open import Class.MonadError.Instances
open import Class.MonadReader.Instances
open import Class.MonadTC.Instances
open import Class.Traversable
 
module _ where
  open import Class.Monad hiding (Monad-List) public
 
import Class.Monad
 
record ClauseBuilder (M : Set → Set) : Set₁ where
  field
    Base : Set → Set
    liftBase : Base A → M A
    addPattern : List (String × Arg Type) → Arg Pattern → M A → M A
    toClause : M (Maybe Term) → Base Clause
 
open ClauseBuilder {{...}} public
 
SinglePattern : Set
SinglePattern = List (String × Arg Type) × Arg Pattern
 
typedVarSinglePattern : String → Arg Type → SinglePattern
typedVarSinglePattern n t@(arg i _) = ([ n , t ] , arg i (` ))
 
varSinglePattern : Arg String → SinglePattern
varSinglePattern (arg i n) = ([ n , arg i unknown ] , arg i (` ))
 
multiSinglePattern : List String → Arg Pattern → SinglePattern
multiSinglePattern s p = ((_, vArg unknown) <$> s) , p
 
findIndexDefault : List ℕ → ℕ → ℕ → ℕ
findIndexDefault l d a with filter (λ where (i , x) → x ≟ a) (zipWithIndex _,_ l)
... | []          = d
... | (i , _) ∷ _ = i
 
singlePatternFromPattern : Arg Pattern → SinglePattern
singlePatternFromPattern (arg i p) =
  replicate (length (appearingIndices p)) ("" , vArg unknown) , arg i (replacePatternIndex p)
  where
    appearingIndices : Pattern → List ℕ
    appearingIndicesHelper : List (Arg Pattern) → List ℕ
 
    appearingIndices (Pattern.con c ps) = appearingIndicesHelper ps
    appearingIndices (Pattern.dot t) = [] -- TODO: this is probably wrong
    appearingIndices (` x) = [ x ]
    appearingIndices (Pattern.lit l) = []
    appearingIndices (Pattern.proj f) = []
    appearingIndices (Pattern.absurd x) = []
 
    appearingIndicesHelper [] = []
    appearingIndicesHelper (arg _ p ∷ ps) = appearingIndices p ++ appearingIndicesHelper ps
 
    normalisedIndexList : List ℕ
    normalisedIndexList = sort $ deduplicate _≟_ $ appearingIndices p
 
    replacePatternIndex : Pattern → Pattern
    replacePatternIndexHelper : List (Arg Pattern) → List (Arg Pattern)
 
    replacePatternIndex (Pattern.con c ps) = Pattern.con c (replacePatternIndexHelper ps)
    replacePatternIndex (Pattern.dot t) = Pattern.dot t
    replacePatternIndex (` x) = ` findIndexDefault normalisedIndexList  x
    replacePatternIndex (Pattern.lit l) = Pattern.lit l
    replacePatternIndex (Pattern.proj f) = Pattern.proj f
    replacePatternIndex (Pattern.absurd x) = Pattern.absurd x
 
    replacePatternIndexHelper [] = []
    replacePatternIndexHelper (arg i p ∷ ps) = (arg i (replacePatternIndex p)) ∷ (replacePatternIndexHelper ps)
 
module _ {M : ∀ {a} → Set a → Set a} ⦃ _ : Monad M ⦄ ⦃ me : MonadError (List ErrorPart) M ⦄ ⦃ mre : MonadReader TCEnv M ⦄ ⦃ _ : MonadTC M ⦄ where
 
  ctxSinglePatterns : M (List SinglePattern)
  ctxSinglePatterns = do
    ctx ← getContext
    return (singlePatternFromPattern <$> zipWithIndex (λ where k (_ , (arg i _)) → arg i (` k)) ctx)
 
  -- TODO: add the ability to match initial hidden arguments
  -- TODO: add dot patterns
  constrToPattern : Name → Type → M SinglePattern
  constrToPattern n ty = do
    (introTel , _) ← viewTy <$> (runAndReset $ local (λ env → record env { normalisation = true }) $ inferType (n ◆))
    let patternTel = zipWith (λ where (abs _ (arg i _)) k → arg i (` k)) introTel $ downFrom $ length introTel
    return (((λ where (abs s (arg i t)) → (s , arg i unknown)) <$> introTel) , (vArg $ Pattern.con n patternTel))
 
  -- like constrToPattern, but applied to parameters
  constrToPatternTyped : Name → List Term → M SinglePattern
  constrToPatternTyped n ps = do
    t ← applyWithVisibility n ps
    debugLogᵐ (t ᵛⁿ ∷ᵈᵐ []ᵐ)
    (introTel , _) ← viewTy <$> (local (λ env → record env { normalisation = true }) $ inferType t)
    let patternTel = zipWith (λ where (abs _ (arg i _)) k → arg i (` k)) introTel $ downFrom $ length introTel
    return (((λ where (abs s (arg i t)) → (s , arg i t)) <$> introTel) , (vArg $ Pattern.con n patternTel))
 
  -- all possible patterns for an inductive type
  constructorPatterns : Arg Type → M (List SinglePattern)
  constructorPatterns (arg i ty) = do
    constrs ← getConstrsForType ty
    return $ constrs <&> λ where
      (n , t) → let argInfo = proj₁ $ viewTy t
        in (((λ where (abs n' t') → n' , t') <$> argInfo)
           , arg i (Pattern.con n (zipWithIndex (λ where k (abs _ (arg i _)) → arg i (` (length argInfo ∸ ℕ.suc k))) argInfo)))
 
  -- all possible patterns for an inductive type
  constructorPatterns' : Type → M (List SinglePattern)
  constructorPatterns' ty = do
    constrs ← getConstrsForType ty
    traverse (λ (n , _) → constrToPattern n ty) constrs
 
  -- all possible patterns for an inductive type
  constructorPatternsTyped : Type → List Type → M (List SinglePattern)
  constructorPatternsTyped ty ps = do
    constrs ← getConstrsForType ty
    traverse (λ (n , _) → constrToPatternTyped n ps) constrs
 
ClauseInfo : Set
ClauseInfo = List SinglePattern
 
data ClauseExpr : Set where
  MatchExpr : List (SinglePattern × (ClauseExpr ⊎ Maybe Term)) → ClauseExpr
 
multiClauseExpr : List (NE.List⁺ SinglePattern × (ClauseExpr ⊎ Maybe Term)) → ClauseExpr
multiClauseExpr = MatchExpr ∘ Data.List.map helper
  where
    helper' : (List SinglePattern × (ClauseExpr ⊎ Maybe Term)) → ClauseExpr ⊎ Maybe Term
    helper' ([]     , e) = e
    helper' (p ∷ ps , e) = inj₁ (MatchExpr [ p , helper' (ps , e) ])
 
    helper : (NE.List⁺ SinglePattern × (ClauseExpr ⊎ Maybe Term)) → SinglePattern × (ClauseExpr ⊎ Maybe Term)
    helper (p NE.∷ ps , e) = p , helper' (ps , e)
 
clauseExprToClauseInfo : ClauseExpr → List (ClauseInfo × Maybe Term)
clauseExprToClauseInfo (MatchExpr []) = []
clauseExprToClauseInfo (MatchExpr ((p , inj₁ e) ∷ xs)) =
  (map₁ (p ∷_) <$> clauseExprToClauseInfo e) ++ clauseExprToClauseInfo (MatchExpr xs)
clauseExprToClauseInfo (MatchExpr ((p , inj₂ t) ∷ xs)) = ([ p ] , t) ∷ clauseExprToClauseInfo (MatchExpr xs)
 
clauseInfoToClauseArgs : ClauseInfo → List (String × Arg Type) × List (Arg Pattern)
clauseInfoToClauseArgs i =
  case helper  i of λ where (tel , ps , _) → (tel , ps)
  where
    helper : ℕ → ClauseInfo → List (String × Arg Type) × List (Arg Pattern) × ℕ
    helper k [] = ([] , [] , k)
    helper k ((tel' , arg h p) ∷ i) with helper k i
    ... | (tel , ps , k') = (tel' ++ tel , arg h (mapVariables (_+ k') p) ∷ ps , length tel' + k')
 
clauseInfoToClause : ClauseInfo → Maybe Term → Clause
clauseInfoToClause i t with clauseInfoToClauseArgs i | t
... | tel , ps | just x  = Clause.clause tel ps x
... | tel , ps | nothing = Clause.absurd-clause tel ps
 
clauseExprToClauses : ClauseExpr → List Clause
clauseExprToClauses e = (λ { (i , t) → clauseInfoToClause i t }) <$> clauseExprToClauseInfo e
 
nonBindingClause : Term → Clause
nonBindingClause = Clause.clause [] []
 
clauseExprToPatLam : ClauseExpr → Term
clauseExprToPatLam e = pat-lam (clauseExprToClauses e) []
 
record ContextMonad (M : ∀ {a} → Set a → Set a) ⦃ _ : Monad M ⦄ : Setω where
  field
    introPatternM : SinglePattern → M A → M A
 
  extendContextM : Arg Type → M A → M A
  extendContextM l x = introPatternM (typedVarSinglePattern "" l) x
 
open ContextMonad ⦃...⦄
 
-- No context
module _ where
  open import Class.Monad.Id
 
  ContextMonad-Id : ContextMonad id
  ContextMonad-Id .introPatternM _ a = a
 
module _ {M : ∀ {a} → Set a → Set a} ⦃ _ : Monad M ⦄ ⦃ me : MonadError (List ErrorPart) M ⦄ ⦃ mre : MonadReader TCEnv M ⦄ ⦃ _ : MonadTC M ⦄ where
 
  refineWithSingle : (Term → Term) → M Term → M Term
  refineWithSingle ref x = do
    goalTy ← goalTy
    m ← newMeta unknown
    res ← checkType (ref m) goalTy
    x' ← runWithHole m x
    unify m x'
    return $ ref x'
 
  caseMatch : Term → M ClauseExpr → M Term
  caseMatch t expr = debugLog ("Match" ∷ᵈ t ∷ᵈ []) >> (refineWithSingle (`case_of_ t) $
    (λ expr' → pat-lam (clauseExprToClauses expr') []) <$> expr)
 
  currentTyConstrPatterns : M (List SinglePattern)
  currentTyConstrPatterns = do
    (ty ∷ _ , _) ← viewTy′ <$> goalTy
      where _ → error1 "currentTyConstrPatterns: Goal type is not a forall!"
    constructorPatterns' (unArg ty)
 
stripMetaLambdas : Term → Term
stripMetaLambdas = helper 
  where
    helper : ℕ → Term → Term
    helper k (lam _ (abs _ t)) = helper (k + ) t
    helper k (meta x as) = meta x $ map-Args (mapVars (_∸ k)) $ take (length as ∸ k) as
    helper _ _ = unknown
 
module _ {M : ∀ {a} → Set a → Set a} ⦃ _ : Monad M ⦄ ⦃ me : MonadError (List ErrorPart) M ⦄ ⦃ mre : MonadReader TCEnv M ⦄ ⦃ _ : MonadTC M ⦄ where
 
  isProj : Pattern → Bool
  isProj (Pattern.proj _) = true
  isProj _                = false
 
  -- if the goal is of type (a : A) → B, return the type of the branch of pattern p and new context
  specializeType : SinglePattern → Type → M (Type × Telescope)
  specializeType p@(t , arg i p') goalTy = inDebugPath "specializeType" $ runAndReset do
    debugLog ("Goal type to specialize: " ∷ᵈ goalTy ∷ᵈ [])
    cls@((Clause.clause tel _ _) ∷ _) ← return $ clauseExprToClauses $ MatchExpr $
        (p , inj₂ (just unknown)) ∷
        (if isProj p' then [] else [ varSinglePattern (arg i "_") , inj₂ (just unknown) ])
      where _ → error1 "BUG"
    debugLog ("With pattern: " ∷ᵈ cls ∷ᵈ [])
    (pat-lam (cl@(Clause.clause tel' _ (meta x ap)) ∷ _) []) ← checkType (pat-lam cls []) goalTy
      where t → debugLog ("BUG in specializeType:" ∷ᵈ t ∷ᵈ "\nWith pattern:" ∷ᵈ cls
                  ∷ᵈ "\nWith type:" ∷ᵈ goalTy ∷ᵈ "\nSinglePattern:" -- ∷ᵈ {!p!}
                  ∷ᵈ []) >> error1 "BUG"
    let varsToUnbind = 
    let newCtx = tel'
    let m = meta x (map-Args (mapVars (_∸ varsToUnbind)) $ take (length ap ∸ varsToUnbind) ap)
    debugLog1 "New context:"
    logContext newCtx
    goalTy' ← extendContext' newCtx $ inferType m
    return (goalTy' , newCtx)
 
  ContextMonad-MonadTC : ContextMonad M
  ContextMonad-MonadTC .introPatternM pat@(_ , arg _ p) x = do
    goalTy ← goalTy
    (newGoalTy , newContext) ← specializeType pat goalTy
    extendContext' newContext (runWithGoalTy newGoalTy x)
 
module ClauseExprM {M : ∀ {a} → Set a → Set a} ⦃ _ : Monad M ⦄ ⦃ _ : ContextMonad M ⦄ where
 
  -- Construct a ClauseExpr in M and extend the context appropriately
  matchExprM : List (SinglePattern × M (ClauseExpr ⊎ Maybe Term)) → M ClauseExpr
  matchExprM = _<$>_ MatchExpr ∘ traverse (λ (a , b) → (a ,_) <$> introPatternM a b)
 
  multiMatchExprM : List (NE.List⁺ SinglePattern × M (ClauseExpr ⊎ Maybe Term)) → M ClauseExpr
  multiMatchExprM = matchExprM ∘ Data.List.map helper
    where
      helper' : (List SinglePattern × M (ClauseExpr ⊎ Maybe Term)) → M (ClauseExpr ⊎ Maybe Term)
      helper' ([]     , e) = e
      helper' (p ∷ ps , e) = inj₁ <$> (matchExprM [ p , helper' (ps , e) ])
 
      helper : (NE.List⁺ SinglePattern × M (ClauseExpr ⊎ Maybe Term)) → SinglePattern × M (ClauseExpr ⊎ Maybe Term)
      helper (p NE.∷ ps , e) = (p , helper' (ps , e))
 
  singleMatchExpr : SinglePattern → M (ClauseExpr ⊎ Maybe Term) → M ClauseExpr
  singleMatchExpr p x = matchExprM [ p , x ]
 
  singleTelescopeMatchExpr : NE.List⁺ SinglePattern → M (ClauseExpr ⊎ Maybe Term) → M ClauseExpr
  singleTelescopeMatchExpr (p NE.∷ ps) x = helper p ps x
    where
      helper : SinglePattern → List SinglePattern → M (ClauseExpr ⊎ Maybe Term) → M ClauseExpr
      helper p [] x = singleMatchExpr p x
      helper p (p' ∷ ps) x = singleMatchExpr p $ inj₁ <$> helper p' ps x
 
  introExpr : Arg String → M (ClauseExpr ⊎ Maybe Term) → M ClauseExpr
  introExpr n x = singleMatchExpr (varSinglePattern n) x
 
  introsExpr : NE.List⁺ (Arg String) → M (ClauseExpr ⊎ Maybe Term) → M ClauseExpr
  introsExpr (p NE.∷ ps) x = helper p ps x
    where
      helper : Arg String → List (Arg String) → M (ClauseExpr ⊎ Maybe Term) → M ClauseExpr
      helper n [] x = introExpr n x
      helper n (n' ∷ ns) x = introExpr n $ inj₁ <$> helper n ns x
 
  contMatch : M ClauseExpr → M (ClauseExpr ⊎ Maybe Term)
  contMatch expr = inj₁ <$> expr
 
  finishMatch : M Term → M (ClauseExpr ⊎ Maybe Term)
  finishMatch t = (inj₂ ∘ just) <$> t
 
  bindCtxMatchExpr : ⦃ me : MonadError (List ErrorPart) M ⦄ ⦃ mre : MonadReader TCEnv M ⦄ ⦃ _ : MonadTC M ⦄
                   → M ClauseExpr → M ClauseExpr
  bindCtxMatchExpr x = do
    e ← ctxSinglePatterns
    case NE.fromList e of λ where
      (just e') → singleTelescopeMatchExpr e' $ contMatch x
      nothing → x
 
clauseTelescope : Clause → List (String × Arg Type)
clauseTelescope (Clause.clause tel _ _) = tel
clauseTelescope (Clause.absurd-clause tel _) = tel
 
module _ where
  open import Class.Monad.Id
 
  open ClauseExprM ⦃ Monad-Id ⦄ ⦃ ContextMonad-Id ⦄
 
  instanciatePattern : SinglePattern → List (Arg Type)
  instanciatePattern p = proj₂ <$>
    (clauseTelescope $ from-just $ head $ clauseExprToClauses $ singleMatchExpr p $ finishMatch unknown)
 
  instanciatePatterns : List SinglePattern → Term → List Clause
  instanciatePatterns [] t = [ Clause.clause [] [] t ]
  instanciatePatterns (x ∷ ps) t =
    clauseExprToClauses $ singleTelescopeMatchExpr (x NE.∷ ps) (finishMatch t)
 
ctxBindingClause : Term → TC Clause
ctxBindingClause t = do
  pats ← ctxSinglePatterns
  (c ∷ _) ← return $ instanciatePatterns (reverse pats) t
    where _ → error1 "Bug in ctxBindingClause"
  return c