"""Locate the SECOND period-doubling of the cascade (branch switching).

Just above tau*_1 the period-1 orbit is unstable and the period-doubled
orbit is born stable, so forward simulation lands on it. We capture it,
Newton-polish it at period ~2T, then continue it upward in tau and
root-solve its dominant multiplier + 1 = 0 to get tau*_2.
"""

import json
import os
import time

import numpy as np
from scipy.optimize import least_squares

from bif_common import (M, dde_residual, derivative_samples, fundamentality_gap,
                        richardson_multiplier, rhs, spectral_tail)

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "..", "reference_outputs")
TAU1 = 12.795726635219


def integrate(tau, n_steps, dt, x0=0.9):
    nd = int(round(tau / dt))
    x = np.empty(n_steps + nd + 1)
    x[: nd + 1] = x0

    def delayed(idx):
        i = min(max(int(np.floor(idx)) - 1, 0), x.size - 4)
        w = idx - (i + 1)
        p = x[i:i + 4]
        return (-w * (w - 1) * (w - 2) / 6 * p[0]
                + (w + 1) * (w - 1) * (w - 2) / 2 * p[1]
                - (w + 1) * w * (w - 2) / 2 * p[2]
                + (w + 1) * w * (w - 1) / 6 * p[3])

    for s in range(nd, nd + n_steps):
        xi = x[s]
        k1 = rhs(xi, x[s - nd])
        mid = delayed(s + 0.5 - nd)
        k2 = rhs(xi + 0.5 * dt * k1, mid)
        k3 = rhs(xi + 0.5 * dt * k2, mid)
        k4 = rhs(xi + dt * k3, x[s + 1 - nd])
        x[s + 1] = xi + dt / 6.0 * (k1 + 2 * k2 + 2 * k3 + k4)
    return x[nd:]


def polish(m, x0, T0, tau, max_nfev=80000):
    def pack(u):
        return np.concatenate([dde_residual(u[:m], u[m], tau),
                               [derivative_samples(u[:m], u[m])[0]]])
    res = least_squares(pack, np.concatenate([x0, [T0]]), method="lm",
                        xtol=1e-15, ftol=1e-15, gtol=1e-15, max_nfev=max_nfev)
    return res.x[:m], float(res.x[m]), float(np.abs(res.fun[:-1]).max())


