"""Post-hoc plotting for abiogenesis experiment results."""

import json
from pathlib import Path

import numpy as np

try:
    import matplotlib.pyplot as plt
except ImportError:
    plt = None


def _require_matplotlib():
    """Raise a helpful error if matplotlib is not installed."""
    if plt is None:
        raise ImportError(
            "matplotlib is required for plotting. "
            "Install it with: pip install abiogenesis[viz]"
        )


def load_run(checkpoint_path: str) -> tuple[np.ndarray, dict]:
    """Load a checkpoint's arrays and metadata."""
    data = np.load(checkpoint_path)
    meta_path = checkpoint_path.replace(".npz", "_meta.json")
    with open(meta_path) as f:
        meta = json.load(f)
    return data, meta


def plot_ops_scatter(ops_log: np.ndarray, output_path: str | None = None,
                     subsample: int = 1):
    """Phase transition plot: ops count per interaction over time.

    This is the signature plot — a dense cloud near zero that erupts
    into high-ops interactions when replicators emerge.
    """
    _require_matplotlib()
    fig, ax = plt.subplots(figsize=(14, 6))

    x = np.arange(len(ops_log))
    if subsample > 1:
        x = x[::subsample]
        ops = ops_log[::subsample]
    else:
        ops = ops_log

    ax.scatter(x, ops, s=0.1, alpha=0.3, c="black", rasterized=True)
    ax.set_xlabel("Interaction")
    ax.set_ylabel("Operations executed")
    ax.set_title("Operations per interaction (phase transition)")
    ax.set_ylim(bottom=0)
    fig.tight_layout()

    if output_path:
        fig.savefig(output_path, dpi=150)
        print(f"Saved: {output_path}")
    else:
        plt.show()
    plt.close(fig)


def plot_entropy(snapshot_interactions: list[int],
                 snapshot_compression: list[float],
                 output_path: str | None = None):
    """Compression ratio over time — measures structure emergence."""
    _require_matplotlib()
    fig, ax = plt.subplots(figsize=(12, 5))

    ax.plot(snapshot_interactions, snapshot_compression, linewidth=0.8, color="steelblue")
    ax.set_xlabel("Interaction")
    ax.set_ylabel("Compression ratio (lower = more structure)")
    ax.set_title("Entropy (compression ratio) over time")
    ax.axhline(y=1.0, color="gray", linestyle="--", alpha=0.5, label="random baseline")
    ax.legend()
    fig.tight_layout()

    if output_path:
        fig.savefig(output_path, dpi=150)
        print(f"Saved: {output_path}")
    else:
        plt.show()
    plt.close(fig)


def plot_population_histogram(soup_state: np.ndarray, top_n: int = 30,
                               output_path: str | None = None):
    """Bar chart of the most common tapes at a given snapshot."""
    _require_matplotlib()
    from .metrics import top_sequences

    top = top_sequences(soup_state, n=top_n)
    if not top:
        print("No data for population histogram.")
        return

    labels = [f"tape_{i}" for i in range(len(top))]
    counts = [c for _, c in top]

    fig, ax = plt.subplots(figsize=(12, 5))
    ax.bar(range(len(counts)), counts, color="coral")
    ax.set_xticks(range(len(labels)))
    ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=7)
    ax.set_ylabel("Count")
    ax.set_title(f"Top {top_n} most common tapes")
    fig.tight_layout()

    if output_path:
        fig.savefig(output_path, dpi=150)
        print(f"Saved: {output_path}")
    else:
        plt.show()
    plt.close(fig)


def plot_replicator_timeline(snapshot_interactions: list[int],
                              snapshot_replicator_counts: list[dict[str, int]],
                              output_path: str | None = None):
    """Stacked area chart of replicator types over time."""
    _require_matplotlib()
    if not snapshot_replicator_counts:
        print("No replicator data to plot.")
        return

    x = np.array(snapshot_interactions)
    inanimate = np.array([s.get("inanimate", 0) for s in snapshot_replicator_counts])
    viral = np.array([s.get("viral", 0) for s in snapshot_replicator_counts])
    cellular = np.array([s.get("cellular", 0) for s in snapshot_replicator_counts])

    fig, ax = plt.subplots(figsize=(12, 5))
    ax.stackplot(
        x, inanimate, viral, cellular,
        labels=["Inanimate", "Viral", "Cellular"],
        colors=["#aaaaaa", "#e8a838", "#38a832"],
        alpha=0.8,
    )
    ax.set_xlabel("Interaction")
    ax.set_ylabel("Number of distinct replicators")
    ax.set_title("Replicator emergence timeline")
    ax.legend(loc="upper left")
    fig.tight_layout()

    if output_path:
        fig.savefig(output_path, dpi=150)
        print(f"Saved: {output_path}")
    else:
        plt.show()
    plt.close(fig)


