{-# OPTIONS --safe --no-exact-split #-}
module Meta.Regularity where

open import Meta.Prelude
open import Meta.Reflection.Base
open import Meta.Reflection.Subst

open import Data.Bool.Base
open import Data.List.Base
open import Data.List.Instances.FromProduct
open import Data.Maybe.Base
open import Data.Maybe.Instances.Alt
open import Data.Reflection.Abs
open import Data.Reflection.Argument
open import Data.Reflection.Error
open import Data.Reflection.Fixity
open import Data.Reflection.Instances.FromString
open import Data.Reflection.Literal
open import Data.Reflection.Meta
open import Data.Reflection.Name
open import Data.Reflection.Term
open import Data.Unit.Base


{-
A tactic for reducing "transport refl x" in other-wise normal terms. The
implementation is actually surprisingly simple: A term of the form (e.g.)

    transp (λ _ → A) i0 (f (transp (λ _ → B) i0 x))

is already a blueprint for how to normalise it. We simply have to turn it into

    λ i → transp (λ _ → A) i (f (transp (λ _ → B) i x))

so that the constant transports reduce away when (i = i1). Abstracting over i,
this gives a path from the initial term to its "regular normal form", which
may be the worst name ever.

More generally, we replace terms of the form `transp Al φ x` with `transp Al (φ ∨ i) x`
recursively (inside-out), on the condition that Al is constant when i = i1.
-}

private
  open import Agda.Primitive.Cubical using ( primHComp )
  -- We have a double composition operator that doesn't use the
  -- fancy hcomp syntax in its definition. This has better type
  -- inference for one of the macros since it guarantees that the base
  -- (q i) is independent of j without any reduction.
  double-comp
    : ∀ {ℓ} {A : Type ℓ} {w z : A} (x y : A)
    → w ＝ x → x ＝ y → y ＝ z
    → w ＝ z
  double-comp x y p q r i = primHComp
    (λ { j (i = i0) → p (~ j) ; j (i = i1) → r j }) (q i)

  -- The regularity tactic can operate in two modes: `precise` will work
  -- with the type-checker to identify which `transp`s are along refl,
  -- and which should be preserved. The `fast` mode says YOLO and
  -- assumes that **every** application of `transp` is one that would
  -- reduce by regularity. Needless to say, only use `fast` when you're
  -- sure that's the case (e.g. the fibres of a displayed category over
  -- Sets)
  data Regularity-precision : Type where
    precise fast : Regularity-precision
  -- As the name implies, `precise` is `precise`, while `fast` is
  -- `fast`. The reason is that `fast` will avoid traversing many of the
  -- terms involved in a `transp`: It doesn't care about the level, it
  -- doesn't care about the line, and it doesn't care about the
  -- cofibration.

  -- The core of the tactic is this triad of mutually recursive
  -- functions. In all three of them, the `Nat` argument indicates how
  -- many binders we've gone under: it is the dimension variable we
  -- abstracted over at the start.
  refl-transport : ℕ → Term → TC Term
  -- ^ Determines whether an application of `transp` is a case of
  -- regularity, and if so, replaces it by `regular`. Precondition: its
  -- subterms must already be reduced.
  go  : Regularity-precision → ℕ → Term → TC Term
  -- ^ Reduces all the subterms, and finds applications of `transp` to
  -- hand off to `refl-transport`.
  go* : Regularity-precision → ℕ → List (Arg Term) → TC (List (Arg Term))
  -- ^ Isn't the termination checker just lovely?

  refl-transport n tm@(def (quote transp) (ℓ h∷ Al v∷ φ v∷ x v∷ [])) =
    -- This match might make you wonder: Can't Al be a line of
    -- functions, so that the transport will have more arguments? No:
    -- The term is in normal form.
    (do
      debug-print "tactic.regularity" 10 $ "Checking regularity of " ∷ term-err tm ∷ []
      let φ′ = def (quote _∨_) (φ v∷ var n [] v∷ [])
      let tm′ = def (quote transp) (ℓ h∷ Al v∷ φ′ v∷ x v∷ [])
      -- We simply ask Agda to check that the newly constructed term `transp Al (φ ∨ i) x`
      -- is correct, i.e. that Al is constant on (i = i1).
      -- If it isn't, we backtrack and leave the term unchanged.
      -- Note that if Al itself contains constant transports, we have already processed those,
      -- so they reduce away when (i = i1).
      check-type tm′ unknown -- inferType doesn't trigger the constancy check https://github.com/agda/agda/issues/6585
      pure tm′) <|>
    (do
      debug-print "tactic.regularity" 10 $ "NOT a (transport refl): " ∷ term-err tm ∷ []
      pure tm)
  refl-transport _ tm = pure tm

  -- Boring term traversal.
  go pre n (var x args) = var x <$> go* pre n args
  go pre n (con c args) = con c <$> go* pre n args
  go fast n (def (quote transp) (ℓ h∷ Al v∷ φ v∷ x v∷ [])) = do
    x ← go fast n x
    let φ′ = def (quote _∨_) (φ v∷ var n [] v∷ [])
    pure $ def (quote transp) (ℓ h∷ Al v∷ φ′ v∷ x v∷ [])
  go pre n (def f args) = do
    as ← go* pre n args
    refl-transport n (def f as)
  go pre k t@(lam v (abs nm b)) = lam v ∘ abs nm <$> under-abs t (go pre (suc k) b)
  go pre n (pat-lam cs args) = type-error $ "regularity: Can not deal with pattern lambdas" -- TODO can we do something about this?
  go pre n t@(pi (arg i a) (abs nm b)) = do
    a ← go pre n a
    b ← under-abs t (go pre (suc n) b)
    pure (pi (arg i a) (abs nm b))
  go pre n (agda-sort s) = pure (agda-sort s)
  go pre n (lit l) = pure (lit l)
  go pre n (meta x args) = meta x <$> go* pre n args
  go pre n unknown = pure unknown

  go* pre n [] = pure []
  go* pre n (arg i a ∷ as) = do
    a ← go pre n a
    (arg i a ∷_) <$> go* pre n as

  -- To turn a term into a regularity path, given a level of precision,
  -- all we have to do is raise the term by one, do the procedure above,
  -- then wrap it in a lambda. Nice!
  to-regularity-path : Regularity-precision → Term → TC Term
  to-regularity-path pre tm = do
    tm ← raiseTC 1 tm
    -- Since we'll be comparing terms, Agda really wants them to be
    -- well-scoped. Since we shifted eeeverything up by one, we have to
    -- grow the context, too.
    tm ← resetting $ extend-context "i" (argN (quoteTerm I)) $ go pre 0 tm
    pure $ vlam "i" tm

  -- Extend a path x ＝ y to a path x′ ＝ y′, where x′ --> x and y′ --> y
  -- under the given regularity precision. Shorthand for composing
  --    regularity! ∙ p ∙ sym regularity!.
  regular!-worker :
    ∀ {ℓ} {A : Type ℓ} {x y : A}
    → Regularity-precision
    → x ＝ y
    → Term
    → TC ⊤
  regular!-worker {x = x} {y} pre p goal = do
    gt ← infer-type goal
    `x ← quoteTC x
    `y ← quoteTC y
    `p ← quoteTC p
    just (_ , l , r) ← unapply-path =<< infer-type goal
      where _ → type-error []
    l ← normalise =<< wait-for-type l
    r ← normalise =<< wait-for-type r
    reg ← to-regularity-path pre l
    reg′ ← to-regularity-path pre r
    unify-loudly goal $ def (quote double-comp) $
         `x v∷ `y v∷ reg
      v∷ `p
      v∷ def (quote sym) (reg′ v∷ [])
      v∷ []

