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 Mathlib.Probability.Kernel.Composition.CompProd
import Mathlib.MeasureTheory.MeasurableSpace.Embedding
import Lean.Elab.Tactic.Location
import Mathlib.Probability.Kernel.Deterministic
import Mathlib.Combinatorics.Quiver.ReflQuiver
import Mathlib.Probability.Kernel.Category.SFinKer
import Mathlib.MeasureTheory.Integral.Lebesgue.Countable

/-! # Standalone extraction for `KernelEquality`
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
namespace ProbabilityTheory.Kernel
end ProbabilityTheory.Kernel

-- ═══ 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

/-- Extract the universe level from the left side of an equality expression. -/
def getLevelFromEq (eq : Expr) : MetaM Level := do
  let eq ← whnf (← zetaReduce (← instantiateMVars eq))
  let eq := eq.consumeMData
  let some (_, lhs, _) := eq.eq? | throwError "Expected an equality, got: {eq}."
  getDecLevel (← inferType lhs)

end

-- ═══ Tactic.EqLift.Tactic.Kernel.Utils ═══
section
public meta section
open Lean Meta ProbabilityTheory Elab Term

/-- Extract `(X, Y, u, v)` from an expression of type `Kernel X Y`. -/
def getTypesFromKernel (κ : Expr) : MetaM (Expr × Expr × Level × Level) := do
  let κType ← inferType κ
  match κType.getAppFn with
  | Expr.const ``Kernel univs =>
    let args := κType.getAppArgs
    if args.size < 2 then
      throwError "Kernel type with insufficient arguments: {κType}."
    let X := args[0]!
    let Y := args[1]!
    let xLevel := univs[0]!
    let yLevel := univs[1]!
    return (X, Y, xLevel, yLevel)
  | _ => throwError "Expected a kernel type, got: {κType}."

end
end

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

/-- Unlifts an equality expression that has been lifted to a common universe level using the
registered unlifting functions and finisher functions. -/
def unliftEquality := transformEquality getLevelFromEq unliftImplRef unliftFinisherRef

end
end

-- ═══ Tactic.KernelHom.Tactic.Utils ═══
section
public meta section
open Lean Meta ProbabilityTheory

/-- Returns the application `constName` `xs` with `n_impls` last arguments as implicit. -/
def Lean.Meta.mkAppMInst (constName : Name) (xs : Array Expr) (n_impls : Nat) : MetaM Expr := do
  let e ← mkAppM constName xs
  let nones : Array (Option Expr) := Array.replicate n_impls none
  mkAppOptM' e nones

end

-- ═══ Tactic.KernelHom.Tactic.KernelHom ═══
section
public meta section
open Lean Elab Tactic Meta CategoryTheory Parser.Tactic ProbabilityTheory MonoidalCategory
open ProbabilityTheory.Kernel

/-- Recursively decompose a product type into `SFinKer` objects with monoidal tensor structure. -/
partial def decomposeProductToSFinker (X : Expr) (xLvl : Level) : MetaM Expr := do
  match X.getAppFn with
  | Expr.const ``Prod _ =>
    let args := X.getAppArgs
    let t1 ← decomposeProductToSFinker args[0]! xLvl
    let t2 ← decomposeProductToSFinker args[1]! xLvl
    mkAppM ``tensorObj #[t1, t2]
  | _ =>
    mkAppOptM ``SFinKer.of #[X, none]

/-- Compute the `SFinKer` object corresponding to a measurable space. -/
def computeSFinkerOf (X : Expr) (xLvl : Level) : MetaM Expr := do
  match X with
  | Expr.const ``PUnit _ | Expr.const ``Unit _ =>
    let tensorunit := mkConst ``tensorUnit [xLvl, xLvl.succ]
    let sfinker := mkConst ``SFinKer [xLvl]
    mkAppOptM' tensorunit #[sfinker, none, none]
  | _ =>
    decomposeProductToSFinker X xLvl

/-- Compute a measurable equivalence between a type and itself by recursively decomposing
products. -/
partial def idME (X : Expr) : MetaM Expr := do
  match X.getAppFn with
  | Expr.const ``Prod _ =>
    let args := X.getAppArgs
    let id1 ← idME args[0]!
    let id2 ← idME args[1]!
    mkAppM ``MeasurableEquiv.prodCongr #[id1, id2]
  | Expr.const ``PUnit [xLvl] | Expr.const ``Unit [xLvl] =>
    let xLvl ← match xLvl with
      | Level.succ l => pure l
      | _ => throwError "Expected a successor level for PUnit/Unit, got: {xLvl}."
    let punitME := mkConst ``MeasurableEquiv.punit [xLvl, xLvl]
    mkAppM' punitME #[]
  | _ =>
    mkAppOptM ``MeasurableEquiv.refl #[X, none]

