# src/tank_pipe/output.py
"""
Output helpers: persistence (.npz), static plots (.png/.html), animation (.gif).
Uses matplotlib's Agg backend so it works in headless environments.
The PillowWriter is used for GIF output to avoid an ffmpeg dependency.
"""
import html
import os

import numpy as np
import matplotlib
matplotlib.use("Agg")  # headless-safe; must be set before pyplot import
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter


def _pipe_primitives(history, gamma, R_gas):
    W_hist = history['W_hist']
    rho = W_hist[:, 0, :]
    u = W_hist[:, 1, :] / rho
    P = (gamma - 1) * (W_hist[:, 2, :] - 0.5 * rho * u ** 2)
    T = P / (rho * R_gas)
    a = np.sqrt(gamma * P / rho)
    Ma = u / a
    return rho, u, P, T, a, Ma


def save_history(history, pipe, path, gamma, R_gas):
    """
    Persist the full simulation history + pipe geometry to a .npz file.

    Loadable later with:
        d = np.load("results/history.npz")
        rho = d['W_hist'][:, 0, :]
        u   = d['W_hist'][:, 1, :] / rho
        P   = (d['gamma'] - 1) * (d['W_hist'][:, 2, :] - 0.5 * rho * u**2)
    """
    dirname = os.path.dirname(path)
    if dirname:
        os.makedirs(dirname, exist_ok=True)
    np.savez_compressed(
        path,
        t=history['t'],
        P1=history['P1'], T1=history['T1'],
        P2=history['P2'], T2=history['T2'],
        W_hist=history['W_hist'],
        x=pipe.x_centers,
        dx=pipe.dx,
        area=pipe.area,
        gamma=gamma,
        R_gas=R_gas,
    )


def plot_tank_pressure(history, path):
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(history['t'], history['P1'] / 1e6, label="Tank 1 (high pressure)")
    ax.plot(history['t'], history['P2'] / 1e6, label="Tank 2 (low pressure)")
    ax.set_xlabel("Time [s]")
    ax.set_ylabel("Pressure [MPa]")
    ax.set_title("Tank pressures vs time")
    ax.grid(True)
    ax.legend()
    fig.tight_layout()
    fig.savefig(path, dpi=120)
    plt.close(fig)


def plot_tank_temperature(history, path):
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(history['t'], history['T1'], label="Tank 1 (high pressure)")
    ax.plot(history['t'], history['T2'], label="Tank 2 (low pressure)")
    ax.set_xlabel("Time [s]")
    ax.set_ylabel("Temperature [K]")
    ax.set_title("Tank temperatures vs time")
    ax.grid(True)
    ax.legend()
    fig.tight_layout()
    fig.savefig(path, dpi=120)
    plt.close(fig)


def make_pipe_animation(history, pipe, path, gamma, R_gas, stride=10):
    """
    Render a GIF of the pipe's P(x), u(x), T(x) evolution over time.
    Uses PillowWriter so no ffmpeg is needed.
    """
    t = history['t']
    x = pipe.x_centers
    n_steps = history['W_hist'].shape[0]

    # Frame indices: every `stride`th snapshot, plus the final one
    frames = list(range(0, n_steps, stride))
    if frames[-1] != n_steps - 1:
        frames.append(n_steps - 1)

    # Precompute primitives for all frames in one vectorized pass
    _, u, P, T, _, Ma = _pipe_primitives(history, gamma, R_gas)

    fig, axes = plt.subplots(4, 1, figsize=(10, 12), sharex=True)

    # Pressure subplot
    line_P, = axes[0].plot(x, P[0] / 1e6)
    axes[0].set_ylabel("P [MPa]")
    axes[0].set_ylim(P.min() / 1e6 * 0.95, P.max() / 1e6 * 1.05)
    axes[0].grid(True)

    # Velocity subplot
    line_u, = axes[1].plot(x, u[0])
    axes[1].set_ylabel("u [m/s]")
    u_min, u_max = float(u.min()), float(u.max())
    pad = max(1.0, 0.05 * (u_max - u_min) if u_max > u_min else 1.0)
    axes[1].set_ylim(u_min - pad, u_max + pad)
    axes[1].grid(True)

    # Temperature subplot
    line_T, = axes[2].plot(x, T[0])
    axes[2].set_ylabel("T [K]")
    axes[2].set_ylim(T.min() * 0.95, T.max() * 1.05)
    axes[2].grid(True)

    # Mach number subplot
    line_Ma, = axes[3].plot(x, Ma[0])
    axes[3].set_ylabel("Mach [-]")
    axes[3].set_xlabel("x [m]")
    Ma_min, Ma_max = float(Ma.min()), float(Ma.max())
    pad_Ma = max(0.05, 0.05 * (Ma_max - Ma_min) if Ma_max > Ma_min else 0.05)
    axes[3].set_ylim(Ma_min - pad_Ma, Ma_max + pad_Ma)
    axes[3].grid(True)

    title = fig.suptitle("")

    def update(frame_idx):
        line_P.set_ydata(P[frame_idx] / 1e6)
        line_u.set_ydata(u[frame_idx])
        line_T.set_ydata(T[frame_idx])
        line_Ma.set_ydata(Ma[frame_idx])
        title.set_text(
            f"t = {t[frame_idx]:.5f} s   (step {frame_idx + 1}/{n_steps})"
        )
        return line_P, line_u, line_T, line_Ma, title

    anim = FuncAnimation(fig, update, frames=frames, interval=50, blit=False)

    dirname = os.path.dirname(path)
    if dirname:
        os.makedirs(dirname, exist_ok=True)
    anim.save(path, writer=PillowWriter(fps=20))
    plt.close(fig)


