"""Core numerics for the DDE period-doubling bifurcation task.

System (tau is now a FREE PARAMETER, not fixed):
    x'(t) = beta * x(t-tau) / (1 + x(t-tau)^n) - gamma * x(t)
    beta = 0.21, gamma = 0.105, n = 9.8.

Task: locate the critical delay tau* at which the periodic orbit undergoes
a PERIOD-DOUBLING (flip) bifurcation, i.e. its dominant nontrivial Floquet
multiplier equals exactly -1.

Self-certifying: the verifier recomputes, from the submitted (tau*, orbit)
pair alone, (a) the DDE residual of the orbit at that tau*, and (b) the
dominant nontrivial Floquet multiplier, and requires it to equal -1 to
tolerance. No hidden reference value is needed -- "mu = -1" is an intrinsic,
checkable condition.
"""

import numpy as np

BETA = 0.21
GAMMA = 0.105
EXP_N = 9.8
M = 4096                      # samples per period in the submission


def rhs(x, xd):
    return BETA * xd / (1.0 + xd**EXP_N) - GAMMA * x


def fprime(u):
    """d/du [ beta u / (1 + u^n) ]"""
    return BETA * (1 + (1 - EXP_N) * u**EXP_N) / (1 + u**EXP_N) ** 2


def shift_samples(x, shift, T):
    m = x.size
    k = np.fft.rfftfreq(m, d=1.0 / m)
    X = np.fft.rfft(x)
    return np.fft.irfft(X * np.exp(-2j * np.pi * k * (shift / T)), n=m)


def derivative_samples(x, T):
    m = x.size
    k = np.fft.rfftfreq(m, d=1.0 / m)
    dX = np.fft.rfft(x) * (2j * np.pi * k / T)
    if m % 2 == 0:
        dX[-1] = 0.0
    return np.fft.irfft(dX, n=m)


def dde_residual(x, T, tau):
    xd = shift_samples(x, tau % T if T > 0 else tau, T)
    return derivative_samples(x, T) - rhs(x, xd)


def spectral_tail(x):
    X = np.abs(np.fft.rfft(x))
    cut = int(len(X) * 7 / 8)
    denom = np.linalg.norm(X)
    return float(np.linalg.norm(X[cut:]) / denom) if denom > 0 else 1.0


def fundamentality_gap(x, T):
    return float(min(np.abs(x - shift_samples(x, T / k, T)).max()
                     for k in range(2, 7)))


