from __future__ import annotations

import argparse
from pathlib import Path
import pandas as pd
from .data import check_nulls, imbalance_report, stratified_split
from .metrics import threshold_sweep
from .reproducibility import snapshot_environment
from .eval_viz import roc_curve_data, pr_curve_data, save_curve_plot


def main() -> None:
    """CLI entrypoint."""
    parser = argparse.ArgumentParser(prog="tidyfit")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p1 = sub.add_parser("data-summary")
    p1.add_argument("csv")

    p2 = sub.add_parser("imbalance-report")
    p2.add_argument("csv")
    p2.add_argument("--target", required=True)

    p3 = sub.add_parser("threshold-sweep")
    p3.add_argument("csv")
    p3.add_argument("--label-col", default="y_true")
    p3.add_argument("--score-col", default="y_score")

    p4 = sub.add_parser("env-snapshot")
    p4.add_argument("--out", required=True)

    p5 = sub.add_parser("stratified-split")
    p5.add_argument("csv")
    p5.add_argument("--target", required=True)

    p6 = sub.add_parser("curve-data")
    p6.add_argument("csv")
    p6.add_argument("--label-col", default="y_true")
    p6.add_argument("--score-col", default="y_score")
    p6.add_argument("--out-dir", default=None)

    args = parser.parse_args()

    if args.cmd == "data-summary":
        df = pd.read_csv(args.csv)
        print(check_nulls(df).to_string())
    elif args.cmd == "imbalance-report":
        df = pd.read_csv(args.csv)
        print(imbalance_report(df[args.target]).to_string(index=False))
    elif args.cmd == "threshold-sweep":
        df = pd.read_csv(args.csv)
        out = threshold_sweep(df[args.label_col], df[args.score_col])
        print(out.to_string(index=False))
    elif args.cmd == "env-snapshot":
        p = snapshot_environment(args.out)
        print(str(p))
    elif args.cmd == "stratified-split":
        df = pd.read_csv(args.csv)
        train, val, test = stratified_split(df, target=args.target, random_state=42)
        print(
            f"train={len(train[0])} val={len(val[0])} test={len(test[0])} "
            f"target={args.target}"
        )
    elif args.cmd == "curve-data":
        df = pd.read_csv(args.csv)
        roc = roc_curve_data(
            df[args.label_col], df[args.score_col], thresholds=[0.2, 0.5, 0.8]
        )
        pr = pr_curve_data(
            df[args.label_col], df[args.score_col], thresholds=[0.2, 0.5, 0.8]
        )

        if args.out_dir:
            out_dir = Path(args.out_dir)
            out_dir.mkdir(parents=True, exist_ok=True)
            pd.DataFrame(roc, columns=["threshold", "fpr", "tpr"]).to_csv(
                out_dir / "roc_points.csv", index=False
            )
            pd.DataFrame(pr, columns=["threshold", "precision", "recall"]).to_csv(
                out_dir / "pr_points.csv", index=False
            )
            save_curve_plot(roc, "fpr", "tpr", "ROC", out_dir / "roc.png")
            save_curve_plot(pr, "recall", "precision", "PR", out_dir / "pr.png")
            print(f"saved={out_dir}")
        else:
            print(f"roc_points={len(roc)} pr_points={len(pr)}")