def plot_pipe_final_profiles(history, pipe, path, gamma, R_gas):
    """Plot final pipe P(x), u(x), T(x), and Mach(x) as a static PNG."""
    x = pipe.x_centers
    _, u, P, T, _, Ma = _pipe_primitives(history, gamma, R_gas)

    fig, axes = plt.subplots(2, 2, figsize=(12, 8), sharex=True)
    axes = axes.ravel()

    axes[0].plot(x, P[-1] / 1e6)
    axes[0].set_ylabel("P [MPa]")
    axes[0].set_title("Final pressure profile")
    axes[0].grid(True)

    axes[1].plot(x, u[-1])
    axes[1].set_ylabel("u [m/s]")
    axes[1].set_title("Final velocity profile")
    axes[1].grid(True)

    axes[2].plot(x, T[-1])
    axes[2].set_xlabel("x [m]")
    axes[2].set_ylabel("T [K]")
    axes[2].set_title("Final temperature profile")
    axes[2].grid(True)

    axes[3].plot(x, Ma[-1])
    axes[3].axhline(1.0, color="r", linestyle="--", linewidth=1, label="Mach 1")
    axes[3].set_xlabel("x [m]")
    axes[3].set_ylabel("Mach [-]")
    axes[3].set_title("Final Mach profile")
    axes[3].grid(True)
    axes[3].legend()

    fig.tight_layout()
    fig.savefig(path, dpi=120)
    plt.close(fig)


