"""
Plot numerical vs exact solution from solution.csv produced by MonitorSolutionCSV.
Works for both 1D and 2D output.

Usage:
    python plot_solution.py                    # plot all timesteps
    python plot_solution.py --step 0 10 50     # plot specific steps
    python plot_solution.py --animate          # animate over time
"""

import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# ── Load ─────────────────────────────────────────────────────────────────────

def load(path="solution.csv"):
    df = pd.read_csv(path)
    df = df.sort_values(["time", "x"] + (["y"] if "y" in df.columns else []))
    return df

# ── Detect dimensionality ─────────────────────────────────────────────────────

def is_2d(df):
    return "y" in df.columns

# ── 1-D plots ─────────────────────────────────────────────────────────────────

def plot_1d_steps(df, steps=None):
    times = sorted(df["time"].unique())
    if steps is not None:
        times = [times[i] for i in steps if i < len(times)]

    n = len(times)
    ncols = min(4, n)
    nrows = (n + ncols - 1) // ncols
    fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 3.5 * nrows),
                             squeeze=False)
    axes_flat = axes.flatten()

    for ax, t in zip(axes_flat, times):
        sub = df[np.isclose(df["time"], t)]
        ax.plot(sub["x"], sub["phi_exact"], "k--", lw=1.5, label="Exact")
        ax.plot(sub["x"], sub["phi_num"],   "b-",  lw=1.5, label="DG")
        ax.set_title(f"t = {t:.4g}")
        ax.set_xlabel("x")
        ax.set_ylabel("φ")
        ax.legend(fontsize=7)
        ax.grid(alpha=0.3)

    for ax in axes_flat[n:]:
        ax.set_visible(False)

    fig.suptitle("DG vs Exact — 1D Advection", fontsize=13, y=1.01)
    fig.tight_layout()
    plt.savefig("solution_1d.png", dpi=150, bbox_inches="tight")
    print("Saved solution_1d.png")
    plt.show()


def animate_1d(df):
    times = sorted(df["time"].unique())
    fig, ax = plt.subplots(figsize=(7, 4))

    sub0 = df[np.isclose(df["time"], times[0])]
    line_ex, = ax.plot(sub0["x"], sub0["phi_exact"], "k--", lw=1.5, label="Exact")
    line_dg, = ax.plot(sub0["x"], sub0["phi_num"],   "b-",  lw=1.5, label="DG")
    title = ax.set_title("")
    ax.set_xlabel("x"); ax.set_ylabel("φ")
    ax.legend(); ax.grid(alpha=0.3)

    # Fix axis limits over all time
    ax.set_ylim(df["phi_exact"].min() - 0.05, df["phi_exact"].max() + 0.1)

    def update(frame):
        t = times[frame]
        sub = df[np.isclose(df["time"], t)]
        line_ex.set_data(sub["x"], sub["phi_exact"])
        line_dg.set_data(sub["x"], sub["phi_num"])
        title.set_text(f"t = {t:.4g}")
        return line_ex, line_dg, title

    ani = animation.FuncAnimation(fig, update, frames=len(times),
                                  interval=80, blit=True)
    ani.save("solution_1d.gif", writer="pillow", fps=12)
    print("Saved solution_1d.gif")
    plt.show()

# ── 2-D plots ─────────────────────────────────────────────────────────────────

def plot_2d_steps(df, steps=None):
    from scipy.interpolate import griddata

    times = sorted(df["time"].unique())
    if steps is not None:
        times = [times[i] for i in steps if i < len(times)]

    xi = np.linspace(df["x"].min(), df["x"].max(), 200)
    yi = np.linspace(df["y"].min(), df["y"].max(), 200)
    Xi, Yi = np.meshgrid(xi, yi)

    for t in times:
        sub = df[np.isclose(df["time"], t)]
        pts = sub[["x", "y"]].values

        fig, axes = plt.subplots(1, 3, figsize=(14, 4))
        vmin = sub["phi_exact"].min()
        vmax = sub["phi_exact"].max()
        kw = dict(vmin=vmin, vmax=vmax, cmap="viridis")

        for ax, col, label in zip(axes,
                                   ["phi_num", "phi_exact",
                                    None],          # third panel = error
                                   ["DG Numerical", "Exact", "Error |DG − Exact|"]):
            if col is not None:
                Zi = griddata(pts, sub[col].values, (Xi, Yi), method="linear")
                im = ax.pcolormesh(Xi, Yi, Zi, **kw)
            else:
                err = np.abs(sub["phi_num"].values - sub["phi_exact"].values)
                Zi = griddata(pts, err, (Xi, Yi), method="linear")
                im = ax.pcolormesh(Xi, Yi, Zi, cmap="hot_r",
                                   vmin=0, vmax=err.max())
            fig.colorbar(im, ax=ax)
            ax.set_title(label); ax.set_xlabel("x"); ax.set_ylabel("y")

        fig.suptitle(f"t = {t:.4g}", fontsize=13)
        fig.tight_layout()
        fname = f"solution_2d_t{t:.4f}.png"
        plt.savefig(fname, dpi=150, bbox_inches="tight")
        print(f"Saved {fname}")
        plt.show()

# ── L2 error over time ────────────────────────────────────────────────────────

def plot_l2_error(df):
    records = []
    for t, grp in df.groupby("time"):
        err = np.sqrt(((grp["phi_num"] - grp["phi_exact"])**2).mean())
        records.append((t, err))
    records = sorted(records)
    ts, errs = zip(*records)

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.semilogy(ts, errs, "o-", lw=1.5)
    ax.set_xlabel("time"); ax.set_ylabel("L² error (RMS over cells)")
    ax.set_title("L² error over time"); ax.grid(alpha=0.3)
    fig.tight_layout()
    plt.savefig("l2_error.png", dpi=150, bbox_inches="tight")
    print("Saved l2_error.png")
    plt.show()

# ── CLI ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv",     default="solution.csv")
    parser.add_argument("--step",    nargs="*", type=int, default=None,
                        help="Indices of timesteps to plot (default: all)")
    parser.add_argument("--animate", action="store_true",
                        help="Produce animated GIF (1D only)")
    parser.add_argument("--error",   action="store_true",
                        help="Also plot L² error over time")
    args = parser.parse_args()

    df = load(args.csv)
    print(f"Loaded {len(df)} rows, {df['time'].nunique()} timesteps, "
          f"{'2D' if is_2d(df) else '1D'} data.")

    if args.error:
        plot_l2_error(df)

    if is_2d(df):
        plot_2d_steps(df, steps=args.step)
    else:
        if args.animate:
            animate_1d(df)
        else:
            plot_1d_steps(df, steps=args.step)