import Base
import ./nat.bend as N

# Series.bend - exact Nat sums over 1..n (triangular and square-pyramidal).
#
# All sums are exact over Nat (no division): triangular uses the doubled
# form 2*T(n) = n*(n+1); squares use the sixfold form
# 6*S(n) = n*(n+1)*(2*n+1). Both loops are tail-recursive accumulators,
# structural on the first (Nat) argument.
#
# Geometric series are DEFERRED until rat.bend lands: closed forms need
# exact Nat division (r^n - 1)/(r - 1), which is unavailable.
#
# Proof status:
# - PROVEN: trisum_go_spec, tri2_succ, tri_formula (triangular side);
#   sqsum_go_spec, sqW_grow, sq62_distr (square stepping stones);
#   closed instances tri_10, sq_5 (definitional).
# - NOTE-DROPPED (see NOTE at end): sq62_succ, sq_formula.

# tri2: doubled triangular number, 2*T(n) = n*(n+1).
def tri2(+n: Nat) -> Nat:
  Nat.mul(n, Nat.add(n, 1n))

# trisum_go: tail loop, acc holds the partial sum.
def trisum_go(n: Nat, acc: Nat) -> Nat:
  match n:
    case 0n:
      acc
    case 1n++p:
      trisum_go(p, Nat.add(acc, 1n+p))

def trisum(n: Nat) -> Nat:
  trisum_go(n, 0n)

# sq62: sixfold square sum, 6*S(n) = n*(n+1)*(2*n+1).
def sq62(+n: Nat) -> Nat:
  Nat.mul(Nat.mul(n, Nat.add(n, 1n)), Nat.add(n, Nat.add(n, 1n)))

# sqsum_go: tail loop accumulating (1+p)^2 each step.
def sqsum_go(n: Nat, acc: Nat) -> Nat:
  match n:
    case 0n:
      acc
    case 1n++p:
      sqsum_go(p, Nat.add(acc, Nat.mul(1n+p, 1n+p)))

def sqsum(n: Nat) -> Nat:
  sqsum_go(n, 0n)

# Closed instances (definitional: the loops normalize on literals).
# tri(10) = 10*11/2 = 55.
law tri_10:
  {trisum(10n) == 55n : Nat}

def tri_10():
  {==}

# sqsum(5) = 1+4+9+16+25 = 55 = 5*6*11/6.
law sq_5:
  {sqsum(5n) == 55n : Nat}

def sq_5():
  {==}

# trisum_go_spec: the accumulator holds the partial sum.
# Proof is induction on n; the step chains the IH (at the extended
# accumulator) with N.add_assoc and the IH (at 1+p), mirroring
# nat.bend's mul_succ_r nesting.
law trisum_go_spec:
  for n: Nat
  for +acc: Nat
  {trisum_go(n, acc) == Nat.add(acc, trisum(n)) : Nat}

def trisum_go_spec(n, acc):
  match n:
    case 0n:
      Equal.sym(Nat, Nat.add(acc, trisum(0n)), acc, N.add_zero_r(acc))
    case 1n++p:
      Equal.trans(Nat,
        trisum_go(p, Nat.add(acc, 1n+p)),
        Nat.add(Nat.add(acc, 1n+p), trisum(p)),
        Nat.add(acc, trisum_go(p, 1n+p)),
        trisum_go_spec(p, Nat.add(acc, 1n+p)),
        Equal.trans(Nat,
          Nat.add(Nat.add(acc, 1n+p), trisum(p)),
          Nat.add(acc, Nat.add(1n+p, trisum(p))),
          Nat.add(acc, trisum_go(p, 1n+p)),
          N.add_assoc(acc, 1n+p, trisum(p)),
          Equal.cong(Nat, Nat, w => Nat.add(acc, w),
            Nat.add(1n+p, trisum(p)), trisum_go(p, 1n+p),
            Equal.sym(Nat, trisum_go(p, 1n+p), Nat.add(1n+p, trisum(p)),
              trisum_go_spec(p, 1n+p)))))

# tri2_succ (ATTEMPT 1): one-step growth of the doubled triangular number,
# tri2(1+p) == tri2(p) + 2*(1+p). Pure rewriting, no induction: distr +
# one + mul-unfold + succ + assoc + comm, nested like nat.bend mul_succ_r.
law tri2_succ:
  for +p: Nat
  {tri2(1n+p) == Nat.add(tri2(p), Nat.add(1n+p, 1n+p)) : Nat}

