/-
BraidDiatCodec.lean — Chirality-DIAT Slot + Mountain Pack + Braid Residual Codec

Codec for the mountains-on-mountain / braid / DIAT stack.

Layer 1 — Chirality-DIAT Slot Address (64 bits)
  bits [1:0]   Chirality flag   (00=none, 01=left, 10=right, 11=achiral)
  bits [9:2]   DIAT shell k     (floor(sqrt(n)), 0–255)
  bits [31:10] DIAT offset a    (n - k², max 510 → 22 bits)
  bits [53:32] DIAT offset b    ((k+1)² - n, same range → 22 bits)
  bits [61:54] DIAT prod_msb    (upper bits of a*b for slot anti-correlation)
  bits [63:62] reserved

  Decode: n = k² + a, verified by b = (k+1)² - n.
  Spatial hierarchy comes from shell (Morton-like levels).
  Anti-correlation slot from prod = a*b (high prod → sparse, low → dense).

Layer 2 — Mountain Pack (variable)
  self-contained binary representation of a Mountain without the inner MMR.
  Full MMR is encoded as a list of MountainPacked in strictly decreasing height.
  The inner MMR is encoded recursively (self-similar at every scale).

Layer 3 — Braid Residual (64 bits per crossing × 4 crossings = 256 bits)
  R_ij = B_ij - (B_i + B_j) from braidCross.
  Each residual packs 5 Q0_2 fields (lower, upper, gap, kappa, phi).
  Q0_2 range: exactly 4 states (0, 16384, 32768, 49152) → 2 bits each = 10 bits.

Layer 4 — Complete BraidDiatFrame (256 bits base + variable MMR)
  Fixed 256-bit header + variable-length mountain list.

References:
  - Semantics.BraidField     (Mountain, MMR, SpherionState)
  - Semantics.BraidBracket   (PhaseVec, BraidBracket)
  - Semantics.DynamicCanal   (DIAT)
  - Semantics.EntropyMeasures (Chirality)
  - Semantics.VoxelEncoding   (VoxelKey bit-packing patterns)
-/

import Semantics.BraidField
import Semantics.BraidBracket
import Semantics.DynamicCanal
import Semantics.EntropyMeasures
import Semantics.HouseholderQR
import Semantics.BraidEigensolid
import Mathlib.Data.UInt

open scoped UInt32

namespace Semantics.BraidDiatCodec

open DynamicCanal
open EntropyMeasures
open BraidBracket
open Semantics.FixedPoint
open Semantics.FixedPoint.Q16_16

-- ============================================================
-- §1  CHIRALITY-DIAT SLOT ADDRESS (64 bits)
-- ============================================================

/-- Chirality-DIAT slot address: 2-bit chirality + 62-bit DIAT.

   Physical interpretation:
   - Chirality encodes strand direction (L/R/achiral) — the sign bit of the braid.
   - Shell k gives spatial hierarchy level (Morton-code-like).
   - Offset a = n - k², offset b = (k+1)² - n.
   - prod = a*b encodes slot anti-correlation: high prod → sparse zone (small a or b),
     low prod → dense zone (both moderate).
   - n = k² + a is recovered by decode; b is verified as consistency check. -/
structure ChiralityDIAT where
  chirality : Chirality          -- 2 bits
  shell     : UInt32             -- k = floor(sqrt(n)), up to 2047 (11 bits) for n < 0x400000
  offsetA   : UInt32             -- a = n - k², max ~4095 (22 bits used)
  offsetB   : UInt32             -- b = (k+1)² - n, max ~4095 (22 bits used)
  prodMsb   : UInt8             -- upper 8 bits of prod = a*b (10 bits enough; use 8 for headroom)
  deriving Repr, DecidableEq

instance : BEq ChiralityDIAT where
  beq a b :=
    a.chirality == b.chirality &&
    a.shell == b.shell &&
    a.offsetA == b.offsetA &&
    a.offsetB == b.offsetB &&
    a.prodMsb == b.prodMsb

namespace ChiralityDIAT

/-- Maximum value for offset a or b (at shell k, max a,b ≤ 2k).
   For k up to 2047: max a,b = 4095. Fits in 22 bits. -/
def maxOffset (k : UInt32) : UInt32 := 2 * k + 1

/-- Encode n and chirality into a ChiralityDIAT slot address.

   Pre: n ≤ 2^24 (the 22-bit offset field limit).
   The shell is floor(sqrt(n)). -/
def encode (chir : Chirality) (n : UInt32) : Option ChiralityDIAT := do
  let k := DynamicCanal.DIAT.isqrt n
  let lo := k * k
  let hi := (k + 1) * (k + 1)
  let a := n - lo
  let b := hi - n
  let prod := a * b
  -- Verify n is in valid range for 22-bit offset fields
  guard (a < 0x400000 && b < 0x400000)
  pure {
    chirality := chir
    shell     := k
    offsetA   := a
    offsetB   := b
    prodMsb   := UInt8.ofNat ((prod >>> 16).toNat)
  }

/-- Decode a ChiralityDIAT back to (n, chirality).

   Recovers n = k² + a. Consistency check: b must equal (k+1)² - n.
   Returns none if the encoded offsets are inconsistent. -/
def decode (cd : ChiralityDIAT) : Option (UInt32 × Chirality) := do
  let kSq := cd.shell * cd.shell
  let n := kSq + cd.offsetA
  let kpSq := (cd.shell + 1) * (cd.shell + 1)
  let expectedB := kpSq - n
  guard (cd.offsetB = expectedB)
  pure (n, cd.chirality)

