import itertools, json
import numpy as np

SEED = 123
ORDERS = list(itertools.permutations(range(4)))

def rho(a):
    return float(np.max(np.abs(np.linalg.eigvals(a))))

def block_maps(lam, gamma):
    # A local error model: each update changes one coordinate block. The directed
    # sensitivities are deliberately nonsymmetric, as neural block Jacobians need not
    # be symmetric. At lam=0 all blocks are decoupled and order must be irrelevant.
    c = np.array([[0, 1.0, -.6, .3], [-.8, 0, .9, -.4],
                  [.5, -.7, 0, 1.1], [-.3, .8, -.9, 0]], dtype=float)
    out = []
    for a in range(4):
        t = np.eye(4)
        t[a, :] -= gamma * (np.eye(4)[a, :] + lam * c[a, :])
        out.append(t)
    return out

def composed(lam, gamma, order):
    j = np.eye(4)
    for b in order:                         # order[0] acts first
        j = block_maps(lam, gamma)[b] @ j
    return j

def metrics(lam, gamma):
    vals = np.array([rho(composed(lam, gamma, o)) for o in ORDERS])
    return vals

def observed_ratio(j, e0=None, n=100):
    e = np.ones(j.shape[0]) if e0 is None else np.asarray(e0, float).copy()
    rs = []
    for _ in range(n):
        en = j @ e
        rs.append(np.linalg.norm(en) / np.linalg.norm(e))
        e = en
    return float(np.mean(rs[-20:]))

def crossing(lam, order, criterion, lo=.001, hi=2.0):
    # Find the gamma boundary predicted by rho(J)=1. For this family rho is
    # monotonic over the relevant interval; return None if no crossing exists.
    f = lambda g: criterion(composed(lam, g, order)) - 1.0
    if f(lo) >= 0: return lo
    if f(hi) < 0: return None
    for _ in range(60):
        mid = (lo + hi) / 2
        if f(mid) < 0: lo = mid
        else: hi = mid
    return (lo + hi) / 2

def toy_verification():
    # Prediction 1: at zero coupling, all order radii coincide exactly.
    # Prediction 2: increasing directed coupling increases order spread.
    coupling = []
    for lam in [0, .1, .2, .3, .4, .5, .6, .7, .8]:
        vals = metrics(lam, .42)
        coupling.append({'lambda': lam, 'best_rho': float(vals.min()),
                         'worst_rho': float(vals.max()),
                         'order_gap': float(vals.max()-vals.min()),
                         'best_order': list(ORDERS[int(vals.argmin())])})
    # Prediction 3: rho=1 is the stability boundary; compare predicted crossing
    # against direct finite-trajectory growth classification.
    stability = []
    for lam in [.2, .4, .6, .8]:
        vals = metrics(lam, .42)
        order = ORDERS[int(vals.argmin())]
        pred = crossing(lam, order, rho)
        test_gammas = np.linspace(.05, 1.5, 30)
        observed = None
        for g in test_gammas:
            j = composed(lam, g, order)
            e = np.ones(4)
            for _ in range(80): e = j @ e
            if np.linalg.norm(e) > 1.05: # finite-run growth, empirical boundary
                observed = float(g); break
        stability.append({'lambda': lam, 'order': list(order),
                          'predicted_rho1_gamma': pred,
                          'observed_first_growth_gamma': observed})
    # Direct asymptotic-ratio check for the selected composition.
    lam, gamma = .5, .42
    vals = metrics(lam, gamma); order = ORDERS[int(vals.argmin())]
    j = composed(lam, gamma, order)
    ratio = {'lambda': lam, 'gamma': gamma, 'order': list(order),
             'predicted_rho': rho(j),
             'observed_tail_ratio': observed_ratio(j, [.8, -.4, .6, 1.1])}
    return {'coupling_sweep': coupling, 'stability_sweep': stability,
            'ratio_check': ratio}

def neural_experiment():
    import torch
    torch.manual_seed(SEED); np.random.seed(SEED)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        x = torch.linspace(-2, 2, 256, device=device).unsqueeze(1)
        y = torch.sin(3*x) + .15*x*x
        def run(mode):
            torch.manual_seed(SEED)
            net = torch.nn.Sequential(torch.nn.Linear(1,16), torch.nn.Tanh(),
                                      torch.nn.Linear(16,1)).to(device)
            blocks = [[net[0].weight, net[0].bias], [net[2].weight, net[2].bias]]
            order = [0, 1]; losses=[]; orders=[]; lr=.035
            for step in range(180):
                if mode == 'adaptive' and step % 15 == 0:
                    # A cheap empirical cross-block proxy: test each candidate by
                    # one block update and score the resulting loss.
                    scores=[]
                    for candidate in ([0,1],[1,0]):
                        saved=[p.detach().clone() for p in net.parameters()]
                        bi=candidate[0]; net.zero_grad(); q=((net(x)-y)**2).mean(); q.backward()
                        with torch.no_grad():
                            for p in blocks[bi]: p -= lr*p.grad
                        scores.append(float(((net(x)-y)**2).mean()))
                        with torch.no_grad():
                            for p,s in zip(net.parameters(),saved): p.copy_(s)
                    order = [0,1] if scores[0] <= scores[1] else [1,0]
                if mode == 'sim':
                    net.zero_grad(); loss=((net(x)-y)**2).mean(); loss.backward()
                    with torch.no_grad():
                        for p in net.parameters(): p -= lr*p.grad
                else:
                    for bi in order:
                        net.zero_grad(); loss=((net(x)-y)**2).mean(); loss.backward()
                        with torch.no_grad():
                            for p in blocks[bi]: p -= lr*p.grad
                orders.append(tuple(order)); losses.append(float(((net(x)-y)**2).mean()))
            return losses, orders
        out={m:run(m) for m in ['sim','fixed','adaptive']}
        return {'device':device, 'final_loss':{k:v[0][-1] for k,v in out.items()},
                'loss_at_30':{k:v[0][29] for k,v in out.items()},
                'adaptive_order_counts':{str(o):out['adaptive'][1].count(o) for o in [(0,1),(1,0)]}}
    except Exception as e:
        return {'device':'cpu-fallback','error':repr(e)}

def main():
    result={'seed':SEED, 'toy':toy_verification(), 'neural':neural_experiment()}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__ == '__main__': main()