#!/usr/bin/env python3
"""Reproduce numerical illustrations for the fixed linear predator--prey model.

This is numerical corroboration, not a validated computation or a proof.
Analytic existence, shape preservation and cyclicity require the manuscript's
theorems. No global extension of the local growth profile is implemented here.

Run: python reproduce_cycles.py
Dependencies: numpy, scipy, mpmath, matplotlib.
"""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import mpmath as mp
import numpy as np
import scipy
from scipy.integrate import simpson, solve_ivp
from scipy.optimize import brentq, minimize_scalar
from scipy.special import hyp2f1, lambertw


CASES = {
    3: ("0.05", "0.10", "0.15"),
    4: ("0.04", "0.08", "0.12", "0.16"),
}
SOLVER_CONTROLS = {
    "standard": {"rtol": 2e-12, "atol": 2e-14, "max_step": 0.075},
    "refined": {"rtol": 5e-14, "atol": 2e-15, "max_step": 0.02},
}
STABLE = "#0072B2"
UNSTABLE = "#D55E00"
INK = "#263238"


def moment_mp(j: int, h: mp.mpf) -> mp.mpf:
    """J_j(h), using the analytic hypergeometric expression, 0<h<1/2."""
    cj = mp.power(2, -j - 1) * mp.binomial(2*j + 4, j + 2)
    return (
        mp.pi * cj * h**(j + 2)
        * mp.hyp2f1(mp.mpf(2*j + 3)/4, mp.mpf(2*j + 5)/4,
                     j + 3, 2*h)
    )


def moment_np(j: int, h: np.ndarray) -> np.ndarray:
    cj = float(mp.power(2, -j - 1) * mp.binomial(2*j + 4, j + 2))
    return np.pi*cj*h**(j + 2)*hyp2f1(
        (2*j + 3)/4, (2*j + 5)/4, j + 3, 2*h
    )


def interpolate(target_strings: tuple[str, ...], digits: int) -> dict:
    n = len(target_strings)
    targets = [mp.mpf(h) for h in target_strings]
    matrix = mp.matrix([[moment_mp(j, h) for j in range(n)] for h in targets])
    rhs = mp.matrix([moment_mp(n, h) for h in targets])
    coeff = list(mp.lu_solve(matrix, rhs)) + [mp.mpf(-1)]
    residuals = [sum(coeff[j]*moment_mp(j, h) for j in range(n+1)) for h in targets]
    slopes = [mp.diff(lambda z: sum(coeff[j]*moment_mp(j, z)
                                    for j in range(n+1)), h) for h in targets]
    # Independent quadrature check at an ordinary point, away from the endpoints.
    h_check = mp.mpf("0.117")
    r = mp.sqrt(2*h_check)
    moment_check_errors = []
    for j in range(n+1):
        exponent = mp.mpf(j) + mp.mpf("1.5")
        quad_value = 2*r**(2*j+4)*mp.quad(
            lambda t: ((1-t*t)/(1+r*t))**exponent, [-1, 0, 1]
        )
        moment_check_errors.append(abs((quad_value-moment_mp(j, h_check))/quad_value))
    return {
        "coefficients": [mp.nstr(c, digits) for c in coeff],
        "target_energies": list(target_strings),
        "melnikov_derivatives": [mp.nstr(s, digits) for s in slopes],
        "interpolation_residual_max": mp.nstr(max(map(abs, residuals)), 8),
        "quadrature_relative_error_max": mp.nstr(max(moment_check_errors), 8),
        "quadrature_check_energy": str(h_check),
    }


def s_of_x(x):
    """Stable signed energy coordinate s(x) for a=m=1."""
    x = np.asarray(x, dtype=float)
    u = x-1.0
    # The series avoids cancellation in u-log1p(u) close to the equilibrium.
    series = sum((-u)**j/(j+2) for j in range(14))
    near = u*np.sqrt(np.maximum(0.0, 2*series))
    far = np.sign(u)*np.sqrt(np.maximum(0.0, 2*(u-np.log1p(u))))
    return np.where(np.abs(u) < 1e-3, near, far)


