"""Inference on raw (height, width, time-bin) HDF5 histograms."""
import argparse
import json
from pathlib import Path

import h5py
import numpy as np
import torch
from model import HistCoarseMultiScaleMLP


def preprocess(histogram, stats):
    data = np.asarray(histogram, dtype=np.float32)
    if data.shape != (32, 32, 64):
        raise ValueError(f'Expected (32, 32, 64), got {data.shape}')
    if not np.isfinite(data).all() or (data < 0).any():
        raise ValueError('Histogram counts must be finite and nonnegative')
    if stats['mode'] != 'photon_count_10000_log1p_standardize':
        raise ValueError('Unsupported preprocessing mode')
    total = float(data.sum(dtype=np.float64))
    if total > 0:
        data = data * (10000.0 / total)
    data = np.log1p(data)
    data = (data - stats['mean']) / max(stats['std'], 1e-6)
    return torch.from_numpy(np.transpose(data, (2, 0, 1)).copy()).unsqueeze(0).unsqueeze(0)


def load_model(model_dir=None, device='cpu'):
    root = Path(model_dir) if model_dir else Path(__file__).resolve().parent
    config = json.loads((root / 'config.json').read_text())
    stats = json.loads((root / 'preprocess_stats.json').read_text())
    checkpoint = torch.load(root / 'best.pt', map_location='cpu', weights_only=True)
    if stats != checkpoint['preprocess_stats']:
        raise ValueError('Preprocessing statistics differ from checkpoint')
    model = HistCoarseMultiScaleMLP(config['hidden_features'], config['dropout'])
    model.load_state_dict(checkpoint['model_state'], strict=True)
    model.to(device).eval()
    return model, config, stats


def predict(histogram, model, config, stats):
    device = next(model.parameters()).device
    with torch.inference_mode():
        logit = model(preprocess(histogram, stats).to(device)).item()
        probability = torch.sigmoid(torch.tensor(logit)).item()
    label = int(probability >= config['threshold'])
    return {'logit': logit, 'probability_positive': probability,
            'label': label, 'label_name': config['id2label'][str(label)]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('histogram', type=Path)
    parser.add_argument('--model-dir', type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument('--device', default='cpu')
    args = parser.parse_args()
    model, config, stats = load_model(args.model_dir, args.device)
    with h5py.File(args.histogram, 'r') as f:
        histogram = f[config['h5_key']][()]
    print(json.dumps(predict(histogram, model, config, stats), indent=2))


if __name__ == '__main__':
    main()