"""Plotting helpers for model diagnostics.

Four figures:
    * :func:`plot_parity`           - test-set pred vs true per target.
    * :func:`plot_r2_bars`          - per-target R^2 across models.
    * :func:`plot_training_curves`  - train + val concrete loss vs epoch.
    * :func:`plot_rve_graph`        - one generated 2D concrete RVE.

All functions accept a ``save_path``; if provided the figure is written and
the matplotlib ``Figure`` is also returned for further customisation.
"""

from __future__ import annotations

import os
from typing import Dict, Iterable, List, Optional, Sequence

import matplotlib

matplotlib.use("Agg")  # safe headless backend
import matplotlib.pyplot as plt
import numpy as np
import torch

from .graph_generator import ConcreteGraph
from .schema import (
    CONCRETE_TARGETS,
    EDGE_TYPE_AGGREGATE_AGGREGATE,
    EDGE_TYPE_ITZ,
    EDGE_TYPE_MORTAR_MORTAR,
    NODE_TYPE_AGGREGATE,
    NODE_TYPE_MORTAR,
)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _ensure_dir(path: str) -> None:
    parent = os.path.dirname(os.path.abspath(path))
    os.makedirs(parent, exist_ok=True)


def _r2(true: np.ndarray, pred: np.ndarray) -> float:
    ss_res = float(np.sum((true - pred) ** 2))
    ss_tot = float(np.sum((true - true.mean()) ** 2))
    if ss_tot < 1e-20:
        return float("nan")
    return 1.0 - ss_res / ss_tot


# ---------------------------------------------------------------------------
# 1. Parity plots
# ---------------------------------------------------------------------------


def plot_parity(
    predictions: Dict[str, np.ndarray],
    truths: np.ndarray,
    model_order: Sequence[str],
    save_path: Optional[str] = None,
    target_names: Sequence[str] = CONCRETE_TARGETS,
):
    """One subplot per target, one row per model.

    Parameters
    ----------
    predictions: dict
        Mapping ``model_name -> (N, T)`` test predictions (numpy).
    truths: ndarray
        Test ground-truth array of shape ``(N, T)``.
    model_order: sequence of str
        Which models (and in which order) to draw.
    """

    n_models = len(model_order)
    n_targets = len(target_names)
    fig, axes = plt.subplots(
        n_models,
        n_targets,
        figsize=(2.6 * n_targets, 2.6 * n_models),
        squeeze=False,
    )
    for i, name in enumerate(model_order):
        pred = predictions[name]
        for t in range(n_targets):
            ax = axes[i, t]
            y_true = truths[:, t]
            y_pred = pred[:, t]
            r2 = _r2(y_true, y_pred)
            lo = float(min(y_true.min(), y_pred.min()))
            hi = float(max(y_true.max(), y_pred.max()))
            span = hi - lo if hi > lo else max(1.0, abs(hi))
            lo -= 0.05 * span
            hi += 0.05 * span

            ax.plot([lo, hi], [lo, hi], "k--", lw=0.8, alpha=0.6)
            ax.scatter(y_true, y_pred, s=22, alpha=0.7, edgecolor="none")
            ax.set_xlim(lo, hi)
            ax.set_ylim(lo, hi)
            ax.set_aspect("equal", adjustable="box")
            ax.tick_params(labelsize=7)
            ax.text(
                0.04,
                0.93,
                f"$R^2$={r2:.2f}",
                transform=ax.transAxes,
                fontsize=8,
                va="top",
            )
            if i == n_models - 1:
                ax.set_xlabel(target_names[t], fontsize=8)
            if t == 0:
                ax.set_ylabel(f"{name}\npred", fontsize=8)
    fig.suptitle("Test-set parity (true on x, predicted on y)", fontsize=11)
    fig.tight_layout(rect=(0, 0, 1, 0.97))
    if save_path is not None:
        _ensure_dir(save_path)
        fig.savefig(save_path, dpi=160)
    return fig


# ---------------------------------------------------------------------------
# 2. Per-target R^2 bars
# ---------------------------------------------------------------------------


def plot_r2_bars(
    r2_by_model: Dict[str, np.ndarray],
    model_order: Sequence[str],
    save_path: Optional[str] = None,
    target_names: Sequence[str] = CONCRETE_TARGETS,
):
    n_targets = len(target_names)
    n_models = len(model_order)
    x = np.arange(n_targets)
    width = 0.8 / max(1, n_models)

    fig, ax = plt.subplots(figsize=(1.2 * n_targets + 2.0, 4.0))
    colours = plt.cm.tab10(np.linspace(0, 1, max(3, n_models)))
    for i, name in enumerate(model_order):
        r2 = np.asarray(r2_by_model[name])
        offset = (i - (n_models - 1) / 2.0) * width
        ax.bar(x + offset, r2, width=width * 0.95, label=name, color=colours[i])
    ax.axhline(0.0, color="k", lw=0.6, alpha=0.5)
    ax.set_xticks(x)
    ax.set_xticklabels(target_names, rotation=30, ha="right", fontsize=8)
    ax.set_ylabel("Test-set $R^2$")
    ax.set_title("Per-target $R^2$ by model")
    ax.legend(fontsize=8, frameon=False)
    ax.set_ylim(min(-0.2, min((r2_by_model[n].min() for n in model_order))) - 0.05, 1.05)
    fig.tight_layout()
    if save_path is not None:
        _ensure_dir(save_path)
        fig.savefig(save_path, dpi=160)
    return fig


# ---------------------------------------------------------------------------
# 3. Training curves
# ---------------------------------------------------------------------------


