from __future__ import annotations

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

import numpy as np
import torch
import yaml


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from model.ho_fno_pipe import HOFNOPipe, count_parameters  # noqa: E402


class UnitTransformer:
    def __init__(self, tensor: torch.Tensor, axes: tuple[int, ...] = (0, 1), eps: float = 1e-8) -> None:
        self.axes = axes
        self.eps = eps
        self.mean = tensor.mean(dim=axes, keepdim=True)
        self.std = tensor.std(dim=axes, keepdim=True) + eps

    def encode(self, tensor: torch.Tensor) -> torch.Tensor:
        return (tensor - self.mean) / self.std

    def decode(self, tensor: torch.Tensor) -> torch.Tensor:
        return tensor * self.std + self.mean

    def to(self, device: torch.device) -> "UnitTransformer":
        self.mean = self.mean.to(device)
        self.std = self.std.to(device)
        return self

    def state_dict(self) -> dict[str, Any]:
        return {"mean": self.mean.cpu(), "std": self.std.cpu(), "axes": self.axes, "eps": self.eps}

    @classmethod
    def from_state_dict(cls, state: dict[str, Any]) -> "UnitTransformer":
        obj = cls.__new__(cls)
        obj.mean = state["mean"]
        obj.std = state["std"]
        obj.axes = tuple(state.get("axes", (0, 1)))
        obj.eps = float(state.get("eps", 1e-8))
        return obj


def load_config(path: str | os.PathLike[str] | None = None) -> dict[str, Any]:
    cfg_path = Path(path) if path else ROOT / "config" / "config.yaml"
    with cfg_path.open("r", encoding="utf-8") as handle:
        cfg = yaml.safe_load(handle)
    cfg["_config_path"] = str(cfg_path)
    return cfg


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def resolve(path: str | os.PathLike[str], base: Path | None = None) -> Path:
    p = Path(path)
    if p.is_absolute():
        return p
    return (base or ROOT) / p


def select_device(require_accelerator: bool = False) -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if require_accelerator:
        raise RuntimeError("accelerator is required, but torch.cuda.is_available() is False")
    return torch.device("cpu")


def load_pipe_tensors(cfg: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    data_cfg = cfg["data"]
    data_dir = Path(cfg["paths"]["data_dir"])
    train_size = int(data_cfg["train_size"])
    test_size = int(data_cfg["test_size"])
    total_used = int(data_cfg["total_used"])
    r1, r2 = (int(v) for v in data_cfg.get("downsample", [1, 1]))
    s1 = int(((int(data_cfg["resolution"][0]) - 1) / r1) + 1)
    s2 = int(((int(data_cfg["resolution"][1]) - 1) / r2) + 1)

    x_raw = np.load(data_dir / data_cfg["x_file"], mmap_mode="r")[:total_used, ::r1, ::r2][:, :s1, :s2]
    y_raw = np.load(data_dir / data_cfg["y_file"], mmap_mode="r")[:total_used, ::r1, ::r2][:, :s1, :s2]
    q_raw = np.load(data_dir / data_cfg["q_file"], mmap_mode="r")[
        :total_used, int(data_cfg["output_channel_index"]), ::r1, ::r2
    ][:, :s1, :s2]

    coords = np.stack((np.asarray(x_raw, dtype=np.float32), np.asarray(y_raw, dtype=np.float32)), axis=-1)
    target = np.asarray(q_raw, dtype=np.float32)[..., None]
    x_all = torch.from_numpy(coords.copy())
    y_all = torch.from_numpy(target.copy())
    return x_all[:train_size], y_all[:train_size], x_all[-test_size:], y_all[-test_size:]


def make_normalized_loaders(cfg: dict[str, Any]):
    x_train, y_train, x_test, y_test = load_pipe_tensors(cfg)
    axes = tuple(int(v) for v in cfg["data"].get("normalizer_axes", [0, 1]))
    x_normalizer = UnitTransformer(x_train, axes=axes)
    y_normalizer = UnitTransformer(y_train, axes=axes)
    x_train_n = x_normalizer.encode(x_train)
    x_test_n = x_normalizer.encode(x_test)
    y_train_n = y_normalizer.encode(y_train)
    y_test_n = y_normalizer.encode(y_test)

    batch_size = int(cfg["training"]["batch_size"])
    loader_kwargs = {
        "batch_size": batch_size,
        "num_workers": int(cfg["data"].get("num_workers", 0)),
        "pin_memory": bool(cfg["data"].get("pin_memory", False)),
    }
    train_ds = torch.utils.data.TensorDataset(x_train_n, y_train_n)
    test_ds = torch.utils.data.TensorDataset(x_test_n, y_test_n)
    train_loader = torch.utils.data.DataLoader(train_ds, shuffle=True, **loader_kwargs)
    test_loader = torch.utils.data.DataLoader(test_ds, shuffle=False, **loader_kwargs)
    return train_loader, test_loader, x_normalizer, y_normalizer


def build_model(cfg: dict[str, Any]) -> HOFNOPipe:
    m = cfg["model"]
    return HOFNOPipe(
        input_channels=int(m["input_channels"]),
        output_channels=int(m["output_channels"]),
        order=int(m["order"]),
        layers=int(m["layers"]),
        width=int(m["width"]),
        modes=tuple(int(v) for v in m["modes"]),
        add_unit_grid=bool(m.get("add_unit_grid", True)),
        mlp_expansion=int(m.get("mlp_expansion", 4)),
        mlp_dropout=float(m.get("mlp_dropout", 0.0)),
        drop_path_rate=float(m.get("drop_path_rate", 0.0)),
    )


def relative_l2_per_sample(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    pred_f = pred.reshape(pred.shape[0], -1)
    target_f = target.reshape(target.shape[0], -1)
    denom = torch.linalg.vector_norm(target_f, dim=1).clamp_min(1e-12)
    return torch.linalg.vector_norm(pred_f - target_f, dim=1) / denom


def relative_l2_sum(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return relative_l2_per_sample(pred, target).sum()


def save_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2)


__all__ = [
    "ROOT",
    "UnitTransformer",
    "build_model",
    "count_parameters",
    "load_config",
    "make_normalized_loaders",
    "relative_l2_per_sample",
    "relative_l2_sum",
    "resolve",
    "save_json",
    "select_device",
    "set_seed",
]