"""
Step 2: Reed-Solomon erasure coding on the verified GF(256) core.

  (a) neural coding equivalence -> parity computed via the NEURAL gf_mul equals
      the golden RS parity (the verified units drive the code)
  (b) any-k-of-n recovery       -> RS(k=4, m=2): reconstruct from EVERY possible
      2-shard-loss pattern, bit-exact
  (c) shards are bit-verified    -> each shard addressed by SHA-256
  (d) honesty                    -> n shards total (k+m)/k x the data (redundancy,
      not shrinkage)
"""
import os
import hashlib
import itertools
import torch
from storage import gf256, rs
from storage.common import verify

torch.manual_seed(0)

print("=" * 62)
print("STEP 2  --  Reed-Solomon erasure coding (verified GF core)")
print("=" * 62)

# load or train the verified GF units
try:
    ck = torch.load("GF256.pt", map_location="cpu")
    from storage.common import mlp
    net_log = mlp(8, 8, 256, 2); net_log.load_state_dict(ck["log"])
    net_exp = mlp(8, 8, 256, 2); net_exp.load_state_dict(ck["exp"])
    print("loaded verified GF units from GF256.pt")
except Exception:
    net_log, net_exp = gf256.train_units()

# (a) neural coding equivalence
neq = rs.neural_parity_equiv(net_log, net_exp)
print(f"(a) neural gf_mul parity == golden RS parity  -> {'PASS' if neq else 'FAIL'}")

# (b) any-k-of-n recovery over a real 'drive image'
K, M = 4, 2
N = K + M
image = os.urandom(300 * 1024)               # 300 KB payload
shards, L, orig = rs.encode(image, K, M)
print(f"    RS({K}+{M}): {N} shards of {L//1024} KB each  (tolerates up to {M} lost)")

fails = 0
patterns = list(itertools.combinations(range(N), M))    # every 2-shard-loss case
for lost in patterns:
    present = {i: bytes(shards[i]) for i in range(N) if i not in lost}
    if rs.decode(present, K, M, L, orig) != image:
        fails += 1
print(f"(b) recover from every {M}-loss pattern ({len(patterns)} of them): "
      f"{'PASS' if fails == 0 else f'FAIL({fails})'}")

# (c) shards bit-verified by SHA-256
ids = [hashlib.sha256(bytes(s)).hexdigest()[:12] for s in shards]
verified = all(hashlib.sha256(bytes(shards[i])).hexdigest()[:12] == ids[i] for i in range(N))
print(f"(c) shard hashes (bit-verified): {'PASS' if verified else 'FAIL'}  "
      f"e.g. shard0={ids[0]}")

# (d) honesty: storage overhead is exactly the redundancy, nothing shrinks
overhead = N * L / len(image)
print(f"(d) storage: {len(image)//1024} KB -> {N*L//1024} KB ({overhead:.2f}x)  "
      f"-> {'PASS (redundancy, no shrink)' if overhead >= 1.0 else 'FAIL'}")

allpass = neq and fails == 0 and verified and overhead >= 1.0
print("=" * 62)
print(f"OVERALL: {'ALL PASS' if allpass else 'some checks failed'}")