import json, random
import numpy as np
import torch

SEED = 7
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
try:
    torch.zeros(1, device=DEVICE)
except Exception:
    DEVICE = 'cpu'

# A 10D system with an exactly rank-2 invariant state manifold.
# x = c + Bq, S(x) = c + B M q.  Thus the faithful latent dimension is 2.
def make_map(gamma):
    rng = np.random.RandomState(23)
    B, _ = np.linalg.qr(rng.randn(10, 2))
    B = B.astype(np.float32)
    c = np.array([.4,-.2,.3,.1,-.1,.2,.05,-.15,.12,.08], np.float32)
    M = np.diag([gamma, .72 * gamma]).astype(np.float32)
    A = B @ M @ B.T
    b = c - A @ c
    return A.astype(np.float32), b.astype(np.float32), B, c, M

def apply(A, b, x):
    return x @ A.T + b

def manifold_states(B, c, n, seed=11):
    q = np.random.RandomState(seed).randn(n, 2).astype(np.float32) * 1.4
    return c[None, :] + q @ B.T

def train_latent(A, b, B, c, m, epochs=1300):
    x = manifold_states(B, c, 1800, 11 + m)
    y = apply(A, b, x)
    E = torch.nn.Linear(10, m, device=DEVICE)
    D = torch.nn.Linear(m, 10, device=DEVICE)
    T = torch.nn.Linear(m, m, device=DEVICE)
    opt = torch.optim.Adam(list(E.parameters()) + list(D.parameters()) + list(T.parameters()), lr=.025)
    xt, yt = torch.tensor(x, device=DEVICE), torch.tensor(y, device=DEVICE)
    for _ in range(epochs):
        idx = torch.randint(0, len(x), (128,), device=DEVICE)
        xv, yv = xt[idx], yt[idx]
        z, zn = E(xv), T(E(xv))
        loss = ((E(yv)-zn)**2).mean() + ((D(zn)-yv)**2).mean() + .5*((D(z)-xv)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return E, D, T

def learned_metrics(A, b, B, c, E, D, T, gamma):
    x = torch.tensor(manifold_states(B, c, 300, 99), device=DEVICE)
    y = torch.tensor(apply(A, b, x.cpu().numpy()), device=DEVICE)
    with torch.no_grad():
        enc = ((E(y)-T(E(x)))**2).mean().sqrt().item()
        dec = ((D(T(E(x)))-y)**2).mean().sqrt().item()
        rec = ((D(E(x))-x)**2).mean().sqrt().item()
        z = E(x); gaps=[]
        for _ in range(12):
            gaps.append(torch.linalg.vector_norm(D(T(z))-D(z), dim=1).mean().item())
            z = T(z)
        ratio = float(np.median(np.array(gaps[6:]) / np.array(gaps[5:11])))
        fp = c + np.zeros(10, np.float32)  # exact fixed point is c
        fp_t = torch.tensor(fp[None, :], device=DEVICE)
        fp_res = torch.linalg.vector_norm(D(T(E(fp_t)))-fp_t, dim=1).item()
    return {'enc_residual':enc, 'decode_residual':dec, 'reconstruction':rec,
            'decoded_fixed_point_residual':fp_res, 'decoded_rate':ratio,
            'predicted_rate':gamma}

def exact_checks():
    # Prediction 1: scalar contraction boundary is gamma*Lambda=gamma < 1.
    boundary=[]
    for g in [.6,.9,.99,1.,1.01,1.2]:
        d=1.;
        for _ in range(12): d=g*d
        boundary.append({'gamma':g, 'predicted_contracts':g < 1,
                         'observed_contracts':d < 1, '12_step_ratio':d})
    # Prediction 2: an m-dimensional faithful representation needs m >= rank(M)=2.
    A,b,B,c,M=make_map(.82)
    x=manifold_states(B,c,500,41)
    # exact best rank-m linear reconstruction of centered manifold states
    _,s,_=np.linalg.svd((x-c).T, full_matrices=False)
    rank=[]
    for m in [1,2,3]:
        residual=np.sqrt(np.maximum(0., np.sum(s[m:]**2) / x.shape[0])) if m < len(s) else 0.
        rank.append({'m':m, 'predicted_zero_residual':m>=2,
                     'observed_reconstruction_rmse':float(residual)})
    # Prediction 3: ideal decoded trajectory contracts at dominant latent rate gamma.
    rates=[]
    for g in [.55,.70,.82,.92]:
        _,_,_,_,M=make_map(g); q=1.; vals=[]
        for _ in range(12): vals.append(abs(q)); q=g*q
        rates.append({'gamma':g, 'predicted_asymptotic_rate':g,
                      'observed_rate':float(np.median(np.array(vals[6:]) / np.array(vals[5:11])))})
    return boundary,rank,rates

def main():
    boundary, rank, rates = exact_checks()
    A,b,B,c,_=make_map(.82)
    learned=[]
    for m in [1,2,3]:
        E,D,T=train_latent(A,b,B,c,m)
        row=learned_metrics(A,b,B,c,E,D,T,.82); row['m']=m; learned.append(row)
    # Same manifold initial states and 12 update evaluations.
    x=manifold_states(B,c,300,123); fixed=c
    direct=x.copy()
    for _ in range(12): direct=apply(A,b,direct)
    E,D,T=train_latent(A,b,B,c,2)
    with torch.no_grad():
        z=E(torch.tensor(x,device=DEVICE))
        for _ in range(12): z=T(z)
        idea=D(z).cpu().numpy()
    out={'device':DEVICE, 'exact_predictions':{'contraction_boundary':boundary,
          'rank_threshold':rank, 'rate_transfer':rates}, 'learned_sweep':learned,
          'comparison':{'baseline_S_evals':12,'idea_S_evals':0,'idea_T_evals':12,
            'baseline_fixed_point_rmse':float(np.sqrt(np.mean((direct-fixed)**2))),
            'idea_fixed_point_rmse':float(np.sqrt(np.mean((idea-fixed)**2))),
            'baseline_output_residual':float(np.sqrt(np.mean((apply(A,b,direct)-direct)**2))),
            'idea_output_residual':float(np.sqrt(np.mean((apply(A,b,idea)-idea)**2)))}}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()