from __future__ import annotations

import argparse
from pathlib import Path

from sepsis_mcp.runner import RunConfig, run_experiments


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="sepsis-mcp")
    subparsers = parser.add_subparsers(dest="command", required=True)

    run_parser = subparsers.add_parser("run")
    run_parser.add_argument("--data-root", type=Path, required=True)
    run_parser.add_argument("--train-hospital", choices=["A", "B"], default="A")
    run_parser.add_argument("--test-mode", choices=["aa", "ab", "both"], default="both")
    run_parser.add_argument(
        "--model-type",
        choices=["sklearn_gbdt", "xgboost", "logistic_regression"],
        default="sklearn_gbdt",
    )
    run_parser.add_argument("--train-patients", type=int, default=128)
    run_parser.add_argument("--selection-patients", type=int, default=0)
    run_parser.add_argument("--calibration-patients", type=int, default=64)
    run_parser.add_argument("--test-patients", type=int, default=128)
    run_parser.add_argument("--lookback-hours", type=int, default=6)
    run_parser.add_argument("--horizon-hours", type=int, default=6)
    run_parser.add_argument("--alpha", type=float, default=0.1)
    run_parser.add_argument("--random-state", type=int, default=0)
    run_parser.add_argument("--weighted-shrinkage-lambda", type=float, default=0.5)
    run_parser.add_argument("--mondrian-tilted-shrinkage-lambda", type=float, default=0.5)
    run_parser.add_argument("--mondrian-tilted-bandwidth", type=float, default=None)
    run_parser.add_argument(
        "--missingness-grouping-strategy",
        choices=["coverage_gap_variable", "mask_cluster"],
        default=None,
    )
    run_parser.add_argument("--enable-external-baselines", action="store_true")
    run_parser.add_argument("--enable-learned-partition", action="store_true")
    run_parser.add_argument(
        "--mask-strategy",
        choices=["none", "random_drop", "block_missing", "selective_missing"],
        default="none",
    )
    run_parser.add_argument("--mask-rate", type=float, default=0.3)
    run_parser.add_argument("--block-length", type=int, default=4)
    run_parser.add_argument("--mask-random-state", type=int, default=0)
    run_parser.add_argument("--selective-feature", action="append", dest="selective_features")
    run_parser.add_argument("--output-dir", type=Path, required=True)
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command == "run":
        config = RunConfig(
            data_root=args.data_root,
            train_hospital=args.train_hospital,
            test_mode=args.test_mode,
            model_type=args.model_type,
            train_patients=args.train_patients,
            selection_patients=args.selection_patients,
            calibration_patients=args.calibration_patients,
            test_patients=args.test_patients,
            lookback_hours=args.lookback_hours,
            horizon_hours=args.horizon_hours,
            alpha=args.alpha,
            random_state=args.random_state,
            weighted_shrinkage_lambda=args.weighted_shrinkage_lambda,
            mondrian_tilted_shrinkage_lambda=args.mondrian_tilted_shrinkage_lambda,
            mondrian_tilted_bandwidth=args.mondrian_tilted_bandwidth,
            missingness_grouping_strategy=args.missingness_grouping_strategy,
            enable_external_baselines=args.enable_external_baselines,
            enable_learned_partition=args.enable_learned_partition,
            mask_strategy=args.mask_strategy,
            mask_rate=args.mask_rate,
            block_length=args.block_length,
            mask_random_state=args.mask_random_state,
            selective_features=args.selective_features,
            output_dir=args.output_dir,
        )
        run_experiments(config)
        return 0

    parser.error(f"unsupported command: {args.command}")
    return 2


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