import json
from pathlib import Path
import numpy as np
from scipy.optimize import minimize


def mu_factors(N, m, q, weights=None):
    weights = np.ones(m) if weights is None else np.asarray(weights)
    ell = np.arange(1, m + 1)
    cq = np.cos(2 * np.pi * q * ell / N)
    return np.array([np.sum(weights * cq * (1 - np.cos(2 * np.pi * k * ell / N))) for k in range(N)])


def eigenvalues(N, m, q, kappa=1.0, weights=None):
    weights = np.ones(m) if weights is None else np.asarray(weights)
    ell = np.arange(1, m + 1)
    cq = np.cos(2 * np.pi * q * ell / N)
    return np.array([kappa * np.sum(weights * cq * (np.exp(2j * np.pi * k * ell / N) - 1)) for k in range(N)])


def rhs(theta, m, kappa=1.0, weights=None):
    weights = np.ones(m) if weights is None else np.asarray(weights)
    y = np.zeros_like(theta)
    for l, a in enumerate(weights, 1):
        y += a * np.sin(np.roll(theta, -l) - theta)
    return kappa * y


def mode_amplitude(x, profile, k):
    N = len(x)
    u = np.exp(2j * np.pi * k * np.arange(N) / N)
    return abs(np.vdot(u, x - profile) / N)


def simulate(profile, m, k, kappa, h, steps, weights=None, eps=1e-7):
    N = len(profile)
    u = np.exp(2j * np.pi * k * np.arange(N) / N)
    x = profile + eps * np.real(u)
    amps = [mode_amplitude(x, profile, k)]
    for _ in range(steps):
        x = x + h * rhs(x, m, kappa, weights)
        amps.append(mode_amplitude(x, profile, k))
    return np.asarray(amps)


def fit_slope(a, h):
    t = np.arange(len(a)) * h
    return float(np.polyfit(t, np.log(np.maximum(a, 1e-30)), 1)[0])


def main():
    N, m, q, k = 32, 4, 2, 1
    profile = 2 * np.pi * q * np.arange(N) / N
    mu = mu_factors(N, m, q)
    lam = eigenvalues(N, m, q)[k]
    # Prediction 1: infinitesimal slope equals -mu.
    a = simulate(profile, m, k, 1.0, .002, 1000)
    slope = fit_slope(a, .002)
    # Prediction 2: slope scales linearly with coupling.
    scaling = []
    for kap in [.25, .5, 1., 1.5]:
        aa = simulate(profile, m, k, kap, .001, 500)
        scaling.append({'kappa': kap, 'predicted': float(-kap * mu[k]), 'observed': fit_slope(aa, .001)})
    # Prediction 3: explicit-Euler boundary |1+h lambda|=1.
    lb = eigenvalues(N, m, 0)[5]
    hcrit = float(-2 * lb.real / abs(lb) ** 2)
    hs = np.linspace(.2 * hcrit, 1.8 * hcrit, 65)
    stable = np.abs(1 + hs * lb) <= 1
    crossings = np.where(stable[:-1] & ~stable[1:])[0]
    observed_hcrit = float((hs[crossings[0]] + hs[crossings[0] + 1]) / 2)
    # Negative-factor prediction.
    candidates = [(qq, kk, float(mu_factors(N, m, qq)[kk])) for qq in range(N) for kk in range(1, N)
                  if mu_factors(N, m, qq)[kk] < -.05]
    qq, kk, negmu = candidates[0]
    prof_neg = 2 * np.pi * qq * np.arange(N) / N
    an = simulate(prof_neg, m, kk, 1., .001, 1000)
    neg_slope = fit_slope(an, .001)
    # Baseline versus Fourier-shaped coupling on the same unstable profile/mode.
    baseline_w = np.ones(m)
    target = .25
    objective = lambda w: (mu_factors(N, m, qq, w)[kk] - target) ** 2 + .02 * np.sum((w - 1.) ** 2)
    opt = minimize(objective, baseline_w, method='L-BFGS-B', bounds=[(-3., 3.)] * m)
    shaped_w = opt.x
    base_mu = float(mu_factors(N, m, qq, baseline_w)[kk])
    shaped_mu = float(mu_factors(N, m, qq, shaped_w)[kk])
    base_amp = simulate(prof_neg, m, kk, 1., .001, 1000, baseline_w)
    shaped_amp = simulate(prof_neg, m, kk, 1., .001, 1000, shaped_w)
    comparison = {'profile_q': qq, 'mode': kk, 'target_mu': target,
      'baseline': {'weights': baseline_w.tolist(), 'mu': base_mu, 'slope': fit_slope(base_amp, .001),
                   'amplitude_ratio': float(base_amp[-1] / base_amp[0])},
      'fourier_shaped': {'weights': shaped_w.tolist(), 'mu': shaped_mu, 'slope': fit_slope(shaped_amp, .001),
                         'amplitude_ratio': float(shaped_amp[-1] / shaped_amp[0])}}
    report = {'config': {'N': N, 'm': m, 'q': q, 'mode': k},
      'growth_rate_test': {'predicted': float(-mu[k]), 'observed': slope, 'relative_error': abs(slope + mu[k]) / abs(mu[k])},
      'coupling_scaling_test': scaling,
      'euler_boundary_test': {'lambda': [float(lb.real), float(lb.imag)], 'predicted_hcrit': hcrit,
          'observed_hcrit': observed_hcrit, 'relative_error': abs(observed_hcrit - hcrit) / hcrit},
      'negative_factor_test': {'q': qq, 'mode': kk, 'mu': negmu, 'predicted_growth': -negmu,
          'observed_growth': neg_slope, 'relative_error': abs(neg_slope + negmu) / abs(negmu)},
      'baseline_vs_shaped': comparison}
    Path('results.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()