import copy
from pathlib import Path

import yaml

from train import train
from evaluate import evaluate


def _load_config_dict(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as handle:
        return yaml.safe_load(handle) or {}


def _dump_config_dict(cfg: dict, path: Path) -> None:
    with path.open("w", encoding="utf-8") as handle:
        yaml.safe_dump(cfg, handle, sort_keys=False)


def _apply_overrides(cfg: dict, overrides: dict) -> dict:
    updated = copy.deepcopy(cfg)
    for key_path, value in overrides.items():
        levels = key_path.split(".")
        node = updated
        for level in levels[:-1]:
            node = node.setdefault(level, {})
        node[levels[-1]] = value
    return updated


def main(config_path: str = "config.yaml") -> None:
    base_path = Path(config_path)
    base_cfg = _load_config_dict(base_path)

    scenarios = [
        {
            "name": "simple_with_attention",
            "overrides": {
                "model.backbone": "simple_unet",
                "model.attention.enabled": True,
            },
        },
        {
            "name": "ddpm_with_attention",
            "overrides": {
                "model.backbone": "ddpm_unet",
                "model.attention.enabled": True,
            },
        },
        {
            "name": "ddpm_no_attention",
            "overrides": {
                "model.backbone": "ddpm_unet",
                "model.attention.enabled": False,
            },
        },
        {
            "name": "simple_no_attention",
            "overrides": {
                "model.backbone": "simple_unet",
                "model.attention.enabled": False,
            },
        },
    ]

    for scenario in scenarios:
        scenario_cfg = _apply_overrides(base_cfg, scenario["overrides"])
        scenario_cfg = _apply_overrides(
            scenario_cfg,
            {
                "training.epochs": 300, # Number of epochs
                "training.out_dir": f"outputs/{scenario['name']}",
                    "training.preview_eval_interval": 50, # Preview evaluation every 50 epochs
                    "training.preview_eval_samples": 1, # Preview 1 sample
                    "training.preview_save_dir": scenario["name"],
                    "evaluation.save_dir": scenario["name"],
            },
        )

        temp_path = base_path.parent / f"{scenario['name']}_config.yaml"
        _dump_config_dict(scenario_cfg, temp_path)

        print(f"\n[Pipeline] Running scenario: {scenario['name']}")
        train(str(temp_path))
        evaluate(str(temp_path))

        print(f"[Pipeline] Completed scenario: {scenario['name']}")


if __name__ == "__main__":
    main()