def x_of_s(s):
    """Invert s(x) with the two real Lambert-W branches.

    W_0 gives x<1 for s<0; W_-1 gives x>1 for s>0. A local inverse
    series avoids the numerical branch-point loss of precision near s=0.
    """
    s = np.asarray(s, dtype=float)
    flat = s.reshape(-1)
    result = np.empty_like(flat)
    near = np.abs(flat) < 1e-4
    z = flat[near]
    result[near] = 1+z+z*z/3+z**3/36-z**4/270+z**5/4320
    for positive, branch in ((False, 0), (True, -1)):
        mask = ~near & ((flat > 0) if positive else (flat < 0))
        argument = -np.exp(-1-flat[mask]**2/2)
        result[mask] = -lambertw(argument, branch).real
    return result.reshape(s.shape)


def growth(x, coeff, delta):
    s = s_of_x(x)
    return 1-s*s/2+delta*s**3*np.polynomial.polynomial.polyval(s*s, coeff)


def growth_derivatives(x, coeff, delta):
    """Return alpha' and alpha'', avoiding a removable singularity at x=1."""
    x = np.asarray(x, dtype=float)
    s = s_of_x(x)
    sp = np.ones_like(s)
    np.divide(x-1, x*s, out=sp, where=np.abs(s) > 1e-12)
    first = -(x-1)/x
    second = -1/(x*x)
    for j, bj in enumerate(coeff):
        degree = 2*j+3
        first += delta*degree*bj*s**(degree-1)*sp
        second += delta*degree*bj*s**(degree-2)*((degree-2)*sp*sp+1/(x*x))
    return first, second


def integrate_return(h, coeff, delta, controls, sample=False):
    dcoeff = np.arange(1, len(coeff))*coeff[1:]

    def rhs(_time, state):
        q, p, _log_multiplier = state
        if 1+q <= 0:
            raise RuntimeError("Trajectory left the positive-predator coordinate domain.")
        z = p*p/(1+q)
        poly = np.polynomial.polynomial.polyval(z, coeff)
        dpoly = np.polynomial.polynomial.polyval(z, dcoeff)
        perturb = delta*p**3/(1+q)**1.5*poly
        div = delta*p*p/(1+q)**1.5*(3*poly+2*z*dpoly)
        return [p, -q+perturb, div]

    def section(_time, state):
        return state[0]

    section.direction = 1
    section.terminal = False
    sol = solve_ivp(
        rhs, (0, 8), [0, np.sqrt(2*h), 0], method="DOP853",
        events=section, dense_output=sample, **controls
    )
    crossings = np.flatnonzero(sol.t_events[0] > 1.0)
    if not sol.success or len(crossings) != 1:
        raise RuntimeError("A unique full return was not found: " + sol.message)
    i = crossings[0]
    period = float(sol.t_events[0][i])
    endpoint = sol.y_events[0][i]
    answer = {
        "displacement": float(endpoint[1]**2/2-h),
        "multiplier": float(np.exp(endpoint[2])),
        "log_multiplier": float(endpoint[2]),
        "period_sigma": period,
        "return_function_evaluations": int(sol.nfev),
    }
    if sample:
        times = np.linspace(0, period, 6001)
        trajectory = sol.sol(times)
        q, p = trajectory[:2]
        s = p/np.sqrt(1+q)
        x = x_of_s(s)
        y = 1+q
        dt_dsigma = np.ones_like(x)
        np.divide(s, x-1, out=dt_dsigma, where=np.abs(s) > 1e-10)
        dt_dsigma /= np.sqrt(y)
        alpha_prime, _ = growth_derivatives(x, coeff, delta)
        independent_log_mu = simpson(x*alpha_prime*dt_dsigma, x=times)
        answer.update({
            "period_original_t": float(simpson(dt_dsigma, x=times)),
            "independent_log_multiplier_original_coordinates": float(independent_log_mu),
            "coordinate_inversion_error_max": float(np.max(np.abs(s_of_x(x)-s))),
            "min_x": float(np.min(x)), "max_x": float(np.max(x)),
            "min_y": float(np.min(y)), "max_y": float(np.max(y)),
            "max_abs_s": float(np.max(np.abs(s))),
            "sample": {"sigma": times, "q": q, "p": p, "s": s, "x": x, "y": y},
        })
    return answer


