import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import pandas as pd
from sklearn.metrics import classification_report, confusion_matrix


def evaluate_classification(predictions: str | Path) -> dict:
    df = pd.read_csv(predictions)
    y_true = df["label_id"].astype(int)
    y_pred = df["predicted_label_id"].astype(int)
    diff = (y_true - y_pred).abs()
    report = {
        "exact_accuracy": float((diff == 0).mean()),
        "within_1_accuracy": float((diff <= 1).mean()),
        "large_error_rate": float((diff >= 2).mean()),
        "confusion_matrix": confusion_matrix(y_true, y_pred).tolist(),
        "label_report": classification_report(y_true, y_pred, output_dict=True, zero_division=0),
    }
    if "store_id" in df.columns:
        report["store_accuracy"] = (df.assign(correct=diff == 0).groupby("store_id")["correct"].mean().to_dict())
    if "steam_level" in df.columns:
        report["steam_level_accuracy"] = (df.assign(correct=diff == 0).groupby("steam_level")["correct"].mean().to_dict())
    return report


def main() -> int:
    parser = argparse.ArgumentParser(description="Evaluate density or steam classification predictions.")
    parser.add_argument("--predictions", required=True)
    parser.add_argument("--output", default="reports/density/classification_report.json")
    args = parser.parse_args()
    report = evaluate_classification(args.predictions)
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    Path(args.output).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())