"""Configuration, exact MAT contract, sampling and checkpoint utilities."""
from __future__ import annotations

import hashlib
import json
import random
from pathlib import Path

import numpy as np
import torch
import yaml
from scipy.io import loadmat

ROOT = Path(__file__).resolve().parents[1]


def read_config(path=None):
    path = Path(path) if path else ROOT / "config/config.yaml"
    with path.open() as handle:
        cfg = yaml.safe_load(handle)
    if cfg["model"]["hidden_layers"] != [50] * 4:
        raise ValueError("Formal experiment requires exactly four 50-unit hidden layers")
    if cfg["training"]["iterations"] != 4000:
        raise ValueError("Formal experiment requires exactly 4000 iterations")
    return cfg


def resolve(path):
    p = Path(path)
    return p if p.is_absolute() else ROOT / p


def save_json(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False))
    tmp.replace(path)


def setup(cfg, device=None):
    seed = int(cfg["seed"])
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.set_num_threads(int(cfg["runtime"]["cpu_threads"]))
    dtype = getattr(torch, cfg["model"]["dtype"])
    torch.set_default_dtype(dtype)
    device = torch.device(device or cfg["runtime"]["device"])
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("Requested accelerator is unavailable; run inside SLURM allocation")
    return device, dtype


def build_model(cfg, device):
    from model.TSAPINN import TSAPINN
    m = cfg["model"]
    return TSAPINN(**{k: m[k] for k in ["hidden_layers", "sigma", "zeta_sin", "zeta_cos", "reynolds", "bias_initialization"]}).to(device)


def load_dns(cfg):
    path = Path(cfg["data"]["path"])
    if not path.is_file():
        raise FileNotFoundError(f"Required author DNS data absent: {path}")
    raw = loadmat(path)
    for key, shape in cfg["data"]["expected_shapes"].items():
        if raw[key].shape != tuple(shape) or not np.isfinite(raw[key]).all():
            raise ValueError(f"Invalid DNS field {key}: expected finite {shape}")
    xy = np.asarray(raw["X_star"], dtype=np.float64)
    times = np.asarray(raw["t"], dtype=np.float64).ravel()
    if not np.all(np.diff(times) > 0):
        raise ValueError("DNS time must be strictly increasing")
    return {"xy": xy, "t": times, "uv": np.asarray(raw["U_star"]).transpose(2, 0, 1),
            "p": np.asarray(raw["p_star"]).T, "path": str(path),
            "sha256": hashlib.sha256(path.read_bytes()).hexdigest()}


def sample_indices(cfg, dns):
    """Literal paper per-time counts, geometric boundary; fixed throughout training.

    Only 296 unique boundary nodes exist. Balanced sampling with replacement
    provides 100 direct DNS observations per side without inventing new labels.
    Boundary corner duplication and replacement are explicit implementation choices.
    """
    rng = np.random.RandomState(cfg["seed"])
    xy, times = dns["xy"], dns["t"]
    xmin, ymin = xy.min(0)
    xmax, ymax = xy.max(0)
    sides = [np.flatnonzero(np.isclose(xy[:, 0], xmin)), np.flatnonzero(np.isclose(xy[:, 0], xmax)),
             np.flatnonzero(np.isclose(xy[:, 1], ymin)), np.flatnonzero(np.isclose(xy[:, 1], ymax))]
    boundary = np.unique(np.concatenate(sides))
    interior = np.setdiff1d(np.arange(len(xy)), boundary)
    nt, tr = len(times), cfg["training"]
    nb, nr = int(tr["boundary_points_per_time"]), int(tr["residual_points_per_time"])
    if nb % 4:
        raise ValueError("Balanced side sampling needs boundary count divisible by 4")
    bc = np.stack([np.concatenate([rng.choice(side, nb // 4, replace=True) for side in sides]) for _ in times])
    res = np.stack([rng.choice(interior, nr, replace=False) for _ in times])
    val = np.stack([rng.choice(np.setdiff1d(interior, res[k]), tr["validation_points_per_time"], replace=False) for k in range(nt)])
    return {"boundary": bc, "residual": res, "validation": val}


def gather(dns, indices):
    nt, n = indices.shape
    xy = dns["xy"][indices]
    time = np.broadcast_to(dns["t"][:, None, None], (nt, n, 1))
    coordinates = np.concatenate((xy, time), 2).reshape(-1, 3)
    labels = dns["uv"][np.arange(nt)[:, None], indices].reshape(-1, 2)
    return coordinates, labels


def tensor(array, device, dtype):
    return torch.as_tensor(np.ascontiguousarray(array), dtype=dtype, device=device)


def predict(model, coordinates, device, dtype, batch_size):
    values = []
    model.eval()
    with torch.no_grad():
        for start in range(0, len(coordinates), batch_size):
            q = tensor(coordinates[start:start + batch_size], device, dtype)
            values.append(model.flow_taylor(q).cpu().numpy())
    return np.concatenate(values)


def relative_l2(prediction, reference):
    difference = np.asarray(prediction, dtype=np.float64) - reference
    numerator = np.sum(difference * difference, axis=0, dtype=np.float64)
    denominator = np.sum(np.asarray(reference, dtype=np.float64)**2, axis=0, dtype=np.float64)
    return np.sqrt(numerator / denominator)


def atomic_checkpoint(path, payload):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    torch.save(payload, tmp)
    tmp.replace(path)


def load_checkpoint(cfg, device, path=None):
    path = resolve(path or cfg["outputs"]["checkpoint"])
    checkpoint = torch.load(path, map_location=device, weights_only=False)
    saved_config = checkpoint["config"].copy()
    saved_config["data"] = saved_config["data"].copy()
    saved_config["data"]["path"] = cfg["data"]["path"]
    if saved_config != cfg:
        raise ValueError("Checkpoint configuration does not match requested configuration")
    model = build_model(cfg, device)
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    model.eval()
    return model, checkpoint