"""Small endpoint-vs-envelope optimizer comparison.

Uses a two-layer MLP on a fixed synthetic regression batch. Both methods form
exactly the same gradient proposal; the baseline checks constraints only at
the endpoint, while the envelope method scans interpolation points and halves
until all sampled points are feasible.
"""
import json
import random
from pathlib import Path
import numpy as np
import torch

SEED = 584
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
try:
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(SEED)
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
except Exception:
    DEVICE = torch.device("cpu")


def make_data():
    g = torch.Generator(device="cpu").manual_seed(SEED)
    x = torch.randn(96, 2, generator=g)
    y = (torch.sin(2.0 * x[:, :1]) + 0.5 * x[:, 1:2] ** 2).float()
    return x.to(DEVICE), y.to(DEVICE)


class MLP(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(2, 16)
        self.l2 = torch.nn.Linear(16, 1)

    def forward(self, x):
        h = torch.tanh(self.l1(x))
        return self.l2(h), h


def flat_params(model):
    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])


def set_flat(model, flat):
    pos = 0
    with torch.no_grad():
        for p in model.parameters():
            n = p.numel()
            p.copy_(flat[pos:pos+n].reshape_as(p))
            pos += n


def evaluate(model, x, y, flat=None):
    old = flat_params(model)
    if flat is not None:
        set_flat(model, flat)
    with torch.no_grad():
        pred, h = model(x)
        loss = torch.mean((pred - y) ** 2)
        rms = torch.sqrt(torch.mean(h * h) + 1e-12)
    set_flat(model, old)
    return float(loss.cpu()), float(rms.cpu())


def constraints(model, x, y, flat, loss_limit, rms_limit):
    loss, rms = evaluate(model, x, y, flat)
    return max(loss - loss_limit, rms - rms_limit), loss, rms


def train(mode, lr, steps=45, K=8):
    torch.manual_seed(SEED)
    model = MLP().to(DEVICE)
    x, y = make_data()
    initial = flat_params(model)
    base_loss, base_rms = evaluate(model, x, y, initial)
    # Small safety margins make this a loss-growth/activation envelope,
    # rather than requiring the current point to be exactly on a boundary.
    loss_margin, rms_margin = 0.015, 0.10
    accepted, violations, losses, step_sizes = 0, [], [], []
    for _ in range(steps):
        model.zero_grad(set_to_none=True)
        pred, _ = model(x)
        loss = torch.mean((pred - y) ** 2)
        loss.backward()
        grad = torch.cat([p.grad.detach().reshape(-1) for p in model.parameters()])
        current = flat_params(model)
        proposal = current - lr * grad
        cur_loss, cur_rms = evaluate(model, x, y, current)
        loss_limit = cur_loss + loss_margin
        rms_limit = max(base_rms + rms_margin, cur_rms + rms_margin)

        def scan(scale):
            vals = []
            for s in torch.linspace(0, 1, K + 1, device=DEVICE):
                v, pl, pr = constraints(model, x, y, current + s * scale * (proposal-current), loss_limit, rms_limit)
                vals.append((v, pl, pr))
            return max(z[0] for z in vals), max(z[1] for z in vals), max(z[2] for z in vals)

        endpoint_v, endpoint_peak_loss, endpoint_peak_rms = scan(1.0) if mode == "envelope" else constraints(model, x, y, proposal, loss_limit, rms_limit)
        scale = 1.0
        if mode == "endpoint":
            accepted_now = endpoint_v <= 1e-8
            observed_v, peak_loss, peak_rms = endpoint_v, endpoint_peak_loss if isinstance(endpoint_peak_loss, float) else endpoint_v, endpoint_peak_rms if isinstance(endpoint_peak_rms, float) else 0.0
        else:
            observed_v, peak_loss, peak_rms = scan(scale)
            while observed_v > 1e-8 and scale > 2.0 ** -12:
                scale *= 0.5
                observed_v, peak_loss, peak_rms = scan(scale)
            accepted_now = observed_v <= 1e-8
        if accepted_now:
            set_flat(model, current + scale * (proposal-current))
            accepted += 1
        step_sizes.append(scale if accepted_now else 0.0)
        violations.append(float(observed_v))
        losses.append(evaluate(model, x, y)[0])
    return {
        "mode": mode, "lr": lr, "K": K, "final_loss": losses[-1],
        "best_loss": min(losses), "accepted_steps": accepted,
        "mean_accepted_fraction": float(np.mean([z for z in step_sizes if z > 0])) if accepted else 0.0,
        "max_recorded_constraint_violation": max(violations),
        "positive_violation_steps": sum(v > 1e-7 for v in violations),
        "loss_curve": losses,
    }


def main():
    results = []
    for lr in (0.25, 0.5, 1.0, 2.0):
        results += [train("endpoint", lr), train("envelope", lr, K=8), train("envelope", lr, K=16)]
    out = {"seed": SEED, "device": str(DEVICE), "results": results}
    Path("mlp_comparison_results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()