module Regularity where
  open Regularity-precision public
  -- The reflection interface: Regularity.reduce! will, well, reduce a
  -- term. There's a lot of blocking involved in making this work.
  macro
    reduce! : Term → TC ⊤
    reduce! goal = do
      just (_ , l , r) ← unapply-path =<< infer-type goal
        where _ → type-error []
      reg ← to-regularity-path precise =<< (wait-for-type =<< normalise l)
      unify-loudly goal reg

    -- We then have wrappers that reduce on one side, and expand on the
    -- other, depending on how precise you want the reduction to be.
    precise! : ∀ {ℓ} {A : Type ℓ} {x y : A} → x ＝ y → Term → TC ⊤
    fast! : ∀ {ℓ} {A : Type ℓ} {x y : A} → x ＝ y → Term → TC ⊤

    precise! p goal = regular!-worker precise p goal
    fast! p goal = regular!-worker fast p goal

    -- For debugging purposes, this macro will take a term and output
    -- its (transport refl)-normal form, according to the given level of
    -- precision.
    reduct : Regularity-precision → Term → Term → TC ⊤
    reduct pres tm _ = do
      orig ← wait-for-type =<< normalise tm
      tm ← to-regularity-path pres orig
      red ← applyTC tm (argN (con (quote i1) [])) >>= normalise
      `pres ← quoteTC pres
      type-error
        [ "The term\n\n  " , term-err orig , "\n\nreduces modulo "
        , term-err `pres , " regularity to\n\n  "
        , term-err red , "\n" ]

-- Test cases.
module
  _ {ℓ ℓ′} {A : Type ℓ} {B : Type ℓ′} (f g : A → B) (x : A)
    (a-loop : (i : I) → Type ℓ [ (i ∨ ~ i) ↦ (λ ._ → A) ])
  where private

  p : (h : ∀ x → f x ＝ g x)
    → transport (λ i → A → B)
        (λ x → transport refl (transport refl (f (transport refl x)))) ＝ g
  p h = Regularity.reduce! ∙ fun-ext h

  q : transport refl (transp (λ i → outS (a-loop i)) i0 x) ＝ transp (λ i → outS (a-loop i)) i0 x
  q = Regularity.reduce!

  -- Imprecise/fast reduction: According to it, the normal form of the
  -- transport below is refl. That's.. not the case, at least we don't
  -- know so. Precise regularity handles it, though.
  -- q′ : ⊤
  -- q′ = {! Regularity.reduct Regularity.fast (transp (λ i → outS (a-loop i)) i0 x) !}

  r : (h : ∀ x → f x ＝ g x) → transport refl (f x) ＝ transport refl (g x)
  r h = Regularity.precise! (h x)

  s : (h : ∀ x → f x ＝ g x) → transport refl (g x) ＝ transport refl (f x)
  s h = sym $ Regularity.fast! (h x)
