"""Reproduce the dense, continuously tuned Part 5 sensor curves.

Only compact plotted summaries are written. No sampled rack archives are needed
by the browser or committed to the repository. Requires Python, NumPy and SciPy.
"""
from pathlib import Path
import argparse
import hashlib
import heapq
import json
import math
import time

import numpy as np
from scipy.optimize import brentq, minimize_scalar

HERE = Path(__file__).resolve().parent
DEFAULT_OUTPUT = HERE.parent.parent / "assets" / "sensor-ridge-sweeps.js"
MASTER_SEED = 2026100405
RACKS = 1024
PRESETS = {"n5s4": (5, 4), "n20s1": (20, 1), "n5s25": (5, 25), "n20s40": (20, 40)}
WIDTHS = sorted(set(range(1, 101)) | {int(round(x)) for x in np.geomspace(101, 4000, 97)} | {1000, 4000})


def seed_for(key):
    text = f"sensor-ridge-v2:{MASTER_SEED}:{key}"
    return int.from_bytes(hashlib.sha256(text.encode("ascii")).digest()[:8], "little")


class SpectralRacks:
    """Thin-SVD formula for each rack's exact physical and empirical risk.

    X = U diag(sqrt(d)) V^T, z = U^T y, b = V^T a. The ridge coefficients
    along V are sqrt(d) z / (d + f d_min). The population error is exactly
    (b . coefficients - 1)^2 + variance * ||coefficients||^2.
    """

    def __init__(self, d, z, b, label_norm2, n, variance):
        if np.any(d <= 0) or not np.isfinite(d).all():
            raise ArithmeticError("Nonpositive singular value; abort, never discard a rack")
        self.r = d / d[:, :1]
        self.q = np.sqrt(d) * z / d[:, :1]
        self.b, self.z = b, z
        self.label_norm2 = label_norm2
        self.residual = np.maximum(0, label_norm2 - np.sum(z*z, axis=1))
        self.n, self.variance = n, variance
        self.count = d.shape[0]

    def values(self, parameter, positive=False, derivative=False):
        if positive:
            # t=f/(1+f) compactifies the ENTIRE nonnegative branch. t=1 is
            # exactly f=+infinity, the zero predictor, not a finite search cap.
            t = float(parameter)
            denominator = self.r * (1-t) + t
            coefficients = self.q * (1-t) / denominator
            residual_coefficients = self.z * t / denominator
            differential = -self.q / (denominator*denominator)
        else:
            f = float(parameter)
            denominator = self.r + f
            coefficients = self.q / denominator
            residual_coefficients = self.z * f / denominator
            differential = -self.q / (denominator*denominator)
        survival = np.sum(self.b * coefficients, axis=1)
        risk = (survival-1)**2 + self.variance*np.sum(coefficients*coefficients, axis=1)
        train = (self.residual + np.sum(residual_coefficients*residual_coefficients, axis=1))/self.n
        if derivative:
            slope = 2*(survival-1)*np.sum(self.b*differential, axis=1) + 2*self.variance*np.sum(coefficients*differential, axis=1)
            return risk, train, slope
        return risk, train

    def median(self, parameter, positive=False, derivative=False):
        values = self.values(parameter, positive, derivative)
        risk = values[0]
        middle = [(self.count-1)//2, self.count//2]
        active = np.argpartition(risk, middle)[middle]
        value = float(np.mean(risk[active]))
        return (value, float(np.mean(values[2][active]))) if derivative else value

    def interval_lower_bound(self, lower, upper, positive):
        """Certified lower bound for the median risk throughout an interval.

        Each singular coefficient keeps its sign and decreases in magnitude
        along either branch. Thus its endpoint values bound signal survival,
        while the upper endpoint gives the smallest possible noise energy.
        Taking the median preserves this pointwise lower-bound ordering.
        """
        if positive:
            high_coeff = self.q*(1-upper)/(self.r*(1-upper)+upper)
            low_coeff = self.q*(1-lower)/(self.r*(1-lower)+lower)
        else:
            high_coeff = self.q/(self.r+upper)
            if lower == -1:
                return float(np.median(self.variance*np.sum(high_coeff*high_coeff, axis=1)))
            low_coeff = self.q/(self.r+lower)
        low_signal, high_signal = self.b*low_coeff, self.b*high_coeff
        minimum_survival = np.sum(np.minimum(low_signal, high_signal), axis=1)
        maximum_survival = np.sum(np.maximum(low_signal, high_signal), axis=1)
        closest = np.maximum(0, np.maximum(1-maximum_survival, minimum_survival-1))
        bound = closest*closest + self.variance*np.sum(high_coeff*high_coeff, axis=1)
        return float(np.median(bound))


def branch_minimum(racks, positive, resolution=65, tolerance=5e-7):
    """Solve stationary points and median-switch cusps, using analytic slopes.

    Chebyshev brackets resolve both ends of the compact domain. Brent solves
    derivative sign changes; bounded refinement also handles a cusp where the
    median's active pair changes. Grid points bracket roots, never select a
    displayed finite-grid penalty. The positive endpoint is the exact zero rule.
    """
    unit = (1-np.cos(np.linspace(0, math.pi, resolution)))/2
    grid = unit if positive else -1+unit
    if not positive:
        # The first Gram pole is excluded mathematically. Its representable
        # neighbor is used only to bracket roots, never as an admissible result.
        grid[0] = np.nextafter(-1., 0.)
    observations = np.array([racks.median(x, positive, True) for x in grid])
    candidates = [(float(grid[-1]), float(observations[-1, 0]))]
    if positive:
        candidates.append((0., float(observations[0, 0])))
    for i in range(len(grid)-1):
        if observations[i, 1] < 0 < observations[i+1, 1]:
            root = brentq(lambda x: racks.median(x, positive, True)[1], grid[i], grid[i+1], xtol=1e-13)
            candidates.append((float(root), racks.median(root, positive)))
    for i in range(1, len(grid)-1):
        if observations[i, 0] <= min(observations[i-1, 0], observations[i+1, 0]):
            fit = minimize_scalar(lambda x: racks.median(x, positive), bounds=(grid[i-1], grid[i+1]), method="bounded", options={"xatol": 1e-13})
            candidates.append((float(fit.x), float(fit.fun)))
    best = min(candidates, key=lambda item: item[1])
    # Analytic interval bounds prevent a narrow median-switch cusp from being
    # missed between root brackets. Refine only intervals that can improve the
    # current answer by more than tolerance. This covers the entire domain and
    # provides a numerical global-optimality tolerance, rather than assuming
    # that the median curve is smooth or unimodal.
    bounds_grid = grid.copy()
    if not positive:
        bounds_grid[0] = -1.
    queue = []
    for lower, upper in zip(bounds_grid[:-1], bounds_grid[1:]):
        bound = racks.interval_lower_bound(lower, upper, positive)
        if bound < best[1]-tolerance:
            heapq.heappush(queue, (bound, float(lower), float(upper)))
    while queue:
        bound, lower, upper = heapq.heappop(queue)
        if bound >= best[1]-tolerance:
            break
        middle = (lower+upper)/2
        if middle == lower or middle == upper:
            raise ArithmeticError("Ridge tuning exhausted floating-point resolution")
        value = racks.median(middle, positive)
        if value < best[1]:
            best = (middle, value)
        for a, b in [(lower, middle), (middle, upper)]:
            child_bound = racks.interval_lower_bound(a, b, positive)
            if child_bound < best[1]-tolerance:
                heapq.heappush(queue, (child_bound, a, b))
    if not positive and best[0] <= -1:
        raise ArithmeticError("The Gram pole is not an admissible predictor")
    return best


def tune(racks, resolution=65):
    positive = branch_minimum(racks, True, resolution)
    negative = branch_minimum(racks, False, resolution)
    if positive[1] <= negative[1]:
        t, risk = positive
        return (None if t == 1 else t/(1-t)), risk
    return negative


def summarize(racks, p, resolution=65):
    positive = branch_minimum(racks, True, resolution)
    negative = branch_minimum(racks, False, resolution)
    t, positive_risk = positive
    positive_f = None if t == 1 else t/(1-t)
    f = positive_f if positive_risk <= negative[1] else negative[0]
    interpolation_risk, interpolation_train = racks.values(0)
    if f is None:
        ridge_risk, ridge_train = racks.values(1, True)
    else:
        ridge_risk, ridge_train = racks.values(f)
    positive_values, positive_train = racks.values(t, True)
    if not all(np.isfinite(values).all() for values in (ridge_risk, ridge_train, positive_values, positive_train)):
        raise ArithmeticError("Nonfinite risk; abort, never filter a rack")
    return dict(lo=float(np.quantile(interpolation_risk, .25)), interp=float(np.median(interpolation_risk)),
                hi=float(np.quantile(interpolation_risk, .75)), ridge=float(np.median(ridge_risk)),
                lambda_=f, interpTrain=0. if p >= racks.n else float(np.median(interpolation_train)),
                ridgeTrain=0. if f == 0 and p >= racks.n else float(np.median(ridge_train)), positiveRidge=float(np.median(positive_values)),
                positiveLambda=positive_f, positiveRidgeTrain=0. if positive_f == 0 and p >= racks.n else float(np.median(positive_train)))


def training_ratios(racks, p, signed_f, positive_f):
    """Take ratios within each rack, then summarize; never divide medians."""
    baseline = racks.label_norm2 / racks.n
    def median_ratio(f):
        _, train = racks.values(1, True) if f is None else racks.values(f)
        return 0. if f == 0 and p >= racks.n else float(np.median(train / baseline))
    return dict(interpTrainRatio=median_ratio(0), ridgeTrainRatio=median_ratio(signed_f),
                positiveRidgeTrainRatio=median_ratio(positive_f))


def preset(key, count=RACKS, widths=WIDTHS, resolution=65, progress=True, existing=None):
    """Nested sensor columns, same ordinary Gaussian labels at every width.

    Rotation invariance lets a=(1,...,1), whose squared norm is exactly p. Each
    new sensor contributes an independent Gaussian noise column; all existing
    columns remain fixed. Racks are independent and targets are not normalized.
    """
    n, variance = PRESETS[key]
    rng = np.random.Generator(np.random.PCG64(seed_for(key)))
    labels = rng.standard_normal((count, n))
    label_norm2 = np.sum(labels*labels, axis=1)
    noise = rng.standard_normal((max(widths), count, n))
    columns = labels[None, :, :] + math.sqrt(variance)*noise
    gram = np.zeros((count, n, n))
    response_projection = np.zeros((count, n))
    rows = []
    last = 0
    started = time.perf_counter()
    for ip, p in enumerate(widths):
        block = np.transpose(columns[last:p], (1, 2, 0))
        gram += block @ np.swapaxes(block, 1, 2)
        response_projection += np.sum(block, axis=2)
        if p <= n:
            X = np.transpose(columns[:p], (1, 2, 0))
            u, singular, vt = np.linalg.svd(X, full_matrices=False)
            d = singular*singular
            # Sort ascending, matching the Gram eigensystem in wider racks.
            d, u, vt = d[:, ::-1], u[:, :, ::-1], vt[:, ::-1, :]
            z = np.einsum("rik,ri->rk", u, labels)
            b = np.sum(vt, axis=2)
        else:
            d, u = np.linalg.eigh(gram)
            z = np.einsum("rik,ri->rk", u, labels)
            b = np.einsum("rik,ri->rk", u, response_projection)/np.sqrt(d)
        spectral = SpectralRacks(d, z, b, label_norm2, n, variance)
        if existing is None:
            row = summarize(spectral, p, resolution)
            row.update(training_ratios(spectral, p, row["lambda_"], row["positiveLambda"]))
        else:
            # Reuse the stored penalties: this migration never retunes or changes
            # the raw summaries. Verify the replay before adding ratio fields.
            for field, f in [("interpTrain", 0), ("ridgeTrain", existing["lambda"][ip]),
                             ("positiveRidgeTrain", existing["positiveLambda"][ip])]:
                _, train = spectral.values(1, True) if f is None else spectral.values(f)
                raw = 0. if f == 0 and p >= n else float(np.median(train))
                if not np.isclose(raw, existing[field][ip], rtol=3e-6, atol=2e-8):
                    raise ArithmeticError(f"Replay mismatch at {key}, p={p}, {field}: {raw} != {existing[field][ip]}")
            row = training_ratios(spectral, p, existing["lambda"][ip], existing["positiveLambda"][ip])
        rows.append(row)
        last = p
        if progress and (ip % 25 == 0 or ip == len(widths)-1):
            print(f"{key} p={p}: {ip+1}/{len(widths)} widths; {time.perf_counter()-started:.1f}s", flush=True)
    out = dict(n=n, variance=variance, p=list(widths), racks=count)
    for field in rows[0]:
        out["lambda" if field == "lambda_" else field] = [None if row[field] is None else float(f"{row[field]:.9g}") for row in rows]
    return out


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--racks", type=int, default=RACKS)
    parser.add_argument("--resolution", type=int, default=65)
    parser.add_argument("--preset", choices=PRESETS, action="append")
    parser.add_argument("--add-training-ratios", action="store_true",
                        help="Replay the saved fixture and add per-rack training ratios, preserving all existing fields")
    args = parser.parse_args()
    if args.add_training_ratios:
        source = args.output.read_text()
        data = json.loads(source.split("const data=", 1)[1].split(";\n", 1)[0])
        metadata = json.loads(source.split("{value:", 1)[1].split(",enumerable:false", 1)[0])
        for key, cell in data.items():
            replay = preset(key, cell["racks"], cell["p"], existing=cell)
            for field in ["interpTrainRatio", "ridgeTrainRatio", "positiveRidgeTrainRatio"]:
                cell[field] = replay[field]
        metadata["trainingRatios"] = "Median of each rack's training MSE divided by mean(y^2) on that same rack; raw fields and stored penalties preserved"
        metadata["trainingRatioGeneratorSha256"] = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
    else:
        data = {key: preset(key, args.racks, resolution=args.resolution) for key in args.preset or PRESETS}
        metadata = dict(schema=4, masterSeed=MASTER_SEED, racks=args.racks,
                    model="X=y a^T+sqrt(v) Z; ordinary Gaussian y and Z; ||a||^2=p",
                    pairing="Nested sensor columns and fixed Gaussian labels within each independent rack",
                    tuning="Common Gram-minimum multiplier minimizing median analytic physical test risk, stored separately for positive-only and signed ridge; stationary points and median-switch cusps; analytic interval bounds certify risk to5e-7; full positive branch includes zero predictor at +infinity; negative branch excludes first pole",
                    numericalTolerance=5e-7,
                    seed="first eight SHA256 bytes, little-endian, of sensor-ridge-v2:master:preset; NumPy PCG64",
                    retained="All racks; no rejection, filtering, replacement, clipping or seed search",
                    quantiles="NumPy linear quartiles and median; nine significant digits",
                    trainingRatios="Median of each rack's training MSE divided by mean(y^2) on that same rack; raw fields retained",
                    generatorSha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest())
    source = "(function(root){\"use strict\";const data=" + json.dumps(data, separators=(",", ":"), allow_nan=False) + ";\n"
    source += "Object.defineProperty(data,\"provenance\",{value:" + json.dumps(metadata, separators=(",", ":")) + ",enumerable:false});\n"
    source += "if(typeof module!==\"undefined\"&&module.exports)module.exports=data;if(root)root.SENSOR_RIDGE_SWEEPS=data;})(typeof window!==\"undefined\"?window:globalThis);\n"
    args.output.write_text(source)
    print(f"Wrote {args.output}: {args.output.stat().st_size:,} bytes", flush=True)


if __name__ == "__main__":
    main()