def generate_all_plots(checkpoint_path: str, output_dir: str | None = None):
    """Generate all four plots from a checkpoint file."""
    _require_matplotlib()
    data, meta = load_run(checkpoint_path)

    if output_dir is None:
        output_dir = str(Path(checkpoint_path).parent / "plots")
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    ops_log = data["ops_log"]
    tapes = data["tapes"]

    # Subsample if > 1M points for the scatter plot
    subsample = max(1, len(ops_log) // 1_000_000)

    plot_ops_scatter(
        ops_log, output_path=f"{output_dir}/ops_scatter.png",
        subsample=subsample,
    )
    plot_entropy(
        meta["snapshot_interactions"],
        meta["snapshot_compression"],
        output_path=f"{output_dir}/entropy.png",
    )
    plot_population_histogram(
        tapes,
        output_path=f"{output_dir}/population_histogram.png",
    )
    plot_replicator_timeline(
        meta["snapshot_interactions"],
        meta["snapshot_replicator_counts"],
        output_path=f"{output_dir}/replicator_timeline.png",
    )

    # Depth timeline if data available
    if "snapshot_max_depth" in meta and meta["snapshot_max_depth"]:
        plot_depth_timeline(
            meta["snapshot_interactions"],
            meta["snapshot_max_depth"],
            meta["snapshot_mean_depth"],
            output_path=f"{output_dir}/depth_timeline.png",
        )

    print(f"\nAll plots saved to {output_dir}/")


def plot_depth_timeline(snapshot_interactions: list[int],
                        snapshot_max_depth: list[int],
                        snapshot_mean_depth: list[float],
                        output_path: str | None = None):
    """Plot max and mean symbiogenesis depth over time."""
    _require_matplotlib()
    if not snapshot_max_depth:
        print("No depth data to plot.")
        return

    fig, ax = plt.subplots(figsize=(12, 5))

    ax.plot(snapshot_interactions, snapshot_max_depth,
            linewidth=1.0, color="firebrick", label="Max depth")
    ax.plot(snapshot_interactions, snapshot_mean_depth,
            linewidth=1.0, color="steelblue", label="Mean depth")
    ax.set_xlabel("Interaction")
    ax.set_ylabel("Symbiogenesis depth")
    ax.set_title("Symbiogenesis tree depth over time")
    ax.legend()
    fig.tight_layout()

    if output_path:
        fig.savefig(output_path, dpi=150)
        print(f"Saved: {output_path}")
    else:
        plt.show()
    plt.close(fig)


def plot_phase_diagram(run_configs: list[dict],
                       output_path: str | None = None):
    """Plot phase diagram from multiple sweep runs.

    Each entry in run_configs should have:
      - "label": str
      - "mutation_rate": float
      - "snapshot_interactions": list[int]
      - "snapshot_compression": list[float]
    """
    _require_matplotlib()
    if not run_configs:
        print("No data for phase diagram.")
        return

    fig, ax = plt.subplots(figsize=(12, 6))

    colors = ["#e41a1c", "#377eb8", "#4daf4a", "#984ea3", "#ff7f00"]
    for i, run in enumerate(run_configs):
        color = colors[i % len(colors)]
        ax.plot(
            run["snapshot_interactions"],
            run["snapshot_compression"],
            linewidth=1.0, color=color,
            label=run.get("label", f"run_{i}"),
        )

    ax.set_xlabel("Interaction")
    ax.set_ylabel("Compression ratio")
    ax.set_title("Phase diagram — compression vs mutation rate")
    ax.axhline(y=1.0, color="gray", linestyle="--", alpha=0.5, label="random baseline")
    ax.legend()
    fig.tight_layout()

    if output_path:
        fig.savefig(output_path, dpi=150)
        print(f"Saved: {output_path}")
    else:
        plt.show()
    plt.close(fig)
