import json, math, os
import numpy as np
from scipy.linalg import eigvals, expm, svdvals


def norm2(a):
    return float(svdvals(a)[0])


def make_case(p=32, n=128, seed=0):
    rng = np.random.default_rng(seed)
    X = rng.normal(size=(n, p))
    H = X.T @ X / n
    K = rng.normal(size=(p, p))
    np.fill_diagonal(K, 0.0)
    # explicitly non-reciprocal, with a substantial antisymmetric component
    K = K - K.T * 0.35 + 0.35 * (K + K.T) / 2
    K /= norm2(K)
    return H, K


def susceptibility(J, gamma=0.02):
    p = J.shape[0]
    R = np.linalg.inv(gamma * np.eye(p) - J)
    return float(np.mean(np.abs(R) ** 2))


def transient_gain(J, times):
    return max(norm2(expm(J * t)) for t in times)


def run_mechanism():
    p, n = 32, 128
    q = p / n
    mp_delta = (1 - math.sqrt(q)) ** 2
    H, K = make_case(p, n, 7)
    observed_delta = float(np.min(np.linalg.eigvalsh(H)))
    k_norm = norm2(K)
    pred_ec = mp_delta / k_norm
    empirical_ec = observed_delta / k_norm
    eps_grid = np.linspace(0, 1.5, 31)
    rows = []
    for eps in eps_grid:
        J = -H + eps * K
        ev = eigvals(J)
        rows.append({
            'eps': float(eps),
            'ratio': float(eps * k_norm / observed_delta),
            'chi': susceptibility(J),
            'max_real_eig': float(np.max(ev.real)),
            'gain': float(transient_gain(J, np.linspace(0, 20, 41)))
        })
    # For the claimed crossover, report the first point whose chi differs
    # from passive chi by 25 percent (either suppression or amplification).
    chi0 = rows[0]['chi']
    changed = [r for r in rows if abs(r['chi'] / chi0 - 1) >= .25]
    observed_chi_cross = changed[0]['eps'] if changed else float('nan')
    # instability boundary is a separate falsifiable prediction: eigenvalues
    # should remain stable for purely dissipative perturbations, but not in general.
    unstable = [r for r in rows if r['max_real_eig'] >= 0]
    instability_eps = unstable[0]['eps'] if unstable else float('nan')
    return {
        'q': q, 'mp_delta': mp_delta, 'observed_delta': observed_delta,
        'k_norm': k_norm, 'predicted_ec_mp': pred_ec,
        'predicted_ec_empirical_gap': empirical_ec,
        'observed_chi_25pct_eps': observed_chi_cross,
        'observed_instability_eps': instability_eps,
        'rows': rows
    }


def train_linear_demo(seed=11, steps=500, penalty=False):
    # Same tiny supervised long-horizon state-space problem for both methods.
    rng = np.random.default_rng(seed)
    p, T, batch = 12, 35, 64
    torch = __import__('torch')
    torch.manual_seed(seed)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        dtype = torch.float32
        Atrue = torch.tensor(rng.normal(size=(p,p)), dtype=dtype, device=device)
        Atrue = Atrue / (1.15 * torch.linalg.matrix_norm(Atrue))
        Btrue = torch.tensor(rng.normal(size=(p,1)), dtype=dtype, device=device)
        W = torch.randn(p,p, device=device) * .08
        U = torch.randn(p,1, device=device) * .08
        W.requires_grad_(); U.requires_grad_()
        opt = torch.optim.Adam([W,U], lr=.025)
        last_loss = None
        for step in range(steps):
            x = torch.randn(batch,T,1,device=device)
            with torch.no_grad():
                h = torch.zeros(batch,p,device=device)
                ys=[]
                for t in range(T):
                    h = torch.tanh(h @ Atrue.T + x[:,t] @ Btrue.T)
                    ys.append(h)
                y = torch.stack(ys,1)
            h = torch.zeros(batch,p,device=device)
            pred=[]
            for t in range(T):
                h = torch.tanh(h @ W.T + x[:,t] @ U.T)
                pred.append(h)
            pred=torch.stack(pred,1)
            loss=((pred-y)**2).mean()
            if penalty:
                # Exact small-batch local resolvent proxy at a fixed damping.
                # Use W's symmetric Gram passive estimate and residual activity.
                H = (W.T @ W) / p
                J = W
                R = torch.linalg.inv(.08*torch.eye(p,device=device)-J)
                chi=(R*R).mean()
                loss = loss + .002 * torch.relu(torch.log1p(chi)-torch.tensor(math.log(3.0),device=device))**2
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_([W,U], 5.0); opt.step()
            last_loss=float(loss.detach().cpu())
        with torch.no_grad():
            ev=torch.linalg.eigvals(W).real.max().item()
            # sampled transient gain using torch matrix exponentials
            gain=max(torch.linalg.matrix_norm(torch.matrix_exp(W*t),2).item() for t in np.linspace(0,8,17))
            R=torch.linalg.inv(.08*torch.eye(p,device=device)-W)
            chi=(R*R).mean().item()
        return {'loss':last_loss,'max_real_eig':ev,'gain':gain,'chi':chi,'device':device}
    except Exception as exc:
        # CUDA failures and numerical backend failures fall back to CPU.
        if device == 'cuda':
            os.environ['CUDA_VISIBLE_DEVICES']=''
            return train_linear_demo(seed, steps, penalty)
        raise exc


if __name__ == '__main__':
    out={'mechanism':run_mechanism(), 'baseline':train_linear_demo(penalty=False), 'idea':train_linear_demo(penalty=True)}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    m=out['mechanism']
    print(json.dumps({'mechanism_summary':{k:v for k,v in m.items() if k!='rows'}, 'baseline':out['baseline'],'idea':out['idea']},indent=2))