def tri2_succ(p):
  Equal.trans(Nat,
    Nat.mul(1n+p, Nat.add(1n+p, 1n)),
    Nat.add(Nat.mul(1n+p, 1n+p), Nat.mul(1n+p, 1n)),
    Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
    N.mul_add_distr(1n+p, 1n+p, 1n),
    Equal.trans(Nat,
      Nat.add(Nat.mul(1n+p, 1n+p), Nat.mul(1n+p, 1n)),
      Nat.add(Nat.mul(1n+p, 1n+p), 1n+p),
      Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
      Equal.cong(Nat, Nat, w => Nat.add(Nat.mul(1n+p, 1n+p), w),
        Nat.mul(1n+p, 1n), 1n+p, N.mul_one_r(1n+p)),
      Equal.trans(Nat,
        Nat.add(Nat.mul(1n+p, 1n+p), 1n+p),
        Nat.add(Nat.add(1n+p, Nat.mul(p, 1n+p)), 1n+p),
        Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
        Equal.cong(Nat, Nat, w => Nat.add(w, 1n+p),
          Nat.mul(1n+p, 1n+p), Nat.add(1n+p, Nat.mul(p, 1n+p)),
          {==}),
        Equal.trans(Nat,
          Nat.add(Nat.add(1n+p, Nat.mul(p, 1n+p)), 1n+p),
          Nat.add(Nat.add(1n+p, Nat.add(Nat.mul(p, p), p)), 1n+p),
          Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
          Equal.cong(Nat, Nat, w => Nat.add(Nat.add(1n+p, w), 1n+p),
            Nat.mul(p, 1n+p), Nat.add(Nat.mul(p, p), p),
            N.mul_succ_r(p, p)),
          Equal.trans(Nat,
            Nat.add(Nat.add(1n+p, Nat.add(Nat.mul(p, p), p)), 1n+p),
            Nat.add(1n+p, Nat.add(Nat.add(Nat.mul(p, p), p), 1n+p)),
            Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
            N.add_assoc(1n+p, Nat.add(Nat.mul(p, p), p), 1n+p),
            Equal.trans(Nat,
              Nat.add(1n+p, Nat.add(Nat.add(Nat.mul(p, p), p), 1n+p)),
              Nat.add(1n+p, Nat.add(1n+p, Nat.add(Nat.mul(p, p), p))),
              Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
              Equal.cong(Nat, Nat, w => Nat.add(1n+p, w),
                Nat.add(Nat.add(Nat.mul(p, p), p), 1n+p),
                Nat.add(1n+p, Nat.add(Nat.mul(p, p), p)),
                N.add_comm(Nat.add(Nat.mul(p, p), p), 1n+p)),
              Equal.trans(Nat,
                Nat.add(1n+p, Nat.add(1n+p, Nat.add(Nat.mul(p, p), p))),
                Nat.add(Nat.add(1n+p, 1n+p), Nat.add(Nat.mul(p, p), p)),
                Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
                Equal.sym(Nat,
                  Nat.add(Nat.add(1n+p, 1n+p), Nat.add(Nat.mul(p, p), p)),
                  Nat.add(1n+p, Nat.add(1n+p, Nat.add(Nat.mul(p, p), p))),
                  N.add_assoc(1n+p, 1n+p, Nat.add(Nat.mul(p, p), p))),
                Equal.trans(Nat,
                  Nat.add(Nat.add(1n+p, 1n+p), Nat.add(Nat.mul(p, p), p)),
                  Nat.add(Nat.add(Nat.mul(p, p), p), Nat.add(1n+p, 1n+p)),
                  Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
                  N.add_comm(Nat.add(1n+p, 1n+p), Nat.add(Nat.mul(p, p), p)),
                  Equal.cong(Nat, Nat, w => Nat.add(w, Nat.add(1n+p, 1n+p)),
                    Nat.add(Nat.mul(p, p), p),
                    Nat.mul(p, Nat.add(p, 1n)),
                    Equal.trans(Nat,
                      Nat.add(Nat.mul(p, p), p),
                      Nat.mul(p, 1n+p),
                      Nat.mul(p, Nat.add(p, 1n)),
                      Equal.sym(Nat, Nat.mul(p, 1n+p),
                        Nat.add(Nat.mul(p, p), p), N.mul_succ_r(p, p)),
                      Equal.cong(Nat, Nat, w => Nat.mul(p, w),
                        1n+p, Nat.add(p, 1n),
                        Equal.sym(Nat, Nat.add(p, 1n), 1n+p,
                          N.add_comm(p, 1n)))))))))))))

