import os
from pathlib import Path
from typing import Optional, Tuple

import torch
import yaml


def load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def save_checkpoint(model, optimizer: torch.optim.Optimizer, iteration: int, path: str):
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    torch.save(
        {
            "iteration": iteration,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
        },
        path,
    )
    print(f"[Checkpoint] saved → {path}")


def load_checkpoint(model, path: str, device=None) -> int:
    ckpt = torch.load(path, map_location=device or "cpu")
    model.load_state_dict(ckpt["model_state_dict"], strict=False)
    if "optimizer_state_dict" in ckpt and hasattr(model, "optimizer"):
        try:
            model.optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        except Exception:
            pass
    it = ckpt.get("iteration", 0)
    print(f"[Checkpoint] loaded ← {path} (iter {it})")
    return it


def setup_helmholtz_config(config: dict, k_test=None, center=None, radius=None, eps_test=None) -> float:
    """Lock scatterer/helmholtz config for test/viz. Returns wave number k."""
    hc = config.setdefault("environment", {}).setdefault("mesh_refinement", {}).setdefault("fem", {}).setdefault("helmholtz", {})
    sc = hc.setdefault("scatterer", {})
    sc["mode"] = "fixed"
    if center is not None:
        sc["cx"], sc["cy"] = center[0], center[1]
    if radius is not None:
        sc["radius"] = radius
    if eps_test is not None:
        sc["eps_r"] = eps_test
    if k_test is not None:
        hc["wave_number_mode"] = "fixed"
        hc["wave_number"] = k_test
    return hc.get("wave_number", 6.0)


def parse_center(center_str: Optional[str]) -> Optional[Tuple[float, float]]:
    if center_str is None:
        return None
    parts = center_str.split(",")
    if len(parts) != 2:
        raise ValueError(f"Invalid --center format (expected 'cx,cy'): {center_str}")
    return (float(parts[0].strip()), float(parts[1].strip()))