import json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

SEED = 1380
np.random.seed(SEED)
random.seed(SEED)
torch.manual_seed(SEED)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception:
    device = torch.device('cpu')

# Scalar residual map h_next=(1+a)h. Its exact finite-time Lyapunov rate is log|1+a|.
def lyapunov_autodiff(a, L=12, dim=8):
    h = torch.zeros(dim, device=device, requires_grad=True)
    v = torch.ones(dim, device=device)
    v = v / torch.linalg.vector_norm(v)
    rates = []
    for _ in range(L):
        def f(x):
            return x + a * x
        _, jv = torch.func.jvp(f, (h,), (v,))
        gain = float(torch.linalg.vector_norm(jv).detach().cpu())
        rates.append(math.log(max(gain, 1e-30)))
        v = jv / max(gain, 1e-30)
        h = (h + a * h).detach().requires_grad_(True)
    return float(np.mean(rates))

def controller(rhat, target=0.0, kappa=0.7, dt=0.1, qmax=2.0):
    return float(np.clip(2*kappa*max(rhat-target, 0.0)*dt, 0, qmax))

def mechanism_sweeps():
    # A: rhat=log|1+a| and the boundary rhat=0 occurs at a=0.
    avals = np.array([-0.45, -0.35, -0.25, -0.15, -0.05, 0.0, 0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65])
    boundary = []
    for a in avals:
        pred = math.log(abs(1+a))
        obs = lyapunov_autodiff(float(a))
        boundary.append({'a': float(a), 'pred_r': pred, 'observed_r': obs,
                         'abs_error': abs(pred-obs)})

    # B: q=2*kappa*dt*[r-r_target]_+, so fitted positive-side slope is predicted.
    kappa, dt, target = 0.7, 0.1, 0.0
    rs = np.linspace(-0.4, 0.8, 25)
    qrows = [{'r': float(r), 'q_pred': controller(float(r), target, kappa, dt)} for r in rs]
    positive = [x for x in qrows if x['r'] > target]
    slope = float(np.polyfit([x['r'] for x in positive], [x['q_pred'] for x in positive], 1)[0])

    # C: directly sample gate noise; conditional variance should be q*u*(1-u/K).
    rng = np.random.default_rng(SEED)
    K, uvals = 1.0, [0.2, 0.5, 0.8]
    samples = []
    for r in [0.1, 0.3, 0.6]:
        q = controller(r, target, kappa, dt)
        for u in uvals:
            z = 1.0 + math.sqrt(q*u*(1-u/K))*rng.standard_normal(200000)
            samples.append({'r': r, 'u': u, 'q': q, 'emp_var': float(np.var(z, ddof=1)),
                            'pred_var': q*u*(1-u/K)})
    x = np.array([s['pred_var'] for s in samples])
    y = np.array([s['emp_var'] for s in samples])
    var_slope = float(np.dot(x, y) / np.dot(x, x))
    return {'boundary': boundary, 'q_sweep': qrows, 'q_slope': slope,
            'q_slope_pred': 2*kappa*dt, 'gate_variance': samples,
            'variance_slope': var_slope}

class ToyNet(nn.Module):
    def __init__(self, adaptive=False, fixed_q=0.08, depth=8):
        super().__init__()
        self.adaptive, self.fixed_q, self.depth = adaptive, fixed_q, depth
        self.a = nn.Parameter(torch.tensor(0.18))
        self.readout = nn.Linear(1, 1)
    def forward(self, x, return_r=False):
        h = x
        rates = []
        for _ in range(self.depth):
            r = torch.log(torch.abs(1 + self.a) + 1e-8)
            rates.append(r)
            if self.adaptive:
                q = torch.clamp(2*0.7*torch.relu(r.detach())*0.1, 0, 0.5)
            else:
                q = torch.tensor(self.fixed_q, device=x.device)
            eps = torch.randn_like(h)
            z = 1 + torch.sqrt(q * 0.5 * (1-0.5)) * eps
            h = h + z * self.a * h
        out = self.readout(h)
        return (out, torch.stack(rates).mean()) if return_r else out

def mini_training():
    g = torch.Generator(device=device); g.manual_seed(SEED)
    x = torch.linspace(-1, 1, 96, device=device).reshape(-1, 1)
    y = 0.7*x + 0.1*torch.sin(5*x)
    results = {}
    for name, adaptive in [('fixed_noise', False), ('lyapunov_noise', True)]:
        torch.manual_seed(SEED)
        model = ToyNet(adaptive=adaptive).to(device)
        opt = torch.optim.Adam(model.parameters(), lr=0.025)
        losses, rates, qs = [], [], []
        for step in range(180):
            opt.zero_grad(set_to_none=True)
            pred, r = model(x, return_r=True)
            loss = ((pred-y)**2).mean()
            loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
            losses.append(float(loss.detach().cpu())); rates.append(float(r.detach().cpu()))
            qs.append(controller(rates[-1]) if adaptive else 0.08)
        results[name] = {'final_loss': losses[-1], 'best_loss': min(losses),
                         'initial_loss': losses[0], 'mean_r': float(np.mean(rates)),
                         'mean_q': float(np.mean(qs))}
    return results

def main():
    out = {'device': str(device), 'mechanism': mechanism_sweeps(), 'training': mini_training()}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps({'device': out['device'], 'q_slope': out['mechanism']['q_slope'],
                      'q_slope_pred': out['mechanism']['q_slope_pred'],
                      'variance_slope': out['mechanism']['variance_slope'],
                      'training': out['training']}, indent=2))

if __name__ == '__main__':
    main()