"""Configuration, supplied MAT data, and portable checkpoint utilities."""

import hashlib
import json
import os
import random
from pathlib import Path

import numpy as np
import scipy.io
import torch
import yaml
from torch.nn import functional as F

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


def package_path(value):
    path = Path(value).expanduser()
    return path if path.is_absolute() else ROOT / path


def load_config(path):
    with open(package_path(path)) as stream:
        return yaml.safe_load(stream)


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


def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def compute_device(name):
    device = torch.device(name)
    if device.type == "cuda":
        if not os.environ.get("SLURM_JOB_ID"):
            raise RuntimeError("Accelerator execution requires a Slurm allocation")
        if not torch.cuda.is_available():
            raise RuntimeError("No accelerator visible inside the Slurm job")
    return device


def read_fields(config):
    data = config["data"]
    path = package_path(data["directory"]) / data["filename"]
    arrays = scipy.io.loadmat(path, variable_names=[data["input_key"], data["target_key"]])
    fields = []
    for key in [data["input_key"], data["target_key"]]:
        array = arrays[key]
        if list(array.shape) != data["expected_shape"] or array.dtype != np.float64:
            raise ValueError(f"Unexpected supplied data {key}: {array.shape}, {array.dtype}")
        if not np.isfinite(array).all():
            raise ValueError(f"Nonfinite values in {key}")
        tensor = torch.from_numpy(array.astype(np.float32))
        # Paper leaves resizing unspecified: author's linear interpolation,
        # align_corners=False, evaluates at source centers 15.5,47.5,... .
        tensor = F.interpolate(tensor[:, None], size=data["resolution"],
                               mode=data["interpolation"], align_corners=False)
        fields.append(tensor.transpose(1, 2).contiguous())
    return fields


def split_indices(config):
    data = config["data"]
    permutation = torch.randperm(data["expected_shape"][0], generator=torch.Generator().manual_seed(config["seed"]))
    n, m = data["n_train"], data["n_val"]
    if n + m != len(permutation):
        raise ValueError("The formal experiment must use all 2048 supplied samples")
    return permutation[:n], permutation[n:n + m]


def fit_normalizers(initial, final, train_ids, eps):
    result = {"eps": eps}
    for prefix, field in [("a", initial), ("u", final)]:
        training = field[train_ids]
        result[prefix + "_mean"] = training.mean()
        result[prefix + "_std"] = training.std(correction=1)
    return result


def encode(field, stats, prefix):
    return (field - stats[prefix + "_mean"].to(field)) / (stats[prefix + "_std"].to(field) + stats["eps"])


def decode(field, stats, prefix="u"):
    return field * (stats[prefix + "_std"].to(field) + stats["eps"]) + stats[prefix + "_mean"].to(field)


def make_inputs(initial, stats):
    physical = encode(initial, stats, "a")
    grid = torch.linspace(0, 1, initial.shape[1], dtype=initial.dtype).view(1, -1, 1)
    return torch.cat([physical, grid.expand(initial.shape[0], -1, -1)], dim=-1)


def relative_l2(prediction, target):
    difference = (prediction - target).flatten(1)
    return torch.linalg.vector_norm(difference, dim=1) / (torch.linalg.vector_norm(target.flatten(1), dim=1) + 1e-12)


def atomic_checkpoint(path, state):
    path = package_path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_name(path.name + ".tmp")
    torch.save(state, temp)
    temp.replace(path)


def load_checkpoint(path):
    # Tensor/primitive-only checkpoint; no pickled custom class or path equality.
    state = torch.load(package_path(path), map_location="cpu", weights_only=True)
    if state.get("format") != "lino-burgers-v1":
        raise ValueError("Unsupported LiNO checkpoint")
    return state


def file_digest(path):
    digest = hashlib.sha256()
    with open(path, "rb") as stream:
        for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()