#!/usr/bin/env python3
"""Export the actual saved best states and check safe-weight inference on compute."""
import argparse
import hashlib
import json
import os
import sys
import tempfile
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
from safetensors.torch import load_file, save_file

HERE = Path(__file__).resolve().parent
GEMINI = Path('/e/scratch/reformo/schuhmann1_moss/whisper_score_regression/gemini_finetune_preparation_20261005')
RELEASE = GEMINI.parent / 'humaneness_ears_release_20261007/model'


def digest(path):
    h = hashlib.sha256()
    with Path(path).open('rb') as stream:
        while part := stream.read(8 * 1024 * 1024):
            h.update(part)
    return h.hexdigest()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', choices=('base', 'small'), required=True)
    args = parser.parse_args()
    size = args.model
    release_size = 'medium' if size == 'small' else size
    torch.set_num_threads(2)
    source = GEMINI / 'training' / ('whisper_' + size) / 'best.pt'
    checkpoint = torch.load(source, map_location='cpu', weights_only=False)
    assert checkpoint['epoch_completed'] == 2 and checkpoint['update'] == 914
    norm = json.loads((RELEASE / 'training_normalization.json').read_text())
    for key in ['score_names', 'score_mean', 'score_std', 'cps_mean', 'cps_std']:
        assert checkpoint['normalization'][key] == norm[key], key
    state = {k: v.detach().contiguous() for k, v in checkpoint['model'].items()}
    assert all(not k.startswith(('decoder.', 'model.decoder.')) for k in state)
    output = RELEASE / release_size / 'model.safetensors'
    save_file(state, output, metadata={'format': 'pt', 'author': 'Christoph Schuhmann',
              'license': 'cc-by-4.0', 'checkpoint_sha256': digest(source),
              'epochs': '2', 'update': '914', 'teacher': 'gemini-3.8-flash'})
    restored = load_file(output)
    assert set(restored) == set(state)
    assert all(torch.equal(state[k], restored[k]) for k in state)
    sys.path.insert(0, str(RELEASE))
    from inference import Predictor
    from model import LayeredMultiTaskWhisper
    device = 'cuda:' + os.environ.get('LOCAL_RANK', '0')
    predictor = Predictor(RELEASE, size, device, threads=2)
    original = LayeredMultiTaskWhisper(RELEASE / release_size, 53, initialize_pretrained=False).to(device).eval()
    original.load_state_dict(checkpoint['model'], strict=True)
    assert sum(p.numel() for p in original.parameters()) == json.loads(
        (RELEASE / 'training/configs' / (size + '_recorded_run.json')).read_text())['trainable_total_parameters']
    # Exact saved-weights parity, including layer-mixture and event heads.
    generator = torch.Generator().manual_seed(20261007)
    mel = torch.randn((1, 80, 100), generator=generator).to(device)
    mask = torch.ones((1, 100), dtype=torch.long, device=device)
    with torch.inference_mode():
        a = original(mel, mask, predict_events=True, max_pred_events=32)
        b = predictor.model(mel, mask, predict_events=True, max_pred_events=32)
    assert set(a) == set(b)
    assert all(torch.equal(a[k], b[k]) for k in a), 'Export altered inference outputs'
    del original, checkpoint, state, restored, a, b
    # End-to-end new CLI/API on an actual held-out waveform, including resampling.
    row = None
    with (GEMINI / 'prepared/targets.jsonl').open() as stream:
        for line in stream:
            candidate = json.loads(line)
            if candidate['split'] == 'test' and candidate['ready_for_whisper_training']:
                row = candidate
                break
    sys.path.insert(0, str(RELEASE / 'training/gemini_finetune'))
    from common import TarReader
    wave = TarReader().audio(row, verify=True)
    with tempfile.TemporaryDirectory(prefix='whisper-release-') as directory:
        path = Path(directory) / 'heldout.wav'
        sf.write(path, wave, 16000, subtype='FLOAT')
        result = predictor.predict(path, include_frames=True)
        assert len(result['scores_raw']) == len(result['scores_z']) == len(result['scores_training_z']) == 192
        assert len(result['orange_timbre_128d']) == 128 and len(result['orange_identity_250d']) == 250
        assert len(result['burst_probability_per_20ms_frame']) > 0
        json.dumps(result, allow_nan=False)
        for event in result['binary_burst_regions'] + result['burst_event_proposals']:
            assert 0 <= event['start_s'] < event['end_s'] <= result['duration_s'] + 1e-6
        for event in result['burst_event_proposals']:
            assert len(event['top3_classes']) == 3
        # The advertised one-core CPU entry point also runs on real audio.
        cpu = Predictor(RELEASE, size, 'cpu', threads=1)
        short = Path(directory) / 'heldout-short.wav'
        sf.write(short, wave[:16000], 16000, subtype='FLOAT')
        cpu_result = cpu.predict(short)
        json.dumps(cpu_result, allow_nan=False)
        assert len(cpu_result['scores_raw']) == 192
    evidence = {'model': size, 'checkpoint_epoch': 2, 'updates': 914,
                'source_checkpoint_sha256': digest(source), 'safe_weights_sha256': digest(output),
                'safe_weights_bytes': output.stat().st_size,
                'all_tensors_bitwise_equal': True, 'all_model_outputs_bitwise_equal': True,
                'heldout_audio_sha256': row['sha256_mp3'], 'heldout_duration_s': result['duration_s'],
                'targets': 192, 'speaker_dimensions': [128, 250], 'finite_outputs': True,
                'timestamps_have_start_and_end': True, 'device': device,
                'one_thread_cpu_inference_passed': True,
                'torch': torch.__version__, 'slurm_job': os.environ.get('SLURM_JOB_ID')}
    path = RELEASE / 'verification' / (size + '.json')
    path.parent.mkdir(exist_ok=True)
    path.write_text(json.dumps(evidence, indent=2) + '\n')
    print(json.dumps(evidence), flush=True)


if __name__ == '__main__':
    main()