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

/-! # Standalone extraction for `elabKernelDiagramCmd`
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
namespace Mathlib.Tactic.Widget
end Mathlib.Tactic.Widget
namespace Mathlib.Tactic
end Mathlib.Tactic
namespace MeasureTheory
end MeasureTheory
namespace Mathlib.Tactic.Widget.StringDiagram
end Mathlib.Tactic.Widget.StringDiagram
namespace KernelDiagram
end KernelDiagram

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

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

/-- Unfold kernel operations in an expression. -/
def unfoldKernelOp (e : Expr) : MetaM Expr := do
  let names := (.empty |> NameSet.insert <| ``Kernel.prod) |> NameSet.insert <| ``Kernel.compProd
  transform e (post := fun e => do
    let e' ← deltaExpand e names.contains
    let e' ← Core.betaReduce e'
    return .done e')

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

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

/-- Get the original type and its universe from a `SFinKer.of` expression. -/
partial def getTypeFromSFinKer (e : Expr) : MetaM Expr := do
  match e.getAppFn with
  | Expr.const ``tensorUnit [eLvl, _] =>
    return mkConst ``PUnit [eLvl.succ]
  | Expr.const ``SFinKer.of _ =>
    let args := e.getAppArgs
    return args[0]!
  | Expr.const ``MonoidalCategory.tensorObj _ =>
    let args := e.getAppArgs
    let SY := args[args.size - 1]!
    let SX := args[args.size - 2]!
    let Y ← getTypeFromSFinKer SY
    let X ← getTypeFromSFinKer SX
    mkAppOptM ``Prod #[X, Y]
  | _ => throwError "Expected a SFinKer.of expression, got: {e}."

/-- 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}."

end

-- ═══ Tactic.KernelHom.Tactic.KernelDiagram ═══
section
public meta section
open Lean Meta Elab Command ProofWidgets Mathlib.Tactic.Widget
open Mathlib.Tactic BicategoryLike Penrose Server
open MeasureTheory ProbabilityTheory CategoryTheory
open CategoryTheory
open scoped MonoidalCategory ComonObj
namespace Mathlib.Tactic.Widget.StringDiagram

/-- The kernelized penrose variable associated with a node. -/
def Node.toPenroseVar_kernel (n : Node) : MetaM PenroseVar := do
  let expr ←
    try
      match n.e.getAppFn with
      | Expr.const ``SFinKer.of _ => do
        let res ← getTypeFromSFinKer n.e
        pure res
      | _ => do
        let (expr, _) ← transformHomToKernel n.e []
        pure expr
    catch _ =>
      pure n.e
  return ⟨"E", [n.vPos, n.hPosSrc, n.hPosTar], expr⟩

open scoped Jsx in
/-- Construct a kernelized string diagram from a Penrose `sub`stance program and
expressions `embeds` to display as labels in the diagram. -/
def mkKernelDiagram (nodes : List (List Node)) (strands : List (List Strand)) :
    DiagramBuilderM PUnit := do
  /- Add 2-morphisms. -/
  for x in nodes.flatten do
    match x with
    | .atom _ => do addPenroseVar "Atom" (← x.toPenroseVar_kernel)
    | .id _ => do StringDiagram.addPenroseVar "Id" (← x.toPenroseVar_kernel)
  /- Add constraints. -/
  for l in nodes do
    for (x₁, x₂) in l.consecutivePairs do
      DiagramBuilderM.addInstruction
        s!"Left({← x₁.toPenroseVar_kernel}, {← x₂.toPenroseVar_kernel})"
  /- Add constraints. -/
  for (l₁, l₂) in nodes.consecutivePairs do
    if let some x₁ := l₁.head? then
      if let some x₂ := l₂.head? then
        DiagramBuilderM.addInstruction
          s!"Above({← x₁.toPenroseVar_kernel}, {← x₂.toPenroseVar_kernel})"
  /- Add 1-morphisms as strings. -/
  for l in strands do
    for s in l do
      StringDiagram.addConstructor "Mor1" s.toPenroseVar
        "MakeString" [← s.startPoint.toPenroseVar_kernel, ← s.endPoint.toPenroseVar_kernel]

