import argparse
import json
from pathlib import Path
from typing import Iterable

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    precision_recall_fscore_support,
)

from utils import ensure_dir


def compute_metrics(y_true: Iterable[int], y_pred: Iterable[int]) -> dict:
    y_true = np.asarray(list(y_true))
    y_pred = np.asarray(list(y_pred))
    precision, recall, f1, _ = precision_recall_fscore_support(
        y_true,
        y_pred,
        average="macro",
        zero_division=0,
    )
    return {
        "accuracy": float(accuracy_score(y_true, y_pred)),
        "precision_macro": float(precision),
        "recall_macro": float(recall),
        "f1_macro": float(f1),
        "support": int(len(y_true)),
    }


def plot_confusion(cm: np.ndarray, label_names: list[str], output_path: str | Path, title: str) -> None:
    plt.figure(figsize=(6, 5))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=label_names, yticklabels=label_names)
    plt.xlabel("Predicted")
    plt.ylabel("True")
    plt.title(title)
    plt.tight_layout()
    plt.savefig(output_path, dpi=200)
    plt.close()


def save_evaluation_bundle(
    output_dir: str | Path,
    split_name: str,
    y_true: Iterable[int],
    y_pred: Iterable[int],
    label_names: list[str],
    dataframe: pd.DataFrame | None = None,
    probabilities: np.ndarray | None = None,
) -> dict:
    output_dir = ensure_dir(output_dir)
    y_true = np.asarray(list(y_true))
    y_pred = np.asarray(list(y_pred))
    metrics = compute_metrics(y_true, y_pred)
    report_dict = classification_report(
        y_true,
        y_pred,
        target_names=label_names,
        zero_division=0,
        output_dict=True,
    )
    report_text = classification_report(
        y_true,
        y_pred,
        target_names=label_names,
        zero_division=0,
    )
    cm = confusion_matrix(y_true, y_pred)

    with (Path(output_dir) / f"{split_name}_metrics.json").open("w", encoding="utf-8") as file:
        json.dump(
            {
                "metrics": metrics,
                "classification_report": report_dict,
                "confusion_matrix": cm.tolist(),
            },
            file,
            indent=2,
            ensure_ascii=False,
        )

    with (Path(output_dir) / f"{split_name}_cls_report.txt").open("w", encoding="utf-8") as file:
        file.write(report_text)

    plot_confusion(cm, label_names, Path(output_dir) / f"{split_name}_confusion_matrix.png", f"{split_name} confusion matrix")

    if split_name == "test":
        with (Path(output_dir) / "cls_report.txt").open("w", encoding="utf-8") as file:
            file.write(report_text)
        plot_confusion(cm, label_names, Path(output_dir) / "confusion_matrix.png", "Test confusion matrix")

    if dataframe is not None:
        predictions = dataframe.copy()
        predictions["y_true"] = y_true
        predictions["y_pred"] = y_pred
        predictions["y_pred_label"] = [label_names[index] for index in y_pred]
        if probabilities is not None:
            for class_index, label_name in enumerate(label_names):
                predictions[f"prob_{label_name}"] = probabilities[:, class_index]
        split_path = Path(output_dir) / f"pred_{split_name}.csv"
        predictions.to_csv(split_path, index=False)
        if split_name == "test":
            predictions.to_csv(Path(output_dir) / "pred_test.csv", index=False)

    return metrics


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate a predictions CSV file.")
    parser.add_argument("--input", required=True, help="CSV file with y_true and y_pred columns.")
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--label_names", nargs="+", default=["non-clickbait", "clickbait"])
    args = parser.parse_args()

    frame = pd.read_csv(args.input)
    save_evaluation_bundle(
        output_dir=args.output_dir,
        split_name="custom",
        y_true=frame["y_true"].tolist(),
        y_pred=frame["y_pred"].tolist(),
        label_names=args.label_names,
        dataframe=frame,
    )


if __name__ == "__main__":
    main()