from __future__ import annotations

import json
import math
import os
import random
from pathlib import Path
from typing import Any

import numpy as np
import yaml


def ensure_dir(path: str | Path) -> Path:
    """Create a directory if needed and return it as a Path."""
    out = Path(path)
    out.mkdir(parents=True, exist_ok=True)
    return out


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    try:
        import torch

        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except Exception:
        pass


def load_yaml(path: str | Path) -> dict[str, Any]:
    with Path(path).open("r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def save_json(obj: Any, path: str | Path) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8") as f:
        json.dump(to_jsonable(obj), f, indent=2, sort_keys=True)


def to_jsonable(obj: Any) -> Any:
    if isinstance(obj, dict):
        return {str(k): to_jsonable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [to_jsonable(v) for v in obj]
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, np.generic):
        return obj.item()
    if isinstance(obj, Path):
        return str(obj)
    return obj


def project_root() -> Path:
    return Path(__file__).resolve().parents[1]


def resolve_path(path: str | Path, base: str | Path | None = None) -> Path:
    p = Path(path)
    if p.is_absolute():
        return p
    return (Path(base) if base is not None else project_root()) / p


def configure_matplotlib_cache(output_dir: str | Path) -> None:
    cache_dir = ensure_dir(Path(output_dir) / "mplconfig")
    os.environ.setdefault("MPLCONFIGDIR", str(cache_dir))


def rmse(mean_state: np.ndarray, truth: np.ndarray) -> float:
    return float(np.sqrt(np.mean((mean_state - truth) ** 2)))


def ensemble_spread(ensemble: np.ndarray) -> float:
    if ensemble.shape[0] < 2:
        return float("nan")
    variance = np.var(ensemble, axis=0, ddof=1)
    return float(np.sqrt(np.mean(variance)))


def rescale_to_rmse(error: np.ndarray, target_rmse: float) -> np.ndarray:
    current = np.sqrt(np.mean(error**2))
    if current == 0:
        raise ValueError("Cannot rescale a zero error vector.")
    return error * (target_rmse / current)


def cyclic_distances(points: np.ndarray, ref: float, domain_size: int) -> np.ndarray:
    distances = np.abs(points.astype(float) - float(ref))
    return np.minimum(distances, domain_size - distances)


def gaspari_cohn(distances: np.ndarray, radius: float) -> np.ndarray:
    """Gaspari-Cohn taper with compact support at ``radius``.

    The usual Gaspari-Cohn polynomial has support at 2c. Here the public
    parameter is the zero-cutoff radius, so c = radius / 2.
    """
    distances = np.asarray(distances, dtype=float)
    if radius <= 0:
        return np.zeros_like(distances)

    r = distances / (radius / 2.0)
    out = np.zeros_like(r)

    mask1 = r <= 1.0
    x = r[mask1]
    out[mask1] = (
        1.0
        - (5.0 / 3.0) * x**2
        + (5.0 / 8.0) * x**3
        + 0.5 * x**4
        - 0.25 * x**5
    )

    mask2 = (r > 1.0) & (r < 2.0)
    x = r[mask2]
    out[mask2] = (
        4.0
        - 5.0 * x
        + (5.0 / 3.0) * x**2
        + (5.0 / 8.0) * x**3
        - 0.5 * x**4
        + (1.0 / 12.0) * x**5
        - (2.0 / (3.0 * x))
    )
    return np.clip(out, 0.0, 1.0)


def summarize_cycles(
    prior_rmse: np.ndarray,
    analysis_rmse: np.ndarray,
    prior_spread: np.ndarray,
    analysis_spread: np.ndarray,
    discard_cycles: int,
) -> dict[str, float]:
    sl = slice(discard_cycles, None)
    return {
        "discard_cycles": int(discard_cycles),
        "mean_prior_rmse": float(np.mean(prior_rmse[sl])),
        "mean_analysis_rmse": float(np.mean(analysis_rmse[sl])),
        "mean_prior_spread": float(np.mean(prior_spread[sl])),
        "mean_analysis_spread": float(np.mean(analysis_spread[sl])),
        "prior_spread_rmse_ratio": float(np.mean(prior_spread[sl]) / np.mean(prior_rmse[sl])),
        "analysis_spread_rmse_ratio": float(
            np.mean(analysis_spread[sl]) / np.mean(analysis_rmse[sl])
        ),
        "stable": bool(np.all(np.isfinite(analysis_rmse))),
    }


def exact_grid_std_anomalies(rng: np.random.Generator, n_members: int, k: int, std: float) -> np.ndarray:
    anomalies = rng.normal(0.0, 1.0, size=(n_members, k))
    anomalies -= anomalies.mean(axis=0, keepdims=True)
    sample_std = anomalies.std(axis=0, ddof=1, keepdims=True)
    sample_std = np.where(sample_std == 0.0, 1.0, sample_std)
    return anomalies / sample_std * std


def best_device(requested: str = "auto") -> str:
    if requested != "auto":
        return requested
    try:
        import torch

        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "cpu"


def format_seconds(seconds: float) -> str:
    if seconds < 60:
        return f"{seconds:.1f}s"
    minutes = math.floor(seconds / 60)
    return f"{minutes}m {seconds - 60 * minutes:.1f}s"
