"""Shared configuration, raw Burgers data and metric contracts."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import random
import socket
import sys

import numpy as np
from scipy.io import loadmat
import torch
import yaml

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))
from model.ronom import RONOM


def parser(description):
    p = argparse.ArgumentParser(description=description)
    p.add_argument('--config', default='config/config.yaml')
    return p


def load_config(filename):
    with open(filename, encoding='utf-8') as handle:
        cfg = yaml.safe_load(handle)
    return cfg


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


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


def update_run_manifest(cfg, **changes):
    path = resolve(cfg['paths']['results']) / 'run_manifest.json'
    if path.exists():
        value = json.loads(path.read_text())
        value.update(changes)
        save_json(path, value)


def seed_everything(seed, threads=8):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.set_num_threads(threads)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False


def device_for(name):
    if name in ('cuda', 'dcu'):
        if not os.environ.get('SLURM_JOB_ID'):
            raise RuntimeError('Accelerator execution must run inside a Slurm allocation.')
        if not torch.cuda.is_available():
            raise RuntimeError('The allocated compute node has no accessible accelerator.')
        if torch.cuda.device_count() != 1:
            raise RuntimeError('Expected one Slurm-visible DCU; check allocation device visibility. '
                               f'CUDA_VISIBLE_DEVICES={os.environ.get("CUDA_VISIBLE_DEVICES")} '
                               f'HIP_VISIBLE_DEVICES={os.environ.get("HIP_VISIBLE_DEVICES")} '
                               f'ROCR_VISIBLE_DEVICES={os.environ.get("ROCR_VISIBLE_DEVICES")} '
                               f'SLURM_JOB_GPUS={os.environ.get("SLURM_JOB_GPUS")}')
        print(f'Slurm job={os.environ["SLURM_JOB_ID"]} host={socket.gethostname()} '
              f'accelerator={torch.cuda.get_device_name()} count={torch.cuda.device_count()}', flush=True)
        return torch.device('cuda')
    return torch.device(name)


def load_splits(cfg, names):
    """Read only the supplied .mat. No downloads, augmentation, or normalization.

    scipy decompresses the raw MATLAB array; only requested split tensors survive
    this function. train.py requests train/validation and never constructs test.
    """
    d = cfg['data']
    path = resolve(d['directory']) / d['filename']
    raw = loadmat(path, variable_names=['input', 'output', 'tspan', 'visc'])
    if 'visc' not in raw or not np.isclose(float(raw['visc'].squeeze()), 0.01):
        raise ValueError('Raw Burgers viscosity must equal 0.01')
    initial, target = raw['input'], raw['output']
    times = raw['tspan'].reshape(-1)
    if initial.shape != (1000, 1025) or target.shape != (1000, 101, 1025):
        raise ValueError(f'Unexpected raw Burgers shapes: {initial.shape}, {target.shape}')
    if times.shape != (101,) or not np.allclose(times, np.linspace(0, 1, 101), atol=1e-12):
        raise ValueError('Raw tspan does not match the paper interval [0,1].')
    counts = (d['train_samples'], d['validation_samples'], d['test_samples'])
    if counts != (600, 200, 200):
        raise ValueError('Formal Burgers protocol requires exactly 600/200/200 samples.')
    if (d['space_stride'], d['time_stride']) != (32, 10):
        raise ValueError('This experiment requires strides 32 in space and 10 in time.')
    spans = {'train': (0, 600), 'validation': (600, 800), 'test': (800, 1000)}
    splits = {}
    for name in names:
        begin, end = spans[name]
        u = np.ascontiguousarray(initial[begin:end, ::32], dtype=np.float32)
        y = np.ascontiguousarray(target[begin:end, ::10, ::32], dtype=np.float32)
        if not np.isfinite(u).all() or not np.isfinite(y).all():
            raise ValueError(f'Nonfinite data in {name}')
        mismatch = float(np.max(np.abs(u - y[:, 0])))
        if mismatch > 1e-5:
            raise ValueError(f'input and output(t=0) disagree in {name}: {mismatch}')
        splits[name] = (torch.from_numpy(u), torch.from_numpy(y), np.arange(begin, end))
        print(f'Data split={name} indices=[{begin},{end}) input={u.shape} '
              f'target={y.shape} initial_consistency_max={mismatch:.3e}', flush=True)
    return splits, torch.tensor(times[::10], dtype=torch.float32), path


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


def build_model(cfg, device):
    model = RONOM(**cfg['model']).to(device)
    if cfg['training'].get('jit_flow', False):
        model.latent_flow = torch.jit.script(model.latent_flow)
    count = sum(p.numel() for p in model.parameters())
    if count != 414217:
        raise ValueError(f'Model parameter count {count} differs from paper Appendix B: 414217')
    print(f'RONOM parameters={count}, latent=4x33=132', flush=True)
    return model


def checkpoint_save(path, value):
    path = resolve(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + '.tmp')
    torch.save(value, tmp)
    tmp.replace(path)


def load_best(cfg, device):
    path = resolve(cfg['paths']['checkpoint'])
    checkpoint = torch.load(path, map_location='cpu', weights_only=False)
    # Only architecture is compared: data/results paths never gate loading.
    if checkpoint['model_config'] != cfg['model']:
        raise ValueError('Model architecture in config differs from checkpoint.')
    model = build_model(cfg, device)
    model.load_state_dict(checkpoint['model_state_dict'], strict=True)
    model.eval()
    print(f'Loaded best checkpoint epoch={checkpoint["epoch"]} '
          f'validation_mse={checkpoint["validation_mse"]:.9e}', flush=True)
    return model, checkpoint, path


def metrics(prediction, target):
    err = prediction.double() - target.double()
    mse_i = err.square().flatten(1).mean(1)
    rmse_i = mse_i.sqrt()
    relative = err.flatten(1).norm(dim=1) / target.double().flatten(1).norm(dim=1).clamp_min(1e-30)
    return {
        'samples': len(rmse_i),
        'mse': mse_i.mean().item(),
        'global_rmse': mse_i.mean().sqrt().item(),
        'trajectory_rmse_mean': rmse_i.mean().item(),
        'trajectory_rmse_std_ddof1': rmse_i.std(correction=1).item() if len(rmse_i) > 1 else 0.,
        'relative_l2_mean': relative.mean().item(),
        'rmse_by_time': err.square().mean(dim=(0, 2)).sqrt().cpu().tolist(),
    }, rmse_i.cpu().numpy(), relative.cpu().numpy()


@torch.no_grad()
def evaluate(model, initial, target, times, batch_size, collect_latent=False):
    predictions, latents = [], []
    model.eval()
    for begin in range(0, len(initial), batch_size):
        output = model(initial[begin:begin + batch_size], times, return_latent=collect_latent)
        if collect_latent:
            output, zt = output
            latents.append(zt.cpu())
        predictions.append(output.cpu())
    predicted = torch.cat(predictions)
    if not torch.isfinite(predicted).all():
        raise RuntimeError('Evaluation produced nonfinite predictions')
    summary, per_sample, relative = metrics(predicted, target.cpu())
    return summary, predicted, per_sample, relative, torch.cat(latents) if latents else None