import Mathlib.MeasureTheory.Order.Lattice
import Lean.Meta.Tactic.Replace
import Lean.Meta.Tactic.Rewrite
import Lean.Meta.DecLevel
import Lean.Meta.Transform
import Lean.Util.Recognizers
import Mathlib.Probability.Kernel.Composition.Prod
import Lean.Elab.Tactic.Location

/-! # Standalone extraction for `liftEquality`
Definitions are copied verbatim; theorem proofs are replaced by `sorry`.
Auto-generated by Referee. -/

set_option quotPrecheck false

-- Namespace stubs (so later `open`s resolve).
namespace Finset
end Finset
namespace Lean
end Lean
namespace Lean.Meta
end Lean.Meta
namespace ProbabilityTheory
end ProbabilityTheory

-- ═══ ForMathlib.MeasureTheory.Order.Lattice ═══
section
open Finset
variable {α δ : Type*} [MeasurableSpace δ] [SemilatticeInf α] {m : MeasurableSpace α} [MeasurableInf₂ α]
attribute [to_dual existing] MeasurableInf₂
end

-- ═══ Tactic.EqLift.Tactic.Utils ═══
section
public meta section
open Lean Elab Tactic Meta Parser.Tactic

/-- A type alias for lifting/unlifting functions. -/
abbrev liftMetadata := Expr → Level → List Expr → MetaM (Expr × List Expr)

/-- A type alias for finisher functions that construct the final proof of equality after lifting/
unlifting inner expressions. -/
abbrev finisherMetadata := Expr → Expr → Expr → Expr → Level → MetaM Expr

/-- Transforms an expression using the registered lifting/unlifting functions given in `impl_ref`. Returns the first successful transformation along with the updated list of proofs. -/
def transformExpr (e : Expr) (maxLvl : Level) (proofs : List Expr)
    (impl_ref : IO.Ref (Array (liftMetadata))) : MetaM (Expr × List Expr) := do
  let handlers ← impl_ref.get
  let (lift_expr, proofs) ← handlers.firstM (fun h => h e maxLvl proofs) <|> do
    throwError "No transform handler found for {e}."
  return (lift_expr, proofs)

/-- Rewrites the type of `mvarId` at the `n`-th occurrence using `heq`.-/
def Lean.MVarId.nthRewrite (mvarId : MVarId) (n : Nat) (heq : Expr) : MetaM MVarId := do
  let r ← mvarId.rewrite (← mvarId.getType) heq (config := { occs := .pos [n] })
  mvarId.replaceTargetEq r.eNew r.eqProof

/-- Constructs a proof of equality between the original and transformed expressions using the provided proofs and finisher functions. -/
def constructProof (eqProofType lhs rhs lhs_t rhs_t : Expr) (maxLvl : Level) (proofs : List Expr)
    (finisher_ref : IO.Ref (Array finisherMetadata)) : MetaM Expr := do
  let mvar ← mkFreshExprSyntheticOpaqueMVar eqProofType
  let mvarId := mvar.mvarId!
  let propext := mkConst ``propext
  match ← mvarId.apply propext with
  | [mvarId] =>
    let proofs := proofs.reverse
    let mut mvarId := mvarId
    for proof in proofs do
      mvarId ← mvarId.nthRewrite 1 proof
    let handlers ← finisher_ref.get
    let e ← handlers.firstM (fun h => do
      let finisher ← h lhs rhs lhs_t rhs_t maxLvl
      unless ← isDefEq (← mvarId.getType) (← inferType finisher) do
        throwError "Type mismatch: expected {← mvarId.getType}, got {← inferType finisher}."
      mvarId.assign finisher
      instantiateMVars mvar
    ) <|> do
      throwError m!"No finisher found for {eqProofType}."
    return e
  | _ =>
    throwError "Failed to apply propext while building kernel_lift equivalence proof for
      {eqProofType}."

/-- Lifts or unlifts an equality expression by transforming both sides using the registered lifting/
unlifting functions. Returns the transformed equality and a proof of equality between the original
and transformed expressions. -/
def transformEquality (getLvl : Expr → MetaM Level) (lift_ref : IO.Ref (Array liftMetadata))
    (finisher_ref : IO.Ref (Array finisherMetadata)) (eq : Expr) : MetaM (Expr × Expr) := do
  let e ← whnfR <| ← zetaReduce <| ← instantiateMVars eq
  let e := e.consumeMData
  let lvl ← getLvl eq
  let some (_, lhs, rhs) := e.eq? | throwError "Expected an equality, got: {e}."
  let (lhs_transformed, proofs) ← transformExpr lhs lvl [] lift_ref
  let (rhs_transformed, proofs) ← transformExpr rhs lvl proofs lift_ref
  let eq_transformed ← mkEq lhs_transformed rhs_transformed
  let eq_proof_type ← mkEq eq eq_transformed
  let proof ← constructProof
    eq_proof_type
    lhs rhs
    lhs_transformed rhs_transformed
    lvl
    proofs
    finisher_ref
  return (eq_transformed, proof)

end
end

-- ═══ Tactic.EqLift.Tactic.Universe ═══
section
public meta section
open Lean Meta ProbabilityTheory

/-- Recursively traverses an expression and collects all universe levels. -/
def collectExprUniverses.aux (e : Expr) : List Level :=
  match e with
  | Expr.const _ univs => univs
  | Expr.sort u => [u]
  | Expr.app f a => aux f ++ aux a
  | Expr.lam _ t b _ => aux t ++ aux b
  | Expr.forallE _ t b _ => aux t ++ aux b
  | Expr.letE _ t v b _ => aux t ++ aux v ++ aux b
  | Expr.mdata _ b => aux b
  | Expr.proj _ _ b => aux b
  | Expr.bvar _ | Expr.fvar _ | Expr.mvar _ | Expr.lit _ => []

/-- Recursively traverse an expression and collect universe levels found.
Returns a list of all unique universe levels encountered. -/
def collectExprUniverses (e : Expr) : MetaM (List Level) := do
  let e ← instantiateMVars e
  let e ← zetaReduce e
  return (collectExprUniverses.aux e).eraseDups

/-- Compute the maximum universe level from a list of levels. -/
def computeMaxLevel (levels : List Level) : MetaM Level :=
  match levels with
    | [] => throwError "Expected at least one universe level, got an empty list."
    | head :: tail => pure (tail.foldl Level.max head)

end

-- ═══ Tactic.EqLift.Tactic.Lift ═══
section
public meta section
open Lean Elab Tactic Meta Parser.Tactic

/-- Gets the maximum universe level from an equality expression by collecting all universe levels
from the left-hand side and right-hand side of the equality. -/
def getMaxLvl (eq : Expr) : MetaM Level := do
  let univs ← collectExprUniverses eq
  computeMaxLevel univs

/-- Lifts an equality expression to a common universe level using the registered lifting functions
and finisher functions. -/
def liftEquality := transformEquality getMaxLvl liftImplRef liftFinisherRef

end
end