def compute_case(n, target_strings, delta, digits, output):
    interpolation = interpolate(target_strings, digits)
    coeff = np.array([float(c) for c in interpolation["coefficients"]])
    targets = np.array([float(h) for h in target_strings])
    gap = float(np.min(np.diff(targets)))
    brackets = [(max(0.001, h-gap/2), h+gap/2) for h in targets]
    runs = {}
    samples = []
    for accuracy, controls in SOLVER_CONTROLS.items():
        rows = []
        for i, (target, bracket) in enumerate(zip(targets, brackets), start=1):
            displacements = [integrate_return(h, coeff, delta, controls)["displacement"]
                             for h in bracket]
            if displacements[0]*displacements[1] >= 0:
                raise RuntimeError(f"N={n}, cycle={i}: chosen bracket has no sign change.")
            root = brentq(
                lambda h: integrate_return(h, coeff, delta, controls)["displacement"],
                *bracket, xtol=5e-13, rtol=5e-13, maxiter=100
            )
            result = integrate_return(root, coeff, delta, controls,
                                      sample=(accuracy == "refined"))
            if "sample" in result:
                samples.append(result.pop("sample"))
            result.update({
                "cycle_index": i, "target_energy": float(target),
                "actual_section_energy": float(root),
                "energy_shift": float(root-target),
                "stability": "stable" if result["multiplier"] < 1 else "unstable",
                "bracket": list(bracket), "bracket_displacements": displacements,
            })
            rows.append(result)
        runs[accuracy] = rows
    comparisons = []
    for standard, refined in zip(runs["standard"], runs["refined"]):
        comparisons.append({
            "cycle_index": standard["cycle_index"],
            "energy_difference": abs(standard["actual_section_energy"]-refined["actual_section_energy"]),
            "multiplier_difference": abs(standard["multiplier"]-refined["multiplier"]),
        })
    projection_min = min(row["min_x"] for row in runs["refined"])
    projection_max = max(row["max_x"] for row in runs["refined"])
    width = projection_max-projection_min
    interval = [projection_min-0.02*width, projection_max+0.02*width]
    grid = np.unique(np.r_[np.linspace(*interval, 50001), 1.0])
    first, second = growth_derivatives(grid, coeff, delta)
    i_max = int(np.argmax(second))
    candidate_values = [float(second[i_max])]
    if 0 < i_max < len(grid)-1:
        optimization = minimize_scalar(
            lambda xx: -float(growth_derivatives(xx, coeff, delta)[1]),
            bounds=(grid[i_max-1], grid[i_max+1]), method="bounded",
            options={"xatol": 1e-12}
        )
        candidate_values.append(-float(optimization.fun))
    concavity = {
        "method": "50001-point grid plus local scalar refinement; not interval arithmetic",
        "cycle_projection_x": [projection_min, projection_max],
        "checked_interval_x": interval,
        "sample_count": int(len(grid)),
        "estimated_max_alpha_second": max(candidate_values),
        "estimated_min_alpha_second": float(np.min(second)),
        "min_alpha_on_checked_interval": float(np.min(growth(grid, coeff, delta))),
        "alpha_at_peak": 1.0, "alpha_prime_at_peak": 0.0,
        "alpha_second_at_peak": -1.0,
    }
    if concavity["estimated_max_alpha_second"] >= 0:
        raise RuntimeError(f"Sampled strict concavity check failed for N={n}.")
    with (output/f"coefficients_N{n}.csv").open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["power_j", "coefficient_b_j"])
        writer.writerows(enumerate(interpolation["coefficients"]))
    for i, sample in enumerate(samples, start=1):
        columns = list(sample)
        np.savetxt(output/f"cycle_N{n}_{i}.csv",
                   np.column_stack([sample[c] for c in columns]),
                   delimiter=",", header=",".join(columns), comments="", fmt="%.16e")
    return {
        "number_of_selected_cycles": n, "delta": delta,
        "interpolation": interpolation, "return_maps": runs,
        "solver_comparisons": comparisons, "concavity_check": concavity,
    }, coeff, samples