# tri_formula (ATTEMPT 1): doubled triangular formula, by induction on n.
# Base is definitional (both sides unfold to 0n). Step chains the
# accumulator spec (both copies), N.add_shuffle, the IH, and tri2_succ.
law tri_formula:
  for n: Nat
  {Nat.add(trisum(n), trisum(n)) == tri2(n) : Nat}

def tri_formula(n):
  match n:
    case 0n:
      {==}
    case 1n++p:
      Equal.trans(Nat,
        Nat.add(trisum(1n+p), trisum(1n+p)),
        Nat.add(Nat.add(1n+p, trisum(p)), Nat.add(1n+p, trisum(p))),
        tri2(1n+p),
        Equal.trans(Nat,
          Nat.add(trisum_go(p, 1n+p), trisum_go(p, 1n+p)),
          Nat.add(Nat.add(1n+p, trisum(p)), trisum_go(p, 1n+p)),
          Nat.add(Nat.add(1n+p, trisum(p)), Nat.add(1n+p, trisum(p))),
          Equal.cong(Nat, Nat, w => Nat.add(w, trisum_go(p, 1n+p)),
            trisum_go(p, 1n+p), Nat.add(1n+p, trisum(p)),
            trisum_go_spec(p, 1n+p)),
          Equal.cong(Nat, Nat, w => Nat.add(Nat.add(1n+p, trisum(p)), w),
            trisum_go(p, 1n+p), Nat.add(1n+p, trisum(p)),
            trisum_go_spec(p, 1n+p))),
        Equal.trans(Nat,
          Nat.add(Nat.add(1n+p, trisum(p)), Nat.add(1n+p, trisum(p))),
          Nat.add(Nat.add(1n+p, 1n+p), Nat.add(trisum(p), trisum(p))),
          tri2(1n+p),
          N.add_shuffle(1n+p, trisum(p), 1n+p, trisum(p)),
          Equal.trans(Nat,
            Nat.add(Nat.add(1n+p, 1n+p), Nat.add(trisum(p), trisum(p))),
            Nat.add(Nat.add(1n+p, 1n+p), tri2(p)),
            tri2(1n+p),
            Equal.cong(Nat, Nat, w => Nat.add(Nat.add(1n+p, 1n+p), w),
              Nat.add(trisum(p), trisum(p)), tri2(p), tri_formula(p)),
            Equal.trans(Nat,
              Nat.add(Nat.add(1n+p, 1n+p), tri2(p)),
              Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
              tri2(1n+p),
              N.add_comm(Nat.add(1n+p, 1n+p), tri2(p)),
              Equal.sym(Nat, tri2(1n+p),
                Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
                tri2_succ(p))))))

# sqsum_go_spec: the accumulator holds the partial square sum.
# Same shape as trisum_go_spec (only the accumulated term differs).
law sqsum_go_spec:
  for n: Nat
  for +acc: Nat
  {sqsum_go(n, acc) == Nat.add(acc, sqsum(n)) : Nat}

def sqsum_go_spec(n, acc):
  match n:
    case 0n:
      Equal.sym(Nat, Nat.add(acc, sqsum(0n)), acc, N.add_zero_r(acc))
    case 1n++p:
      Equal.trans(Nat,
        sqsum_go(p, Nat.add(acc, Nat.mul(1n+p, 1n+p))),
        Nat.add(Nat.add(acc, Nat.mul(1n+p, 1n+p)), sqsum(p)),
        Nat.add(acc, sqsum_go(p, Nat.mul(1n+p, 1n+p))),
        sqsum_go_spec(p, Nat.add(acc, Nat.mul(1n+p, 1n+p))),
        Equal.trans(Nat,
          Nat.add(Nat.add(acc, Nat.mul(1n+p, 1n+p)), sqsum(p)),
          Nat.add(acc, Nat.add(Nat.mul(1n+p, 1n+p), sqsum(p))),
          Nat.add(acc, sqsum_go(p, Nat.mul(1n+p, 1n+p))),
          N.add_assoc(acc, Nat.mul(1n+p, 1n+p), sqsum(p)),
          Equal.cong(Nat, Nat, w => Nat.add(acc, w),
            Nat.add(Nat.mul(1n+p, 1n+p), sqsum(p)),
            sqsum_go(p, Nat.mul(1n+p, 1n+p)),
            Equal.sym(Nat, sqsum_go(p, Nat.mul(1n+p, 1n+p)),
              Nat.add(Nat.mul(1n+p, 1n+p), sqsum(p)),
              sqsum_go_spec(p, Nat.mul(1n+p, 1n+p))))))