def write_summary_report(history, pipe, path, gamma, R_gas, config):
    """Write an HTML summary report for the latest simulation outputs."""
    t = history['t']
    P1 = history['P1']
    T1 = history['T1']
    P2 = history['P2']
    T2 = history['T2']
    W_hist = history['W_hist']
    x = pipe.x_centers

    _, u, P, T, _, Ma = _pipe_primitives(history, gamma, R_gas)

    V1 = config['V1']
    V2 = config['V2']

    m1_init = P1[0] * V1 / (R_gas * T1[0])
    m1_final = P1[-1] * V1 / (R_gas * T1[-1])
    m2_init = P2[0] * V2 / (R_gas * T2[0])
    m2_final = P2[-1] * V2 / (R_gas * T2[-1])
    mpipe_init = float(np.sum(W_hist[0, 0, :] * pipe.area * pipe.dx))
    mpipe_final = float(np.sum(W_hist[-1, 0, :] * pipe.area * pipe.dx))
    m_total_init = m1_init + m2_init + mpipe_init
    m_total_final = m1_final + m2_final + mpipe_final
    rel_err_m = abs(m_total_final - m_total_init) / m_total_init

    U1_init = P1[0] * V1 / (gamma - 1)
    U1_final = P1[-1] * V1 / (gamma - 1)
    U2_init = P2[0] * V2 / (gamma - 1)
    U2_final = P2[-1] * V2 / (gamma - 1)
    Upipe_init = float(np.sum(W_hist[0, 2, :] * pipe.area * pipe.dx))
    Upipe_final = float(np.sum(W_hist[-1, 2, :] * pipe.area * pipe.dx))
    U_total_init = U1_init + U2_init + Upipe_init
    U_total_final = U1_final + U2_final + Upipe_final
    rel_err_U = abs(U_total_final - U_total_init) / U_total_init

    idx_ma = np.unravel_index(np.argmax(Ma), Ma.shape)
    idx_p = np.unravel_index(np.argmax(P), P.shape)

    sections = [
        ("Simulation setup", [
            ("Gamma", f"{gamma:.3f}"),
            ("R_gas", f"{R_gas:.3f} J/(kg K)"),
            ("Tank 1", f"V={V1:.3f} m^3, P0={config['P1_INIT']/1e6:.6f} MPa, T0={config['T1_INIT']:.3f} K"),
            ("Tank 2", f"V={V2:.3f} m^3, P0={config['P2_INIT']/1e6:.6f} MPa, T0={config['T2_INIT']:.3f} K"),
            ("Pipe", f"L={config['L']:.3f} m, D={config['D']*1e3:.3f} mm, N={config['N_CELLS']}"),
            ("Run control", f"t_end={config['T_END']:.6f} s, CFL={config['CFL']:.3f}, steps={len(t)}"),
        ]),
        ("Tank states", [
            ("Tank 1 pressure", f"{P1[0]/1e6:.6f} -> {P1[-1]/1e6:.6f} MPa"),
            ("Tank 2 pressure", f"{P2[0]/1e6:.6f} -> {P2[-1]/1e6:.6f} MPa"),
            ("Tank 1 temperature", f"{T1[0]:.6f} -> {T1[-1]:.6f} K"),
            ("Tank 2 temperature", f"{T2[0]:.6f} -> {T2[-1]:.6f} K"),
        ]),
        ("Pipe extrema", [
            ("Max pressure", f"{P[idx_p]/1e6:.6f} MPa at t={t[idx_p[0]]:.6e} s, x={x[idx_p[1]]:.6f} m"),
            ("Max velocity", f"{u.max():.6f} m/s"),
            ("Min / max temperature", f"{T.min():.6f} / {T.max():.6f} K"),
            ("Max Mach", f"{Ma[idx_ma]:.6f} at t={t[idx_ma[0]]:.6e} s, x={x[idx_ma[1]]:.6f} m"),
            ("Final Mach range", f"{Ma[-1].min():.6f} -> {Ma[-1].max():.6f}"),
            ("Final supersonic cells", f"{int(np.sum(Ma[-1] > 1.0))} / {Ma.shape[1]}"),
        ]),
        ("Conservation check", [
            ("Total mass", f"{m_total_init:.12e} -> {m_total_final:.12e} kg (rel err {rel_err_m:.3e})"),
            ("Total energy", f"{U_total_init:.12e} -> {U_total_final:.12e} J (rel err {rel_err_U:.3e})"),
        ]),
    ]

    parts = [
        "<!doctype html>",
        "<html lang='en'>",
        "<head>",
        "<meta charset='utf-8'>",
        "<title>Pipe system simulation summary</title>",
        "<style>",
        "body { font-family: Arial, sans-serif; margin: 24px; line-height: 1.45; }",
        "h1, h2 { margin-bottom: 0.3em; }",
        "table { border-collapse: collapse; width: 100%; margin: 12px 0 24px; }",
        "th, td { border: 1px solid #ccc; padding: 8px 10px; text-align: left; vertical-align: top; }",
        "th { background: #f5f5f5; width: 28%; }",
        "img { max-width: 100%; height: auto; border: 1px solid #ddd; margin: 8px 0 24px; }",
        "code { background: #f5f5f5; padding: 1px 4px; }",
        "</style>",
        "</head>",
        "<body>",
        "<h1>Pipe system simulation summary</h1>",
        f"<p>Generated from <code>results/history.npz</code>. Final simulation time: {t[-1]:.6f} s.</p>",
    ]

    for title, rows in sections:
        parts.append(f"<h2>{html.escape(title)}</h2>")
        parts.append("<table>")
        for key, value in rows:
            parts.append(
                f"<tr><th>{html.escape(str(key))}</th><td>{html.escape(str(value))}</td></tr>"
            )
        parts.append("</table>")

    parts.extend([
        "<h2>Figures</h2>",
        "<p><img src='tank_pressure.png' alt='Tank pressure history'></p>",
        "<p><img src='tank_temperature.png' alt='Tank temperature history'></p>",
        "<p><img src='pipe_final_profiles.png' alt='Final pipe profiles'></p>",
        "<p>Animation: <a href='pipe_animation.gif'>pipe_animation.gif</a></p>",
        "</body>",
        "</html>",
    ])

    dirname = os.path.dirname(path)
    if dirname:
        os.makedirs(dirname, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write("\n".join(parts))