def set_plot_style():
    plt.rcParams.update({
        "font.family": "serif", "font.size": 10.95,
        "axes.labelsize": 10.95, "axes.titlesize": 10.95,
        "xtick.labelsize": 10.95, "ytick.labelsize": 10.95,
        "legend.fontsize": 10.95,
        "pgf.texsystem": "pdflatex", "pgf.rcfonts": False,
        "pgf.preamble": r"\usepackage[T1]{fontenc}\usepackage{lmodern}",
        "axes.spines.top": False, "axes.spines.right": False,
        "axes.edgecolor": "#546E7A", "text.color": INK,
        "axes.labelcolor": INK, "xtick.color": INK, "ytick.color": INK,
        "pdf.fonttype": 42, "ps.fonttype": 42,
        "savefig.facecolor": "white", "figure.facecolor": "white",
    })


def save_figure(fig, output, stem):
    fig.savefig(output/f"{stem}.pgf", backend="pgf")
    fig.savefig(output/f"{stem}.pdf", backend="pgf")
    import subprocess
    subprocess.run(["pdftoppm", "-singlefile", "-r", "160", "-png",
                    str(output/f"{stem}.pdf"), str(output/stem)], check=True,
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    plt.close(fig)


def plot_melnikov(results, coeffs, output):
    fig, axes = plt.subplots(1, 2, figsize=(156/25.4, 3.05), constrained_layout=True)
    for n, ax in zip((3, 4), axes):
        targets = np.array([float(h) for h in CASES[n]])
        gap = np.min(np.diff(targets))
        h = np.linspace(targets[0]-0.35*gap, targets[-1]+0.35*gap, 1401)
        m = sum(bj*moment_np(j, h) for j, bj in enumerate(coeffs[n]))
        curve = m/(h*h)
        normalization = np.max(np.abs(curve))
        ax.plot(h, curve/normalization, color=INK, lw=1.8)
        ax.axhline(0, color="#B0BEC5", lw=0.9)
        for row in results[n]["return_maps"]["refined"]:
            ht = row["target_energy"]
            stable = row["stability"] == "stable"
            color = STABLE if stable else UNSTABLE
            ax.axvline(ht, color=color, alpha=0.19, lw=0.8)
            ax.plot(ht, 0, "o", color=color, markersize=5.5,
                    markerfacecolor=color if stable else "white", markeredgewidth=1.4)
        ax.set(xlabel=r"Energy $h$", ylabel=r"Normalized $M(h)/h^2$",
               title=rf"$N={n}$", ylim=(-1.1, 1.1))
        ax.set_xticks(targets)
        ax.grid(axis="y", color="#ECEFF1", lw=0.65)
    save_figure(fig, output, "melnikov_prescribed_zeros")


def plot_phase(result, samples, output):
    fig, ax = plt.subplots(figsize=(121.68/25.4, 4.6), constrained_layout=True)
    for row, sample in zip(result["return_maps"]["refined"], samples):
        stable = row["stability"] == "stable"
        color = STABLE if stable else UNSTABLE
        i = row["cycle_index"]
        ax.plot(sample["x"], sample["y"], color=color, lw=1.8,
                ls="-" if stable else "--",
                label=rf"$\Gamma_{i}$: {'attracting' if stable else 'repelling'}, $\mu={row['multiplier']:.7f}$")
        idx = int((0.11+0.11*(i-1))*(len(sample["x"])-1))
        offset = 65
        ax.annotate("", xy=(sample["x"][idx+offset], sample["y"][idx+offset]),
                    xytext=(sample["x"][idx], sample["y"][idx]),
                    arrowprops={"arrowstyle": "-|>", "color": color, "lw": 1.4,
                                "mutation_scale": 12})
    ax.plot(1, 1, "o", color=INK, markersize=5)
    ax.annotate(r"$E=(1,1)$", (1, 1), xytext=(9, 8), textcoords="offset points", fontsize=10.95)
    ax.set(xlabel=r"Prey density $x$", ylabel=r"Predator density $y$",
           title=rf"Three nested cycles ($\delta={result['delta']:g}$)")
    ax.set_aspect("equal", adjustable="datalim")
    ax.margins(x=0.08, y=0.10)
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.14), frameon=False, fontsize=10.95)
    save_figure(fig, output, "three_cycles_original_coordinates")


