From Stdlib Require Import Arith Lia.
From Stdlib Require Import List.
Import ListNotations.

(* ----- sigma3: sum of cubes of divisors ----- *)

Fixpoint divisors_aux (n d : nat) : list nat :=
  match d with
  | 0 => []
  | S d' =>
    if Nat.eqb (Nat.modulo n (S d')) 0
    then (S d') :: divisors_aux n d'
    else divisors_aux n d'
  end.

Definition divisors (n : nat) : list nat :=
  match n with
  | 0 => []
  | _ => divisors_aux n n
  end.

Definition sigma3 (n : nat) : nat :=
  fold_right (fun d acc => d * d * d + acc) 0 (divisors n).

(* ----- hachimoji_letter: sigma3 mod 8 ----- *)

Definition hachimoji_letter (s3 : nat) : nat :=
  Nat.modulo s3 8.

(* ----- cartan_weight: 273 if equal, 256 if complementary, 0 otherwise ----- *)
(* Complementary hachimoji letters sum to a multiple of 8.                    *)

Definition cartan_weight (a b : nat) : nat :=
  if Nat.eqb a b then 273
  else if Nat.eqb (Nat.modulo (a + b) 8) 0 then 256
  else 0.

(* ----- Helper: C(k,2) = k*(k-1)/2 via tail recursion ----- *)

Fixpoint tail_pair_count (k : nat) : nat :=
  match k with
  | 0 | 1 => 0
  | S k' => k' + tail_pair_count k'
  end.

(* ----- angrysphinx_gate ----- *)
(* Collisions = number of unordered pairs among tail elements (i.e., among   *)
(* elements after the first). Energy = 273 - 239*collisions, saturating at   *)
(* zero. The gate passes iff energy > 0. The constant 239 is 273 - 34, where *)
(* 34 is the single-collision residual energy.                                *)

Definition angrysphinx_gate (els : list nat) : bool * nat * nat :=
  let n := length els in
  let collisions :=
    match n with
    | 0 | 1 => 0
    | S n' => tail_pair_count n'
    end
  in
  let penalty := collisions * 239 in
  let energy := Nat.sub 273 penalty in
  let passed := negb (Nat.eqb energy 0) in
  (passed, collisions, energy).

(* ----- hachimoji_encode ----- *)

Definition hachimoji_encode (n : nat) : nat * nat :=
  let s3 := sigma3 n in
  (s3, hachimoji_letter s3).

(* ===== COMPUTE DIRECTIVES ===== *)

(* sigma3 verification *)
Compute sigma3 1.   (* = 1 *)
Compute sigma3 2.   (* = 9 *)
Compute sigma3 3.   (* = 28 *)
Compute sigma3 10.  (* = 1134 *)

(* angrysphinx_gate verification *)
Compute angrysphinx_gate [1;2].      (* = (true, 0, 273) *)
Compute angrysphinx_gate [1;2;3].    (* = (true, 1, 34)  *)
Compute angrysphinx_gate [1;2;3;4].  (* = (false, 3, 0)  *)

(* Test sigma3 for all n=1..10 *)
Compute sigma3 1.
Compute sigma3 2.
Compute sigma3 3.
Compute sigma3 4.
Compute sigma3 5.
Compute sigma3 6.
Compute sigma3 7.
Compute sigma3 8.
Compute sigma3 9.
Compute sigma3 10.

(* ===== THEOREMS ===== *)

Theorem gate_12_passes : angrysphinx_gate [1;2] = (true, 0, 273).
Proof. reflexivity. Qed.

Theorem gate_1234_fails : angrysphinx_gate [1;2;3;4] = (false, 3, 0).
Proof. reflexivity. Qed.