"""Run two classic chaotic systems and draw their evolution and characteristic plots.

Usage: python3 chaos_homework.py
"""

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 numpy as np
from scipy.integrate import solve_ivp


ROOT = Path(__file__).resolve().parent
LORENZ = {"sigma": 10.0, "rho": 28.0, "beta": 8.0 / 3.0}
LORENZ_Y0 = np.array([1.0, 1.0, 1.0])
LOGISTIC_R = 3.9
LOGISTIC_X0 = 0.2


def lorenz_rhs(_t: float, y: np.ndarray, *, sigma: float, rho: float, beta: float) -> np.ndarray:
    x, yy, z = y
    return np.array([sigma * (yy - x), x * (rho - z) - yy, x * yy - beta * z])


def lorenz_trajectory(t_end: float, dt: float, y0: np.ndarray = LORENZ_Y0):
    t = np.arange(0.0, t_end + dt / 2, dt)
    solution = solve_ivp(
        lambda tt, yy: lorenz_rhs(tt, yy, **LORENZ),
        (0.0, t_end),
        np.asarray(y0, dtype=float),
        t_eval=t,
        method="DOP853",
        rtol=1e-10,
        atol=1e-12,
    )
    if not solution.success or not np.isfinite(solution.y).all():
        raise RuntimeError(f"Lorenz integration failed: {solution.message}")
    return solution


def logistic_series(r: float, x0: float, n: int) -> np.ndarray:
    if not 0.0 <= r <= 4.0 or not 0.0 <= x0 <= 1.0 or n < 1:
        raise ValueError("Require 0 <= r <= 4, 0 <= x0 <= 1 and n >= 1")
    x = np.empty(n + 1)
    x[0] = x0
    for k in range(n):
        x[k + 1] = r * x[k] * (1.0 - x[k])
    return x


def smoke() -> None:
    sol = lorenz_trajectory(1.0, 0.05)
    logistic = logistic_series(LOGISTIC_R, LOGISTIC_X0, 20)
    assert np.max(np.abs(sol.y[:, 0] - LORENZ_Y0)) < 1e-14
    assert np.isfinite(sol.y).all()
    assert np.all((0.0 <= logistic) & (logistic <= 1.0))
    result = {
        "lorenz_points": len(sol.t),
        "lorenz_final": sol.y[:, -1].tolist(),
        "logistic_points": len(logistic),
        "logistic_final": float(logistic[-1]),
    }
    out = ROOT / "results" / "smoke.json"
    out.parent.mkdir(exist_ok=True)
    out.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")
    print(json.dumps(result, ensure_ascii=False))


def save_csv(path: Path, header: tuple[str, ...], rows) -> None:
    path.parent.mkdir(exist_ok=True)
    with path.open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(header)
        writer.writerows(rows)


def save_plot(fig, name: str) -> None:
    path = ROOT / "figures" / name
    path.parent.mkdir(exist_ok=True)
    fig.savefig(path, dpi=220, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved {path.relative_to(ROOT)}")


def bifurcation_data(n_r: int = 1200, burn: int = 800, keep: int = 120):
    """Vectorized logistic iteration; retain only post-transient states."""
    r = np.linspace(2.5, 4.0, n_r)
    x = np.full_like(r, LOGISTIC_X0)
    for _ in range(burn):
        x = r * x * (1.0 - x)
    retained = np.empty((keep, n_r))
    for k in range(keep):
        x = r * x * (1.0 - x)
        retained[k] = x
    if not np.isfinite(retained).all() or np.any((retained < 0) | (retained > 1)):
        raise RuntimeError("Logistic scan left the invariant interval [0, 1]")
    return r, retained


def run_full() -> None:
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 10,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "figure.facecolor": "white",
    })

    # System 1: continuous Lorenz equations.
    sol = lorenz_trajectory(t_end=40.0, dt=0.01)
    save_csv(ROOT / "results" / "lorenz_trajectory.csv", ("t", "x", "y", "z"),
             zip(sol.t, *sol.y))

    fig, ax = plt.subplots(figsize=(8, 3.5))
    ax.plot(sol.t, sol.y[0], color="#0072B2", linewidth=0.8)
    ax.set(xlabel="Time t", ylabel="x(t)", title="Lorenz system: time evolution")
    ax.grid(alpha=0.2)
    save_plot(fig, "01_lorenz_evolution.png")

    fig = plt.figure(figsize=(7, 6))
    ax = fig.add_subplot(projection="3d")
    stable = sol.t >= 5.0  # Remove initial transient only for the attractor view.
    ax.plot(*sol.y[:, stable], color="#0072B2", linewidth=0.45)
    ax.set(xlabel="x", ylabel="y", zlabel="z", title="Lorenz attractor")
    ax.view_init(elev=25, azim=-60)
    save_plot(fig, "02_lorenz_attractor.png")

    # System 2: discrete logistic map.
    logistic = logistic_series(LOGISTIC_R, LOGISTIC_X0, 300)
    save_csv(ROOT / "results" / "logistic_trajectory.csv", ("n", "x_n"),
             enumerate(logistic))

    fig, ax = plt.subplots(figsize=(8, 3.5))
    n = np.arange(logistic.size)
    ax.plot(n, logistic, color="#D55E00", linewidth=0.8)
    ax.set(xlabel="Iteration n", ylabel=r"$x_n$", ylim=(0, 1),
           title=f"Logistic map: time evolution (r = {LOGISTIC_R})")
    ax.grid(alpha=0.2)
    save_plot(fig, "03_logistic_evolution.png")

    r, retained = bifurcation_data()
    fig, ax = plt.subplots(figsize=(9, 5))
    ax.plot(np.broadcast_to(r, retained.shape).ravel(), retained.ravel(),
            ",", color="black", alpha=0.32)
    ax.axvline(3.0, color="#D55E00", linestyle="--", linewidth=1,
               label="First period doubling: r = 3")
    ax.set(xlabel="Growth parameter r", ylabel="Long-run x", xlim=(2.5, 4),
           ylim=(0, 1), title="Logistic map: bifurcation diagram")
    ax.legend(frameon=False, loc="upper left", fontsize=8)
    save_plot(fig, "04_logistic_bifurcation.png")

    summary = {
        "lorenz": {"parameters": LORENZ, "initial_state": LORENZ_Y0.tolist(),
                   "time_range": [0, 40], "step": 0.01, "trajectory_points": len(sol.t)},
        "logistic": {"r_for_evolution": LOGISTIC_R, "initial_state": LOGISTIC_X0,
                     "trajectory_iterations": 300, "bifurcation_r_range": [2.5, 4.0],
                     "bifurcation_r_count": len(r), "transient_iterations": 800,
                     "retained_iterations_per_r": retained.shape[0]},
    }
    (ROOT / "results" / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n")
    print("Done. Figures and numerical data are in figures/ and results/.")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", help="run a small numerical check")
    args = parser.parse_args()
    if args.smoke:
        smoke()
    else:
        run_full()


if __name__ == "__main__":
    main()