import Erdos1135.KrasikovLagarias.EliminationPolicy
import Erdos1135.KrasikovLagarias.Retarded

/-!
# Annotated critical trees for KL elimination

This module records only the regular finite tree shape underlying the next
semantic elimination layer.  It does not assert that splitting preserves an
invariant or that any deletion policy terminates.
-/

namespace Erdos1135
namespace KrasikovLagarias
namespace EliminationCriticalTree

open EliminationPolicy
open EliminationResidue

mutual

/-- A finite expansion tree whose root and every descendant carry an exact
principal-index/shift label. -/
inductive Tree (k : Nat) where
  | terminal (label : Label k)
  | principalOnly (label : Label k) (principal : Tree k)
  | principalMin (label : Label k) (principal : Tree k)
      (auxiliary : NonemptyFamily k)

/-- A structurally nonempty finite family of auxiliary child trees. -/
inductive NonemptyFamily (k : Nat) where
  | one (tree : Tree k)
  | cons (tree : Tree k) (tail : NonemptyFamily k)

end

namespace NonemptyFamily

/-- Membership in a finite nonempty auxiliary family. -/
def Member {k : Nat} (child : Tree k) : NonemptyFamily k → Prop
  | .one tree => child = tree
  | .cons tree tail => child = tree ∨ Member child tail

end NonemptyFamily

namespace Tree

/-- The label attached to the root node. -/
def rootLabel {k : Nat} : Tree k → Label k
  | .terminal label => label
  | .principalOnly label _ => label
  | .principalMin label _ _ => label

/-- The principal residue index attached to the root node. -/
def rootIndex {k : Nat} (tree : Tree k) : PrincipalIndex k :=
  tree.rootLabel.index

/-- The exact shift code attached to the root node. -/
def rootShift {k : Nat} (tree : Tree k) : EliminationShift :=
  tree.rootLabel.shift

end Tree

mutual

/-- Evaluate an annotated tree at base time `y`. -/
noncomputable def Tree.eval {k : Nat} (tree : Tree k)
    (values : PrincipalIndex k → Real → Real) (y : Real) : Real :=
  match tree with
  | .terminal label => values label.index (y + label.value)
  | .principalOnly _ principal => principal.eval values y
  | .principalMin _ principal auxiliary =>
      principal.eval values y + auxiliary.evalMinimum values y

/-- Minimum evaluation of a structurally nonempty auxiliary family. -/
noncomputable def NonemptyFamily.evalMinimum {k : Nat}
    (family : NonemptyFamily k) (values : PrincipalIndex k → Real → Real)
    (y : Real) : Real :=
  match family with
  | .one tree => tree.eval values y
  | .cons tree tail => Min.min (tree.eval values y) (tail.evalMinimum values y)

end

mutual

/-- Evaluate the coefficient expression represented by an annotated tree. -/
noncomputable def Tree.coefficientValue {k : Nat} (tree : Tree k)
    (coefficients : PrincipalIndex k → Real) (lambda : Real) : Real :=
  match tree with
  | .terminal label => coefficients label.index * lambda ^ label.value
  | .principalOnly _ principal => principal.coefficientValue coefficients lambda
  | .principalMin _ principal auxiliary =>
      principal.coefficientValue coefficients lambda +
        auxiliary.coefficientMinimum coefficients lambda

/-- Minimum coefficient value of a structurally nonempty auxiliary family. -/
noncomputable def NonemptyFamily.coefficientMinimum {k : Nat}
    (family : NonemptyFamily k) (coefficients : PrincipalIndex k → Real)
    (lambda : Real) : Real :=
  match family with
  | .one tree => tree.coefficientValue coefficients lambda
  | .cons tree tail =>
      Min.min (tree.coefficientValue coefficients lambda)
        (tail.coefficientMinimum coefficients lambda)

end

mutual

/-- Forget annotations except at leaves, producing the existing retarded
expression.  A terminal keeps its exact label shift. -/
noncomputable def Tree.toExpr {k : Nat} (tree : Tree k) :
    Retarded.Expr (PrincipalIndex k) :=
  match tree with
  | .terminal label => .leaf label.index label.value
  | .principalOnly _ principal => principal.toExpr
  | .principalMin _ principal auxiliary =>
      .add principal.toExpr auxiliary.toExprMinimum

/-- Convert a nonempty auxiliary family to a nested binary minimum. -/
noncomputable def NonemptyFamily.toExprMinimum {k : Nat}
    (family : NonemptyFamily k) : Retarded.Expr (PrincipalIndex k) :=
  match family with
  | .one tree => tree.toExpr
  | .cons tree tail => .min tree.toExpr tail.toExprMinimum

