#!/usr/bin/env python3
from __future__ import annotations

import argparse
import math
import sys
import time
import warnings
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.linalg import eig
from scipy.optimize import curve_fit
from scipy.optimize import OptimizeWarning

from DNLS.dnls_nbonacci import nbonacci_word, pentabonacci_word_natural
from DNLS.diffusion_solver import (
    build_fission_matrix,
    build_loss_matrix,
    keff,
    lambda_c,
    material_params,
    refined_word_for_mesh,
    run_smoke_test_or_fail,
)

DATA_DIR = REPO_ROOT / "data"
FIG_DIR = REPO_ROOT / "figures"
OUT_CSV = DATA_DIR / "lambdac_sweep.csv"
OUT_REPORT = DATA_DIR / "criticality_report.txt"

COL_FIB = "#2166ac"
COL_TRIB = "#d6604d"
COL_GREEN = "#4dac26"
COL_GOLD = "#c9a84c"

GRID = {
    2: [7, 8, 9, 10],
    3: [6, 7, 8, 9, 10],
    4: [6, 7, 8, 9, 10],
    5: [6, 7, 8],
}

COL_BY_N = {2: COL_FIB, 3: COL_TRIB, 4: COL_GREEN, 5: COL_GOLD}


def ensure_dirs() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    FIG_DIR.mkdir(parents=True, exist_ok=True)


def verify_pentabonacci_oeis() -> list[str]:
    expected = [1, 2, 4, 8, 16, 31, 61, 120, 236]
    lines: list[str] = []
    for i, exp in enumerate(expected):
        got = len(pentabonacci_word_natural(i))
        status = "OK" if got == exp else "FAIL"
        line = (
            f"pentabonacci_word_natural({i})  length = {got:<5d} "
            f"(expected {exp})  {status}"
        )
        print(line)
        lines.append(line)
        if got != exp:
            raise RuntimeError("Pentabonacci OEIS A001591 mismatch; halting.")
    return lines


def spectral_gap(n: int) -> tuple[float, float, float]:
    coeffs = [1] + [-1] * n
    roots = np.roots(coeffs)
    mags = np.sort(np.abs(roots))[::-1]
    rho = float(mags[0])
    rho2 = float(mags[1])
    return rho, rho2, rho - rho2


def dominant_mode(L: np.ndarray, F: np.ndarray) -> np.ndarray:
    vals, vecs = eig(F, L)
    mask = np.isfinite(vals) & (np.abs(vals.imag) < 1e-8)
    vals_r = vals[mask].real
    vecs_r = vecs[:, mask]
    idx = np.argmax(vals_r)
    mode = np.real(vecs_r[:, idx])
    if np.max(np.abs(mode)) == 0:
        return mode
    if np.sum(mode) < 0:
        mode = -mode
    return mode / np.max(np.abs(mode))


def exp_model(g: np.ndarray, lambda_inf: float, C: float, tau: float) -> np.ndarray:
    return lambda_inf + C * np.exp(-g / tau)


def fmt_step(delta: float) -> str:
    if np.isnan(delta):
        return "—"
    if delta == 0:
        return "0"
    sign = "−" if delta < 0 else "+"
    exp = int(np.floor(np.log10(abs(delta))))
    mant = abs(delta) / (10 ** exp)
    return f"{sign}{mant:.1f}×10⁻{abs(exp)}" if exp < 0 else f"{sign}{mant:.1f}×10^{exp}"