# sqW_grow (TRY 1 on the sq-formula track): split of the odd factor,
# (2(1+p)+1) == (2p+1)+2. Pure rewriting: unfold + succ + assoc + comm.
# Needed by the sq62 growth bridge (NOTE below on why the bridge itself
# is deferred).
law sqW_grow:
  for +p: Nat
  {Nat.add(1n+p, Nat.add(1n+p, 1n))
    == Nat.add(Nat.add(p, Nat.add(p, 1n)), 2n) : Nat}

def sqW_grow(p):
  Equal.trans(Nat,
    Nat.add(1n+p, Nat.add(1n+p, 1n)),
    Nat.add(1n+p, Nat.add(1n, Nat.add(p, 1n))),
    Nat.add(Nat.add(p, Nat.add(p, 1n)), 2n),
    Equal.cong(Nat, Nat, w => Nat.add(1n+p, w),
      Nat.add(1n+p, 1n), Nat.add(1n, Nat.add(p, 1n)),
      {==}),
    Equal.trans(Nat,
      Nat.add(1n+p, Nat.add(1n, Nat.add(p, 1n))),
      Nat.add(1n, Nat.add(p, Nat.add(1n, Nat.add(p, 1n)))),
      Nat.add(Nat.add(p, Nat.add(p, 1n)), 2n),
      {==},
      Equal.trans(Nat,
        Nat.add(1n, Nat.add(p, Nat.add(1n, Nat.add(p, 1n)))),
        Nat.add(1n, Nat.add(1n, Nat.add(p, Nat.add(p, 1n)))),
        Nat.add(Nat.add(p, Nat.add(p, 1n)), 2n),
        Equal.cong(Nat, Nat, w => Nat.add(1n, w),
          Nat.add(p, 1n+Nat.add(p, 1n)),
          1n+Nat.add(p, Nat.add(p, 1n)),
          N.add_succ_r(p, Nat.add(p, 1n))),
        Equal.trans(Nat,
          Nat.add(1n, Nat.add(1n, Nat.add(p, Nat.add(p, 1n)))),
          Nat.add(Nat.add(1n, 1n), Nat.add(p, Nat.add(p, 1n))),
          Nat.add(Nat.add(p, Nat.add(p, 1n)), 2n),
          Equal.sym(Nat,
            Nat.add(Nat.add(1n, 1n), Nat.add(p, Nat.add(p, 1n))),
            Nat.add(1n, Nat.add(1n, Nat.add(p, Nat.add(p, 1n)))),
            N.add_assoc(1n, 1n, Nat.add(p, Nat.add(p, 1n)))),
          Equal.trans(Nat,
            Nat.add(Nat.add(1n, 1n), Nat.add(p, Nat.add(p, 1n))),
            Nat.add(2n, Nat.add(p, Nat.add(p, 1n))),
            Nat.add(Nat.add(p, Nat.add(p, 1n)), 2n),
            Equal.cong(Nat, Nat, w => Nat.add(w, Nat.add(p, Nat.add(p, 1n))),
              Nat.add(1n, 1n), 2n,
              {==}),
            N.add_comm(2n, Nat.add(p, Nat.add(p, 1n))))))))

# sq62_distr (TRY 2 on the sq-formula track): distr skeleton of the
# sixfold-square growth, sq62(1+p) == tri2(p)*W + (2(1+p))*W where
# W = (2(1+p)+1). Nodes mirror the tri2_succ opening (tri2_succ cong +
# comm-flip + distr + double flip-back); all rewrites are syntactic.
law sq62_distr:
  for +p: Nat
  {sq62(1n+p)
    == Nat.add(Nat.mul(tri2(p), Nat.add(1n+p, Nat.add(1n+p, 1n))),
        Nat.mul(Nat.add(1n+p, 1n+p),
          Nat.add(1n+p, Nat.add(1n+p, 1n)))) : Nat}

