import Mathlib.MeasureTheory.Order.Lattice
import Lean.Elab.Tactic.Location
import Mathlib.Probability.Kernel.Composition.Prod
import Mathlib.Probability.Kernel.Composition.CompProd
import Mathlib.MeasureTheory.MeasurableSpace.Embedding
import Lean.Meta.DecLevel
import Lean.Meta.Transform
import Lean.Util.Recognizers
import Lean.Meta.Tactic.Replace
import Lean.Meta.Tactic.Rewrite
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 `transformKernelToHom`
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.KernelHom.Tactic.KernelHom ═══
section
public meta section
open Lean Elab Tactic Meta CategoryTheory Parser.Tactic ProbabilityTheory MonoidalCategory
open ProbabilityTheory.Kernel

/-- Recursive transformation from kernel expressions to morphism expressions in the `SFinKer`
category. -/
partial def transformKernelToHom (e : Expr) (proofs : List Expr) :
    MetaM (Expr × List Expr) := do
  match e.getAppFn with
  | Expr.const ``Kernel.comp _ =>
    let args := e.getAppArgs
    let η := args[args.size - 2]!
    let κ := args[args.size - 1]!
    let (X, Y, xLvl, yLvl) ← getTypesFromKernel η
    let (Z, _, tLvl, _) ← getTypesFromKernel κ
    let SX ← computeSFinkerOf X xLvl
    let SY ← computeSFinkerOf Y yLvl
    let SZ ← computeSFinkerOf Z tLvl
    let comp_hom_proof ← mkAppMInst ``comp_hom #[SX, SY, SZ, ← idME X, ← idME Y, ← idME Z, η, κ] 2
    let (κ', proofs_κ) ← transformKernelToHom κ proofs
    let (η', proofs_η) ← transformKernelToHom η proofs_κ
    return (← mkAppM ``CategoryStruct.comp #[κ', η'], comp_hom_proof :: proofs_η)
  | Expr.const ``Kernel.parallelComp _ =>
    if ← checkWhiskerLeft e then
      let (X, Y, _, _) ← getTypesFromKernel e
      let (SZ, Z, SX, X, SY, Y, κ) ← constructWhiskersArgs e X Y false
      let (κ', proofs_κ) ← transformKernelToHom κ proofs
      let whisker_left_hom_proof ← mkAppMInst ``Kernel.whiskerLeft
          #[SX, SY, SZ, ← idME X, ← idME Y, ← idME Z, κ] 1
      let whiskerleft ← mkAppM ``MonoidalCategory.whiskerLeft #[SZ, κ']
      return (whiskerleft, whisker_left_hom_proof :: proofs_κ)
    else if ← checkWhiskerRight e then
      let (X, Y, _, _) ← getTypesFromKernel e
      let (SZ, Z, SX, X, SY, Y, κ) ← constructWhiskersArgs e X Y true
      let (κ', proofs_κ) ← transformKernelToHom κ proofs
      let whiskerright ← mkAppM ``MonoidalCategory.whiskerRight #[κ', SZ]
      let whiskerright_hom_proof ← mkAppMInst ``Kernel.whiskerRight
        #[SX, SY, SZ, ← idME X, ← idME Y, ← idME Z, κ] 1
      return (whiskerright, whiskerright_hom_proof :: proofs_κ)
    else
      let args := e.getAppArgs
      let κ := args[args.size - 2]!
      let η := args[args.size - 1]!
      let (X, Y, xLvl, yLvl) ← getTypesFromKernel κ
      let (Z, T, zLvl, tLvl) ← getTypesFromKernel η
      let SX ← computeSFinkerOf X xLvl
      let SY ← computeSFinkerOf Y yLvl
      let SZ ← computeSFinkerOf Z zLvl
      let ST ← computeSFinkerOf T tLvl
      let parallelComp_hom_proof ← mkAppMInst ``parallelComp_hom
        #[SX, SY, SZ, ST, ← idME X, ← idME Y, ← idME Z, ← idME T, κ, η] 2
      let (κ', proofs_κ) ← transformKernelToHom κ proofs
      let (η', proofs_η) ← transformKernelToHom η proofs_κ
      return (← mkAppM ``tensorHom #[κ', η'], parallelComp_hom_proof :: proofs_η)
  | Expr.const ``Kernel.id [xLvl] =>
    let X := e.getAppArgs[0]!
    let SX ← computeSFinkerOf X xLvl
    let id_hom_proof ← mkAppM ``id_hom #[SX, ← idME X]
    return (← mkAppM ``CategoryStruct.id #[SX], id_hom_proof :: proofs)
  | Expr.const ``Kernel.discard [xLvl, punitLvl] =>
    let X := e.getAppArgs[0]!
    let SX ← computeSFinkerOf X xLvl
    let discard_const := mkConst ``counit [xLvl, xLvl, punitLvl]
    let discard_hom_proof ← mkAppM' discard_const #[SX, ← idME X]
    return (← mkAppOptM ``ComonObj.counit #[none, none, none, SX, none],
      discard_hom_proof :: proofs)
  | Expr.const ``Kernel.copy [xLvl] =>
    let X := e.getAppArgs[0]!
    let SX ← computeSFinkerOf X xLvl
    let copy_hom_proof ← mkAppM ``comul #[SX, ← idME X]
    return (← mkAppOptM ``ComonObj.comul #[none, none, none, SX, none], copy_hom_proof :: proofs)
  | Expr.const ``Kernel.swap [xLvl, yLvl] =>
    let X := e.getAppArgs[0]!
    let Y := e.getAppArgs[1]!
    let SX ← computeSFinkerOf X xLvl
    let SY ← computeSFinkerOf Y yLvl
    let swap_hom_proof ← mkAppM ``braiding_hom #[SX, SY, ← idME X, ← idME Y]
    let braiding ← mkAppM ``Iso.hom #[← mkAppM ``BraidedCategory.braiding #[SX, SY]]
    return (braiding, swap_hom_proof :: proofs)
  | Expr.const ``Kernel.lift [_, y₀Lvl, _] =>
    let (X, Y, xLvl, yLvl) ← getTypesFromKernel e
    let args := e.getAppArgs
    let κ := args[args.size - 1]!
    if ← checkLeftUnitor κ then
      let punitLvl ← match args[0]!.getAppFn with
        | Expr.const ``Prod [punitLvl, _] => pure punitLvl
        | _ => throwError "Expected a product with PUnit as the first component, got {args[0]!}."
      let ey₀ := args[args.size - 2]!
      let (leftUnitorExpr, left_unitor_hom_proof) ← constructUnitors Y ey₀ yLvl y₀Lvl punitLvl 0
      return (leftUnitorExpr, left_unitor_hom_proof :: proofs)
    else if ← checkRightUnitor κ then
      let punitLvl ← match args[0]!.getAppFn with
        | Expr.const ``Prod [_, punitLvl] => pure punitLvl
        | _ => throwError "Expected a product with PUnit as the first component, got {args[0]!}."
      let ey₀ := args[args.size - 2]!
      let (rightUnitorExpr, right_unitor_hom_proof) ← constructUnitors Y ey₀ yLvl y₀Lvl punitLvl 1
      return (rightUnitorExpr, right_unitor_hom_proof :: proofs)
    else if ← checkAssociatorHom κ then
      let (ex₀, ey₀, ez₀) ← getMEFromThreeProds args[args.size - 2]!
      let (associatorExpr, associator_hom_proof) ← constructAssociatorHom X Y ex₀ ey₀ ez₀
      return (associatorExpr, associator_hom_proof :: proofs)
    else if ← checkAssociatorInv κ then
      let (ex₀, ey₀, ez₀) ← getMEFromThreeProds args[args.size - 3]!
      let (associatorInvExpr, associator_inv_hom_proof) ← constructAssociatorInv X Y ex₀ ey₀ ez₀
      return (associatorInvExpr, associator_inv_hom_proof :: proofs)
    else
      let SX ← computeSFinkerOf X xLvl
      let SY ← computeSFinkerOf Y yLvl
      let homExpr ← mkAppOptM ``ProbabilityTheory.Kernel.hom
        #[X, Y, none, none, SX, SY, (← idME X), (← idME Y), e, none]
      pure (homExpr, proofs)
  | _ =>
    throwError "Expected a lifted kernel expression, got: {e}."

end