"""MVP verification of robust parameter-update envelopes.

The model is f(x) = (a b)x, a two-layer scalar linear neural network.
At (a,b)=(1,4), the proposed update to (4,1) has zero endpoint loss for
x=1,y=4, but its interpolation has a loss peak at the midpoint.
"""
import json
import math
import random
from pathlib import Path
import numpy as np

SEED = 584
random.seed(SEED)
np.random.seed(SEED)


def scalar_loss(a, b, target=4.0):
    return (a * b - target) ** 2


def path_values(theta0, theta1, grid):
    a0, b0 = theta0
    a1, b1 = theta1
    a = a0 + grid * (a1 - a0)
    b = b0 + grid * (b1 - b0)
    output = a * b
    loss = (output - 4.0) ** 2
    return output, loss


def envelope_accept(theta0, theta1, loss_budget, K):
    grid = np.linspace(0.0, 1.0, K + 1)
    _, losses = path_values(theta0, theta1, grid)
    violation = float(np.max(losses - loss_budget))
    return violation <= 1e-12, violation, float(np.max(losses)), grid[int(np.argmax(losses))]


def largest_feasible_fraction(theta0, theta1, loss_budget, K=16):
    # Scale the proposed direction, then use bisection on the accepted fraction.
    def ok(frac):
        end = np.asarray(theta0) + frac * (np.asarray(theta1) - theta0)
        return envelope_accept(theta0, end, loss_budget, K)[0]
    lo, hi = 0.0, 1.0
    if ok(hi):
        return 1.0
    for _ in range(45):
        mid = (lo + hi) / 2.0
        if ok(mid):
            lo = mid
        else:
            hi = mid
    return lo


def run():
    theta0 = np.array([1.0, 4.0])
    theta1 = np.array([4.0, 1.0])
    budget = scalar_loss(*theta0) + 1e-10  # endpoint-only loss-growth budget

    # Core numerical claim, with a dense reference grid.
    dense = np.linspace(0.0, 1.0, 100001)
    outputs, losses = path_values(theta0, theta1, dense)
    dense_peak = float(np.max(losses))
    peak_s = float(dense[int(np.argmax(losses))])
    endpoint_max = max(float(losses[0]), float(losses[-1]))
    core = {
        "start_loss": float(losses[0]), "end_loss": float(losses[-1]),
        "endpoint_max_loss": endpoint_max, "dense_max_loss": dense_peak,
        "dense_peak_s": peak_s, "interior_violation": dense_peak - budget,
        "analytic_midpoint_output": float(outputs[len(outputs)//2]),
    }

    # Compare endpoint acceptance against the envelope rule at several K values.
    comparisons = {}
    for K in (4, 8, 16, 32):
        accepted, violation, peak, where = envelope_accept(theta0, theta1, budget, K)
        comparisons[str(K)] = {
            "endpoint_only_accepts": True,
            "envelope_accepts": bool(accepted),
            "sampled_max_violation": violation,
            "sampled_max_loss": peak,
            "sampled_peak_s": where,
        }

    # An acceptance-transition measurement: robust halving finds the largest
    # fraction of this direction whose entire sampled path is feasible.
    thresholds = {str(K): largest_feasible_fraction(theta0, theta1, budget, K)
                  for K in (8, 16, 32)}

    # A small repeated-update proxy: each trial has the same candidate update.
    # This isolates the acceptance mechanism rather than confounding it with
    # a changing optimizer state.
    trials = 50
    endpoint_accepted = sum(1 for _ in range(trials)
                            if scalar_loss(*theta1) <= budget)
    robust_accepted = sum(1 for _ in range(trials)
                          if envelope_accept(theta0, theta1, budget, 16)[0])

    result = {
        "seed": SEED,
        "model": "two-layer scalar linear network f(x)=(a*b)x, x=1, target=4",
        "constraint": "loss(theta(s)) <= loss(theta(0)) + 1e-10",
        "core_check": core,
        "K_comparison": comparisons,
        "largest_feasible_fraction": thresholds,
        "50_identical_candidate_trials": {
            "endpoint_only_accepted": endpoint_accepted,
            "robust_envelope_accepted": robust_accepted,
        },
        "interpretation": "Endpoint rule accepts the weight swap, while every envelope grid detects its interior loss spike; K=8,16,32 agree on rejection and similar halving threshold.",
    }
    out = Path("robust_envelope_results.json")
    out.write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    run()