"""Data, Sobol, PDE and metric contracts shared by the three entry points."""
import hashlib
import json
import os
import random
import socket
from pathlib import Path

import h5py
import numpy as np
import torch
import yaml

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


def load_config(path):
    with open(path, encoding="utf-8") as stream:
        config = yaml.safe_load(stream)
    if config["training"]["batch_size"] != 64 or config["training"]["epochs"] != 20000:
        raise ValueError("Formal protocol must retain batch_size=64 and epochs=20000")
    if config["training"]["early_stopping"]:
        raise ValueError("User requires complete formal training; early stopping disabled")
    return config


def resolved(config, name):
    path = Path(config["paths"][name])
    return path if path.is_absolute() else ROOT / path


def atomic_json(path, payload):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False), encoding="utf-8")
    tmp.replace(path)


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


def require_compute_node():
    # A self-submitted allocation also records a job name unique to this case.
    jid, name = os.environ.get("SLURM_JOB_ID"), os.environ.get("SLURM_JOB_NAME", "")
    if not jid or not name.startswith("qcpinn_cavity_") or "login" in socket.gethostname().lower():
        raise RuntimeError("Training/smoke requires a new case-owned SLURM compute allocation")
    return {"job_id": jid, "job_name": name, "hostname": socket.gethostname(),
            "array_job_id": os.environ.get("SLURM_ARRAY_JOB_ID"),
            "array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID")}


def seed_everything(seed, threads=1):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.set_num_threads(threads)


def build_model(config):
    from model.QCPINN import QCPINN
    m = config["model"]
    model = QCPINN(**{key: m[key] for key in ["input_dim", "output_dim", "hidden_dim", "n_qubits", "quantum_layers", "backend", "shots"]})
    model = model.to(dtype=getattr(torch, m["dtype"]), device=m["device"])
    n = sum(p.numel() for p in model.parameters())
    if n != m["expected_parameters"]:
        raise ValueError(f"Parameter count {n} != paper count {m['expected_parameters']}")
    return model


def make_sobol_pools(config, seed):
    s = config["sampling"]
    pools = {}
    for index, kind in enumerate(["interior", "lid", "left", "right", "bottom", "initial"]):
        count = s["interior"] if kind == "interior" else s["initial"] if kind == "initial" else s["per_wall"]
        engine = torch.quasirandom.SobolEngine(3, scramble=s["scramble"], seed=seed*17+index)
        xyz = engine.draw(count).to(dtype=getattr(torch, config["model"]["dtype"]))
        xyz[:, 0] *= config["physics"]["t_max"]
        if kind == "lid": xyz[:, 2] = 1.0
        if kind == "left": xyz[:, 1] = 0.0
        if kind == "right": xyz[:, 1] = 1.0
        if kind == "bottom": xyz[:, 2] = 0.0
        if kind == "initial": xyz[:, 0] = 0.0
        pools[kind] = xyz.to(config["model"]["device"])
    return pools


def sample_batches(pools, size, rng):
    return {key: points[torch.as_tensor(rng.choice(len(points), size=size, replace=False), device=points.device)]
            for key, points in pools.items()}


def derivative(field, coords):
    return torch.autograd.grad(field, coords, torch.ones_like(field), create_graph=True)[0]


def pde_residuals(model, coordinates, physics):
    coords = coordinates.detach().clone().requires_grad_(True)
    uvp = model(coords)
    u, v, p = uvp.unbind(-1)
    du, dv, dp = [derivative(field, coords) for field in (u, v, p)]
    uxx = derivative(du[:, 1], coords)[:, 1]
    uyy = derivative(du[:, 2], coords)[:, 2]
    vxx = derivative(dv[:, 1], coords)[:, 1]
    vyy = derivative(dv[:, 2], coords)[:, 2]
    ru = du[:, 0] + u*du[:, 1] + v*du[:, 2] + dp[:, 1]/physics["rho"] - physics["mu"]*(uxx+uyy)
    rv = dv[:, 0] + u*dv[:, 1] + v*dv[:, 2] + dp[:, 2]/physics["rho"] - physics["mu"]*(vxx+vyy)
    rc = du[:, 1] + dv[:, 2]
    return torch.stack((ru, rv, rc), dim=-1)


def loss_terms(model, batches, config):
    physics = config["physics"]
    if physics["loss_formula"] != "paper_literal":
        raise ValueError("This reproduction implements the literal Appendix A.2 losses")
    residuals = pde_residuals(model, batches["interior"], physics)
    terms = {"physics": residuals.square().mean(0).sum()}
    names = ["lid", "left", "right", "bottom", "initial"]
    sizes = [len(batches[k]) for k in names]
    fields = torch.split(model(torch.cat([batches[k] for k in names])), sizes)
    lid, *rest = fields
    terms["lid"] = (1-lid[:, 0]+lid[:, 1]).square().mean()
    terms["walls"] = sum(out[:, :2].sum(-1).square().mean() for out in rest[:3])
    terms["initial"] = rest[3].sum(-1).square().mean()
    terms["total"] = sum(physics["loss_weights"][key]*value for key, value in terms.items())
    return terms


def validation_indices(config, total):
    size = config["validation"]["size"]
    if size >= total:
        raise ValueError("Validation must leave nonempty holdout")
    return np.sort(np.random.default_rng(config["validation"]["seed"]).choice(total, size, replace=False))


def load_validation(config):
    key = config["validation"]["reference_key"]
    with h5py.File(resolved(config, "reference"), "r") as source:
        data = source[key]
        if data.shape[0] != 6:
            raise ValueError("Expected MATLAB v7.3 dataset of shape (6,N)")
        ids = validation_indices(config, data.shape[1])
        rows = data[:, ids].T
    if not np.isfinite(rows).all():
        raise ValueError("Reference validation contains NaN/Inf")
    return rows[:, :3], rows[:, 3:], ids


def relative_metrics(error_squares, reference_squares):
    metrics = {}
    for i, key in enumerate(("u", "v", "p")):
        den, num = float(reference_squares[i]), float(error_squares[i])
        value = float(np.sqrt(num/den)) if den > 0 else None
        metrics[key] = {"relative_l2": value, "relative_l2_percent": None if value is None else 100*value,
                        "error_sum_squares": num, "reference_sum_squares": den,
                        "reason": None if value is not None else "zero_reference_norm"}
    return metrics


@torch.no_grad()
def predict(model, coords, size):
    param = next(model.parameters())
    return np.concatenate([model(torch.as_tensor(coords[i:i+size], dtype=param.dtype, device=param.device)).cpu().numpy()
                           for i in range(0, len(coords), size)], axis=0)


def evaluate(model, coordinates, reference, size):
    was_training = model.training
    model.eval()
    predicted = predict(model, coordinates, size).astype(np.float64)
    model.train(was_training)
    errors = predicted-reference
    return relative_metrics(np.sum(errors**2, 0), np.sum(reference**2, 0))


def selection_score(metrics):
    values = [metrics[k]["relative_l2"] for k in ("u", "v", "p")]
    if any(v is None or not np.isfinite(v) for v in values):
        raise ValueError("Validation L2 is undefined/nonfinite")
    return float(np.mean(values))


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


def load_checkpoint(path, model):
    checkpoint = torch.load(path, map_location=next(model.parameters()).device, weights_only=False)
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    return checkpoint