#!/usr/bin/env python3
from __future__ import annotations

import argparse
from pathlib import Path
import sys

import pandas as pd

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

from src.config_utils import prepare_output_dirs, read_config, resolve_path


def main() -> None:
    parser = argparse.ArgumentParser(description="Create task CSV for ML/DL bias-correction experiments.")
    parser.add_argument("--config", default="config/experiment.yaml")
    args = parser.parse_args()

    cfg_path = Path(args.config).resolve()
    cfg = read_config(cfg_path)
    dirs = prepare_output_dirs(cfg, base_dir=cfg_path.parent.parent)

    rows = []
    for source_model in cfg["source_models"]:
        for lead in cfg["leads"]:
            for target in cfg["targets"]:
                for spatial_mode in cfg["spatial_modes"]:
                    for predictor_set in cfg["predictor_sets"]:
                        for method in cfg["methods"]:
                            # RAW does not need predictor set; keep only one P2 row to avoid duplicates.
                            if str(method).lower() == "raw" and predictor_set != cfg["predictor_sets"][0]:
                                continue
                            rows.append({
                                "source_model": source_model,
                                "lead": int(lead),
                                "target": target,
                                "spatial_mode": spatial_mode,
                                "predictor_set": predictor_set,
                                "method": method,
                            })
    df = pd.DataFrame(rows)
    task_file = dirs["output"] / "tasks.csv"
    df.to_csv(task_file, index=False)
    print(f"Wrote {len(df)} tasks to {task_file}")


if __name__ == "__main__":
    main()