#!/usr/bin/env python3
"""Generate task manifests for parallel Slurm execution."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from dr_sfc.config import GD1_DIR, LEADS, MODELS

AI_MODELS = ("ridge", "elasticnet", "rf", "gbm", "xgboost")
FEATURE_SETS = {
    "pr": {
        "SP1": ("pr", "q2", "rh2", "w700", "olr"),
        "SP2": ("pr", "q2", "w700", "olr", "u850", "h500"),
        "SP3": ("pr", "olr", "q2", "rh2", "slp", "u850", "u200", "w700", "h500", "h850"),
    },
    "t2": {
        "ST1": ("t2", "h500", "h850", "slp", "q2"),
        "ST2": ("t2", "h500", "h850", "w700", "q2", "rh2", "slp"),
        "ST3": ("t2", "olr", "q2", "rh2", "slp", "u850", "u200", "w700", "h500", "h850"),
    },
}
SPATIAL_UNITS = ("national", "R1", "R2", "R3", "R4", "R5", "R6", "R7")


def stage1_bc_tasks() -> list[dict[str, object]]:
    tasks: list[dict[str, object]] = []
    task_id = 0
    for variable in ("pr", "t2"):
        for lead in LEADS:
            for rcm in MODELS:
                for ai_model in AI_MODELS:
                    for feature_set, features in FEATURE_SETS[variable].items():
                        for spatial_unit in SPATIAL_UNITS:
                            task_id += 1
                            tasks.append(
                                {
                                    "task_id": task_id,
                                    "stage": "stage1_bc",
                                    "variable": variable,
                                    "lead": lead,
                                    "rcm": rcm,
                                    "ai_model": ai_model,
                                    "feature_set": feature_set,
                                    "features": list(features) + ["lon", "lat"],
                                    "spatial_unit": spatial_unit,
                                    "output_dir": str(GD1_DIR / "outputs" / "BC" / variable / f"lead{lead}" / rcm / ai_model / feature_set / spatial_unit),
                                }
                            )
    return tasks


def stage2_ope_tasks(year: int, month: int, day: int) -> list[dict[str, object]]:
    tasks: list[dict[str, object]] = []
    task_id = 0
    for lead in LEADS:
        task_id += 1
        tasks.append(
            {
                "task_id": task_id,
                "stage": "stage2_operational_lead",
                "year": year,
                "month": month,
                "day": day,
                "lead": lead,
            }
        )
    return tasks


def write_jsonl(tasks: list[dict[str, object]], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for task in tasks:
            f.write(json.dumps(task, sort_keys=True) + "\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)

    bc = subparsers.add_parser("stage1-bc", help="Generate Stage 1 bias-correction training task manifest")
    bc.add_argument("--output", type=Path, default=Path("manifests/stage1_bc_tasks.jsonl"))

    ope = subparsers.add_parser("stage2-operational", help="Generate Stage 2 operational lead task manifest")
    ope.add_argument("--year", type=int, required=True)
    ope.add_argument("--month", type=int, required=True)
    ope.add_argument("--day", type=int, required=True)
    ope.add_argument("--output", type=Path, default=None)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "stage1-bc":
        tasks = stage1_bc_tasks()
        output = args.output
    elif args.command == "stage2-operational":
        tasks = stage2_ope_tasks(args.year, args.month, args.day)
        output = args.output or Path(f"manifests/stage2_operational_{args.year:04d}{args.month:02d}{args.day:02d}.jsonl")
    else:
        raise ValueError(args.command)
    write_jsonl(tasks, output)
    print(output)
    print(len(tasks))


if __name__ == "__main__":
    main()