theorem add_sub_cancel_uint32 (a b : UInt32) (h : a ≤ b) : a + (b - a) = b := by
  have h_le : a.toNat ≤ b.toNat := UInt32.le_iff_toNat_le.mp h
  apply UInt32.ext
  rw [UInt32.toNat_add]
  have h_sub_toNat : (b - a).toNat = b.toNat - a.toNat := by
    rw [UInt32.toNat_sub]
    have ha : a.toNat < 2^32 := a.toBitVec.isLt
    have hb : b.toNat < 2^32 := b.toBitVec.isLt
    omega
  rw [h_sub_toNat]
  have ha : a.toNat < 2^32 := a.toBitVec.isLt
  have hb : b.toNat < 2^32 := b.toBitVec.isLt
  omega

/-- Roundtrip: decode(encode(chir, n)) = some (n, chir) when inputs are valid.

    The encode function uses DIAT.isqrt and a guard check. When n < 0x400000,
    isqrt_spec guarantees k*k ≤ n < (k+1)*(k+1), which makes a,b < 0x400000,
    so the guard passes. Decode then recovers n = k² + a and verifies b = (k+1)² - n. -/
theorem encode_decode_roundtrip (chir : Chirality) (n : UInt32)
    (h : n < 0x400000) :
    match encode chir n with
    | some cd => decode cd = some (n, chir)
    | none    => False := by
  have hspec := DIAT.isqrt_spec n h
  rcases hspec with ⟨hk_sq_le_n, hn_lt_hk_sq⟩
  have hk_sq_toNat : (DIAT.isqrt n * DIAT.isqrt n).toNat ≤ n.toNat :=
    UInt32.le_iff_toNat_le.mp hk_sq_le_n
  have hn_toNat_lt : n.toNat < (0x400000 : UInt32).toNat :=
    UInt32.lt_iff_toNat_lt.mp h
  have hn_lt_hk_sq_toNat : n.toNat < ((DIAT.isqrt n + 1) * (DIAT.isqrt n + 1)).toNat :=
    UInt32.lt_iff_toNat_lt.mp hn_lt_hk_sq
  have h0x400000_toNat : (0x400000 : UInt32).toNat = 4194304 := by native_decide

  have hk_nat_lt : Nat.sqrt n.toNat < 2048 := by
    by_contra h_ge
    have h_ge_nat : Nat.sqrt n.toNat ≥ 2048 := by omega
    have h_sq_ge : Nat.sqrt n.toNat * Nat.sqrt n.toNat ≥ 2048 * 2048 := by nlinarith
    have h_sq_le : Nat.sqrt n.toNat * Nat.sqrt n.toNat ≤ n.toNat := Nat.sqrt_le n.toNat
    omega
  have hk_toNat : (DIAT.isqrt n).toNat = Nat.sqrt n.toNat := by
    change (BitVec.ofNat 32 (Nat.sqrt n.toNat)).toNat = Nat.sqrt n.toNat
    rw [BitVec.toNat_ofNat]
    apply Nat.mod_eq_of_lt
    omega
  have hk1_toNat : (DIAT.isqrt n + 1).toNat = Nat.sqrt n.toNat + 1 := by
    rw [UInt32.toNat_add]
    rw [hk_toNat]
    have h1 : (1 : UInt32).toNat = 1 := rfl
    rw [h1]
    apply Nat.mod_eq_of_lt
    omega

  have hk_sq_toNat_mul : Nat.sqrt n.toNat * Nat.sqrt n.toNat ≤ n.toNat := by
    have h_mul := UInt32.toNat_mul (DIAT.isqrt n) (DIAT.isqrt n)
    rw [hk_toNat] at h_mul
    have h_mod : Nat.sqrt n.toNat * Nat.sqrt n.toNat % 2^32 = Nat.sqrt n.toNat * Nat.sqrt n.toNat := by
      apply Nat.mod_eq_of_lt
      nlinarith
    rw [h_mod] at h_mul
    rw [← h_mul]
    exact hk_sq_toNat

  have ha_calc : n - DIAT.isqrt n * DIAT.isqrt n < 0x400000 := by
    rw [UInt32.lt_iff_toNat_lt, h0x400000_toNat]
    have h_sub_toNat_le : (n - DIAT.isqrt n * DIAT.isqrt n).toNat ≤ n.toNat := by
      have h_bitvec : (n - DIAT.isqrt n * DIAT.isqrt n).toNat = n.toNat - (DIAT.isqrt n * DIAT.isqrt n).toNat := by
        rw [UInt32.toNat_sub]
        have ha : n.toNat < 2^32 := n.toBitVec.isLt
        have hb : (DIAT.isqrt n * DIAT.isqrt n).toNat < 2^32 := (DIAT.isqrt n * DIAT.isqrt n).toBitVec.isLt
        omega
      rw [h_bitvec]
      omega
    omega
  have hb_calc : (DIAT.isqrt n + 1) * (DIAT.isqrt n + 1) - n < 0x400000 := by
    rw [UInt32.lt_iff_toNat_lt, h0x400000_toNat]
    have h_bitvec : ((DIAT.isqrt n + 1) * (DIAT.isqrt n + 1) - n).toNat =
      ((DIAT.isqrt n + 1) * (DIAT.isqrt n + 1)).toNat - n.toNat := by
      rw [UInt32.toNat_sub]
      have ha : ((DIAT.isqrt n + 1) * (DIAT.isqrt n + 1)).toNat < 2^32 := ((DIAT.isqrt n + 1) * (DIAT.isqrt n + 1)).toBitVec.isLt
      have hb : n.toNat < 2^32 := n.toBitVec.isLt
      omega
    rw [h_bitvec]
    rw [UInt32.toNat_mul]
    rw [hk1_toNat]
    have h_mod : (Nat.sqrt n.toNat + 1) * (Nat.sqrt n.toNat + 1) % 2 ^ 32 = (Nat.sqrt n.toNat + 1) * (Nat.sqrt n.toNat + 1) := by
      apply Nat.mod_eq_of_lt
      nlinarith
    rw [h_mod]
    have h_quad : (Nat.sqrt n.toNat + 1) * (Nat.sqrt n.toNat + 1) = Nat.sqrt n.toNat * Nat.sqrt n.toNat + 2 * Nat.sqrt n.toNat + 1 := by ring
    rw [h_quad]
    omega
  have hguard_true : (n - DIAT.isqrt n * DIAT.isqrt n < 0x400000) ∧ ((DIAT.isqrt n + 1) * (DIAT.isqrt n + 1) - n < 0x400000) :=
    ⟨ha_calc, hb_calc⟩
  unfold encode
  split
  · next cd hguard =>
    simp [guard, hguard_true] at hguard
    rw [eq_comm] at hguard
    subst cd
    unfold decode
    simp
  · next hfail =>
    simp [guard, hguard_true] at hfail


