"""Data, functional, provenance and computation boundary shared by entrypoints."""
import argparse
import csv
import hashlib
import json
import os
import random
import socket
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if (ROOT / '.deps').exists():
    sys.path.insert(0, str(ROOT / '.deps'))
sys.path.insert(0, str(ROOT))

import numpy as np
import torch
import yaml


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


def load_config(path):
    with open(path) as f:
        cfg = yaml.safe_load(f)
    return cfg


def save_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))
    temporary.replace(path)


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



def model_state_on_cpu(model):
    """Export PyTorch state while handling neuraloperator 1.0.2 metadata.

    BaseModel.state_dict inserts an unprefixed dictionary key '_metadata' even
    when nested in NOTS. That entry is not a parameter/buffer and must not be
    passed to the enclosing nn.Module's strict loader. The model architecture
    is saved separately as checkpoint.model_config. OrderedDict._metadata (the
    PyTorch version information attribute) is preserved.
    """
    state = model.state_dict()
    metadata = state.pop('_metadata', None)
    if metadata is not None and not isinstance(metadata, dict):
        raise TypeError('Unexpected neuraloperator initialization metadata')
    for name, value in state.items():
        if not torch.is_tensor(value):
            raise TypeError(f'Unexpected non-tensor model state: {name}')
        state[name] = value.detach().cpu()
    return state


def write_csv(path, rows):
    path = Path(path)
    if not rows:
        return
    temporary = path.with_suffix('.csv.tmp')
    with temporary.open('w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)
    temporary.replace(path)


def sha256(path):
    h = hashlib.sha256()
    with open(path, 'rb') as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b''):
            h.update(chunk)
    return h.hexdigest()


def config_hash(cfg):
    return hashlib.sha256(json.dumps(cfg, sort_keys=True).encode()).hexdigest()


def case_guard(cfg):
    """Refuse login-node execution and attachment to any pre-existing job."""
    if cfg['runtime']['require_case_slurm_job']:
        actual = os.environ.get('SLURM_JOB_ID', '')
        registry = Path(cfg['paths']['output_dir']) / '.onescience/jobs.json'
        if not actual or not registry.is_file():
            raise RuntimeError('Compute requires a newly submitted case-owned SLURM job')
        jobs = json.loads(registry.read_text())['jobs']
        if actual not in [str(j['job_id']) for j in jobs]:
            raise RuntimeError(f'Refusing unregistered job {actual}')
        if not os.environ.get('SLURM_JOB_NODELIST'):
            raise RuntimeError('Missing allocated compute node list')
    torch.set_num_threads(cfg['runtime']['num_threads'])
    device = torch.device(cfg['runtime']['device'])
    if device.type == 'cuda' and not torch.cuda.is_available():
        raise RuntimeError('Allocated accelerator is unavailable; CPU fallback is disabled')
    return device


def seed_all(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def rng_state():
    return {'python': random.getstate(), 'numpy': np.random.get_state(),
            'torch': torch.get_rng_state(), 'cuda': torch.cuda.get_rng_state_all()}


def restore_rng(state):
    random.setstate(state['python'])
    np.random.set_state(state['numpy'])
    torch.set_rng_state(state['torch'])
    torch.cuda.set_rng_state_all(state['cuda'])


def load_data(cfg, split='train'):
    path = Path(cfg['paths']['data_dir']) / cfg['paths'][split + '_file']
    raw = torch.load(path, map_location='cpu', weights_only=True)
    out = []
    for key in ('x', 'y'):
        a = raw[key].float()
        if a.ndim == 3:
            a = a.unsqueeze(1)
        if a.ndim != 4 or tuple(a.shape[1:]) != (1, 16, 16):
            raise ValueError(f'{path}:{key} invalid shape {a.shape}')
        if not torch.isfinite(a).all():
            raise ValueError(f'{path}:{key} contains nonfinite data')
        out.append(a.contiguous())
    if out[0].shape != out[1].shape:
        raise ValueError('Input and pressure pairing mismatch')
    if split == 'train' and len(out[0]) != cfg['experiment']['pool_size']:
        raise ValueError('The required 1000-candidate official pool is unavailable')
    meta = {'path': str(path), 'sha256': sha256(path), 'n': len(out[0]),
            'shape': list(out[0].shape), 'x_range': [out[0].min().item(), out[0].max().item()],
            'y_range': [out[1].min().item(), out[1].max().item()], 'dtype': 'float32'}
    return out[0], out[1], meta


def negative_total_flow_rates(u, a):
    """Appendix D.1: integral(a grad(u).n), maximize negative physical outflow.

    First-order one-sided normal differences and rectangular edge quadrature.
    Equal 16x16 spacings cancel analytically; corners count once per edge.
    This deliberately uses original permeability values, not a display mask.
    """
    if u.shape != a.shape or u.ndim != 4:
        raise ValueError('Functional requires matching [B,1,H,W] tensors')
    h, w = u.shape[-2:]
    left = -(a[..., :, 0] * (u[..., :, 1] - u[..., :, 0])).sum(-1) * (w / h)
    right = (a[..., :, -1] * (u[..., :, -1] - u[..., :, -2])).sum(-1) * (w / h)
    bottom = -(a[..., 0, :] * (u[..., 1, :] - u[..., 0, :])).sum(-1) * (h / w)
    top = (a[..., -1, :] * (u[..., -1, :] - u[..., -2, :])).sum(-1) * (h / w)
    return (left + right + bottom + top).mean(-1)


@torch.no_grad()
def predict(model, xs, batch_size):
    model.eval()
    return torch.cat([model(xs[i:i + batch_size]) for i in range(0, len(xs), batch_size)])


def hardware_record():
    return {'hostname': socket.gethostname(), 'job_id': os.environ.get('SLURM_JOB_ID'),
            'node_list': os.environ.get('SLURM_JOB_NODELIST'), 'torch': torch.__version__,
            'accelerator': torch.cuda.get_device_name(), 'device_count': torch.cuda.device_count(),
            'device_memory_bytes': torch.cuda.get_device_properties(0).total_memory}