"""Report on model performance."""

from __future__ import annotations

import json
from copy import deepcopy
from pathlib import Path

import numpy as np
from numpy.typing import NDArray
from sklearn.metrics import classification_report

from vito_crop_classification.vito_logger import Logger, bh_logger


def generate_report(
    y_true: NDArray[np.int64],
    y_pred: NDArray[np.int64],
    save_dir: Path | None = None,
    logger: Logger | None = None,
    verbose: bool = True,
) -> None:
    """Generate the classification report."""
    logger = logger if (logger is not None) else bh_logger
    target_names = sorted(set(y_true) | set(y_pred))
    report = classification_report(y_true, y_pred, target_names=target_names, zero_division=0)
    report_dict = classification_report(
        y_true,
        y_pred,
        target_names=target_names,
        zero_division=0,
        output_dict=True,
    )

    # Resolve confusing naming
    report_dict["sample avg"] = deepcopy(report_dict["weighted avg"])
    del report_dict["weighted avg"]
    report_dict["class avg"] = deepcopy(report_dict["macro avg"])
    del report_dict["macro avg"]

    # Write away the report
    if save_dir is not None:
        with open(save_dir / "report.txt", "w") as f:
            f.write(report)
        with open(save_dir / "report.json", "w") as f:
            json.dump(report_dict, f, indent=2)

    # Show the report
    if verbose:
        logger("\n" + report)
