"""Data, normalization, metrics and checkpoint contracts shared by the CLIs."""
from __future__ import annotations

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

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

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.hno import HNO


def configuration(description):
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--config', default='config/config.yaml')
    parser.add_argument('--device', default='cuda', choices=['cuda', 'cpu'])
    return parser


def read_config(path):
    with open(path, encoding='utf-8') as f:
        c = yaml.safe_load(f)
    for section in ['model', 'data', 'training', 'output']:
        if section not in c:
            raise ValueError(f'Missing config section: {section}')
    if c['model']['grid_size'] != (c['data']['raw_shape'][1] - 1) // c['data']['stride'] + 1:
        raise ValueError('Configured grid and strided raw data disagree')
    if c['training']['precision'] != 'float32':
        raise ValueError('This reproduction uses FP32 model parameters')
    return c


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def compute_device(name):
    if name == 'cuda':
        if not os.environ.get('SLURM_JOB_ID'):
            raise RuntimeError('Accelerator computation requires a Slurm allocation')
        if not torch.cuda.is_available():
            raise RuntimeError('No accelerator visible in this Slurm allocation')
    return torch.device(name)


def read_split(c, split):
    d = c['data']
    count = d['n_train' if split == 'train' else 'n_eval']
    path = Path(d['root']) / d['train_file' if split == 'train' else 'eval_file']
    print(f'Reading {split}: {path}, samples={count}', flush=True)
    raw = loadmat(path, variable_names=['coeff', 'sol'])
    arrays = []
    for key in ['coeff', 'sol']:
        a = raw.pop(key)
        if list(a.shape) != d['raw_shape']:
            raise ValueError(f'{path}:{key} shape {a.shape}, expected {d["raw_shape"]}')
        if not 0 < count <= len(a):
            raise ValueError('Invalid sample count')
        selected = np.array(a[:count, ::d['stride'], ::d['stride']], copy=True, order='C')
        if not np.isfinite(selected).all():
            raise ValueError(f'Nonfinite {key} data')
        tensor = torch.from_numpy(selected)
        if key == 'coeff':
            tensor = tensor.float()
        arrays.append(tensor)
        del a
    if arrays[0].shape[1:] != (c['model']['grid_size'], c['model']['grid_size']):
        raise ValueError('Unexpected downsampled grid')
    return tuple(arrays)


def training_statistics(coeff, solution):
    # Keep non-scalar tensor statistics: preserves source solution precision.
    return {name: value.reshape(1, 1, 1) for name, value in {
        'input_mean': coeff.mean(), 'input_std': coeff.std() + 1e-8,
        'target_mean': solution.mean(), 'target_std': solution.std() + 1e-8,
    }.items()}


def encode_input(x, stats):
    return ((x - stats['input_mean']) / stats['input_std']).unsqueeze(-1)


def decode_output(y, stats):
    return y[..., 0] * stats['target_std'] + stats['target_mean']


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


def central_gradient(field, spacing):
    extended = F.pad(field, (1, 1, 1, 1))
    dx = (extended[:, 1:-1, 2:] - extended[:, 1:-1, :-2]) / (2 * spacing)
    dy = (extended[:, 2:, 1:-1] - extended[:, :-2, 1:-1]) / (2 * spacing)
    return dx, dy


def objective(prediction, target, training):
    primary = relative_l2(prediction, target).sum()
    interior_prediction = F.pad(prediction[:, 1:-1, 1:-1], (1, 1, 1, 1))
    predicted_gradients = central_gradient(interior_prediction, training['derivative_spacing'])
    target_gradients = central_gradient(target, training['derivative_spacing'])
    auxiliary = sum(relative_l2(p, t).sum() for p, t in zip(predicted_gradients, target_gradients))
    return primary + training['derivative_weight'] * auxiliary, primary, auxiliary


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


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


def load_checkpoint(path, device):
    checkpoint = torch.load(path, map_location='cpu', weights_only=True)
    if checkpoint['format_version'] != 1:
        raise ValueError('Unsupported checkpoint schema')
    model = HNO(**checkpoint['model_config']).to(device)
    model.load_state_dict(checkpoint['model_state'], strict=True)
    stats = {k: v.to(device) for k, v in checkpoint['normalization'].items()}
    return model, stats, checkpoint


def source_hashes():
    paths = [ROOT/'model/hno.py', *sorted((ROOT/'scripts').glob('*.py'))]
    return {str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest() for p in paths}