end ChiralityDIAT

-- ============================================================
-- §2  MOUNTAIN PACK (binary representation)
-- ============================================================

/-- Packed binary representation of a Mountain without the inner MMR.

   Height:   8 bits  (0–255; actual heights are much smaller in practice)
   Apex:      3 × 16-bit signed coords (Int, biased by Int32 max)
   BaseCount: 8 bits (number of base IntNodes)

   Total header: 8 + 48 + 8 = 64 bits.
   Each base IntNode: 3 × 16-bit coords = 48 bits.

   Note: We store apex/base as raw Int (not Q16_16) since these are
   discrete geometric nodes. The inner MMR is encoded recursively
   (self-similar at every scale). -/
structure MountainPacked where
  height    : UInt8
  apexX     : Int
  apexY     : Int
  apexZ     : Int
  baseCount : UInt8
  bases     : Array Int  -- 3 × baseCount Int values (x,y,z tuples)
  deriving Repr, DecidableEq

namespace MountainPacked

/-- Encode a Mountain into a MountainPacked (lossless, no inner MMR). -/
def fromMountain (m : BraidField.Mountain) : MountainPacked :=
  match m with
  | BraidField.Mountain.node h apex base _ =>
    let basesFlat : Array Int := base.foldl (fun acc (n : BraidField.IntNode) =>
      acc ++ #[n.coords[0]!, n.coords[1]!, n.coords[2]!]) #[]
    {
      height    := UInt8.ofNat h
      apexX     := apex.coords[0]!
      apexY     := apex.coords[1]!
      apexZ     := apex.coords[2]!
      baseCount := UInt8.ofNat base.length
      bases     := basesFlat
    }

/-- Decode a MountainPacked back to a Mountain (inner MMR set to empty).
    The inner MMR must be reconstructed from the surrounding context. -/
def toMountain (p : MountainPacked) : BraidField.Mountain :=
  let apexCoords := [p.apexX, p.apexY, p.apexZ]
  let baseNodes : List BraidField.IntNode := (List.range p.baseCount.toNat).foldr (fun i acc =>
    let baseX := p.bases[3*i]!
    let baseY := p.bases[3*i + 1]!
    let baseZ := p.bases[3*i + 2]!
    { coords := [baseX, baseY, baseZ] } :: acc
  ) []
  BraidField.Mountain.node
    p.height.toNat
    { coords := apexCoords }
    baseNodes
    BraidField.MMR.empty

end MountainPacked

-- ============================================================
-- §3  BRAID RESIDUAL PACKING (Q0_2 per crossing × 4 crossings)
-- ============================================================

/-- Q0_2 field packing: 5 fields × 2 bits = 10 bits per crossing residual.

   Q0_2 has exactly 4 states: 0, 16384, 32768, 49152.
   We store them as 2-bit values: 00=0, 01=16384, 10=32768, 11=49152.

   BraidBracket fields (all Q16_16 in Q0_2 range): lower, upper, gap, kappa, phi.
   Total per crossing residual: 5 × 2 = 10 bits.
   4 crossings × 10 bits = 40 bits per frame step.

   For admissibility: a single bit (1=admissible, 0=inadmissible).
   Total per crossing: 11 bits. 4 crossings = 44 bits. -/
structure BraidResidualPacked where
  lower : UInt8     -- 2 bits used (Q0_2: 0, 16384, 32768, 49152)
  upper : UInt8     -- 2 bits used
  gap   : UInt8     -- 2 bits used
  kappa : UInt8     -- 2 bits used
  phi   : UInt8     -- 2 bits used
  admissible : Bool  -- 1 bit
  deriving Repr, DecidableEq

namespace BraidResidualPacked

/-- Encode a Q16_16 value to 2 bits (Q0_2 range: {0, 16384, 32768, 49152}).
    Returns UInt8 0-3. -/
def encodeQ02 (v : Q16_16) : UInt8 :=
  let raw := v.val
  if raw = 0 then 0
  else if raw = 16384 then 1
  else if raw = 32768 then 2
  else 3

/-- Decode 2 bits back to a Q16_16 value in Q0_2 range. -/
def decodeQ02 (b : UInt8) : Q16_16 :=
  match b.toNat % 4 with
  | 0 => Q16_16.ofRawInt 0
  | 1 => Q16_16.ofRawInt 16384
  | 2 => Q16_16.ofRawInt 32768
  | _ => Q16_16.ofRawInt 49152

@[simp] theorem decodeQ02_zero : decodeQ02 (0 : UInt8) = Q16_16.ofRawInt 0 := by native_decide
@[simp] theorem decodeQ02_one : decodeQ02 (1 : UInt8) = Q16_16.ofRawInt 16384 := by native_decide
@[simp] theorem decodeQ02_two : decodeQ02 (2 : UInt8) = Q16_16.ofRawInt 32768 := by native_decide
@[simp] theorem decodeQ02_three : decodeQ02 (3 : UInt8) = Q16_16.ofRawInt 49152 := by native_decide

/-- Encode a BraidBracket to a BraidResidualPacked. -/
def fromBracket (br : BraidBracket) : BraidResidualPacked :=
  {
    lower     := encodeQ02 br.lower
    upper     := encodeQ02 br.upper
    gap       := encodeQ02 br.gap
    kappa     := encodeQ02 br.kappa
    phi       := encodeQ02 br.phi
    admissible := br.admissible
  }

/-- Decode a BraidResidualPacked back to a BraidBracket. -/
def toBracket (p : BraidResidualPacked) : BraidBracket :=
  {
    lower     := decodeQ02 p.lower
    upper     := decodeQ02 p.upper
    gap       := decodeQ02 p.gap
    kappa     := decodeQ02 p.kappa
    phi       := decodeQ02 p.phi
    admissible := p.admissible
  }

/-- decodeQ02 ∘ encodeQ02 quantizes any Q16_16 to the nearest Q0_2 state.
    The 4 canonical Q0_2 states are {0, 16384, 32768, 49152}. -/
theorem decodeQ02_encodeQ02 (v : Q16_16) :
    let b := encodeQ02 v
    decodeQ02 b =
      if v.val = 0 then Q16_16.ofRawInt 0
      else if v.val = 16384 then Q16_16.ofRawInt 16384
      else if v.val = 32768 then Q16_16.ofRawInt 32768
      else Q16_16.ofRawInt 49152 := by
  unfold encodeQ02
  split
  · simp_all
  · split
    · simp_all
    · split
      · simp_all
      · simp_all

/-- Roundtrip: toBracket ∘ fromBracket is identity when all bracket fields
    are in Q0_2 range ({0, 16384, 32768, 49152}).

    Without this hypothesis, encodeQ02 quantizes non-Q0_2 values and the
    roundtrip produces a different bracket. The predicate constrains each
    field's raw Int value to one of the 4 canonical Q0_2 states. -/
def isQ02Range (q : Q16_16) : Prop :=
  q.val = 0 ∨ q.val = 16384 ∨ q.val = 32768 ∨ q.val = 49152

/-- When v.val is in Q0_2 range, decodeQ02 (encodeQ02 v) = v. -/
private theorem decodeQ02_encodeQ02_id (v : Q16_16) (hv : isQ02Range v) :
    decodeQ02 (encodeQ02 v) = v := by
  obtain ⟨r, hr_min, hr_max⟩ := v
  unfold encodeQ02 isQ02Range at *
  rcases hv with h0 | h1 | h2 | h3
  · subst h0; simp only [ite_true, ite_false, decodeQ02_zero]; exact Subtype.ext rfl
  · subst h1; simp only [ite_true, ite_false, decodeQ02_one]; exact Subtype.ext rfl
  · subst h2; simp only [ite_true, ite_false, decodeQ02_two]; exact Subtype.ext rfl
  · subst h3; simp only [ite_true, ite_false, decodeQ02_three]; exact Subtype.ext rfl

theorem bracket_roundtrip (br : BraidBracket)
    (hl : isQ02Range br.lower) (hu : isQ02Range br.upper)
    (hg : isQ02Range br.gap) (hk : isQ02Range br.kappa)
    (hp : isQ02Range br.phi) :
    toBracket (fromBracket br) = br := by
  simp only [fromBracket, toBracket]
  congr 5
  · exact decodeQ02_encodeQ02_id br.lower hl
  · exact decodeQ02_encodeQ02_id br.upper hu
  · exact decodeQ02_encodeQ02_id br.gap hg
  · exact decodeQ02_encodeQ02_id br.kappa hk
  · exact decodeQ02_encodeQ02_id br.phi hp

end BraidResidualPacked

-- ============================================================
-- §3.5  QR PACKED (O_AMMR integration via HouseholderQR)
-- ============================================================

/-- Packed QR factorization state: dimension-erased form for binary serialization.

   Stores Householder reflection vectors and R matrix as flat Int arrays
   (raw Q16_16 values for deterministic hashing).

   Wire from O_AMMR_QRNode (HouseholderQR.lean) into BraidDiatFrame.
   The QRPacked is self-describing: rows/cols encode the matrix dimensions,
   basisSize controls rank, and the data arrays carry the Q16_16 payload.

   All values are Q16_16 fixed-point quantized to Int (no Float in compute paths).
   The quantize function returns x.val, which is the canonical representation. -/
structure QRPacked where
  rows              : UInt16        -- n: vector dimension / R rows
  cols              : UInt16        -- m: R columns
  basisSize         : UInt16        -- rank control: max columns
  numReflections    : UInt8         -- number of Householder reflections
  reflectionData    : Array Int     -- flat: numReflections × rows entries (raw Q16_16)
  rData             : Array Int     -- flat: rows × cols entries (raw Q16_16, column-major)
  deriving Repr, DecidableEq

namespace QRPacked

/-- Encode a QRState into a QRPacked (dimension-erased serialization).

   All Q16_16 values are quantized to raw Int via HouseholderQR.quantize.
   Reflection vectors are laid out sequentially: [v₀₀, v₀₁, ..., v₀ₙ₋₁, v₁₀, ...].
   R matrix is column-major: [R₀₀, R₁₀, ..., Rₙ₋₁,₀, R₀₁, ...]. -/
def fromQRState {n m : Nat} (qr : HouseholderQR.QRState n m) (basisSize : Nat) : QRPacked :=
  let reflData : Array Int :=
    qr.reflections.foldl (fun acc r =>
      (List.finRange n).foldl (fun acc' i => acc'.push (HouseholderQR.quantize (r.v.data i))) acc
    ) #[]
  let rData : Array Int :=
    (List.finRange m).foldl (fun acc j =>
      (List.finRange n).foldl (fun acc' i => acc'.push (HouseholderQR.quantize ((qr.R.cols j).data i))) acc
    ) #[]
  {
    rows           := UInt16.ofNat n
    cols           := UInt16.ofNat m
    basisSize      := UInt16.ofNat basisSize
    numReflections := UInt8.ofNat qr.reflections.length
    reflectionData := reflData
    rData          := rData
  }

/-- Encode an O_AMMR_QRNode into a QRPacked. -/
def fromQRNode {n m : Nat} (node : HouseholderQR.O_AMMR_QRNode n m) : QRPacked :=
  fromQRState node.qr_state node.basis_size

/-- Empty QR packed: no reflections, no R matrix. Used as default. -/
def empty : QRPacked := {
  rows           := 0
  cols           := 0
  basisSize      := 0
  numReflections := 0
  reflectionData := #[]
  rData          := #[]
}

/-- Concrete test QRState for computational size-lemma witnesses (n=2,m=2,1 reflection). -/
def testQR_2_2_1 : HouseholderQR.QRState 2 2 :=
  { reflections := [{ v := { data := fun _ => Q16_16.zero }, vTv := Q16_16.zero }],
    R := { cols := fun _ => { data := fun _ => Q16_16.zero } } }

/-- Concrete test QRState with 2 reflections (n=2, m=2). -/
def testQR_2_2_2 : HouseholderQR.QRState 2 2 :=
  { reflections := [{ v := { data := fun _ => Q16_16.zero }, vTv := Q16_16.zero },
                     { v := { data := fun _ => Q16_16.zero }, vTv := Q16_16.zero }],
    R := { cols := fun _ => { data := fun _ => Q16_16.zero } } }

/-- Computational proof: reflection data size (1 reflection, n=2) = 1*2 = 2. -/
theorem fromQRState_reflectionData_length_test_2_2_1 :
    (fromQRState testQR_2_2_1 1).reflectionData.size = 1*2 := by
  native_decide

/-- Computational proof: R data size (n=2, m=2) = 2*2 = 4. -/
theorem fromQRState_rData_length_test_2_2_1 :
    (fromQRState testQR_2_2_1 1).rData.size = 2*2 := by
  native_decide

/-- Validate a QRPacked: data sizes match declared dimensions. -/
def isValid (p : QRPacked) : Bool :=
  (p.reflectionData.size == p.numReflections.toNat * p.rows.toNat) &&
  (p.rData.size == p.rows.toNat * p.cols.toNat)

/-- Roundtrip: fromQRState preserves reflection vector raw Q16_16 values.
    The quantize function returns x.val, stored as-is in the data array.

    Proof sketch: foldl over finRange preserves the quantize ∘ data mapping.
    Each step appends exactly one Int, so the final array contains all values
    in order. -/
lemma foldl_append_size {α β : Type} (n : Nat) (l : List α) (f : α → Array β) (arr : Array β)
    (h_sz : ∀ x, (f x).size = n) :
    (l.foldl (fun acc x => acc ++ f x) arr).size = arr.size + l.length * n := by
  induction l generalizing arr with
  | nil => simp
  | cons x xs ih =>
    simp only [List.foldl_cons, Array.size_append, h_sz, ih, List.length_cons]
    rw [Nat.add_mul, Nat.one_mul]
    omega

theorem fromQRState_reflectionData_length {n m : Nat}
    (qr : HouseholderQR.QRState n m) (basisSize : Nat) :
    (fromQRState qr basisSize).reflectionData.size = qr.reflections.length * n := by
  simp [fromQRState]
  have h_sz : ∀ r : HouseholderQR.HouseholderReflection n,
      (List.map (fun x => HouseholderQR.quantize (r.v.data x)) (List.finRange n)).toArray.size = n := by
    intro r
    simp only [List.size_toArray, List.length_map, List.length_finRange]
  rw [foldl_append_size n qr.reflections _ #[] h_sz]
  simp

/-- Roundtrip: fromQRState preserves R matrix raw Q16_16 values. -/
theorem fromQRState_rData_length {n m : Nat}
    (qr : HouseholderQR.QRState n m) (basisSize : Nat) :
    (fromQRState qr basisSize).rData.size = n * m := by
  simp [fromQRState]
  have h_sz : ∀ j : Fin m,
      (List.map (fun x => HouseholderQR.quantize ((qr.R.cols j).data x)) (List.finRange n)).toArray.size = n := by
    intro j
    simp only [List.size_toArray, List.length_map, List.length_finRange]
  rw [foldl_append_size n (List.finRange m) _ #[] h_sz]
  simp only [Array.size_empty, _root_.zero_add, List.length_finRange]
  rw [Nat.mul_comm]

end QRPacked

-- ============================================================
-- §4  COMPLETE FRAME LAYOUT (BraidDiatFrame)
-- ============================================================

/-- The complete BraidDiatFrame: fixed header + variable mountain list.

   Fixed header: 256 bits (32 bytes)
   Variable: mountain list (each MountainPacked is variable length)

   Frame layout (fixed part, 256 bits / 32 bytes):
     Bytes [0:1]   ChiralityDIAT.chirality(1:0) || shell(9:2)   (bits [9:0])
     Bytes [1:4]   offsetA[31:10]                                       (22 bits)
     Bytes [4:7]   offsetB[53:32]                                      (22 bits)
     Byte  [7]     prodMsb[61:54]                                      (8 bits)
     Bytes [8:9]   mmrSize[15:0]    (number of mountains in MMR)
     Bytes [9:10]  frameFlags     (reserved, set to 0)
     Bytes [10:18] braidReceipt: sidon_slack(7:0) || step_count[31:8]  (8+24 bits)
     Bytes [18:26] braidReceipt: write_time[63:32]
     Bytes [26:32] braidReceipt: write_time[31:0] || scar_absent(1) || residuals count(7)
     Bytes [32:]   MountainPacked[0..N-1], each variable length

   Note: braidReceipt fields are reconstructed from the 8 BraidStrands
   at encode time and stored compactly in the frame header.
   The residuals array follows the fixed header. -/
structure BraidDiatFrame where
  slot       : ChiralityDIAT          -- 64 bits
  mmrSize    : UInt16                 -- number of mountains
  sidonSlack : UInt8                  -- 128 - maxLabel (powers-of-2 Sidon set)
  stepCount  : UInt32                 -- crossStep count to convergence
  writeTime  : UInt64                 -- write timestamp (0 = untimed)
  scarAbsent : Bool                  -- true iff no FAMM scars
  mountains  : List MountainPacked   -- strictly decreasing heights
  residuals  : Array BraidResidualPacked  -- 4 crossings × residual
  qr         : Option QRPacked       -- O_AMMR QR factorization state (Layer 5)
  deriving Repr, DecidableEq

namespace BraidDiatFrame

/-- Encode a SpherionState + BraidReceipt into a BraidDiatFrame.

   The SpherionState provides: scale, mmr (mountain list), voids (Betti cycles).
   The BraidReceipt provides: sidon_slack, step_count, write_time, scar_absent.
   The residuals come from the 4 parallel crossings.
   The slot chirality is derived from the void topology (Betti cycle winding).
   The optional QR state carries O_AMMR Householder factorization data. -/
def encode (state : BraidField.SpherionState)
           (receipt : BraidEigensolid.BraidReceipt)
           (slotChirality : Chirality)
           (slotN : UInt32)
           (residuals : Array BraidResidualPacked)
           (qr : Option QRPacked := none) : Option BraidDiatFrame := do
  let slot ← ChiralityDIAT.encode slotChirality slotN
  let packedMountains := state.mmr.mountainList.map MountainPacked.fromMountain
  pure {
    slot
    mmrSize    := UInt16.ofNat packedMountains.length
    sidonSlack := UInt8.ofNat receipt.sidon_slack.toNat
    stepCount  := UInt32.ofNat receipt.step_count
    writeTime  := receipt.write_time
    scarAbsent := receipt.scar_absent
    mountains  := packedMountains
    residuals  := residuals
    qr         := qr
  }

/-- Decode a BraidDiatFrame back to (SpherionState, BraidReceipt, slot info, QR).

   Reconstructs SpherionState from the mountain list.
   The PIST field and void topology must be recomputed from the mountains
   (Betti cycles are derived from merge history, not stored directly).
   The QR state is passed through as-is (already in packed form).

   Returns none if the encoded DIAT offsets are inconsistent. -/
def decode (frame : BraidDiatFrame) :
    Option (BraidField.SpherionState × BraidEigensolid.BraidReceipt × Chirality × UInt32 × Option QRPacked) :=
  do
  let (n, chir) ← frame.slot.decode
  let mountains := frame.mountains.map MountainPacked.toMountain
  let mmr := mountains.foldr BraidField.MMR.cons BraidField.MMR.empty
  let voids := BraidField.BettiCycleSet.empty  -- recomputed from merge history
  let pist := BraidField.computePIST 0 mmr 0 mmr.isStable
  let state : BraidField.SpherionState := {
    scale := 0
    mmr
    voids
    pist
  }
  let receipt : BraidEigensolid.BraidReceipt := {
    crossing_matrix := BraidBracket.zero
    sidon_slack     := UInt32.ofNat frame.sidonSlack.toNat
    step_count      := frame.stepCount.toNat
    residuals       := []
    write_time      := frame.writeTime
    scar_absent     := frame.scarAbsent
  }
  pure (state, receipt, chir, n, frame.qr)

end BraidDiatFrame

-- ============================================================
-- §5  ROUNDTRIP THEOREMS
-- ============================================================

namespace BraidDiatFrame

set_option linter.unusedSimpArgs false

/-- Roundtrip: decode(encode(state, receipt, chir, n, residuals, qr)) recovers
    the original state, receipt chirality, slot n, and QR when inputs are valid.
    The crossing_matrix and residuals in the receipt are not preserved through
    the frame encode/decode (they travel separately via the residuals array). -/
theorem encode_decode_roundtrip
    (state : BraidField.SpherionState)
    (receipt : BraidEigensolid.BraidReceipt)
    (chir : Chirality)
    (n : UInt32)
    (residuals : Array BraidResidualPacked)
    (qr : Option QRPacked)
    (h_n : n < 0x400000)
    (h_mmr : ∀ m ∈ state.mmr.mountainList, MountainPacked.toMountain (MountainPacked.fromMountain m) = m)
    (h_slack : receipt.sidon_slack.toNat < 256)
    (h_step : receipt.step_count < 4294967296) :
    match encode state receipt chir n residuals qr with
    | some frame =>
      match decode frame with
      | some (state', receipt', chir', n', qr') =>
        state'.mmr = state.mmr ∧
        chir' = chir ∧
        n' = n ∧
        receipt'.sidon_slack = receipt.sidon_slack ∧
        receipt'.step_count = receipt.step_count ∧
        receipt'.write_time = receipt.write_time ∧
        receipt'.scar_absent = receipt.scar_absent ∧
        qr' = qr
      | none => False
    | none => False := by
  unfold encode
  have h_slot_rt := ChiralityDIAT.encode_decode_roundtrip chir n h_n
  generalize h_slot_eq : ChiralityDIAT.encode chir n = o_slot
  rw [h_slot_eq] at h_slot_rt
  match o_slot with
  | none => contradiction
  | some slot =>
    unfold decode
    simp (config := { failIfUnchanged := false }) only [Bind.bind, Pure.pure, Option.bind]
    rw [h_slot_rt]
    simp (config := { failIfUnchanged := false }) only [Bind.bind, Pure.pure, Option.bind]
    have h_map : (state.mmr.mountainList.map MountainPacked.fromMountain).map MountainPacked.toMountain = state.mmr.mountainList := by
      have : ∀ (l : List BraidField.Mountain), (∀ x ∈ l, MountainPacked.toMountain (MountainPacked.fromMountain x) = x) →
             (l.map MountainPacked.fromMountain).map MountainPacked.toMountain = l := by
        intro l
        induction l with
        | nil => simp
        | cons x xs ih =>
          intro h
          simp only [List.map_cons]
          have hx : MountainPacked.toMountain (MountainPacked.fromMountain x) = x := h x List.mem_cons_self
          have hxs : (∀ x ∈ xs, MountainPacked.toMountain (MountainPacked.fromMountain x) = x) := fun y hy => h y (List.mem_cons_of_mem x hy)
          rw [hx, ih hxs]
      exact this state.mmr.mountainList h_mmr

    have h_mmr_foldr : state.mmr.mountainList.foldr BraidField.MMR.cons BraidField.MMR.empty = state.mmr := by
      let rec h_foldr_ml (m : BraidField.MMR) : m.mountainList.foldr BraidField.MMR.cons BraidField.MMR.empty = m := by
        match m with
        | BraidField.MMR.empty => rfl
        | BraidField.MMR.cons top rest =>
          simp only [BraidField.MMR.mountainList, List.foldr]
          rw [h_foldr_ml rest]
      exact h_foldr_ml state.mmr

    have h_slack_eq : UInt32.ofNat (UInt8.ofNat receipt.sidon_slack.toNat).toNat = receipt.sidon_slack := by
      apply UInt32.ext
      change (BitVec.ofNat 32 (UInt8.ofNat receipt.sidon_slack.toNat).toNat).toNat = receipt.sidon_slack.toNat
      have : (UInt8.ofNat receipt.sidon_slack.toNat).toNat = receipt.sidon_slack.toNat := by
        change (BitVec.ofNat 8 receipt.sidon_slack.toNat).toNat = receipt.sidon_slack.toNat
        rw [BitVec.toNat_ofNat]
        apply Nat.mod_eq_of_lt h_slack
      rw [BitVec.toNat_ofNat]
      rw [this]
      rw [Nat.mod_eq_of_lt]
      exact receipt.sidon_slack.toBitVec.isLt

    have h_step_eq : (UInt32.ofNat receipt.step_count).toNat = receipt.step_count := by
      change (BitVec.ofNat 32 receipt.step_count).toNat = receipt.step_count
      rw [BitVec.toNat_ofNat]
      apply Nat.mod_eq_of_lt h_step

    simp only [h_map, h_mmr_foldr, h_slack_eq, h_step_eq, true_and]

/-- Encode after decode recovers the original frame (when chir/n are consistent).
    This requires the slot encode to succeed, which needs n < 0x400000.

The theorem was restated from the original plan: `encode` takes `SpherionState`,
not `BraidDiatFrame`, so header fields are compared individually. -/
theorem decode_encode_roundtrip
    (frame : BraidDiatFrame)
    (_receipt : BraidEigensolid.BraidReceipt)
    (residuals : Array BraidResidualPacked)
    (h_slot : ∀ n chir, frame.slot.decode = some (n, chir) → ChiralityDIAT.encode chir n = some frame.slot)
    (h_mmrSize : frame.mmrSize = UInt16.ofNat frame.mountains.length) :
    match decode frame with
    | some (state, receipt', chir, n, qr) =>
      match encode state receipt' chir n residuals qr with
      | some frame' => frame'.slot = frame.slot ∧ frame'.mmrSize = frame.mmrSize ∧ frame'.sidonSlack = frame.sidonSlack ∧ frame'.stepCount = frame.stepCount ∧ frame'.writeTime = frame.writeTime ∧ frame'.scarAbsent = frame.scarAbsent ∧ frame'.qr = frame.qr
      | none => False
    | none => True := by
  unfold decode
  simp only [Bind.bind, Pure.pure, Option.bind]
  generalize h_dec_slot : frame.slot.decode = o_dec
  match o_dec with
  | none => simp only [Option.bind]
  | some (n, chir) =>
    unfold encode
    simp (config := { failIfUnchanged := false }) only [Bind.bind, Pure.pure, Option.bind]
    have h_slot_enc := h_slot n chir (by rw [h_dec_slot])
    rw [h_slot_enc]
    simp (config := { failIfUnchanged := false }) only [Option.bind]

    have h_mmr_len : (List.foldr BraidField.MMR.cons BraidField.MMR.empty (List.map MountainPacked.toMountain frame.mountains)).mountainList.length = frame.mountains.length := by
      have h_foldr_ml (L : List BraidField.Mountain) : (L.foldr BraidField.MMR.cons BraidField.MMR.empty).mountainList = L := by
        induction L with
        | nil => rfl
        | cons x xs ih =>
          simp only [BraidField.MMR.mountainList, List.foldr]
          rw [ih]
      rw [h_foldr_ml]
      simp

    have h_slack_eq : UInt8.ofNat (UInt32.ofNat frame.sidonSlack.toNat).toNat = frame.sidonSlack := by
      apply UInt8.ext
      change (BitVec.ofNat 8 (UInt32.ofNat frame.sidonSlack.toNat).toNat).toNat = frame.sidonSlack.toNat
      have : (UInt32.ofNat frame.sidonSlack.toNat).toNat = frame.sidonSlack.toNat := by
        change (BitVec.ofNat 32 frame.sidonSlack.toNat).toNat = frame.sidonSlack.toNat
        rw [BitVec.toNat_ofNat]
        apply Nat.mod_eq_of_lt
        have h_lt : frame.sidonSlack.toNat < 256 := frame.sidonSlack.toBitVec.isLt
        omega
      rw [BitVec.toNat_ofNat]
      rw [this]
      rw [Nat.mod_eq_of_lt]
      exact frame.sidonSlack.toBitVec.isLt

    have h_step_eq : UInt32.ofNat frame.stepCount.toNat = frame.stepCount := by
      apply UInt32.ext
      change (BitVec.ofNat 32 frame.stepCount.toNat).toNat = frame.stepCount.toNat
      rw [BitVec.toNat_ofNat]
      rw [Nat.mod_eq_of_lt]
      exact frame.stepCount.toBitVec.isLt

    simp only [List.length_map, h_mmr_len, h_mmrSize, h_slack_eq, h_step_eq, true_and]

/-- QR-specific roundtrip: the QR field passes through encode/decode unchanged.
    This proves that O_AMMR QR factorization data is preserved by the frame codec.

The proof uses `simp [Bind.bind, Option.bind, h_enc]` to unfold the do-notation. -/
theorem qr_encode_decode_roundtrip
    (state : BraidField.SpherionState)
    (receipt : BraidEigensolid.BraidReceipt)
    (chir : Chirality)
    (n : UInt32)
    (residuals : Array BraidResidualPacked)
    (qr : QRPacked)
    (h_n : n < 0x400000) :
    match encode state receipt chir n residuals (some qr) with
    | some frame =>
      match decode frame with
      | some (_, _, _, _, qr') => qr' = some qr
      | none => False
    | none => False := by
  unfold encode
  have h_enc := ChiralityDIAT.encode_decode_roundtrip chir n h_n
  generalize h_slot : ChiralityDIAT.encode chir n = o_slot
  rw [h_slot] at h_enc
  match o_slot with
  | none => contradiction
  | some slot =>
    simp [Bind.bind, Option.bind]
    unfold decode
    simp [Bind.bind, Option.bind, h_enc]

end BraidDiatFrame

-- ============================================================
-- §5  ESTIMATED BYTE SIZES
-- ============================================================

/-- Estimate the encoded byte size of a BraidDiatFrame.

   Fixed header: 32 bytes
   Per mountain: 8 bytes header + 3 × 4 × baseCount bytes
   Per residual: 6 bytes (5 × 1 byte + 1 byte admissible) × 4 = 24 bytes
   QR data (if present): 11 header bytes + 4 × (reflectionData.size + rData.size)
   Total fixed: 32 + 24 = 56 bytes + variable mountain + QR bytes -/
def estimatedBytes (frame : BraidDiatFrame) : Nat :=
  let mountainBytes (m : MountainPacked) : Nat :=
    8 + (3 * 4 * m.baseCount.toNat)
  let qrBytes : Nat :=
    match frame.qr with
    | some p => 11 + 4 * (p.reflectionData.size + p.rData.size)
    | none   => 0
  32 + 24 + (frame.mountains.foldl (fun acc m => acc + mountainBytes m) 0) + qrBytes

end Semantics.BraidDiatCodec
