import Mathlib.MeasureTheory.Order.Lattice
import Mathlib.Probability.Kernel.Composition.Prod
import Mathlib.Probability.Kernel.Composition.CompProd
import Mathlib.MeasureTheory.MeasurableSpace.Embedding
import Lean.Elab.Tactic.Location
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

/-! # Standalone extraction for `unliftDiscard`
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.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}."

/-- Build a measurable equivalence for `e` into universe `maxLvl` (recursive on products). -/
partial def constructMeasurableEquiv (e : Expr) (eLevel maxLvl : Level) : MetaM Expr := do
  let ewhnf ← whnf e
  match ewhnf.getAppFn with
  | Expr.const ``PUnit _ | Expr.const ``Unit _ =>
    mkAppOptM' (Expr.const `MeasurableEquiv.punit [maxLvl, eLevel]) #[]
  | Expr.const ``Prod univs =>
    let args := ewhnf.getAppArgs
    let X := args[0]!
    let Y := args[1]!
    let xLevel := univs[0]!
    let yLevel := univs[1]!
    let ex ← constructMeasurableEquiv X xLevel maxLvl
    let ey ← constructMeasurableEquiv Y yLevel maxLvl
    let res ← mkAppOptM' (Expr.const ``MeasurableEquiv.prodCongr [maxLvl, xLevel, maxLvl, yLevel])
      #[none, none, none, none, none, none, none, none, ex, ey]
    return res
  | _ => mkAppOptM' (Expr.const ``MeasurableEquiv.ulift [eLevel, maxLvl]) #[e, none]

/-- Get the original type from a lifted type. -/
partial def getOriginalType (t : Expr) : MetaM (Expr × Level) := do
  let twhnf ← whnf t
  match twhnf.getAppFn with
  | Expr.const ``PUnit _ | Expr.const ``Unit _ =>
    return (mkConst ``Unit [], 0)
  | Expr.const ``ULift univs =>
    return (twhnf.getAppArgs[0]!, univs[1]!)
  | Expr.const ``Prod _ =>
    let args := twhnf.getAppArgs
    let (X, xLvl) ← getOriginalType args[0]!
    let (Y, yLvl) ← getOriginalType args[1]!
    return (← mkAppM ``Prod #[X, Y], .max xLvl yLvl)
  | _ =>
    return (t, ← getDecLevel (← inferType t))

end
end

-- ═══ Tactic.EqLift.Tactic.Kernel.KernelUnlift ═══
section
public meta section
open Lean Meta Parser.Tactic ProbabilityTheory ProbabilityTheory.Kernel

/-- Unlifts the discard kernel by unlifting the carrier type. -/
def unliftDiscard (e : Expr) (eLvl : Level) (proofs : List Expr) :
    MetaM (Expr × List Expr) := do
  unless e.isAppOf ``Kernel.discard do
    throwError "Expected the discard kernel, but got {e}."
  let (X', _, _, _) ← getTypesFromKernel e
  let (X, xLvl) ← getOriginalType X'
  if X == X' then
    return (e, proofs)
  else
    let ex ← constructMeasurableEquiv X xLvl eLvl
    let discard_const := mkConst ``discard_lift [xLvl, eLvl, Level.zero]
    let discard_unlift_proof ← mkAppM' discard_const #[ex]
    let discard_const := mkConst ``Kernel.discard [xLvl, 0]
    return (← mkAppOptM' discard_const #[X, none], discard_unlift_proof :: proofs)

end
end