"""Configuration, read-only validation and unambiguous metric reductions."""
import hashlib
import json
from pathlib import Path
import numpy as np
import torch
import yaml
from model.comphy_tgv import ComPhyTaylorGreenVortex

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


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


def load_config(path):
    with open(path) as f:
        cfg = yaml.safe_load(f)
    if cfg['training']['steps'] != 100000 or cfg['training']['interior_points'] != 1000:
        raise ValueError('Formal experiment requires 100000 steps and 1000 interior points.')
    if cfg['model']['hidden_layers'] != 2 or cfg['model']['hidden_units'] != 64:
        raise ValueError('Table 5 requires two hidden layers of 64 units.')
    return cfg


def build_model(cfg, device):
    dtype = getattr(torch, cfg['training']['dtype'])
    return ComPhyTaylorGreenVortex(**cfg['model']).to(device=device, dtype=dtype)


def load_reference(path):
    """Validation only. No torch tensors with gradients are created here."""
    with np.load(path, allow_pickle=False) as data:
        coordinates = data['coordinates'].copy()
        target = data['target'].copy()
    if coordinates.shape != (36864, 3) or target.shape != coordinates.shape:
        raise ValueError('Expected coordinates and target [36864,3].')
    if coordinates.dtype != np.float64 or target.dtype != np.float64:
        raise ValueError('Reference must be float64.')
    if not np.isfinite(coordinates).all() or not np.isfinite(target).all():
        raise ValueError('Reference has nonfinite values.')
    coordinates.flags.writeable = False
    target.flags.writeable = False
    return coordinates, target


def metrics(prediction, target):
    error = np.asarray(prediction, dtype=np.float64) - target
    if not np.isfinite(error).all():
        raise FloatingPointError('Nonfinite prediction.')
    e2 = float(np.mean(np.sum(error**2, axis=1)))
    return {'l2_error': e2, 'max_error': float(np.abs(error).max()),
            'l2_norm': float(np.sqrt(e2)), 'scalar_mse': e2/3,
            'relative_l2': float(np.linalg.norm(error)/np.linalg.norm(target)),
            'per_variable_mse': dict(zip(('u','v','p'), np.mean(error**2, axis=0).tolist())),
            'per_variable_max': dict(zip(('u','v','p'), np.abs(error).max(axis=0).tolist()))}


@torch.no_grad()
def predict(model, coordinates, batch_size):
    model.eval()
    p = next(model.parameters())
    blocks = []
    for start in range(0, len(coordinates), batch_size):
        q = torch.tensor(coordinates[start:start+batch_size], device=p.device, dtype=p.dtype)
        # forward is strictly the momentum module, never an average/ensemble.
        blocks.append(model(q).cpu().double().numpy())
    return np.concatenate(blocks)


def atomic_json(path, payload):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_suffix(path.suffix + '.tmp')
    temp.write_text(json.dumps(payload, indent=2, ensure_ascii=False, allow_nan=False)+'\n')
    temp.replace(path)


def atomic_checkpoint(path, payload):
    path = Path(path)
    temp = path.with_suffix(path.suffix + '.tmp')
    torch.save(payload, temp)
    temp.replace(path)


def sha256(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()