def eval_periodic(x, T, t):
    """Trig interpolant of the samples evaluated at arbitrary times t."""
    m = x.size
    X = np.fft.rfft(x) / m
    k = np.arange(X.size)
    w = np.where((k == 0) | ((m % 2 == 0) & (k == m // 2)), 1.0, 2.0)
    return (np.exp(2j * np.pi * np.outer(t / T, k)) @ (X * w)).real


def monodromy_matrix(x, T, tau, nh, steps_per_delay):
    """Monodromy of y' = -gamma y + f'(x(t-tau)) y(t-tau) over one period.

    All nh history-basis vectors are propagated simultaneously with RK4 and
    4-point Lagrange interpolation of delayed values.
    """
    dt = tau / steps_per_delay
    n_steps = int(np.ceil(T / dt))
    dt = T / n_steps
    nd = tau / dt

    tg = np.arange(n_steps + 1) * dt
    c_full = fprime(eval_periodic(x, T, tg - tau))
    c_half = fprime(eval_periodic(x, T, tg[:-1] + 0.5 * dt - tau))

    n_hist = int(np.ceil(nd)) + 3
    Y = np.zeros((n_hist + n_steps + 1, nh))
    hist_t = -tau + np.arange(n_hist) * dt - (n_hist - 1 - nd) * dt
    nodes = np.linspace(-tau, 0.0, nh)
    for b in range(nh):
        e = np.zeros(nh)
        e[b] = 1.0
        Y[:n_hist, b] = np.interp(hist_t, nodes, e)
    i0 = n_hist - 1

    def delayed(idx_float):
        i = min(max(int(np.floor(idx_float)) - 1, 0), Y.shape[0] - 4)
        w = idx_float - (i + 1)
        p = Y[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(n_steps):
        row = i0 + s
        y = Y[row]
        k1 = -GAMMA * y + c_full[s] * delayed(row - nd)
        mid = c_half[s] * delayed(row + 0.5 - nd)
        k2 = -GAMMA * (y + 0.5 * dt * k1) + mid
        k3 = -GAMMA * (y + 0.5 * dt * k2) + mid
        k4 = -GAMMA * (y + dt * k3) + c_full[s + 1] * delayed(row + 1 - nd)
        Y[row + 1] = y + dt / 6.0 * (k1 + 2 * k2 + 2 * k3 + k4)

    out_idx = i0 + (nodes + T) / dt
    Mm = np.empty((nh, nh))
    for j, idx in enumerate(out_idx):
        Mm[j] = delayed(idx)
    return Mm


def dominant_multiplier(x, T, tau, nh=320, steps_per_delay=800):
    """Dominant nontrivial multiplier and the trivial-multiplier canary."""
    ev = np.linalg.eigvals(monodromy_matrix(x, T, tau, nh, steps_per_delay))
    order = np.argsort(-np.abs(ev))
    ev = ev[order[:8]]
    i_triv = np.argmin(np.abs(ev - 1.0))
    canary = abs(ev[i_triv] - 1.0)
    rest = np.delete(ev, i_triv)
    return rest[np.argmax(np.abs(rest))], canary


def richardson_multiplier(x, T, tau):
    """Two-resolution Richardson estimate (RK4 monodromy is O(h^2) here)."""
    mu1, can1 = dominant_multiplier(x, T, tau, nh=280, steps_per_delay=700)
    mu2, can2 = dominant_multiplier(x, T, tau, nh=400, steps_per_delay=1400)
    mu_r = (4.0 * mu2 - mu1) / 3.0
    return mu_r, mu2, max(can1, can2), abs(mu2 - mu1)


# ---------------------------------------------------------------------------
# Verification / grading
# ---------------------------------------------------------------------------
# Acceptance thresholds. The bound on |mu + 1| is set by the VERIFIER's own
# certified accuracy, not by wishful precision: the two-resolution Richardson
# estimator shows a spread of ~2.7e-6 and a trivial-multiplier canary of
# ~2.3e-5 at the bifurcation orbit, so demanding more than ~1e-4 would be
# asking for precision the grader cannot itself vouch for.
TAU_LO, TAU_HI = 10.0, 20.0
RESID_TOL = 1e-8
TAIL_TOL = 1e-10
AMP_MIN = 0.4
FUND_MIN = 0.1
MU_TOL = 1e-4
MU_TOL_PARTIAL = 1e-2


def verify_bifurcation(sol):
    """Recompute everything from the submitted (tau*, orbit) pair alone."""
    out = {}
    for key in ("tau_star", "period", "x"):
        if key not in sol:
            return {"structurally_valid": False, "reason": f"missing '{key}'"}
    tau = float(np.asarray(sol["tau_star"]).reshape(()))
    T = float(np.asarray(sol["period"]).reshape(()))
    x = np.asarray(sol["x"], dtype=np.float64)
    if x.shape != (M,) or not np.all(np.isfinite(x)) or not np.isfinite(T) \
            or not np.isfinite(tau):
        return {"structurally_valid": False,
                "reason": f"x must be finite float64 of shape ({M},)"}
    if not (TAU_LO <= tau <= TAU_HI):
        return {"structurally_valid": False,
                "reason": f"tau_star {tau} outside [{TAU_LO}, {TAU_HI}]"}
    if T <= tau:
        return {"structurally_valid": False, "reason": "period must exceed tau"}
    out["structurally_valid"] = True
    out["tau_star"] = tau
    out["period"] = T
    out["amplitude"] = float(x.max() - x.min())
    out["min_x"] = float(x.min())
    out["resid_linf"] = float(np.abs(dde_residual(x, T, tau)).max())
    out["spectral_tail"] = spectral_tail(x)
    out["fundamentality_gap"] = fundamentality_gap(x, T)
    mu_r, mu_fine, canary, spread = richardson_multiplier(x, T, tau)
    out["mu_re"], out["mu_im"] = float(mu_r.real), float(mu_r.imag)
    out["mu_plus_one"] = float(abs(mu_r + 1.0))
    out["richardson_spread"] = float(spread)
    out["canary"] = float(canary)
    return out


def grade_bifurcation(m):
    if not m.get("structurally_valid"):
        return 0.0, m.get("reason", "invalid submission")
    if m["min_x"] <= 0:
        return 0.0, "orbit must stay positive"
    if m["amplitude"] < AMP_MIN:
        return 0.0, f"amplitude {m['amplitude']:.3f} < {AMP_MIN}"
    if m["spectral_tail"] > TAIL_TOL:
        return 0.0, f"spectral tail {m['spectral_tail']:.2e} (under-resolved)"
    if m["fundamentality_gap"] < FUND_MIN:
        return 0.0, "period is not fundamental (k-fold repetition)"
    if m["resid_linf"] > RESID_TOL:
        return 0.0, f"DDE residual {m['resid_linf']:.2e} > {RESID_TOL}"
    if m["mu_plus_one"] <= MU_TOL:
        return 1.0, "orbit at a certified period-doubling point (mu = -1)"
    if m["mu_plus_one"] <= MU_TOL_PARTIAL:
        return 0.5, f"near the bifurcation: |mu+1| = {m['mu_plus_one']:.2e}"
    return 0.0, f"|mu+1| = {m['mu_plus_one']:.2e}: not a period-doubling point"


# ---------------------------------------------------------------------------
# Second deliverable: branch switching evidence
# ---------------------------------------------------------------------------
# The agent must also produce the PERIOD-DOUBLED orbit that is born at the
# bifurcation, at a fixed evidence delay above tau*. This forces an actual
# branch switch rather than a scalar root-solve. The fundamentality check is
# what makes it non-trivial: the period-1 orbit written at double period
# satisfies the delay equation exactly and has multiplier mu_1^2, so it looks
# like a solution in every respect except that its half-period shift
# reproduces it. Only the fundamentality gap separates the two.
TAU_EVIDENCE = 14.0
T2_LO, T2_HI = 70.0, 95.0


def verify_doubled(sol):
    """Check the period-doubled orbit at the fixed evidence delay."""
    out = {}
    for key in ("period2", "x2"):
        if key not in sol:
            return {"structurally_valid": False, "reason": f"missing '{key}'"}
    T2 = float(np.asarray(sol["period2"]).reshape(()))
    x2 = np.asarray(sol["x2"], dtype=np.float64)
    if x2.shape != (M,) or not np.all(np.isfinite(x2)) or not np.isfinite(T2):
        return {"structurally_valid": False,
                "reason": f"x2 must be finite float64 of shape ({M},)"}
    if not (T2_LO <= T2 <= T2_HI):
        return {"structurally_valid": False,
                "reason": f"period2 {T2} outside [{T2_LO}, {T2_HI}]"}
    out["structurally_valid"] = True
    out["period2"] = T2
    out["resid_linf_2"] = float(np.abs(dde_residual(x2, T2, TAU_EVIDENCE)).max())
    out["fundamentality_gap_2"] = fundamentality_gap(x2, T2)
    out["spectral_tail_2"] = spectral_tail(x2)
    out["amplitude_2"] = float(x2.max() - x2.min())
    out["min_x2"] = float(x2.min())
    return out


def grade_doubled(m):
    if not m.get("structurally_valid"):
        return False, m.get("reason", "invalid")
    if m["min_x2"] <= 0:
        return False, "doubled orbit must stay positive"
    if m["amplitude_2"] < AMP_MIN:
        return False, f"doubled-orbit amplitude {m['amplitude_2']:.3f} < {AMP_MIN}"
    if m["spectral_tail_2"] > TAIL_TOL:
        return False, f"doubled orbit under-resolved ({m['spectral_tail_2']:.2e})"
    if m["resid_linf_2"] > RESID_TOL:
        return False, f"doubled-orbit residual {m['resid_linf_2']:.2e} > {RESID_TOL}"
    if m["fundamentality_gap_2"] < FUND_MIN:
        return False, ("submitted 'doubled' orbit is the period-1 orbit written "
                       "at double period (fundamentality gap "
                       f"{m['fundamentality_gap_2']:.3f} < {FUND_MIN})")
    return True, "genuine period-doubled orbit"