def plot_training_curves(
    histories: Dict[str, List[Dict[str, float]]],
    model_order: Sequence[str],
    save_path: Optional[str] = None,
):
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    train_ax, val_ax = axes
    colours = plt.cm.tab10(np.linspace(0, 1, max(3, len(model_order))))

    for i, name in enumerate(model_order):
        history = histories[name]
        epochs = np.arange(1, len(history) + 1)
        train_loss = np.asarray([h["concrete"] for h in history])
        train_ax.plot(epochs, train_loss, label=name, color=colours[i], lw=1.4)
        if "val_concrete" in history[0]:
            val_loss = np.asarray([h.get("val_concrete", np.nan) for h in history])
            val_ax.plot(epochs, val_loss, label=name, color=colours[i], lw=1.4)
            if np.isfinite(val_loss).any():
                best_idx = int(np.nanargmin(val_loss))
                val_ax.scatter(
                    [epochs[best_idx]],
                    [val_loss[best_idx]],
                    color=colours[i],
                    edgecolor="black",
                    zorder=5,
                    s=42,
                )

    for ax, title in zip(axes, ("train concrete loss", "val concrete loss")):
        ax.set_yscale("log")
        ax.set_xlabel("epoch")
        ax.set_ylabel("normalised MSE")
        ax.set_title(title)
        ax.grid(True, alpha=0.3)
        ax.legend(fontsize=8, frameon=False)
    fig.suptitle("Training curves (best-val checkpoint marked)", fontsize=11)
    fig.tight_layout(rect=(0, 0, 1, 0.95))
    if save_path is not None:
        _ensure_dir(save_path)
        fig.savefig(save_path, dpi=160)
    return fig


# ---------------------------------------------------------------------------
# 4. Concrete RVE graph visualization
# ---------------------------------------------------------------------------


def plot_rve_graph(
    graph: ConcreteGraph,
    rve_size_mm: float = 150.0,
    save_path: Optional[str] = None,
):
    """Draw aggregates as circles, mortar nodes as dots, edges coloured by type."""

    pos = graph.pos.detach().cpu().numpy()
    node_type = graph.node_type.detach().cpu().numpy()
    edge_index = graph.edge_index.detach().cpu().numpy()
    edge_type = graph.edge_type.detach().cpu().numpy()
    x = graph.x.detach().cpu().numpy()

    fig, ax = plt.subplots(figsize=(7, 7))

    # Edges first (so nodes draw on top).
    seen = set()
    edge_styles = {
        EDGE_TYPE_ITZ: dict(color="#d62728", alpha=0.6, lw=0.8, ls="-", label="ITZ"),
        EDGE_TYPE_MORTAR_MORTAR: dict(
            color="#7f7f7f", alpha=0.45, lw=0.6, ls="-", label="mortar-mortar"
        ),
        EDGE_TYPE_AGGREGATE_AGGREGATE: dict(
            color="#1f77b4", alpha=0.6, lw=0.7, ls=":", label="agg-agg"
        ),
    }
    legend_seen = set()
    for k in range(edge_index.shape[1]):
        a = int(edge_index[0, k])
        b = int(edge_index[1, k])
        if a == b:
            continue
        key = (min(a, b), max(a, b))
        if key in seen:
            continue
        seen.add(key)
        style = edge_styles[int(edge_type[k])]
        label = style["label"] if int(edge_type[k]) not in legend_seen else None
        legend_seen.add(int(edge_type[k]))
        ax.plot(
            [pos[a, 0], pos[b, 0]],
            [pos[a, 1], pos[b, 1]],
            color=style["color"],
            alpha=style["alpha"],
            lw=style["lw"],
            ls=style["ls"],
            label=label,
        )

    # Aggregate circles (diameter is the first feature channel).
    for i in range(pos.shape[0]):
        if int(node_type[i]) == NODE_TYPE_AGGREGATE:
            diameter = float(x[i, 0])
            radius = max(1.0, 0.5 * diameter)
            ax.add_patch(
                plt.Circle(
                    (pos[i, 0], pos[i, 1]),
                    radius=radius,
                    facecolor="#bdbdbd",
                    edgecolor="black",
                    lw=0.6,
                    alpha=0.85,
                    zorder=2,
                )
            )

    mortar_pts = pos[node_type == NODE_TYPE_MORTAR]
    ax.scatter(
        mortar_pts[:, 0],
        mortar_pts[:, 1],
        s=18,
        c="#2ca02c",
        edgecolor="black",
        lw=0.3,
        zorder=3,
        label="mortar patch",
    )

    ax.set_xlim(0, rve_size_mm)
    ax.set_ylim(0, rve_size_mm)
    ax.set_aspect("equal")
    ax.set_xlabel("x [mm]")
    ax.set_ylabel("y [mm]")
    ax.set_title("Generated 2D concrete RVE (aggregates + mortar + edge types)")
    ax.legend(fontsize=8, loc="upper right", frameon=True)
    fig.tight_layout()
    if save_path is not None:
        _ensure_dir(save_path)
        fig.savefig(save_path, dpi=160)
    return fig


# ---------------------------------------------------------------------------
# Prediction collection helper
# ---------------------------------------------------------------------------


@torch.no_grad()
def collect_predictions(
    model: torch.nn.Module,
    loader: Iterable,
    device: torch.device,
) -> Dict[str, np.ndarray]:
    """Return ``{"pred": (N, T), "true": (N, T)}`` arrays on CPU."""

    model.eval()
    preds, trues = [], []
    for batch in loader:
        batch = batch.to(device)
        out = model(batch)
        preds.append(out["concrete_pred"].cpu())
        trues.append(batch.concrete_target.cpu())
    return {
        "pred": torch.cat(preds, dim=0).numpy(),
        "true": torch.cat(trues, dim=0).numpy(),
    }