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

SEED = 1595
np.random.seed(SEED)
random.seed(SEED)


def fit_log_slope(values, start=5, stop=None):
    v = np.maximum(np.asarray(values, float), 1e-30)
    stop = len(v) if stop is None else stop
    x = np.arange(start, stop)
    return float(np.polyfit(x, np.log(v[start:stop]), 1)[0])


def trajectory(aq, af, c=0.7, n=80, z0=1.0, y0=0.0):
    z = np.zeros(n); y = np.zeros(n)
    z[0], y[0] = z0, y0
    for t in range(n - 1):
        z[t + 1] = aq * z[t]
        y[t + 1] = af * y[t] + c * z[t]
    return z, y


def max_rate_sweep():
    rows = []
    pairs = [(0.98,.30),(.70,.90),(.55,.55),(.90,.40),(.40,.90),(.80,.60)]
    for aq, af in pairs:
        z, y = trajectory(aq, af)
        norm = np.sqrt(z*z + y*y)
        observed = math.exp(fit_log_slope(norm, 15, 55))
        predicted = max(aq, af)
        rows.append(dict(aq=aq, af=af, predicted=predicted, observed=observed,
                         relative_error=abs(observed-predicted)/predicted))
    return rows


def slow_branch_invariance():
    rows = []
    aq = .92
    for af in [.10,.30,.60,.80]:
        z, y = trajectory(aq, af)
        observed = math.exp(fit_log_slope(np.sqrt(z*z+y*y), 15, 55))
        rows.append(dict(aq=aq, af=af, predicted=aq, observed=observed,
                         relative_error=abs(observed-aq)/aq))
    return rows


def stability_sweep():
    # For this triangular linear transition, eigenvalues are aq and af.
    # Thus the asymptotic norm slope is log(max(|aq|,|af|)); its sign changes at 1.
    rows = []
    for aq, af in [(0.95,.5),(.99,.99),(1.0,.5),(1.01,.5),(.5,1.01),(.5,.98)]:
        z, y = trajectory(aq, af, n=300)
        norm = np.sqrt(z*z+y*y)
        slope = fit_log_slope(norm, 100, 250)
        predicted_rate = max(abs(aq), abs(af))
        predicted_stable = predicted_rate < 1
        observed_rate = math.exp(slope)
        observed_stable = slope < -1e-5
        rows.append(dict(aq=aq, af=af, predicted_rate=predicted_rate,
                         observed_rate=observed_rate, predicted_stable=predicted_stable,
                         observed_stable=observed_stable,
                         slope=slope))
    return rows


def learning_comparison():
    # Fixed linear system identification: split model has independently constrained rates;
    # monolithic matched-size model has unconstrained 2x2 transition.
    rng = np.random.default_rng(SEED)
    T, N = 25, 1600
    X = rng.normal(size=(N,T,1))
    Z = np.zeros((N,T)); Y = np.zeros((N,T))
    for i in range(N):
        for t in range(T-1):
            Z[i,t+1] = .96*Z[i,t] + .08*X[i,t,0] + .03*rng.normal()
            Y[i,t+1] = .42*Y[i,t] + .65*Z[i,t] + .08*X[i,t,0] + .03*rng.normal()
    H = np.stack([Z,Y], axis=-1)
    train = np.arange(1200); test = np.arange(1200,N)
    # Least-squares one-step fit with states supplied, then multi-step free rollout.
    def mse(a,b): return float(np.mean((a-b)**2))
    inp = np.concatenate([H[train,:-1], X[train,:-1]], axis=-1).reshape(-1,3)
    target = H[train,1:].reshape(-1,2)
    W = np.linalg.lstsq(inp, target, rcond=None)[0]
    # split: z uses z,x; y uses y,z,x, preserving triangular structure
    wz = np.linalg.lstsq(inp[:,[0,2]], target[:,0], rcond=None)[0]
    wy = np.linalg.lstsq(inp[:,[1,0,2]], target[:,1], rcond=None)[0]
    one_m = mse(inp @ W, target)
    one_s = mse(np.column_stack([inp[:,[0,2]] @ wz, inp[:,[1,0,2]] @ wy]), target)
    def rollout(kind, idx):
        out=[]
        for i in idx:
            h = H[i,0].copy(); pred=[]
            for t in range(T-1):
                q = X[i,t,0]
                if kind=='mono': h = np.array([h[0],h[1],q]) @ W
                else: h = np.array([np.array([h[0],q]) @ wz, np.array([h[1],h[0],q]) @ wy])
                pred.append(h.copy())
            out.append(pred)
        return np.asarray(out)
    pm, ps = rollout('mono', test), rollout('split', test)
    truth = H[test,1:]
    return dict(monolithic_one_step=one_m, split_one_step=one_s,
                monolithic_rollout=mse(pm,truth), split_rollout=mse(ps,truth),
                n_test=len(test))


def main():
    result = dict(seed=SEED,
                  prediction_1_max_rate=max_rate_sweep(),
                  prediction_2_slow_branch_invariance=slow_branch_invariance(),
                  prediction_3_stability_boundary=stability_sweep(),
                  learning_comparison=learning_comparison())
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()