def sq62_distr(p):
  Equal.trans(Nat,
    Nat.mul(Nat.mul(1n+p, Nat.add(1n+p, 1n)),
      Nat.add(1n+p, Nat.add(1n+p, 1n))),
    Nat.mul(Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
      Nat.add(1n+p, Nat.add(1n+p, 1n))),
    Nat.add(Nat.mul(tri2(p), Nat.add(1n+p, Nat.add(1n+p, 1n))),
      Nat.mul(Nat.add(1n+p, 1n+p), Nat.add(1n+p, Nat.add(1n+p, 1n)))),
    Equal.cong(Nat, Nat,
      w => Nat.mul(w, Nat.add(1n+p, Nat.add(1n+p, 1n))),
      tri2(1n+p), Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
      tri2_succ(p)),
    Equal.trans(Nat,
      Nat.mul(Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
        Nat.add(1n+p, Nat.add(1n+p, 1n))),
      Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)),
        Nat.add(tri2(p), Nat.add(1n+p, 1n+p))),
      Nat.add(Nat.mul(tri2(p), Nat.add(1n+p, Nat.add(1n+p, 1n))),
        Nat.mul(Nat.add(1n+p, 1n+p), Nat.add(1n+p, Nat.add(1n+p, 1n)))),
      N.mul_comm(Nat.add(tri2(p), Nat.add(1n+p, 1n+p)),
        Nat.add(1n+p, Nat.add(1n+p, 1n))),
      Equal.trans(Nat,
        Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)),
          Nat.add(tri2(p), Nat.add(1n+p, 1n+p))),
        Nat.add(Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)), tri2(p)),
          Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)),
            Nat.add(1n+p, 1n+p))),
        Nat.add(Nat.mul(tri2(p), Nat.add(1n+p, Nat.add(1n+p, 1n))),
          Nat.mul(Nat.add(1n+p, 1n+p), Nat.add(1n+p, Nat.add(1n+p, 1n)))),
        N.mul_add_distr(Nat.add(1n+p, Nat.add(1n+p, 1n)),
          tri2(p), Nat.add(1n+p, 1n+p)),
        Equal.trans(Nat,
          Nat.add(Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)), tri2(p)),
            Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)),
              Nat.add(1n+p, 1n+p))),
          Nat.add(Nat.mul(tri2(p), Nat.add(1n+p, Nat.add(1n+p, 1n))),
            Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)),
              Nat.add(1n+p, 1n+p))),
          Nat.add(Nat.mul(tri2(p), Nat.add(1n+p, Nat.add(1n+p, 1n))),
            Nat.mul(Nat.add(1n+p, 1n+p), Nat.add(1n+p, Nat.add(1n+p, 1n)))),
          Equal.cong(Nat, Nat,
            w => Nat.add(w,
              Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)),
                Nat.add(1n+p, 1n+p))),
            Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)), tri2(p)),
            Nat.mul(tri2(p), Nat.add(1n+p, Nat.add(1n+p, 1n))),
            N.mul_comm(Nat.add(1n+p, Nat.add(1n+p, 1n)), tri2(p))),
          Equal.cong(Nat, Nat,
            w => Nat.add(Nat.mul(tri2(p),
              Nat.add(1n+p, Nat.add(1n+p, 1n))), w),
            Nat.mul(Nat.add(1n+p, Nat.add(1n+p, 1n)),
              Nat.add(1n+p, 1n+p)),
            Nat.mul(Nat.add(1n+p, 1n+p),
              Nat.add(1n+p, Nat.add(1n+p, 1n))),
            N.mul_comm(Nat.add(1n+p, Nat.add(1n+p, 1n)),
              Nat.add(1n+p, 1n+p)))))))

# NOTE (dropped laws): sq62_succ and sq_formula are NOT stated - proofs
# exceeded the 2-try budget, so they are documented here instead.
# Intended statements:
#   law sq62_succ: for +p: Nat
#     {sq62(1n+p) == Nat.add(sq62(p), Nat.mul(6n, Nat.mul(1n+p, 1n+p))) : Nat}
#   law sq_formula: for n: Nat
#     {Nat.mul(6n, sqsum(n)) == sq62(n) : Nat}
# Why deferred: the growth bridge needs ~40-50 exact trans nodes -
# W-split (sqW_grow, proven) + comm-flips + distr twice + one_r assembly
# + sixfold Q-tower assembly with sym-flipped expansions (the u = 1+p
# factor is consumed by expansion, so Q = (1+p)^2 leaves cannot be
# re-formed without backward shaping). Kept stepping stones toward it:
# sqsum_go_spec, sqW_grow, sq62_distr (all check). Retry from those.