import numpy as np

try:
    from omegaconf import OmegaConf
except ModuleNotFoundError:
    OmegaConf = None

_DIST_KEYS = ("uniform", "normal", "choice", "loguniform")

def _is_dist_node(x) -> bool:
    """True if x is a dict like {'uniform': [a,b]} etc."""
    return isinstance(x, dict) and len(x) == 1 and next(iter(x.keys())) in _DIST_KEYS

def _sample_dist(spec: dict, rng: np.random.Generator):
    k = next(iter(spec.keys()))
    v = spec[k]

    match k:
        case "uniform":
            lo, hi = v
            return float(rng.uniform(lo, hi))

        case "normal":
            mu, sigma = v
            return float(rng.normal(mu, sigma))

        case "choice":
            # v is a list of options
            return v[int(rng.integers(0, len(v)))]

        case "loguniform":
            lo, hi = v
            # require positive bounds
            if lo <= 0 or hi <= 0:
                raise ValueError(f"loguniform bounds must be > 0, got {v}")
            return float(np.exp(rng.uniform(np.log(lo), np.log(hi))))

    raise ValueError(f"Unknown distribution: {k}")

def sample_config(cfg, rng: np.random.Generator, ignore_paths=()):
    """
    Recursively sample a config.

    - cfg can be an OmegaConf object or plain python container.
    - ignore_paths is a tuple/list of dotpaths to leave untouched (copied as-is).
      Example: ('time.frame_durations_s',)
    """
    # Convert OmegaConf -> plain containers once at the top
    if (
        OmegaConf is not None
        and not isinstance(cfg, (dict, list))
        and OmegaConf.is_config(cfg)
    ):
        cfg = OmegaConf.to_container(cfg, resolve=True)

    ignore_paths = tuple(ignore_paths)

    def rec(node, path: str):
        # If this exact path is ignored, return node as-is
        if path in ignore_paths:
            return node

        # Dict
        if isinstance(node, dict):
            # Distribution leaf
            if _is_dist_node(node):
                return _sample_dist(node, rng)

            # Recurse
            out = {}
            for kk, vv in node.items():
                child_path = f"{path}.{kk}" if path else kk
                out[kk] = rec(vv, child_path)
            return out

        # List
        if isinstance(node, list):
            # Keep list shape; allow lists to contain fixed values or dist nodes
            out = []
            for i, vv in enumerate(node):
                child_path = f"{path}[{i}]"
                out.append(rec(vv, child_path))
            return out

        # Scalar leaf (fixed)
        return node

    return rec(cfg, "")