/-- Construct the proof of equivalence between the original equality and the transformed one. -/
def mkKernelHomEqProof (eqProofType lhs rhs : Expr) (proofs : List Expr) : 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 (X, Y, xLvl, yLvl) ← getTypesFromKernel lhs
    let SX ← computeSFinkerOf X xLvl
    let SY ← computeSFinkerOf Y yLvl
    let e ← mkAppMInst ``hom_congr #[SX, SY, ← idME X, ← idME Y, lhs, rhs] 2
    unless ← isDefEq (← mvarId.getType) (← inferType e) do
      throwError "Type mismatch: expected {← mvarId.getType}, got {← inferType e}."
    mvarId.assign e
    instantiateMVars mvar
  | _ =>
    throwError "Failed to apply propext while building kernel_hom equivalence proof for
      {eqProofType}."

end

-- ═══ Tactic.KernelHom.Tactic.HomKernel ═══
section
public meta section
open Lean Elab Tactic Meta CategoryTheory Parser.Tactic ProbabilityTheory MonoidalCategory
open ProbabilityTheory.Kernel

/-- Recursive transformation from morphism expression in `SFinKer` to kernel expression. -/
partial def transformHomToKernel (e : Expr) (proofs : List Expr) :
    MetaM (Expr × List Expr) := do
  match e.getAppFn with
  | Expr.const ``tensorHom _ =>
    let args := e.getAppArgs
    let κ := args[args.size - 2]!
    let η := args[args.size - 1]!
    let ST := args[args.size - 3]!
    let SZ := args[args.size - 4]!
    let SY := args[args.size - 5]!
    let SX := args[args.size - 6]!
    let (κ', proofs_κ) ← transformHomToKernel κ proofs
    let (η', proofs_η) ← transformHomToKernel η proofs_κ
    let (X, Y, _, _) ← getTypesFromKernel κ'
    let (Z, T, _, _) ← getTypesFromKernel η'
    let parallelComp_hom_proof ← mkAppMInst ``parallelComp_hom
        #[SX, SY, SZ, ST, ← idME X, ← idME Y, ← idME Z, ← idME T, κ', η'] 2
    return (← mkAppM ``Kernel.parallelComp #[κ', η'], parallelComp_hom_proof :: proofs_η)
  | Expr.const ``CategoryStruct.comp _ =>
    let args := e.getAppArgs
    let κ := args[args.size - 2]!
    let η := args[args.size - 1]!
    let SY := args[args.size - 3]!
    let SX := args[args.size - 4]!
    let SZ := args[args.size - 5]!
    let (κ', proofs_κ) ← transformHomToKernel κ proofs
    let (η', proofs_η) ← transformHomToKernel η proofs_κ
    let (X, Y, _, _) ← getTypesFromKernel η'
    let (Z, _, _, _) ← getTypesFromKernel κ'
    let comp_hom_proof ← mkAppMInst ``comp_hom
        #[SX, SY, SZ, ← idME X, ← idME Y, ← idME Z, η', κ'] 2
    return (← mkAppM ``Kernel.comp #[η', κ'], comp_hom_proof :: proofs_η)
  | Expr.const ``CategoryStruct.id [xLvl, _] =>
    let args := e.getAppArgs
    let SX := args[args.size - 1]!
    let X ← getTypeFromSFinKer SX
    let mX' ← synthInstance (mkApp (mkConst ``MeasurableSpace [xLvl]) X)
    let id ← mkAppOptM ``Kernel.id #[X, mX']
    let id_hom_proof ← mkAppM ``id_hom #[SX, ← idME X]
    return (id, id_hom_proof :: proofs)
  | Expr.const ``ComonObj.counit [xLvl, _] =>
    let args := e.getAppArgs
    let SX := args[args.size - 2]!
    let X ← getTypeFromSFinKer SX
    let discard_kernel_const := mkConst ``Kernel.discard [xLvl, xLvl]
    let discard_const := mkConst ``counit [xLvl, xLvl, xLvl]
    let discard_hom_proof ← mkAppM' discard_const #[SX, ← idME X]
    return (← mkAppOptM' discard_kernel_const #[X, none], discard_hom_proof :: proofs)
  | Expr.const ``ComonObj.comul [xLvl, _] =>
    let args := e.getAppArgs
    let SX := args[args.size - 2]!
    let X ← getTypeFromSFinKer SX
    let copy_kernel_const := mkConst ``Kernel.copy [xLvl]
    let copy_hom_proof ← mkAppM ``comul #[SX, ← idME X]
    return (← mkAppOptM' copy_kernel_const #[X, none], copy_hom_proof :: proofs)
  | Expr.const ``Kernel.hom _ =>
    let args := e.getAppArgs
    let κ := args[args.size - 2]!
    return (κ, proofs)
  | Expr.const ``MonoidalCategory.whiskerLeft [eLvl, _] =>
    let (κ, kernel_id, SX, SY, SZ, X, Y, Z) ← deconstructWhiskersHomArgs e eLvl true
    let (κ', proofs_κ) ← transformHomToKernel κ proofs
    let whisker_left_hom_proof ← mkAppMInst ``Kernel.whiskerLeft
      #[SX, SY, SZ, ← idME X, ← idME Y, ← idME Z, κ'] 1
    return (← mkAppM ``Kernel.parallelComp #[kernel_id, κ'], whisker_left_hom_proof :: proofs_κ)
  | Expr.const ``MonoidalCategory.whiskerRight [eLvl, _] =>
    let (κ, kernel_id, SX, SY, SZ, X, Y, Z) ← deconstructWhiskersHomArgs e eLvl false
    let (κ', proofs_κ) ← transformHomToKernel κ proofs
    let whisker_right_hom_proof ← mkAppMInst ``Kernel.whiskerRight
      #[SX, SY, SZ, ← idME X, ← idME Y, ← idME Z, κ'] 1
    return (← mkAppM ``Kernel.parallelComp #[κ', kernel_id], whisker_right_hom_proof :: proofs_κ)
  | Expr.const ``Iso.hom _ =>
    let args := e.getAppArgs
    let iso := args[args.size - 1]!
    match iso.getAppFn with
    | Expr.const ``BraidedCategory.braiding _ =>
      let (braiding_expr, swap_hom_proof) ← deconstructBraiding iso
      return (braiding_expr, swap_hom_proof :: proofs)
    | Expr.const ``leftUnitor [eLvl, _] =>
      let (left_unitor_expr, left_unitor_hom_proof) ← deconstructUnitors iso eLvl true true
      return (left_unitor_expr, left_unitor_hom_proof :: proofs)
    | Expr.const ``rightUnitor [eLvl, _] =>
      let (right_unitor_expr, right_unitor_hom_proof) ← deconstructUnitors iso eLvl false true
      return (right_unitor_expr, right_unitor_hom_proof :: proofs)
    | Expr.const ``MonoidalCategory.associator [eLvl, _] =>
      let (associator_expr, associator_hom_proof) ← deconstructAssociator iso eLvl true
      return (associator_expr, associator_hom_proof :: proofs)
    | _ => throwError "Unexpected isomorphism {iso}."
  | Expr.const ``Iso.inv _ =>
    let args := e.getAppArgs
    let iso := args[args.size - 1]!
    match iso.getAppFn with
    | Expr.const ``BraidedCategory.braiding _ =>
      let (braiding_expr, swap_hom_proof) ← deconstructBraiding iso
      return (braiding_expr, swap_hom_proof :: proofs)
    | Expr.const ``leftUnitor [eLvl, _] =>
      let (left_unitor_expr, left_unitor_inv_hom_proof) ← deconstructUnitors iso eLvl true false
      return (left_unitor_expr, left_unitor_inv_hom_proof :: proofs)
    | Expr.const ``rightUnitor [eLvl, _] =>
      let (right_unitor_expr, right_unitor_inv_hom_proof) ← deconstructUnitors iso eLvl false false
      return (right_unitor_expr, right_unitor_inv_hom_proof :: proofs)
    | Expr.const ``MonoidalCategory.associator [eLvl, _] =>
      let (associator_expr, associator_inv_hom_proof) ← deconstructAssociator iso eLvl false
      return (associator_expr, associator_inv_hom_proof :: proofs)
    | _ => throwError "Unexpected isomorphism {iso}."
  | _ => throwError "Expected a hom expression, got: {e}."

/-- Transform a `SFinKer` equality into an equivalent equality of kernels, along with a proof of
equivalence. -/
def KernelEquality (eq : Expr) : MetaM (Expr × Expr) := do
  let eq ← whnfR <| ← instantiateMVars eq
  let some (_, lhs_hom, rhs_hom) := eq.eq? | throwError "Expected an equality, got: {eq}."
  let (lhs, proofs) ← transformHomToKernel lhs_hom []
  let (rhs, proofs) ← transformHomToKernel rhs_hom proofs
  let kernel_expr ← mkEq lhs rhs
  let (unlifted_expr, unlifted_proof) ← unliftEquality kernel_expr
  let kernel_eq_proof_type ← mkEq kernel_expr eq
  let kernel_eq_proof ← mkAppM ``Eq.symm #[← mkKernelHomEqProof kernel_eq_proof_type lhs rhs proofs]
  return (unlifted_expr, ← mkEqTrans kernel_eq_proof unlifted_proof)

end