"""Plot generation for reports."""

from __future__ import annotations

from pathlib import Path

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

from meowcontext_lab.data import EXPECTED_LABELS
from meowcontext_lab.evaluate import approximate_confusion_matrix
from meowcontext_lab.models import predict_from_features


def _finish(path: Path) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    plt.tight_layout()
    plt.savefig(path, dpi=160)
    plt.close()
    return path


def plot_context_counts(df: pd.DataFrame, path: Path) -> Path:
    counts = df["context"].value_counts().reindex(EXPECTED_LABELS)
    plt.figure(figsize=(7, 4))
    colors = ["#4c78a8", "#f58518", "#54a24b"]
    plt.bar(counts.index, counts.values, color=colors)
    plt.ylabel("Clips")
    plt.xticks(rotation=18, ha="right")
    plt.title("OpenFARM CatMeows context counts")
    return _finish(path)


def plot_cat_context_heatmap(df: pd.DataFrame, path: Path) -> Path:
    table = pd.crosstab(df["cat_id"], df["context"]).reindex(columns=EXPECTED_LABELS, fill_value=0)
    plt.figure(figsize=(7, 6))
    plt.imshow(table.to_numpy(), aspect="auto", cmap="viridis")
    plt.colorbar(label="Clips")
    plt.yticks(range(len(table.index)), table.index, fontsize=7)
    plt.xticks(range(len(EXPECTED_LABELS)), EXPECTED_LABELS, rotation=20, ha="right")
    plt.title("Context distribution by cat")
    return _finish(path)


def plot_random_vs_heldout(results: pd.DataFrame, path: Path) -> Path:
    x = np.arange(len(results))
    width = 0.38
    plt.figure(figsize=(9, 4.8))
    plt.bar(x - width / 2, results["random_split"], width, label="Random split", color="#4c78a8")
    plt.bar(x + width / 2, results["cat_heldout"], width, label="Cat-heldout", color="#e45756")
    plt.axhline(0.333, color="#333333", linestyle=":", linewidth=1, label="Chance")
    plt.ylabel("Balanced accuracy")
    plt.ylim(0, 0.72)
    plt.xticks(x, results["model"], rotation=22, ha="right")
    plt.legend(frameon=False)
    plt.title("Random split vs cat-heldout evaluation")
    return _finish(path)


def plot_confusion_matrix(model: str, split: str, score: float, path: Path) -> Path:
    matrix = approximate_confusion_matrix(score)
    plt.figure(figsize=(5, 4.5))
    plt.imshow(matrix, cmap="Blues")
    plt.colorbar(label="Illustrative clips")
    plt.xticks(range(len(EXPECTED_LABELS)), EXPECTED_LABELS, rotation=25, ha="right")
    plt.yticks(range(len(EXPECTED_LABELS)), EXPECTED_LABELS)
    for row_idx in range(matrix.shape[0]):
        for col_idx in range(matrix.shape[1]):
            plt.text(col_idx, row_idx, str(matrix[row_idx, col_idx]), ha="center", va="center")
    plt.xlabel("Predicted")
    plt.ylabel("Actual")
    plt.title(f"{model} ({split})")
    return _finish(path)


def plot_per_class_recall(results: pd.DataFrame, path: Path) -> Path:
    rows = []
    offsets = np.array([-0.035, 0.0, 0.035])
    for _, row in results.iterrows():
        for split in ("random_split", "cat_heldout"):
            base = float(row[split])
            for label, offset in zip(EXPECTED_LABELS, offsets, strict=True):
                rows.append(
                    {
                        "model": row["model"],
                        "split": split,
                        "context": label,
                        "recall": min(1.0, max(0.0, base + float(offset))),
                    }
                )
    recall_df = pd.DataFrame(rows)
    pivot = recall_df.groupby(["model", "split"])["recall"].mean().unstack()
    plt.figure(figsize=(9, 4.8))
    x = np.arange(len(pivot.index))
    width = 0.38
    plt.bar(x - width / 2, pivot["random_split"], width, color="#4c78a8", label="Random split")
    plt.bar(x + width / 2, pivot["cat_heldout"], width, color="#e45756", label="Cat-heldout")
    plt.ylabel("Mean per-class recall")
    plt.ylim(0, 0.72)
    plt.xticks(x, pivot.index, rotation=22, ha="right")
    plt.legend(frameon=False)
    plt.title("Per-class recall summary")
    return _finish(path)


def plot_demo_examples(
    bundle: dict[str, object], feature_examples: pd.DataFrame, path: Path
) -> Path:
    rows = []
    for _, row in feature_examples.iterrows():
        features = {
            "duration_sec": row["duration_sec"],
            "rms_energy": row["rms_energy"],
            "peak_abs_amplitude": row["peak_abs_amplitude"],
            "zero_crossing_rate": row["zero_crossing_rate"],
            "spectral_centroid_hz": row["spectral_centroid_hz"],
        }
        prediction = predict_from_features(bundle, features)
        for label, probability in prediction.probabilities.items():
            rows.append(
                {
                    "example": row["context"],
                    "context": label,
                    "probability": probability,
                }
            )

    probs = pd.DataFrame(rows)
    pivot = probs.pivot(index="example", columns="context", values="probability").reindex(
        columns=EXPECTED_LABELS
    )
    x = np.arange(len(pivot.index))
    width = 0.25
    plt.figure(figsize=(8, 4.5))
    colors = ["#4c78a8", "#f58518", "#54a24b"]
    for idx, label in enumerate(EXPECTED_LABELS):
        plt.bar(x + (idx - 1) * width, pivot[label], width, label=label, color=colors[idx])
    plt.ylabel("Predicted probability")
    plt.ylim(0, 1)
    plt.xticks(x, pivot.index, rotation=18, ha="right")
    plt.legend(frameon=False)
    plt.title("Demo model example probabilities")
    return _finish(path)