"""Configuration, data and checkpoint contracts shared by all entrypoints."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import numpy as np
import torch
import yaml
from scipy.interpolate import interp1d

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


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


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


def config(filename):
    with path(filename).open() as f:
        cfg = yaml.safe_load(f)
    assert cfg['data']['sensors'] == cfg['model']['sensors']
    assert sum(cfg['data'][s] for s in ['train','validation','test']) == cfg['data']['selected_trajectories']
    return cfg


def accelerator(name):
    if name.startswith('cuda'):
        if not os.environ.get('SLURM_JOB_ID'):
            raise RuntimeError('Accelerator work must run inside a Slurm allocation.')
        if not torch.cuda.is_available():
            raise RuntimeError('Configured accelerator is unavailable; CPU fallback forbidden.')
    return torch.device(name)


def load_data(cfg, stats=None):
    d = cfg['data']
    source = path(d['directory']) / d['file']
    with np.load(source, allow_pickle=False) as f:
        p, radius, t = f['del_p'], f['R'], f['t']
    if p.shape != (1500,1000) or radius.shape != p.shape or t.shape != (1000,):
        raise ValueError(f'Unexpected public data schema: {p.shape}, {radius.shape}, {t.shape}')
    if p.dtype != np.float64 or radius.dtype != np.float64:
        raise ValueError('Expected original float64 public trajectories.')
    if not all(np.isfinite(a).all() for a in [p,radius,t]) or not np.all(np.diff(t)>0):
        raise ValueError('Nonfinite or unordered input data.')
    if not np.allclose([t[0],t[-1]],[0,0.0005],rtol=1e-12,atol=1e-14):
        raise ValueError('Unexpected time range.')
    n = d['selected_trajectories']
    pressure_t = np.linspace(t[0],t[-1],d['sensors'])
    output_t = np.linspace(t[0],t[-1],d['output_points'])
    x = interp1d(t,p[:n],kind=d['interpolation'],axis=-1)(pressure_t)
    y = interp1d(t,radius[:n],kind=d['interpolation'],axis=-1)(output_t)
    nt = d['train']
    if stats is None:
        stats = dict(pressure_mean=float(x[:nt].mean()),pressure_std=float(x[:nt].std()),
                     radius_mean=float(y[:nt].mean()),radius_std=float(y[:nt].std()),
                     time_min=float(t[0]),time_span=float(t[-1]-t[0]))
    for key in ['pressure_std','radius_std','time_span']:
        if stats[key] <= 0: raise ValueError(f'Invalid normalization: {key}')
    xn = ((x-stats['pressure_mean'])/stats['pressure_std']).astype(np.float32)
    yn = ((y-stats['radius_mean'])/stats['radius_std']).astype(np.float32)
    tn = ((output_t-stats['time_min'])/stats['time_span']).astype(np.float32)[:,None]
    val_end = nt+d['validation']
    splits = {'train':slice(0,nt),'validation':slice(nt,val_end),'test':slice(val_end,n)}
    return dict(x=xn,y=yn,time=tn,radius=y,pressure=x,output_t=output_t,
                pressure_t=pressure_t,stats=stats,splits=splits,source=source)


def metrics(prediction, target):
    e = np.asarray(prediction,dtype=np.float64)-np.asarray(target,dtype=np.float64)
    rel = np.linalg.norm(e,axis=1)/np.linalg.norm(target,axis=1)
    return dict(relative_l2_global=float(np.linalg.norm(e)/np.linalg.norm(target)),
                relative_l2_mean=float(rel.mean()),relative_l2_std=float(rel.std()),
                relative_l2_max=float(rel.max()),mse=float(np.mean(e*e)),
                rmse=float(np.sqrt(np.mean(e*e))),mae=float(np.mean(np.abs(e))),
                residual_variance_signed=float(np.var(e)),
                residual_variance_absolute=float(np.var(np.abs(e))),
                residual_linf=float(np.abs(e).max()))


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


def save_checkpoint(filename, value):
    filename = Path(filename)
    filename.parent.mkdir(parents=True,exist_ok=True)
    temp = filename.with_suffix('.tmp')
    torch.save(value,temp)
    temp.replace(filename)


def dataset_manifest(data,cfg):
    h=hashlib.sha256()
    with data['source'].open('rb') as f:
        for block in iter(lambda:f.read(1024*1024),b''):h.update(block)
    return dict(source=str(data['source']),sha256=h.hexdigest(),original_shape=[1500,1000],
                selection='first 1000 trajectories in original file order; remaining 500 excluded',
                train_indices=list(range(800)),validation_indices=list(range(800,900)),test_indices=list(range(900,1000)),
                split_provenance='User-specified reproducible convention; authors original indices are unknown.',
                pressure_sensor_times=data['pressure_t'].tolist(),radius_output_times=data['output_t'].tolist(),
                interpolation=cfg['data']['interpolation'],normalization=data['stats'],statistics_fit_on='training split only')