"""Locate tau* where the dominant Floquet multiplier equals -1.

Continuation in tau from the known orbit at tau = 22.7 (mu = -2.676),
computing mu(tau) at each step, then a secant root-solve on mu(tau) + 1.
"""

import json
import os
import sys
import time

import numpy as np
from scipy.optimize import least_squares

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

HERE = os.path.dirname(os.path.abspath(__file__))
SEED_ORBIT = os.path.join(HERE, "..", "..", "dde-periodic-orbit",
                          "reference_outputs", "orbit_1.npz")
OUT = os.path.join(HERE, "..", "reference_outputs")


def polish(m, x0, T0, tau, max_nfev=60000):
    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 orbit_at(tau, x_seed512, T_seed, log):
    x_c, T_c, r_c = polish(512, x_seed512, T_seed, tau)
    if r_c > 1e-9:
        raise RuntimeError(f"coarse stalled at {r_c:.2e}")
    x_f, T_f, r_f = polish(M, upsample(x_c, M), T_c, tau, max_nfev=30000)
    log(f"    tau={tau:.8f}: T={T_f:.8f} resid={r_f:.2e} "
        f"amp={x_f.max()-x_f.min():.3f}")
    return x_c, T_c, x_f, T_f, r_f


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

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

    d = np.load(SEED_ORBIT)
    x0, T0 = d["x"], float(d["period"])
    idx = np.linspace(0, M, 512, endpoint=False).astype(int)
    x_c, T_c = x0[idx], T0

    log("scanning mu(tau) downward from 22.7")
    trace = []
    tau = 22.7
    prev = None
    bracket = None
    while tau > 10.5:
        try:
            x_c, T_c, x_f, T_f, r_f = orbit_at(tau, x_c, T_c, log)
        except RuntimeError as exc:
            log(f"    tau={tau:.4f}: {exc}")
            break
        mu, canary = dominant_multiplier(x_f, T_f, tau)
        log(f"      mu={mu.real:+.6f} (canary {canary:.1e})")
        trace.append({"tau": tau, "T": T_f, "mu_re": float(mu.real),
                      "mu_im": float(mu.imag), "resid": r_f})
        if prev is not None and (prev[1] + 1.0) * (mu.real + 1.0) < 0:
            bracket = (prev[0], tau)
            log(f"  BRACKET for mu=-1: tau in [{tau:.4f}, {prev[0]:.4f}]")
            break
        prev = (tau, mu.real)
        tau -= 0.4

    with open(os.path.join(HERE, "mu_trace.json"), "w") as fh:
        json.dump({"trace": trace, "bracket": bracket}, fh, indent=2)

    if bracket is None:
        log("no crossing found in scanned range")
        return

    lo, hi = min(bracket), max(bracket)
    log(f"secant refinement in [{lo:.6f}, {hi:.6f}]")
    # re-evaluate endpoints with the Richardson estimator for accuracy
    seeds = {}

    def mu_of(tau):
        nonlocal x_c, T_c
        xc, Tc, xf, Tf, rf = orbit_at(tau, x_c, T_c, log)
        x_c, T_c = xc, Tc
        mu_r, mu2, canary, spread = richardson_multiplier(xf, Tf, tau)
        seeds[tau] = (xf, Tf, rf)
        log(f"      mu_rich={mu_r.real:+.9f} (spread {spread:.1e}, "
            f"canary {canary:.1e})")
        return mu_r.real, xf, Tf, rf

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

    tau_star, xf, Tf, rf, fc = best
    os.makedirs(OUT, exist_ok=True)
    np.savez(os.path.join(OUT, "solution.npz"),
             tau_star=tau_star, period=Tf, x=xf)
    mu_r, mu2, canary, spread = richardson_multiplier(xf, Tf, tau_star)
    summary = {"tau_star": tau_star, "period": Tf, "resid_linf": rf,
               "mu_richardson_re": float(mu_r.real),
               "mu_richardson_im": float(mu_r.imag),
               "mu_plus_one": float(abs(mu_r + 1.0)),
               "richardson_spread": float(spread), "canary": float(canary),
               "spectral_tail": spectral_tail(xf),
               "fundamentality_gap": fundamentality_gap(xf, Tf),
               "amplitude": float(xf.max() - xf.min()),
               "min_x": float(xf.min()),
               "wall_seconds": time.time() - t0}
    log("SUMMARY " + json.dumps(summary, indent=2))
    with open(os.path.join(OUT, "solver_report.json"), "w") as fh:
        json.dump(summary, fh, indent=2)


if __name__ == "__main__":
    main()