from __future__ import annotations

import csv
import json
import os
import random
import sys
from pathlib import Path
from typing import Any, Dict, Tuple

import h5py
import numpy as np
import torch
import yaml
from torch.utils.data import DataLoader, TensorDataset


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

from model.crop_navier_stokes import CROPNavierStokes, count_parameters


def load_config(path: str | os.PathLike[str]) -> Dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)
    cfg["_config_path"] = str(Path(path).resolve())
    cfg["_package_root"] = str(PACKAGE_ROOT)
    return cfg


def resolve_output_path(cfg: Dict[str, Any], relative: str) -> Path:
    out = Path(cfg["paths"]["output_dir"])
    path = Path(relative)
    return path if path.is_absolute() else out / path


def set_seed(seed: int, deterministic: bool = False, benchmark: bool = True) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = deterministic
    torch.backends.cudnn.benchmark = benchmark


def get_device(cfg: Dict[str, Any]) -> torch.device:
    requested = cfg.get("runtime", {}).get("device", "cuda")
    if requested == "cuda" and torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def build_model(cfg: Dict[str, Any]) -> CROPNavierStokes:
    m = cfg["model"]
    return CROPNavierStokes(
        in_channels=cfg["data"]["channels"],
        out_channels=cfg["data"]["channels"],
        modes=m["modes"],
        initial_channels=m["initial_channels"],
        levels=m["levels"],
        residual_blocks_per_level=m["residual_blocks_per_level"],
        bottleneck_blocks=m["bottleneck_blocks"],
        kernel_size=m["kernel_size"],
        input_size=m["input_size"],
        latent_size=m["latent_size"],
        use_bn=m["use_bn"],
    )


def load_tensors(cfg: Dict[str, Any]) -> Tuple[torch.Tensor, torch.Tensor]:
    data_path = Path(cfg["paths"]["data_dir"]) / cfg["paths"]["data_file"]
    key = cfg["data"]["key"]
    input_t = cfg["data"]["input_time_index"]
    target_t = cfg["data"]["target_time_index"]
    with h5py.File(data_path, "r") as f:
        if key not in f:
            raise KeyError(f"HDF5 key {key!r} not found in {data_path}")
        dset = f[key]
        expected = tuple(cfg["data"]["raw_shape"])
        if tuple(dset.shape) != expected:
            raise ValueError(f"Expected {key} shape {expected}, got {tuple(dset.shape)}")
        x_np = np.asarray(dset[input_t, :, :, :], dtype=np.float32).transpose(2, 0, 1)
        y_np = np.asarray(dset[target_t, :, :, :], dtype=np.float32).transpose(2, 0, 1)
    x = torch.from_numpy(x_np[:, None, :, :]).contiguous()
    y = torch.from_numpy(y_np[:, None, :, :]).contiguous()
    return x, y


def make_loaders(cfg: Dict[str, Any]) -> Tuple[DataLoader, DataLoader, DataLoader]:
    x, y = load_tensors(cfg)
    n_train = cfg["data"]["train_size"]
    n_val = cfg["data"]["val_size"]
    n_test = cfg["data"]["test_size"]
    n_total = n_train + n_val + n_test
    if x.shape[0] < n_total:
        raise ValueError(f"Need {n_total} samples, found {x.shape[0]}")

    train = TensorDataset(x[:n_train], y[:n_train])
    val = TensorDataset(x[n_train : n_train + n_val], y[n_train : n_train + n_val])
    test = TensorDataset(x[n_train + n_val : n_total], y[n_train + n_val : n_total])

    batch_size = cfg["training"]["batch_size"]
    workers = cfg["training"].get("num_workers", 0)
    return (
        DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=torch.cuda.is_available()),
        DataLoader(val, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=torch.cuda.is_available()),
        DataLoader(test, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=torch.cuda.is_available()),
    )


def relative_normalized_mse(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((pred - target) ** 2) / torch.mean(target**2).clamp_min(1e-12)


@torch.no_grad()
def mean_relative_l2_percent(model: torch.nn.Module, loader: DataLoader, device: torch.device) -> float:
    model.eval()
    total = 0.0
    count = 0
    for x, y in loader:
        x = x.to(device, non_blocking=True)
        y = y.to(device, non_blocking=True)
        pred = model(x)
        diff = torch.linalg.vector_norm((pred - y).reshape(pred.shape[0], -1), ord=2, dim=1)
        denom = torch.linalg.vector_norm(y.reshape(y.shape[0], -1), ord=2, dim=1).clamp_min(1e-12)
        rel = diff / denom
        total += rel.sum().item()
        count += rel.numel()
    return 100.0 * total / max(1, count)


@torch.no_grad()
def collect_predictions(model: torch.nn.Module, loader: DataLoader, device: torch.device) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    model.eval()
    preds, targets, inputs = [], [], []
    for x, y in loader:
        pred = model(x.to(device, non_blocking=True)).cpu()
        preds.append(pred.numpy().astype(np.float32))
        targets.append(y.numpy().astype(np.float32))
        inputs.append(x.numpy().astype(np.float32))
    return np.concatenate(inputs), np.concatenate(preds), np.concatenate(targets)


def save_json(path: Path, payload: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, sort_keys=True)
    tmp.replace(path)


def append_history(path: Path, row: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    exists = path.exists()
    with open(path, "a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(row.keys()))
        if not exists:
            writer.writeheader()
        writer.writerow(row)