#!/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


def main() -> None:
    parser = argparse.ArgumentParser(description="Create stack task CSV for year-blocked stacking 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)
    st = dict(cfg.get("stacking", {}))
    if not st.get("enabled", False):
        print("Stacking disabled in config; no stack tasks created.")
        return

    rows = []
    for lead in cfg["leads"]:
        for target in cfg["targets"]:
            for spatial_mode in cfg["spatial_modes"]:
                for scope in st.get("scopes", ["multi_source"]):
                    for meta_method in st.get("meta_methods", ["stack_mean", "stack_ridge"]):
                        rows.append({
                            "lead": int(lead),
                            "target": target,
                            "spatial_mode": spatial_mode,
                            "stack_scope": scope,
                            "meta_method": meta_method,
                        })
    df = pd.DataFrame(rows)
    task_file = dirs["output"] / "stack_tasks.csv"
    df.to_csv(task_file, index=False)
    print(f"Wrote {len(df)} stack tasks to {task_file}")


if __name__ == "__main__":
    main()