"""Reproduce the sensor dense refinements and correlator whiskers.

No blog edits. All draws are retained; numerical degeneracy aborts the run.
The proposal keeps the page's physical Gaussian-target model and width grids.
Use: python3 generate.py
"""
from pathlib import Path
import hashlib
import json
import math
import platform
import time
import numpy as np

HERE = Path(__file__).resolve().parent
MASTER_SEED = 2026100404
REFINEMENT_RACKS = 360
CORRELATOR_RACKS = 1000
PRESETS = {"n5s4": (5, 4), "n20s1": (20, 1), "n5s25": (5, 25), "n20s40": (20, 40)}
REFINEMENT_GRID = {key: list(range(4, 16)) if n == 5 else list(range(10, 36)) for key, (n, v) in PRESETS.items()}
WHISKER_GRID = {key: [4, 5, 10, 100, 1000, 4000] if n == 5 else [10, 20, 40, 100, 1000, 4000] for key, (n, v) in PRESETS.items()}


def seed_for(*parts):
    """Unsigned little-endian first64bits of SHA256 of explicit ASCII names."""
    name = ":".join(["sensor-fixtures-v1", str(MASTER_SEED), *map(str, parts)])
    return int.from_bytes(hashlib.sha256(name.encode("ascii")).digest()[:8], "little")


def rng_for(*parts):
    return np.random.Generator(np.random.PCG64(seed_for(*parts)))


def raw_rack(group, key, rack_index, p):
    """Same Gaussian y across widths; independent a/noise for each width.

    y is NOT normalized: each physical quantity is N(0,1), as in the page's
    observed Gaussian examples. Only the response pattern a has known norm.
    """
    n, variance = PRESETS[key]
    y = rng_for(group, key, rack_index, "labels").standard_normal(n)
    a = rng_for(group, key, rack_index, p, "response").standard_normal(p)
    raw_norm = np.linalg.norm(a)
    if raw_norm == 0 or y @ y == 0:
        raise ArithmeticError("Degenerate draw: abort; never replace it")
    a *= math.sqrt(p)/raw_norm
    noise = rng_for(group, key, rack_index, p, "noise").standard_normal((n, p))
    X = y[:, None]*a + math.sqrt(variance)*noise
    return X, y, a, variance


def physical_risk(w, a, variance):
    return float((w @ a - 1)**2 + variance*(w @ w))


def interp_risk(X, y, a, variance):
    """Direct NumPy SVD least squares; no production solver or figure code."""
    w = np.linalg.lstsq(X, y, rcond=None)[0]
    return physical_risk(w, a, variance)


def corr_risk(X, y, a, variance):
    p = len(a)
    w = X.T @ y/(p*(y @ y))
    return physical_risk(w, a, variance)


def q(values, probs):
    return np.quantile(values, probs, axis=1, method="linear")


def sig7(values):
    return [float(f"{v:.7g}") for v in values]


def main():
    started = time.perf_counter()
    samples, labels_norms = {}, {}
    dense, bars = {}, {}
    manifest = dict(schema=1, status="reproducible published dense and whisker summaries",
                    model="X=y a^T+sqrt(v) Z; y iid N(0,1); Z iid N(0,1); a normalized Gaussian, norm²=p",
                    master_seed=MASTER_SEED, seed_policy="first8 SHA256 bytes little-endian of ASCII colon-joined sensor-fixtures-v1:master:group:preset:rack[:p]:component",
                    bit_generator="numpy.random.PCG64", normal_sampler="NumPy Generator.standard_normal",
                    refinement_racks=REFINEMENT_RACKS, correlator_racks=CORRELATOR_RACKS,
                    refinement_grid=REFINEMENT_GRID, whisker_grid=WHISKER_GRID,
                    shared_across_widths="Each rack's unnormalized Gaussian y; response/noise use separate named streams per width",
                    selection_policy="No rejection, filtering, resampling, winsorizing, or seed search. All indices0..count-1 retained; invalid arithmetic aborts.",
                    quantile_policy="numpy.quantile(method='linear'); q25/q50/q75 refinements, q10/q90 correlator whiskers; round final fixtures to7significant digits",
                    runtime=dict(python=platform.python_version(), numpy=np.__version__),
                    generator_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), presets={})
    for key, (n, variance) in PRESETS.items():
        values = np.empty((len(REFINEMENT_GRID[key]), REFINEMENT_RACKS))
        norms = np.empty(REFINEMENT_RACKS)
        for ip, p in enumerate(REFINEMENT_GRID[key]):
            for rack in range(REFINEMENT_RACKS):
                X, y, a, v = raw_rack("refinement", key, rack, p)
                values[ip, rack] = interp_risk(X, y, a, v)
                norms[rack] = y @ y
        quantiles = q(values, [.25, .5, .75])
        dense[key] = dict(p=REFINEMENT_GRID[key], lo=sig7(quantiles[0]), interp=sig7(quantiles[1]), hi=sig7(quantiles[2]))
        samples[key+"_refinement_interpolation_risk"] = values
        samples[key+"_refinement_label_norm2"] = norms
        corr_values = np.empty((len(WHISKER_GRID[key]), CORRELATOR_RACKS))
        corr_norms = np.empty(CORRELATOR_RACKS)
        for ip, p in enumerate(WHISKER_GRID[key]):
            for rack in range(CORRELATOR_RACKS):
                X, y, a, v = raw_rack("whiskers", key, rack, p)
                corr_values[ip, rack] = corr_risk(X, y, a, v)
                corr_norms[rack] = y @ y
        whiskers = q(corr_values, [.1, .9])
        bars[key] = [[p, float(f"{lo:.7g}"), float(f"{hi:.7g}")] for p, lo, hi in zip(WHISKER_GRID[key], *whiskers)]
        samples[key+"_correlator_risk"] = corr_values
        samples[key+"_correlator_label_norm2"] = corr_norms
        manifest["presets"][key] = dict(n=n, variance=variance,
            retained_refinement_samples=int(values.size), retained_correlator_samples=int(corr_values.size),
            refinement_first_label_seed=seed_for("refinement", key, 0, "labels"),
            whisker_first_label_seed=seed_for("whiskers", key, 0, "labels"),
            largest_dense_risk=float(values.max()), largest_correlator_risk=float(corr_values.max()))
        print(key, "retained", values.size, corr_values.size, "elapsed", round(time.perf_counter()-started, 2), flush=True)
    if not all(np.isfinite(x).all() for x in samples.values()):
        raise ArithmeticError("Nonfinite output: abort instead of filtering")
    np.savez_compressed(HERE/"samples.npz", **samples)
    fixture = dict(CAPACITY_REFINEMENTS=dense, CORRELATOR_BARS=bars)
    (HERE/"fixtures.json").write_text(json.dumps(fixture, indent=2)+"\n")
    manifest["elapsed_seconds"] = time.perf_counter()-started
    manifest["artifact_hashes"] = {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in [HERE/"fixtures.json", HERE/"samples.npz"]}
    (HERE/"manifest.json").write_text(json.dumps(manifest, indent=2)+"\n")
    print("elapsed", manifest["elapsed_seconds"], flush=True)


if __name__ == "__main__":
    main()
