import json, math, random
from pathlib import Path
import numpy as np

SEED = 2518
np.random.seed(SEED); random.seed(SEED)


def cocycle_run(l1, l2, K=200, noise=0.20, theta0=None):
    """Random diagonal positive cocycle; log singular increments have known means."""
    if theta0 is None:
        theta0 = np.random.uniform(0.25, 1.25, size=2)
    # tangent vectors are represented by slopes y/x; projective distance is angle.
    slopes = np.tan(theta0)
    logs = np.zeros(2)
    angles = [theta0.copy()]
    q = np.eye(2)
    qr_diag = []
    for t in range(K):
        # independent temporal randomness, but the same scalar derivative applies to both vectors
        inc = np.array([l1, l2]) + noise*np.random.randn(2)
        A = np.diag(np.exp(inc))
        # QR iteration estimates both exponents
        z = A @ q
        q, R = np.linalg.qr(z)
        # enforce positive diagonal, without changing Q directions
        signs = np.sign(np.diag(R)); signs[signs == 0] = 1
        q = q @ np.diag(signs); R = np.diag(signs) @ R
        qr_diag.append(np.log(np.abs(np.diag(R))))
        logs += inc
        slopes *= np.exp(inc[1]-inc[0])
        # angles to the top axis (modulo orientation)
        angles.append(np.arctan(np.abs(slopes)))
    return np.asarray(angles), np.asarray(qr_diag), logs/K


def fit_slope(y, start=20, end=None):
    if end is None: end = len(y)
    x = np.arange(start, end)
    return float(np.polyfit(x, np.log(np.maximum(y[start:end], 1e-300)), 1)[0])


def toy_verification():
    # Prediction 1: projective log-angle slope is -gamma.
    gaps = [0.02, 0.05, 0.10, 0.20, 0.40]
    rows = []
    for g in gaps:
        # average many paths to suppress finite-sample noise
        slopes = []; estimates = []
        for _ in range(80):
            a, qr, means = cocycle_run(0.0, -g, K=180, noise=0.20)
            slopes.append(fit_slope(a[:,0], 30, 170))
            estimates.append(np.mean(qr[:,0]-qr[:,1]))
        obs_s = float(np.mean(slopes)); obs_g = float(np.mean(estimates))
        rows.append({'gamma_true':g, 'predicted_slope':-g, 'observed_slope':obs_s,
                     'predicted_gap':g, 'observed_qr_gap':obs_g,
                     'slope_rel_error':abs(obs_s+g)/g})

    # Prediction 2: time to reach angle epsilon scales log(theta0/epsilon)/gamma.
    threshold_rows = []
    eps = 1e-3
    for g in [0.05, 0.10, 0.20, 0.40]:
        measured=[]; predicted=[]
        for _ in range(100):
            a, _, _ = cocycle_run(0.0, -g, K=500, noise=0.20,
                                  theta0=np.array([0.8, 0.8]))
            hit = np.flatnonzero(a[:,0] <= eps)
            measured.append(int(hit[0]) if len(hit) else 500)
            predicted.append(math.log(math.tan(0.8)/eps)/g)
        threshold_rows.append({'gamma':g, 'predicted_steps':float(np.mean(predicted)),
                              'observed_steps':float(np.mean(measured)),
                              'ratio_observed_predicted':float(np.mean(measured)/np.mean(predicted))})

    # Prediction 3: no gap means no projective contraction; with exactly equal rates,
    # the angle remains constant even though the common exponent can be negative.
    zero=[]
    for l in [0.0, -0.2]:
        a, qr, means = cocycle_run(l, l, K=150, noise=0.20)
        zero.append({'common_lambda':l, 'angle_log_slope':fit_slope(a[:,0], 10, 140),
                     'qr_gap':float(np.mean(qr[:,0]-qr[:,1])),
                     'angle_ratio_final_initial':float(a[-1,0]/a[0,0])})
    return {'gap_sweep':rows, 'threshold_sweep':threshold_rows, 'zero_gap':zero}


# Optional practical check: train a tiny tanh RNN, regularizing a finite-window
# Jacobian gap estimated by explicit 2x2 Jacobians (small enough to be transparent).
def rnn_check(steps=180, seq_len=40, batch=32):
    try:
        import torch
        torch.manual_seed(SEED)
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        class RNN(torch.nn.Module):
            def __init__(self):
                super().__init__(); self.W=torch.nn.Parameter(torch.randn(12,12)*.18)
                self.U=torch.nn.Parameter(torch.randn(12,2)*.18); self.b=torch.nn.Parameter(torch.zeros(12))
                self.out=torch.nn.Linear(12,1)
            def forward(self,x, regularize=False):
                z=torch.zeros(x.shape[0],12,device=x.device); gs=[]; l1s=[]
                for t in range(x.shape[1]):
                    z=torch.tanh(z@self.W.T+x[:,t]@self.U.T+self.b)
                    # Exact Jacobian singular values for the hidden map at this state.
                    if regularize:
                        D=1-z*z
                        J=D.unsqueeze(2)*self.W.unsqueeze(0)
                        s=torch.linalg.svdvals(J)
                        l1s.append(torch.log(s[:,0].clamp_min(1e-8)))
                        gs.append(torch.log((s[:,0]/s[:,1]).clamp_min(1e-8)))
                return self.out(z), (torch.stack(l1s,1).mean() if regularize else None), (torch.stack(gs,1).mean() if regularize else None)
        def run(reg):
            torch.manual_seed(SEED+int(reg)); m=RNN().to(device); opt=torch.optim.Adam(m.parameters(),lr=2e-3)
            losses=[]; gaps=[]; exps=[]
            for k in range(steps):
                x=torch.rand(batch,seq_len,2,device=device); target=x[:,:,0].sum(1)+x[:,:,1].sum(1)
                y,l1,g=m(x,reg); loss=((y[:,0]-target)**2).mean()
                if reg: loss=loss+0.5*torch.relu(0.05-g)**2+0.5*l1*l1
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(m.parameters(),1.0); opt.step()
                losses.append(float(((y[:,0]-target)**2).mean().detach().cpu()))
                if reg: gaps.append(float(g.detach().cpu())); exps.append(float(l1.detach().cpu()))
            return {'final_loss':float(np.mean(losses[-20:])), 'gap':float(np.mean(gaps[-20:])) if reg else None,
                    'lambda1':float(np.mean(exps[-20:])) if reg else None, 'device':device}
        return {'baseline':run(False),'regularized':run(True)}
    except Exception as e:
        return {'error':repr(e), 'fallback':'toy verification still valid'}


if __name__ == '__main__':
    result={'seed':SEED, 'toy':toy_verification(), 'rnn':rnn_check()}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))