def plot_concavity(results, coeffs, delta, output):
    fig, axes = plt.subplots(1, 2, figsize=(156/25.4, 3.05), constrained_layout=True)
    for n, color, style in ((3, STABLE, "-"), (4, UNSTABLE, "--")):
        interval = results[n]["concavity_check"]["checked_interval_x"]
        x = np.linspace(*interval, 1501)
        _, second = growth_derivatives(x, coeffs[n], delta)
        axes[0].plot(x, growth(x, coeffs[n], delta), color=color, ls=style, lw=1.7, label=rf"$N={n}$")
        axes[1].plot(x, second, color=color, ls=style, lw=1.7, label=rf"$N={n}$")
    axes[0].set(xlabel=r"Prey density $x$", ylabel=r"Local growth rate $\alpha(x)$",
                title="Local growth")
    axes[0].plot(1, 1, "o", color=INK, markersize=4)
    axes[1].set(xlabel=r"Prey density $x$", ylabel=r"Second derivative $\alpha''(x)$",
                title="Local curvature")
    axes[1].axhline(0, color="#B0BEC5", lw=0.9)
    for ax in axes:
        ax.legend(frameon=False, fontsize=10.95)
        ax.grid(axis="y", color="#ECEFF1", lw=0.65)
    save_figure(fig, output, "local_growth_and_concavity")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path,
                        default=Path(__file__).resolve().parent/"outputs")
    parser.add_argument("--delta", type=float, default=0.1)
    parser.add_argument("--mp-dps", type=int, default=80)
    args = parser.parse_args()
    if args.delta <= 0:
        parser.error("This illustration uses a positive perturbation parameter.")
    if args.mp_dps < 40:
        parser.error("Use at least 40 decimal digits for moment interpolation.")
    mp.mp.dps = args.mp_dps
    output = args.output_dir.resolve()
    output.mkdir(parents=True, exist_ok=True)
    results, coeffs, samples = {}, {}, {}
    for n, targets in CASES.items():
        results[n], coeffs[n], samples[n] = compute_case(
            n, targets, args.delta, args.mp_dps, output
        )
        print(f"N={n}, coefficients: {results[n]['interpolation']['coefficients']}", flush=True)
        for row in results[n]["return_maps"]["refined"]:
            print(f"  cycle {row['cycle_index']}: h={row['actual_section_energy']:.12g}, "
                  f"mu={row['multiplier']:.12g}, {row['stability']}", flush=True)
        print("  concavity: ", results[n]["concavity_check"], flush=True)
    metadata = {
        "status": "Numerical corroboration; not a validated computation or a proof.",
        "model": "x'=x(alpha(x)-y), y'=y(x-1)",
        "local_growth": "alpha(x)=1-s(x)^2/2+delta*s(x)^3*P(s(x)^2)",
        "coordinate": "s=sign(x-1)*sqrt(2*(x-1-log(x))); q=y-1; p=s*sqrt(y)",
        "moment": "J_j=pi*2^(-j-1)*binomial(2j+4,j+2)*h^(j+2)*2F1((2j+3)/4,(2j+5)/4;j+3;2h)",
        "parameter_delta": args.delta, "mpmath_decimal_digits": args.mp_dps,
        "solver": "scipy.solve_ivp DOP853, positive-direction q=0 return section",
        "solver_controls": SOLVER_CONTROLS,
        "software_versions": {"numpy": np.__version__, "scipy": scipy.__version__,
                              "mpmath": mp.__version__, "matplotlib": matplotlib.__version__},
        "cases": results,
    }
    (output/"results.json").write_text(json.dumps(metadata, indent=2)+"\n", encoding="utf-8")
    fields = ["N", "cycle_index", "target_energy", "actual_section_energy", "energy_shift",
              "multiplier", "log_multiplier", "stability", "period_sigma", "period_original_t",
              "displacement", "min_x", "max_x", "min_y", "max_y"]
    with (output/"cycle_summary.csv").open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        for n in CASES:
            for row in results[n]["return_maps"]["refined"]:
                writer.writerow({key: n if key == "N" else row[key] for key in fields})
    set_plot_style()
    plot_melnikov(results, coeffs, output)
    plot_phase(results[3], samples[3], output)
    plot_concavity(results, coeffs, args.delta, output)
    print(f"Wrote numerical data and figures to {output}", flush=True)


if __name__ == "__main__":
    main()