def upsample(x, m_new):
    X = np.fft.rfft(x)
    Xn = np.zeros(m_new // 2 + 1, dtype=complex)
    Xn[:X.size] = X
    return np.fft.irfft(Xn, n=m_new) * (m_new / x.size)


def capture_doubled(tau, log, dt=0.005, t_run=6000.0):
    """Forward-simulate onto the (stable) period-doubled orbit and estimate
    its period from successive maxima."""
    traj = integrate(tau, int(t_run / dt), dt)
    tail = traj[int(len(traj) * 0.7):]
    # peak times
    pk = np.flatnonzero((tail[1:-1] > tail[:-2]) & (tail[1:-1] > tail[2:])) + 1
    if pk.size < 6:
        raise RuntimeError("too few peaks to estimate a period")
    # period-2 pattern: alternating peak heights -> period = 2 * mean spacing
    gaps = np.diff(pk[-9:]) * dt
    T2 = float(gaps[-8:].sum() / 4.0) if gaps.size >= 8 else float(2 * gaps.mean())
    heights = tail[pk[-9:]]
    alt = float(np.abs(heights[::2].mean() - heights[1::2].mean()))
    log(f"    peaks: mean gap {gaps.mean():.4f}, alternation {alt:.4f}, T2~{T2:.4f}")
    i0 = pk[-9]
    seg = tail[i0:i0 + int(T2 / dt)]
    idx = np.linspace(0, len(seg), 512, endpoint=False).astype(int)
    return seg[idx], T2, alt


def main():
    t0 = time.time()

    def log(msg):
        print(f"[{time.time() - t0:7.1f}s] {msg}", flush=True)

    log(f"capturing the period-doubled branch just above tau*_1 = {TAU1:.6f}")
    seed = None
    for tau_try in (14.0, 15.0, 13.5):
        try:
            x512, T2, alt = capture_doubled(tau_try, log)
        except RuntimeError as exc:
            log(f"  tau={tau_try}: {exc}")
            continue
        if alt < 1e-3:
            log(f"  tau={tau_try}: peaks not alternating (still period-1)")
            continue
        x512 = np.roll(x512, -int(np.argmax(x512)))
        xc, Tc, rc = polish(512, x512, T2, tau_try)
        log(f"  tau={tau_try}: coarse polish resid={rc:.2e} T={Tc:.5f}")
        if rc < 1e-9:
            seed = (tau_try, xc, Tc)
            break
    if seed is None:
        log("FAILED to capture the period-doubled branch")
        return

    tau = seed[0]
    trace = []
    prev = None
    bracket = None
    while tau < 22.0:
        try:
            # ALWAYS re-seed from forward simulation: Newton-continuing the
            # previous solution slides back onto the period-1 orbit written
            # at double period (a genuine trap -- it satisfies the equations,
            # has mu = mu_1^2, and only the fundamentality gap detects it).
            x512, T2, alt = capture_doubled(tau, log)
            x512 = np.roll(x512, -int(np.argmax(x512)))
            xc, Tc, rc = polish(512, x512, T2, tau)
            if rc > 1e-9:
                raise RuntimeError(f"coarse stalled {rc:.2e}")
            xf, Tf, rf = polish(M, upsample(xc, M), Tc, tau, max_nfev=30000)
        except RuntimeError as exc:
            log(f"  tau={tau:.4f}: {exc}")
            break
        mu_r, _, canary, spread = richardson_multiplier(xf, Tf, tau)
        fg = fundamentality_gap(xf, Tf)
        log(f"  tau={tau:.4f}: T={Tf:.5f} resid={rf:.1e} mu={mu_r.real:+.6f} "
            f"fund_gap={fg:.3f} (canary {canary:.1e})")
        trace.append({"tau": tau, "T": Tf, "mu_re": float(mu_r.real),
                      "resid": rf, "fund_gap": float(fg)})
        if prev is not None and (prev[1] + 1.0) * (mu_r.real + 1.0) < 0:
            bracket = (prev[0], tau)
            log(f"  BRACKET for the second doubling: [{prev[0]:.4f}, {tau:.4f}]")
            break
        prev = (tau, mu_r.real)
        tau += 0.5

    with open(os.path.join(HERE, "cascade_trace.json"), "w") as fh:
        json.dump({"trace": trace, "bracket": bracket}, fh, indent=2)
    if bracket is None:
        log("no second doubling bracketed in the scanned range")
        return

    a, b = bracket
    store = {}

    def mu_of(tau_v):
        x512, T2, _ = capture_doubled(tau_v, log)
        x512 = np.roll(x512, -int(np.argmax(x512)))
        xc2, Tc2, _ = polish(512, x512, T2, tau_v)
        xf, Tf, rf = polish(M, upsample(xc2, M), Tc2, tau_v, max_nfev=30000)
        mu_r, _, canary, spread = richardson_multiplier(xf, Tf, tau_v)
        store[tau_v] = (xf, Tf, rf, canary, spread)
        log(f"    tau={tau_v:.9f}: mu={mu_r.real:+.9f} (spread {spread:.1e})")
        return mu_r.real + 1.0

    fa, fb = mu_of(a), mu_of(b)
    best = None
    for it in range(10):
        c = b - fb * (b - a) / (fb - fa)
        fc = mu_of(c)
        log(f"  secant it {it}: tau={c:.10f}, mu+1={fc:.3e}")
        best = c
        if abs(fc) < 1e-9:
            break
        a, fa, b, fb = b, fb, c, fc

    tau2 = best
    xf, Tf, rf, canary, spread = store[tau2]
    os.makedirs(OUT, exist_ok=True)
    np.savez(os.path.join(OUT, "cascade.npz"),
             tau_1=TAU1, tau_2=tau2, period_2=Tf, x2=xf)
    mu_r, _, canary, spread = richardson_multiplier(xf, Tf, tau2)
    summary = {"tau_1": TAU1, "tau_2": tau2, "period_2": Tf,
               "resid_linf_2": rf, "mu2_re": float(mu_r.real),
               "mu2_plus_one": float(abs(mu_r + 1.0)),
               "richardson_spread": float(spread), "canary": float(canary),
               "fundamentality_gap_2": fundamentality_gap(xf, Tf),
               "spectral_tail_2": spectral_tail(xf),
               "amplitude_2": float(xf.max() - xf.min()),
               "period_ratio": Tf / 37.732804942870,
               "wall_seconds": time.time() - t0}
    log("SUMMARY " + json.dumps(summary, indent=2))
    with open(os.path.join(OUT, "cascade_report.json"), "w") as fh:
        json.dump(summary, fh, indent=2)


if __name__ == "__main__":
    main()