end

namespace Tree

mutual

@[simp]
theorem toExpr_eval {k : Nat} (tree : Tree k)
    (values : PrincipalIndex k → Real → Real) (y : Real) :
    tree.toExpr.eval values y = tree.eval values y := by
  cases tree with
  | terminal label => rfl
  | principalOnly label principal => exact toExpr_eval principal values y
  | principalMin label principal auxiliary =>
      simp only [toExpr, Retarded.Expr.eval, eval, toExpr_eval,
        NonemptyFamily.toExprMinimum_eval]

@[simp]
theorem NonemptyFamily.toExprMinimum_eval {k : Nat}
    (family : NonemptyFamily k) (values : PrincipalIndex k → Real → Real)
    (y : Real) :
    family.toExprMinimum.eval values y = family.evalMinimum values y := by
  cases family with
  | one tree => exact toExpr_eval tree values y
  | cons tree tail =>
      simp only [NonemptyFamily.toExprMinimum, Retarded.Expr.eval,
        NonemptyFamily.evalMinimum, toExpr_eval,
        NonemptyFamily.toExprMinimum_eval]

end


mutual

@[simp]
theorem toExpr_coefficientValue {k : Nat} (tree : Tree k)
    (coefficients : PrincipalIndex k → Real) (lambda : Real) :
    tree.toExpr.coefficientValue coefficients lambda =
      tree.coefficientValue coefficients lambda := by
  cases tree with
  | terminal label => rfl
  | principalOnly label principal =>
      exact toExpr_coefficientValue principal coefficients lambda
  | principalMin label principal auxiliary =>
      simp only [toExpr, Retarded.Expr.coefficientValue, coefficientValue,
        toExpr_coefficientValue,
        NonemptyFamily.toExprMinimum_coefficientValue]

@[simp]
theorem NonemptyFamily.toExprMinimum_coefficientValue {k : Nat}
    (family : NonemptyFamily k) (coefficients : PrincipalIndex k → Real)
    (lambda : Real) :
    family.toExprMinimum.coefficientValue coefficients lambda =
      family.coefficientMinimum coefficients lambda := by
  cases family with
  | one tree => exact toExpr_coefficientValue tree coefficients lambda
  | cons tree tail =>
      simp only [NonemptyFamily.toExprMinimum, Retarded.Expr.coefficientValue,
        NonemptyFamily.coefficientMinimum, toExpr_coefficientValue,
        NonemptyFamily.toExprMinimum_coefficientValue]

end


/-- Fixed-time criticality matching the source proof shape.  Principal children
remain critical unconditionally; an auxiliary child is required to remain
critical only when it realizes the auxiliary minimum. -/
inductive CriticalInvariant {k : Nat}
    (values : PrincipalIndex k → Real → Real) (y : Real) : Tree k → Prop where
  | terminal (label : Label k) : CriticalInvariant values y (.terminal label)
  | principalOnly (label : Label k) (principal : Tree k)
      (rootBound : principal.eval values y ≤ values label.index (y + label.value))
      (principalInvariant : CriticalInvariant values y principal) :
      CriticalInvariant values y (.principalOnly label principal)
  | principalMin (label : Label k) (principal : Tree k)
      (auxiliary : NonemptyFamily k)
      (rootBound : principal.eval values y + auxiliary.evalMinimum values y ≤
        values label.index (y + label.value))
      (principalInvariant : CriticalInvariant values y principal)
      (auxiliaryInvariant : ∀ child,
        child.eval values y = auxiliary.evalMinimum values y →
          NonemptyFamily.Member child auxiliary →
            CriticalInvariant values y child) :
      CriticalInvariant values y (.principalMin label principal auxiliary)

/-- The root inequality carried by a critical tree. -/
theorem CriticalInvariant.root_bound {k : Nat}
    {values : PrincipalIndex k → Real → Real} {y : Real} {tree : Tree k}
    (invariant : CriticalInvariant values y tree) :
    tree.eval values y ≤
      values tree.rootIndex (y + tree.rootLabel.value) := by
  cases invariant with
  | terminal label => rfl
  | principalOnly label principal rootBound principalInvariant => exact rootBound
  | principalMin label principal auxiliary rootBound principalInvariant
      auxiliaryInvariant => exact rootBound

end Tree
end EliminationCriticalTree
end KrasikovLagarias
end Erdos1135