"""Compare the optional backend with released Hecate on an actual CT crop."""
from __future__ import annotations

import argparse
import gc
import hashlib
import json
from pathlib import Path
import subprocess
import sys
import time
import tempfile
from unittest.mock import patch

import numpy as np
from PIL import Image

import hecate_directml as dml


def sha(path: Path) -> str:
    with path.open('rb') as stream:
        return hashlib.file_digest(stream, 'sha256').hexdigest()


def check_failure_cleanup(source: np.ndarray, folder: Path) -> None:
    class BrokenSession:
        def run(self, *args, **kwargs):
            raise RuntimeError('injected inference failure')

    real_temporary_directory = tempfile.TemporaryDirectory
    with real_temporary_directory(dir=folder) as parent:
        def directory(*args, **kwargs):
            return real_temporary_directory(dir=parent)
        with patch.object(dml.tempfile, 'TemporaryDirectory', side_effect=directory):
            try:
                dml.predict(BrokenSession(), source, np.zeros(source.shape[-2:], np.uint8))
            except RuntimeError as exc:
                assert str(exc) == 'injected inference failure'
            else:
                raise AssertionError('The injected failure did not propagate')
        assert not list(Path(parent).iterdir()), 'Temporary mappings survived the error'
        # The outer PNG buffer must close on errors too.
        try:
            with real_temporary_directory(dir=parent) as tmp:
                with dml.mapped_array(Path(tmp) / 'prediction', 'uint8', source.shape[-2:]):
                    raise RuntimeError('injected writer failure')
        except RuntimeError as exc:
            assert str(exc) == 'injected writer failure'
        assert not list(Path(parent).iterdir()), 'Output mapping survived the error'


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--checkpoint', type=Path, required=True)
    parser.add_argument('--input', type=Path, required=True, help='Real uint8 ZYX crop, at least 16x96x96')
    parser.add_argument('--output', type=Path, required=True, help='New verification directory')
    parser.add_argument('--stage', choices=['all', 'cpu', 'export', 'directml'], default='all')
    args = parser.parse_args()
    if args.stage == 'all':
        args.output.mkdir(parents=True, exist_ok=False)
        # End each process to release the CPU convolution/export allocator before
        # loading the next model. gc.collect() does not release those native arenas.
        for stage in ('cpu', 'export', 'directml'):
            subprocess.run([sys.executable, str(Path(__file__)), '--checkpoint', str(args.checkpoint),
                '--input', str(args.input), '--output', str(args.output), '--stage', stage], check=True)
        return
    import torch
    import hecate
    torch.set_num_threads(4)
    torch.set_num_interop_threads(1)
    graph = args.output / 'hecate.onnx'
    if args.stage == 'export':
        start = time.perf_counter()
        dml.export_model(args.checkpoint, graph)
        (args.output / 'export.json').write_text(json.dumps({'seconds': time.perf_counter() - start}), encoding='utf-8')
        return
    import zarr
    source = np.load(args.input, allow_pickle=False)
    assert source.dtype == np.uint8 and source.ndim == 3 and min(source.shape[-2:]) >= 96 and source.shape[0] >= 16
    source = source[:, :96, :96]
    assert source.any(), 'A blank array is not a real-data demonstration'
    cases = [('overlap', source, False), ('reverse', source, True),
             ('padded_edge', source[:, :45, :53], False)]
    masked = source.copy()
    masked[:, :, :9] = 0
    cases.append(('empty_columns', masked, False))
    cases.append(('blank_control', np.zeros_like(source), False))
    reference, cpu_seconds = {}, {}
    if args.stage == 'cpu':
        model = hecate.load_model(args.checkpoint, 'cpu')
        for name, volume, reverse in cases:
            value = np.zeros(volume.shape[-2:], np.uint8)
            start = time.perf_counter()
            hecate.predict(model, volume, value, reverse=reverse)
            cpu_seconds[name] = time.perf_counter() - start
            Image.fromarray(value).save(args.output / f'{name}_cpu.png')
        (args.output / 'cpu.json').write_text(json.dumps(cpu_seconds), encoding='utf-8')
        return
    cpu_seconds = json.loads((args.output / 'cpu.json').read_text())
    export_seconds = json.loads((args.output / 'export.json').read_text())['seconds']
    reference = {name: np.asarray(Image.open(args.output / f'{name}_cpu.png')) for name, _, _ in cases}
    check_failure_cleanup(source, args.output)
    session = dml.load_session(graph, profile_prefix=str(args.output / 'dispatch'))
    rows = []
    for name, volume, reverse in cases:
        value = np.zeros(volume.shape[-2:], np.uint8)
        start = time.perf_counter()
        stats = dml.predict(session, volume, value, reverse=reverse)
        elapsed = time.perf_counter() - start
        difference = np.abs(value.astype(np.int16) - reference[name].astype(np.int16))
        row = {'case': name, 'shape': list(volume.shape), 'max_uint8_difference': int(difference.max()),
               'different_pixels': int(np.count_nonzero(difference)), 'pixels': int(value.size),
               'cpu_seconds': cpu_seconds[name], 'directml_seconds': elapsed, **stats}
        Image.fromarray(value).save(args.output / f'{name}_directml.png')
        rows.append(row)
        print(json.dumps(row), flush=True)
        assert row['max_uint8_difference'] <= 1, row
        if name == 'empty_columns':
            assert not value[:, :9].any()
        if name == 'blank_control':
            assert not value.any() and stats['model_forwards'] == 0
    repeated = np.zeros(source.shape[-2:], np.uint8)
    dml.predict(session, source, repeated)
    assert np.array_equal(repeated, np.asarray(Image.open(args.output / 'overlap_directml.png')))
    profile = dml.finish_profile(session)
    assert set(profile['executed_providers']) == {'DmlExecutionProvider'}
    del session
    gc.collect()
    volume_path = args.output / 'surface.zarr'
    options = {'zarr_format': 2} if int(zarr.__version__.split('.')[0]) >= 3 else {}
    volume = zarr.open_array(str(volume_path), mode='w', shape=source.shape,
                             chunks=(16, 32, 32), dtype='uint8', **options)
    volume[:] = source
    cli_png = args.output / 'cli.png'
    command = [sys.executable, str(Path(dml.__file__)), 'predict', '--model', str(graph),
               '--input', str(volume_path), '--spacing-um', '9.6', '--output', str(cli_png), '--profile']
    subprocess.run(command, check=True)
    assert np.array_equal(np.asarray(Image.open(cli_png)), repeated)
    cli = json.loads(cli_png.with_suffix('.json').read_text())
    assert set(cli['profile']['executed_providers']) == {'DmlExecutionProvider'}
    assert cli['checkpoint_sha256'] == sha(args.checkpoint)
    # Refuse accidental output replacement and a wrong physical spacing.
    for change in ('existing_output', 'wrong_spacing'):
        bad = command.copy()
        if change == 'wrong_spacing':
            bad[bad.index('--spacing-um') + 1] = '2.4'
            bad[bad.index('--output') + 1] = str(args.output / 'wrong_spacing.png')
        result = subprocess.run(bad, capture_output=True, text=True)
        assert result.returncode != 0, change
    result = {'status': 'PASS', 'checkpoint_sha256': sha(args.checkpoint), 'input_sha256': sha(args.input),
              'hecate_sha256': sha(Path(hecate.__file__)), 'helper_sha256': sha(Path(dml.__file__)),
              'export_seconds': export_seconds, 'rows': rows, 'exact_repeat': True,
              'zarr_cli_exact': True, 'failure_cleanup': True, 'cpu_fallback_disabled': True,
              'executed_providers': profile['executed_providers'],
              'environment': {'python': sys.version.split()[0], 'torch': torch.__version__},
              'scope': 'Backend equivalence on a real exposed crop and explicit edge controls; no new text or model-accuracy claim.'}
    (args.output / 'result.json').write_text(json.dumps(result, indent=2) + '\n', encoding='utf-8')
    print('PASS: real-data parity, reversed depth, padding, empty columns, replay and Zarr CLI', flush=True)


if __name__ == '__main__':
    main()