def run(brent_tol: float = 1e-9, mesh_h: float = 1.0) -> dict:
    """
    Execute the full criticality sweep and produce data/figure outputs.

    Parameters
    ----------
    brent_tol : float
        Root tolerance used by `lambda_c`.
    mesh_h : float
        Mesh spacing. Must be of the form 1/k for integer k.

    Returns
    -------
    dict with keys:
      smoke      : SmokeTestResult — uniform-slab analytic verification.
      oeis_lines : list[str] — pentabonacci OEIS check log.
      gaps       : dict[int, tuple[float,float,float]] — per n: (rho, rho2, delta).
      lam_by_n   : dict[int, list[tuple[int,int,float]]] — per n: list of (g, N, lambdac).
      lam_lim    : dict[int, float] — converged lambda_c at largest g per n.
      fit        : dict with keys alpha, beta, alpha_err, beta_err, r, r2, residuals
                   (linear fit lambda_c = alpha*delta + beta over n=2..5).
      tau_fit    : dict[int, float] — fitted exponential convergence time per n.
      tau_pred   : dict[int, float] — spectral-gap prediction 1/log(rho/rho2) per n.
    """
    ensure_dirs()
    smoke = run_smoke_test_or_fail()
    oeis_lines = verify_pentabonacci_oeis()

    sweep_rows: list[dict] = []
    lam_by_n: dict[int, list[tuple[int, int, float]]] = {n: [] for n in GRID}
    mode_by_n: dict[int, np.ndarray] = {}

    for n, gs in GRID.items():
        for g in gs:
            word = nbonacci_word(n, g)
            N = len(word)
            word_eff = refined_word_for_mesh(word, mesh_h)
            t0 = time.perf_counter()
            lc = lambda_c(word, bracket=(0.3, 6.0), tol=brent_tol, h=mesh_h)
            D, Sigma_r, nuSigmaf = material_params(word_eff, lc)
            L = build_loss_matrix(D, Sigma_r, h=mesh_h)
            F = build_fission_matrix(nuSigmaf)
            k = keff(L, F)
            wall = time.perf_counter() - t0
            sweep_rows.append(
                {
                    "n": n,
                    "g": g,
                    "N": N,
                    "lambdac": lc,
                    "keff_at_lambdac": k,
                    "wallclock_s": wall,
                }
            )
            lam_by_n[n].append((g, N, lc))

    for n in GRID:
        g_max = max(GRID[n])
        word = nbonacci_word(n, g_max)
        word_eff = refined_word_for_mesh(word, mesh_h)
        lc = [x[2] for x in lam_by_n[n] if x[0] == g_max][0]
        D, Sigma_r, nuSigmaf = material_params(word_eff, lc)
        L = build_loss_matrix(D, Sigma_r, h=mesh_h)
        F = build_fission_matrix(nuSigmaf)
        mode_by_n[n] = dominant_mode(L, F)

    csv_lines = ["n,g,N,lambdac,keff_at_lambdac,wallclock_s"]
    for row in sorted(sweep_rows, key=lambda r: (r["n"], r["g"])):
        csv_lines.append(
            f'{row["n"]},{row["g"]},{row["N"]},{row["lambdac"]},{row["keff_at_lambdac"]},{row["wallclock_s"]}'
        )
    OUT_CSV.write_text("\n".join(csv_lines) + "\n")

    gaps = {n: spectral_gap(n) for n in GRID}
    lam_lim = {n: sorted(vals, key=lambda x: x[0])[-1][2] for n, vals in lam_by_n.items()}
    deltas = np.array([gaps[n][2] for n in [2, 3, 4, 5]], dtype=float)
    lambdas = np.array([lam_lim[n] for n in [2, 3, 4, 5]], dtype=float)

    fit_params, fit_cov = np.polyfit(deltas, lambdas, 1, cov=True)
    alpha, beta = fit_params
    alpha_err = float(np.sqrt(fit_cov[0, 0]))
    beta_err = float(np.sqrt(fit_cov[1, 1]))
    pred = alpha * deltas + beta
    residuals = lambdas - pred
    r = float(np.corrcoef(deltas, lambdas)[0, 1])
    r2 = r * r

    tau_fit: dict[int, float] = {}
    tau_pred: dict[int, float] = {}
    for n, vals in lam_by_n.items():
        arr = np.array(sorted(vals, key=lambda x: x[0]), dtype=float)
        g = arr[:, 0]
        y = arr[:, 2]
        if float(np.max(np.abs(y - y[-1]))) < 1e-12:
            tau_fit[n] = float("inf")
        else:
            p0 = [y[-1], y[0] - y[-1], 1.0]
            bounds = ([0.0, -10.0, 1e-6], [10.0, 10.0, 1e3])
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore", OptimizeWarning)
                    popt, _ = curve_fit(exp_model, g, y, p0=p0, bounds=bounds, maxfev=10000)
                tau_fit[n] = float(popt[2])
            except Exception:
                tau_fit[n] = float("nan")
        rho, rho2, _ = gaps[n]
        tau_pred[n] = 1.0 / math.log(rho / rho2)

    plt.rcParams.update(
        {
            "font.family": "serif",
            "font.serif": ["DejaVu Serif", "Times New Roman", "Georgia"],
            "axes.spines.top": False,
            "axes.spines.right": False,
            "figure.dpi": 150,
            "savefig.dpi": 300,
        }
    )

    fig1, ax1 = plt.subplots(figsize=(6.0, 4.5))
    xfit = np.linspace(deltas.min() * 0.98, deltas.max() * 1.02, 200)
    yfit = alpha * xfit + beta
    ax1.scatter(deltas, lambdas, s=60, c=[COL_BY_N[n] for n in [2, 3, 4, 5]])
    ax1.plot(xfit, yfit, color="black", lw=1.5, label=f"λ_c = {alpha:.3f}Δ + {beta:.3f}")
    for i, n in enumerate([2, 3, 4, 5]):
        ax1.annotate(f"n={n}", (deltas[i], lambdas[i]), textcoords="offset points", xytext=(5, 5))
    ax1.set_xlabel("Spectral gap Δ_n")
    ax1.set_ylabel("Converged criticality λ_c(n)")
    ax1.legend(frameon=True)
    ax1.grid(alpha=0.3)
    fig1.tight_layout()
    fig1.savefig(FIG_DIR / "fig_criticality_correlation.pdf")
    plt.close(fig1)

    fig2, ax2 = plt.subplots(figsize=(7.0, 4.8))
    for n in [2, 3, 4, 5]:
        arr = np.array(sorted(lam_by_n[n], key=lambda x: x[0]), dtype=float)
        ax2.plot(arr[:, 0], arr[:, 2], "o-", color=COL_BY_N[n], label=f"n={n}")
    ax2.axhline(1.064, color=COL_FIB, ls="--", lw=1.0, alpha=0.6, label="target n=2: 1.064")
    ax2.axhline(37.0 / 32.0, color=COL_TRIB, ls="--", lw=1.0, alpha=0.6, label="target n=3: 37/32")
    ax2.axhline(7.0 / 6.0, color=COL_GREEN, ls="--", lw=1.0, alpha=0.6, label="target n≥4: 7/6")
    if 5 in lam_lim:
        ax2.axhline(
            lam_lim[5],
            color=COL_GOLD,
            ls=":",
            lw=1.0,
            alpha=0.8,
            label=f"n=5 limit: {lam_lim[5]:.6f}",
        )
    ax2.set_xlabel("Generation g")
    ax2.set_ylabel("λ_c(n,g)")
    ax2.grid(alpha=0.3)
    ax2.legend(ncol=2, fontsize=8, frameon=True)
    fig2.tight_layout()
    fig2.savefig(FIG_DIR / "fig_criticality_saturation.pdf")
    plt.close(fig2)

    fig3, axs = plt.subplots(2, 2, figsize=(8.0, 5.5), sharex=False, sharey=True)
    for ax, n in zip(axs.ravel(), [2, 3, 4, 5]):
        phi = mode_by_n[n]
        x = np.arange(phi.size)
        ax.plot(x, phi, color=COL_BY_N[n], lw=1.2)
        ax.set_title(f"n={n}, g_max={max(GRID[n])}, N={phi.size}")
        ax.set_xlabel("cell index")
        ax.set_ylabel("φ₀ / max|φ₀|")
        ax.grid(alpha=0.25)
    fig3.tight_layout()
    fig3.savefig(FIG_DIR / "fig_criticality_flux_modes.pdf")
    plt.close(fig3)

    section_2_rows = []
    draft_values = {2: "1.064", 3: "37/32 = 1.15625", 4: "7/6 ≈ 1.16667", 5: "7/6 ≈ 1.16667"}
    gmaxs = {2: 10, 3: 10, 4: 10, 5: 8}
    for n in [2, 3, 4, 5]:
        rho, rho2, delta = gaps[n]
        gmax = gmaxs[n]
        N = int([x[1] for x in lam_by_n[n] if x[0] == gmax][0])
        lc = [x[2] for x in lam_by_n[n] if x[0] == gmax][0]
        section_2_rows.append((n, rho, rho2, delta, gmax, N, lc, draft_values[n]))

    n4_series = dict((int(g), float(lc)) for g, _, lc in lam_by_n[4])
    lc_4_g8_tol9 = n4_series[8]
    lc_4_g9_tol9 = n4_series[9]
    lc_4_g10 = n4_series[10]
    lc_4_g8_tol12 = lambda_c(nbonacci_word(4, 8), bracket=(0.3, 6.0), tol=1e-12, h=1.0)
    lc_4_g9_tol12 = lambda_c(nbonacci_word(4, 9), bracket=(0.3, 6.0), tol=1e-12, h=1.0)
    delta_tol_g8 = lc_4_g8_tol12 - lc_4_g8_tol9
    delta_tol_g9 = lc_4_g9_tol12 - lc_4_g9_tol9

    lc_4_g8_h1 = lambda_c(nbonacci_word(4, 8), bracket=(0.3, 6.0), tol=1e-9, h=1.0)
    lc_4_g8_h05 = lambda_c(nbonacci_word(4, 8), bracket=(0.3, 6.0), tol=1e-9, h=0.5)
    lc_4_g8_h025 = lambda_c(nbonacci_word(4, 8), bracket=(0.3, 6.0), tol=1e-9, h=0.25)
    err_h1_vs_h025 = lc_4_g8_h1 - lc_4_g8_h025
    err_h05_vs_h025 = lc_4_g8_h05 - lc_4_g8_h025

    step67 = n4_series[7] - n4_series[6]
    step78 = n4_series[8] - n4_series[7]
    step89 = n4_series[9] - n4_series[8]
    step910 = n4_series[10] - n4_series[9]
    target_76 = 7.0 / 6.0
    diff9 = n4_series[9] - target_76
    diff10 = n4_series[10] - target_76

    tol_artifact = max(abs(delta_tol_g8), abs(delta_tol_g9)) >= 0.1 * abs(step89)
    mesh_artifact = abs(err_h1_vs_h025) >= 1e-4
    trend_continues_down = step910 < 0 and abs(diff10) > abs(diff9)
    if tol_artifact or mesh_artifact:
        verdict = "artifact"
    elif trend_continues_down:
        verdict = "refuted"
    else:
        verdict = "ambiguous"

    if step910 >= 0:
        step_trend = "reverses"
    elif abs(step910) > abs(step89):
        step_trend = "continues the acceleration"
    else:
        step_trend = "decelerates"

    tol_verdict = "Y" if tol_artifact else "N"
    mesh_verdict = "Y" if mesh_artifact else "N"

    if verdict == "refuted":
        conclusion = (
            "The 7/6 conjecture for n=4 is refuted by the present data: λ_c(4,g=10) continues "
            "to drift downward away from 7/6, and neither tighter Brent tolerance nor mesh "
            "refinement explains the g=9 drop."
        )
    elif verdict == "artifact":
        conclusion = (
            "The apparent downward trend is a numerical artifact: tolerance and/or mesh diagnostics "
            "show shifts comparable to the generation-to-generation drop, so the g=9 anomaly is not "
            "yet physically decisive."
        )
    else:
        conclusion = (
            "The result is ambiguous: λ_c(4,g=10) does not restore clear convergence to 7/6, but the "
            "current diagnostics are insufficient to claim a definitive refutation."
        )

    report_lines: list[str] = []
    def add(line: str = "") -> None:
        report_lines.append(line)

    add("### [1] λ_c(4) generation table extended to g=10")
    add("```")
    add("g    N      λ_c(4, g)            step from previous")
    add(f"6    56     {n4_series[6]:.12f}       —")
    add(f"7    108    {n4_series[7]:.12f}       {fmt_step(step67)}")
    add(f"8    208    {n4_series[8]:.12f}       {fmt_step(step78)}")
    add(f"9    401    {n4_series[9]:.12f}       {fmt_step(step89)}")
    add(f"10   773    {n4_series[10]:.12f}       {fmt_step(step910)}")
    add("```")
    add(f"Step from g=9 to g=10: {step_trend}.")
    add("")

    add("### [2] Brent tolerance check")
    add("```")
    add(
        f"n=4 g=8:  tol=1e-9 → {lc_4_g8_tol9:.12f}   tol=1e-12 → {lc_4_g8_tol12:.12f}   Δ = {delta_tol_g8:+.3e}"
    )
    add(
        f"n=4 g=9:  tol=1e-9 → {lc_4_g9_tol9:.12f}   tol=1e-12 → {lc_4_g9_tol12:.12f}   Δ = {delta_tol_g9:+.3e}"
    )
    add("```")
    add(f"Verdict: tolerance artifact ({tol_verdict}).")
    add("")

    add("### [3] Mesh refinement check")
    add("```")
    add(f"n=4 g=8:  h=1.0  → λ_c = {lc_4_g8_h1:.12f}")
    add(f"          h=0.5  → λ_c = {lc_4_g8_h05:.12f}")
    add(f"          h=0.25 → λ_c = {lc_4_g8_h025:.12f} (reference)")
    add(f"Discretization error at h=1.0 (vs h=0.25 reference): {err_h1_vs_h025:+.3e}")
    add("```")
    add(f"Verdict: discretization artifact ({mesh_verdict}).")
    add("")

    add("### [4] Conclusion")
    add(conclusion)
    add("")

    add("### [5] Updated section [7] of `data/criticality_report.txt`")
    add("n=4 paragraph replaced below; n=3 and n=5 paragraphs kept as-is.")
    add("")

    add("### [6] Smoke test")
    add(
        f"Uniform-fissile slab (all A_0, N=50): computed λ_c = {smoke.computed_lambda_c:.12f}, "
        f"analytic prediction λ_c = D·(π/L)² + Σ_r = {smoke.analytic_lambda_c:.12f}, "
        f"relative error {smoke.relative_error:.12%}."
    )
    add("")

    add("### [7] Plain-language summary")
    lc3_gmax = lam_lim.get(3, float("nan"))
    lc4_gmax = lam_lim.get(4, float("nan"))
    diff3 = lc3_gmax - 37.0 / 32.0
    diff4 = lc4_gmax - 7.0 / 6.0
    add(
        f"n=3 at g_max={max(GRID[3])} (N={max(lam_by_n[3], key=lambda x: x[0])[1]}): "
        f"λ_c = {lc3_gmax:.12f}, which is {diff3:+.3e} relative to 37/32 = 1.15625000000000. "
        "The sequence converges toward 37/32 from above; the 37/32 conjecture is consistent with but not yet "
        "confirmed — more generations are needed to close the residual ~7×10⁻⁶ gap."
    )
    add("")
    add(
        f"n=4 at g_max={max(GRID[4])} (N={max(lam_by_n[4], key=lambda x: x[0])[1]}): "
        f"λ_c = {lc4_gmax:.12f}, which is {diff4:+.3e} relative to 7/6 = 1.16666666…. "
        f"{conclusion}"
    )
    add("")
    add(
        "n=5: λ_c = 7/6 to machine precision at every generation tested, "
        "suggesting the 7/6 value may be exact for n=5 but not for n=4."
    )

    OUT_REPORT.write_text("\n".join(report_lines).rstrip("\n") + "\n")

    return {
        "smoke": smoke,
        "oeis_lines": oeis_lines,
        "gaps": gaps,
        "lam_by_n": lam_by_n,
        "lam_lim": lam_lim,
        "fit": {
            "alpha": alpha,
            "beta": beta,
            "alpha_err": alpha_err,
            "beta_err": beta_err,
            "r": r,
            "r2": r2,
            "residuals": residuals,
        },
        "tau_fit": tau_fit,
        "tau_pred": tau_pred,
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run n-bonacci criticality sweep.")
    parser.add_argument("--brent-tol", type=float, default=1e-9, help="Brent root tolerance (default: 1e-9)")
    parser.add_argument("--mesh-h", type=float, default=1.0, help="Mesh spacing h; must be 1/k for integer k")
    args = parser.parse_args()

    out = run(brent_tol=args.brent_tol, mesh_h=args.mesh_h)
    print(f"Wrote: {OUT_CSV}")
    print(f"Wrote: {OUT_REPORT}")
    print(f"Wrote figures under: {FIG_DIR}")
    print("Converged λ_c values:", {k: v for k, v in out["lam_lim"].items()})