end Mathlib.Tactic.Widget.StringDiagram
namespace KernelDiagram

open scoped Jsx in
/-- Given a kernel expression, return a string diagram. Otherwise `none`. -/
def KernelM? (e : Expr) : MetaM (Option Html) := do
  let e ← instantiateMVars e
  try
    let (e, _) ← transformKernelToHom e []
    let k ← StringDiagram.mkKind e
    let x : Option (List (List StringDiagram.Node) × List (List StringDiagram.Strand))
    ← (match k with
      | .monoidal => do
        let some ctx ← BicategoryLike.mkContext? (ρ := Monoidal.Context) e | return none
        CoherenceM.run (ctx := ctx) do
          let e' := (← BicategoryLike.eval k.name (← MkMor₂.ofExpr e)).expr
          return some (← e'.nodes, ← e'.strands)
      | .bicategory => do
        let some ctx ← BicategoryLike.mkContext? (ρ := Bicategory.Context) e | return none
        CoherenceM.run (ctx := ctx) do
          let e' := (← BicategoryLike.eval k.name (← MkMor₂.ofExpr e)).expr
          return some (← e'.nodes, ← e'.strands)
      | .none => return none)
    match x with
    | none => return none
    | some (nodes, strands) => do
      DiagramBuilderM.run do
        StringDiagram.mkKernelDiagram nodes strands
        trace[string_diagram] "Penrose substance: \n{(← get).sub}"
        match ← DiagramBuilderM.buildDiagram StringDiagram.dsl StringDiagram.sty with
        | some html => return html
        | none => return <span>No non-structural morphisms found.</span>
  catch _ => return none

open scoped Jsx in
/-- Help function for displaying two string diagrams in an equality. -/
def mkEqHtml (lhs rhs : Html) : Html :=
  <div className="flex">
    <div className="w-50">
      <details «open»={true}>
        <summary className="mv2 pointer">Kernel diagram for LHS</summary> {lhs}
      </details>
    </div>
    <div className="w-50">
      <details «open»={true}>
        <summary className="mv2 pointer">Kernel diagram for RHS</summary> {rhs}
      </details>
    </div>
  </div>

/-- Given an equality between kernels, return a string diagram of the LHS and RHS.
Otherwise `none`. -/
def kernelEqM? (e : Expr) : MetaM (Option Html) := do
  try
    let e ← unfoldKernelOp <| ← instantiateMVars e
    let (lifted_e, _) ← liftEquality e
    let some (_, lhs, rhs) := lifted_e.eq? | return none
    let some lhs ← KernelM? lhs | return none
    let some rhs ← KernelM? rhs | return none
    return some <| mkEqHtml lhs rhs
  catch _ => return none

/-- Reduce a forall expression and try to display a kernel diagram
if it is an equality of kernels. -/
def kernelEqMReduce? (e : Expr) : MetaM (Option Html) := do
  forallTelescopeReducing (← whnfR <| ← inferType e) fun _ expr => do
    kernelEqM? expr

end KernelDiagram

/--
Display the kernel diagram for a given term.

Example usage:
```
/- Kernel diagram for an equality theorem. -/
#kernel_diagram ProbabilityTheory.Kernel.deterministic_comp_copy
```
-/
syntax (name := kernelDiagram) "#kernel_diagram " term : command

@[command_elab kernelDiagram, inherit_doc kernelDiagram]
def elabKernelDiagramCmd : CommandElab := fun
  | stx@`(#kernel_diagram $t:term) => do
    let html ← runTermElabM fun _ => do
      let e ← try mkConstWithFreshMVarLevels (← realizeGlobalConstNoOverloadWithInfo t)
        catch _ => Term.levelMVarToParam (← instantiateMVars (← Term.elabTerm t none))
      match ← KernelDiagram.kernelEqMReduce? e with
      | some html => return html
      | none => throwError "could not find an equality of kernels: {e}."
    liftCoreM <| Widget.savePanelWidgetInfo
      (hash HtmlDisplay.javascript)
      (return json% { html: $(← Server.RpcEncodable.rpcEncode html) })
      stx
  | stx => throwError "Unexpected syntax {stx}."

end