"""Configuration, reproducible IO and hardware checks shared by entrypoints."""
import argparse
import hashlib
import json
import os
from pathlib import Path

import numpy as np
import torch
import yaml


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


def arguments(description):
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--config', type=Path, default=ROOT / 'config/config.yaml')
    return parser


def load_config(path):
    path = Path(path).resolve()
    with path.open() as handle:
        config = yaml.safe_load(handle)
    return config, path.parents[1]


def output_path(root, value):
    value = Path(value)
    return value if value.is_absolute() else root / value


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 sha256(path):
    digest = hashlib.sha256()
    with Path(path).open('rb') as handle:
        for part in iter(lambda: handle.read(1024 * 1024), b''):
            digest.update(part)
    return digest.hexdigest()


def device_for(config, require_dcu=True):
    device = torch.device(config['training']['device'])
    if require_dcu:
        if not os.environ.get('SLURM_JOB_ID'):
            raise RuntimeError('This experiment must run in a SLURM DCU allocation.')
        if device.type != 'cuda' or not torch.cuda.is_available() or not torch.version.hip:
            raise RuntimeError('A HIP/DCU PyTorch device is required; CPU fallback is disabled.')
    return device


def runtime_info(device):
    return {'hostname': os.uname().nodename, 'slurm_job_id': os.environ.get('SLURM_JOB_ID'),
            'slurm_partition': os.environ.get('SLURM_JOB_PARTITION'),
            'torch': torch.__version__, 'hip': torch.version.hip,
            'device': str(device), 'device_name': torch.cuda.get_device_name(device),
            'device_memory_bytes': torch.cuda.get_device_properties(device).total_memory}


def errors(prediction, target):
    difference = np.asarray(prediction, dtype=np.float64) - np.asarray(target, dtype=np.float64)
    return {'MAE': np.mean(np.abs(difference), axis=0).tolist(),
            'RMSE': np.sqrt(np.mean(difference ** 2, axis=0)).tolist()}


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