"""Shared configuration, analytic sampling, and reference-only evaluation."""
import hashlib
import json
from pathlib import Path

import numpy as np
import torch
import yaml


LOSS_NAMES = ('u_in', 'v_in', 'p_out', 'u_noslip', 'v_noslip', 'u_res', 'v_res', 'p_res')


def load_config(path):
    path = Path(path).resolve()
    cfg = yaml.safe_load(path.read_text())
    root = (path.parent / cfg['paths']['output_dir']).resolve()
    reference = Path(cfg['evaluation']['reference_vtu'])
    if not reference.is_absolute():
        cfg['evaluation']['reference_vtu'] = str(root / reference)
    cfg['_root'] = str(root)
    return cfg, root


def atomic_json(path, obj):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + '.tmp')
    tmp.write_text(json.dumps(obj, indent=2, ensure_ascii=False, allow_nan=False))
    tmp.replace(path)


def atomic_save(path, obj):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + '.tmp')
    torch.save(obj, tmp)
    tmp.replace(path)


def sample_interior(cfg, device, generator=None, count=None):
    d = cfg['domain']
    n = cfg['training']['batch_size'] if count is None else count
    z = torch.rand(n, 2, device=device, generator=generator)
    low = z.new_tensor([d['x_min'], d['y_min']])
    span = z.new_tensor([d['x_max'] - d['x_min'], d['y_max'] - d['y_min']])
    return low + span * z


def analytic_boundaries(cfg, device):
    d, b = cfg['domain'], cfg['boundary']
    yl = torch.linspace(d['y_min'], d['y_max'], b['left_points'], device=device)
    left = torch.stack((torch.full_like(yl, d['x_min']), yl), -1)
    # Group the whole left boundary in the two left-boundary loss terms,
    # as in the author's discretization: its lower half has zero velocity.
    inlet = left
    yo = torch.linspace(d['y_min'], d['y_max'], b['outlet_points'], device=device)
    outlet = torch.stack((torch.full_like(yo, d['x_max']), yo), -1)
    x = torch.linspace(d['x_min'], d['x_max'], b['points_per_horizontal_wall'], device=device)
    walls = torch.cat((torch.stack((x, torch.full_like(x, d['y_min'])), -1),
                       torch.stack((x, torch.full_like(x, d['y_max'])), -1)))
    u_in = torch.where(yl >= 0, 24.0 * yl * (0.5 - yl), torch.zeros_like(yl))
    return inlet, outlet, walls, u_in


def boundary_losses(model, boundary):
    inlet, outlet, walls, target = boundary
    sizes = [len(inlet), len(outlet), len(walls)]
    pi, po, pw = model(torch.cat((inlet, outlet, walls))).split(sizes)
    return torch.stack(((pi[:, 0] - target).square().mean(), pi[:, 1].square().mean(),
                        po[:, 2].square().mean(), pw[:, 0].square().mean(), pw[:, 1].square().mean()))


def read_reference(path, expected_points=None):
    # VTU is a genuine SU2 field. Never called by sample_interior/analytic_boundaries.
    from vtk import vtkXMLUnstructuredGridReader
    from vtk.util.numpy_support import vtk_to_numpy
    p = Path(path)
    reader = vtkXMLUnstructuredGridReader()
    reader.SetFileName(str(p))
    reader.Update()
    mesh = reader.GetOutput()
    n = mesh.GetNumberOfPoints()
    if expected_points is not None and n != expected_points:
        raise ValueError(f'Expected {expected_points} reference points, got {n}')
    vel = mesh.GetPointData().GetArray('Velocity')
    pressure = mesh.GetPointData().GetArray('Pressure')
    if vel is None or pressure is None:
        raise ValueError('VTU must contain Velocity and Pressure PointData')
    xy = vtk_to_numpy(mesh.GetPoints().GetData())[:, :2].copy()
    velocity = vtk_to_numpy(vel).copy()
    uvp = np.column_stack((velocity[:, 0], velocity[:, 1], vtk_to_numpy(pressure)))
    if not np.isfinite(xy).all() or not np.isfinite(uvp).all():
        raise ValueError('Nonfinite reference data')
    metadata = {'source': str(p.resolve()), 'sha256': hashlib.sha256(p.read_bytes()).hexdigest(),
                'points': n, 'target_field': 'Velocity', 'target_component': 0,
                'training_labels_used': False, 'bounds': [xy.min(0).tolist(), xy.max(0).tolist()]}
    return xy, uvp, metadata, mesh


@torch.no_grad()
def predict(model, xy, device, batch_size=4096):
    pred = []
    for start in range(0, len(xy), batch_size):
        x = torch.as_tensor(xy[start:start + batch_size], dtype=torch.float32, device=device)
        pred.append(model(x).cpu().numpy())
    return np.concatenate(pred)


def errors(pred, ref):
    # Accumulate in float64, unweighted over every original reference mesh point.
    delta = pred.astype(np.float64) - ref.astype(np.float64)
    return {'relative_l2_u': float(np.linalg.norm(delta[:, 0]) / np.linalg.norm(ref[:, 0].astype(np.float64))),
            'relative_l2_v': float(np.linalg.norm(delta[:, 1]) / np.linalg.norm(ref[:, 1].astype(np.float64)))}


def build_model(cfg, device):
    from model.PseudoTimePINN import PseudoTimePINN
    d = cfg['domain']
    return PseudoTimePINN(**cfg['model'], length=d['x_max'] - d['x_min'],
                          height=d['y_max'] - d['y_min']).to(device)


def require_compute_node(cfg):
    import os
    if cfg['training']['require_slurm'] and not os.environ.get('SLURM_JOB_ID'):
        raise RuntimeError('DCU execution requires a Slurm allocation; login-node training is forbidden')
    if not torch.cuda.is_available():
        raise RuntimeError('Allocated DCU is not accessible through torch.cuda')