import json
from pathlib import Path
import numpy as np

SEED = 1222
OUT = Path('results.json')


def trajectory(c0, lambdas, eta, gamma=0.0, n=300):
    c0 = np.asarray(c0, dtype=float)
    factors = 1.0 - eta * np.asarray(lambdas, dtype=float)
    factors[0] -= gamma
    t = np.arange(n + 1)[:, None]
    return c0[None, :] * factors[None, :] ** t, factors


def A(c):
    return 0.5 * float(np.sum(np.asarray(c) ** 2))


def crossing_time_discrete(ca, cb, factors):
    # Exact continuous-time interpolation of the discrete recurrence.
    r1, r2 = abs(factors[0]), abs(factors[1])
    numerator = cb[0] ** 2 - ca[0] ** 2
    denominator = ca[1] ** 2 - cb[1] ** 2
    if numerator <= 0 or denominator <= 0 or r1 <= 0 or r2 <= 0 or r1 == r2:
        return np.nan
    return np.log(denominator / numerator) / (2 * np.log(r1 / r2))


def observed_cross(ca, cb, factors, n=100000):
    t = np.arange(n + 1)
    aa = .5 * (ca[0] ** 2 * factors[0] ** (2*t) + ca[1] ** 2 * factors[1] ** (2*t))
    ab = .5 * (cb[0] ** 2 * factors[0] ** (2*t) + cb[1] ** 2 * factors[1] ** (2*t))
    d = aa - ab
    ix = np.where((d[:-1] * d[1:] <= 0) & (t[:-1] > 0))[0]
    return float(ix[0]) if len(ix) else np.nan


def log_slope(values, start):
    t = np.arange(len(values))
    mask = (t >= start) & (values > 1e-250)
    return float(np.polyfit(t[mask], np.log(values[mask]), 1)[0])


def main():
    # Prediction 1: t* scales as 1/(lambda2-lambda1), and the exact discrete formula
    # agrees with the first integer step at which the ordering reverses.
    ca, cb, eta, lam1 = np.array([.20, 1.00]), np.array([.60, .20]), 1e-3, 1.
    crossing = []
    for gap in [.5, 1., 2., 4.]:
        lambdas = np.array([lam1, lam1 + gap])
        factors = 1 - eta * lambdas
        continuous = np.log((ca[1]**2-cb[1]**2)/(cb[0]**2-ca[0]**2)) / (2*gap)
        pred = crossing_time_discrete(ca, cb, factors)
        obs = observed_cross(ca, cb, factors)
        crossing.append({'gap': gap, 'continuous_time': continuous,
                         'discrete_formula_step': pred, 'observed_step': obs,
                         'relative_error': abs(pred-obs)/obs})

    # Prediction 2: for c_{n+1}=(I-eta L-gamma e1e1^T)c_n, the slow-mode
    # stability edge is |1-eta lambda1-gamma|=1, i.e. gamma=2-eta lambda1.
    lambdas, eta2 = np.array([1., 3.]), .2
    gamma_star = 2 - eta2*lambdas[0]
    probes = np.linspace(gamma_star-.3, gamma_star+.3, 13)
    stability = []
    for gamma in probes:
        factors = np.array([1-eta2*lambdas[0]-gamma, 1-eta2*lambdas[1]])
        stable_exact = bool(np.max(np.abs(factors)) < 1)
        c, _ = trajectory([1., .2], lambdas, eta2, gamma=gamma, n=100)
        growth = float(np.max(np.linalg.norm(c, axis=1))/np.linalg.norm(c[0]))
        stability.append({'gamma': float(gamma), 'max_factor_abs': float(np.max(np.abs(factors))),
                          'predicted_stable': stable_exact, 'observed_bounded_100_steps': growth < 1.000001,
                          'norm_growth': growth})

    # Prediction 3: late A2 slope is 2 log|r_j| for the slowest surviving mode.
    eta3, lambdas3 = .05, np.array([1., 2.5])
    slope = []
    for gamma in [0., .2, .6, .95]:
        c, factors = trajectory([.6, 1.], lambdas3, eta3, gamma=gamma, n=300)
        a2 = .5*np.sum(c*c, axis=1)
        dominant = int(np.argmax(np.abs(factors)))
        slope.append({'gamma': gamma, 'factors': factors.tolist(),
                      'predicted_slope': 2*np.log(abs(factors[dominant])),
                      'observed_slope': log_slope(a2, 100), 'dominant_mode': dominant+1})
    c, factors = trajectory([0., 1.], lambdas3, eta3, n=300)
    a2 = .5*np.sum(c*c, axis=1)
    ideal_filter = {'expected_fast_slope': 2*np.log(abs(factors[1])),
                    'observed_fast_slope': log_slope(a2, 50)}

    # Training analogue: diagonal quadratic objective. The filtered run deliberately
    # begins with larger A2, but suppresses its slow coefficient and overtakes baseline.
    train = []
    for name, x0, gamma in [('baseline', [.60, .20], 0.), ('mode_filtered', [.20, 1.00], .75)]:
        x, factors = trajectory(x0, np.array([1., 3.]), .1, gamma=gamma, n=80)
        a2 = .5*np.sum(x*x, axis=1)
        train.append({'name': name, 'initial_A2': float(a2[0]), 'final_A2': float(a2[-1]),
                      'factors': factors.tolist(), 'a2_first_10': a2[:10].tolist()})
    # Find the crossing over full saved trajectories.
    trajectories = []
    for x0, gamma in [([.60,.20],0.),([.20,1.],.75)]:
        x,_ = trajectory(x0, np.array([1.,3.]), .1, gamma=gamma, n=80)
        trajectories.append(.5*np.sum(x*x,axis=1))
    cross = next((i for i in range(1,81) if trajectories[1][i] < trajectories[0][i]), None)
    result = {'seed': SEED, 'predictions': {'crossing': crossing,
        'stability': {'eta': eta2, 'lambda1': 1., 'predicted_upper_gamma': gamma_star, 'sweep': stability},
        'late_slope': slope, 'ideal_zero_slow_mode': ideal_filter},
        'training_analogue': {'runs': train, 'first_filtered_below_baseline_step': cross}}
    OUT.write_text(json.dumps(result, indent=2))
    summary = {'crossing': [{'gap': x['gap'], 'pred': round(x['discrete_formula_step'],3), 'obs': x['observed_step'], 'relerr': round(x['relative_error'],6)} for x in crossing],
               'stability_edge_predicted': gamma_star,
               'stability_edge_observed_bracket': [x['gamma'] for x in stability if x['predicted_stable']][-1:],
               'late_slope': slope, 'ideal_filter': ideal_filter,
               'training': {'initial_A2': [x['initial_A2'] for x in train], 'final_A2': [x['final_A2'] for x in train], 'cross_step': cross}}
    print(json.dumps(summary, indent=2))

if __name__ == '__main__':
    main()