"""Compare paired BF16 and INT4 runs without ranking them."""
import argparse
import csv
import html
import json
import math
import os
import statistics
from pathlib import Path

import numpy as np
from PIL import Image

from scripts.integrity import sha256, validate_pair
from scripts.provenance import validate_roles


def load_records(root):
    rows = [json.loads(line) for line in (Path(root) / 'records.jsonl').read_text(encoding='utf-8').splitlines() if line]
    records = {}
    for row in rows:
        key = (row['case_id'], row['seed'])
        if key in records:
            raise ValueError(f'Duplicate record: {key}')
        if sha256(Path(root) / row['image']) != row['image_sha256']:
            raise ValueError(f'Image hash mismatch: {key}')
        records[key] = row
    if not records:
        raise ValueError('Empty benchmark run')
    return records


def compare_pixels(left, right):
    with Image.open(left) as image_a, Image.open(right) as image_b:
        if image_a.size != image_b.size:
            raise ValueError('Image dimensions differ')
        image_a, image_b = image_a.convert('RGBA'), image_b.convert('RGBA')
        array_a = np.asarray(image_a, dtype=np.float32) / 255
        array_b = np.asarray(image_b, dtype=np.float32) / 255
        composite_a = array_a[..., :3] * array_a[..., 3:] + 1 - array_a[..., 3:]
        composite_b = array_b[..., :3] * array_b[..., 3:] + 1 - array_b[..., 3:]
        mse = float(np.mean((composite_a - composite_b) ** 2))
        return {'rgb_mae_white': float(np.mean(np.abs(composite_a - composite_b))),
                'rgb_psnr_white_db': None if mse == 0 else -10 * math.log10(mse),
                'alpha_mae': float(np.mean(np.abs(array_a[..., 3] - array_b[..., 3])))}


def summary_from(rows, baseline_env, quantized_env, sample):
    return {'pairs': len(rows), 'cases': len({row['case_id'] for row in rows}),
            'bf16_mean_seconds': statistics.mean(row['bf16_seconds'] for row in rows),
            'int4_mean_seconds': statistics.mean(row['int4_seconds'] for row in rows),
            'bf16_max_allocated_gib': max(row['bf16_peak_allocated_gib'] for row in rows),
            'int4_max_allocated_gib': max(row['int4_peak_allocated_gib'] for row in rows),
            'width': sample['width'], 'height': sample['height'], 'steps': sample['steps'],
            'warmup': baseline_env['warmup'],
            'bf16_weight_bytes': sum(item['size'] for item in baseline_env['model_identity']['files']
                                    if item['path'].endswith('.safetensors')),
            'int4_weight_bytes': sum(item['size'] for item in quantized_env['model_identity']['files']
                                    if item['path'].endswith('.safetensors'))}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--baseline', default='artifacts/eval/bf16')
    parser.add_argument('--quantized', default='artifacts/eval/int4')
    parser.add_argument('--output', default='artifacts/eval')
    args = parser.parse_args()
    base, quant, out = Path(args.baseline), Path(args.quantized), Path(args.output)
    baseline, quantized = load_records(base), load_records(quant)
    if baseline.keys() != quantized.keys():
        raise ValueError('Benchmark cases or seeds do not match')
    baseline_env = json.loads((base / 'environment.json').read_text(encoding='utf-8'))
    quantized_env = json.loads((quant / 'environment.json').read_text(encoding='utf-8'))
    validate_roles(baseline_env['model_identity'], quantized_env['model_identity'])
    for key in ('benchmark_sha256', 'runtime_helper_sha256', 'device_helper_sha256'):
        if baseline_env.get(key) != quantized_env.get(key):
            raise ValueError(f'Benchmark implementation mismatch: {key}')
    for key in ('gpu', 'cuda', 'packages', 'offload', 'warmup', 'generator_device', 'cases_sha256'):
        if baseline_env[key] != quantized_env[key]:
            raise ValueError(f'Runtime mismatch: {key}')
    rows = []
    for key in baseline:
        left, right = baseline[key], quantized[key]
        validate_pair(left, right)
        rows.append({'case_id': key[0], 'seed': key[1],
                     'bf16_seconds': left['seconds'], 'int4_seconds': right['seconds'],
                     'bf16_peak_allocated_gib': left['peak_allocated_bytes'] / 2**30,
                     'int4_peak_allocated_gib': right['peak_allocated_bytes'] / 2**30,
                     **compare_pixels(base / left['image'], quant / right['image'])})
    summary = summary_from(rows, baseline_env, quantized_env, next(iter(baseline.values())))
    out.mkdir(parents=True, exist_ok=True)
    (out / 'summary.json').write_text(json.dumps(summary, indent=2), encoding='utf-8')
    with (out / 'comparison.csv').open('w', encoding='utf-8', newline='') as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    lines = ['# Informal release evaluation: BF16 / INT4', '',
             'This community comparison is reference material. It is not an official evaluation.', '',
             f'{len(rows)} paired outputs on {baseline_env["gpu"]}; {summary["width"]}×{summary["height"]}, '
             f'{summary["steps"]} steps, offload={baseline_env["offload"]}, CFG=1, KV cache enabled.', '',
             ('Warmup excluded. ' if baseline_env['warmup'] else 'No warmup performed. ') +
             'Pixel metrics measure drift. They are not a semantic quality score. '
             'The suite is small and does not establish a ranking.', '',
             '| Case | Seed | BF16 s | INT4 s | BF16 peak GiB | INT4 peak GiB | RGB MAE |',
             '|---|---:|---:|---:|---:|---:|---:|']
    gallery = ['<!doctype html><meta charset="utf-8"><title>BF16 / INT4</title>',
               '<style>body{font:16px system-ui;margin:32px;max-width:1500px}'
               '.pair{display:grid;grid-template-columns:1fr 1fr;gap:20px}img{width:100%;background:'
               'repeating-conic-gradient(#ddd 0 25%,#fff 0 50%) 0/24px 24px}pre{white-space:pre-wrap}</style>',
               '<h1>BF16 / INT4</h1><p>Left: BF16. Right: saved INT4. Same settings and seeds.</p>']
    for row in rows:
        lines.append(f'| {row["case_id"]} | {row["seed"]} | {row["bf16_seconds"]:.2f} | '
                     f'{row["int4_seconds"]:.2f} | {row["bf16_peak_allocated_gib"]:.2f} | '
                     f'{row["int4_peak_allocated_gib"]:.2f} | {row["rgb_mae_white"]:.4f} |')
        key = (row['case_id'], row['seed'])
        left = Path(os.path.relpath(base / baseline[key]['image'], out)).as_posix()
        right = Path(os.path.relpath(quant / quantized[key]['image'], out)).as_posix()
        gallery.append(f'<section><h2>{html.escape(row["case_id"])} seed {row["seed"]}</h2>'
                       f'<pre>{html.escape(baseline[key]["prompt"])}</pre><div class="pair">'
                       f'<a href="{html.escape(left, quote=True)}"><img src="{html.escape(left, quote=True)}" alt="BF16"></a>'
                       f'<a href="{html.escape(right, quote=True)}"><img src="{html.escape(right, quote=True)}" alt="INT4"></a></div></section>')
    lines.extend(['', '## Summary', '',
                  f'Mean latency: BF16 {summary["bf16_mean_seconds"]:.2f}s; INT4 {summary["int4_mean_seconds"]:.2f}s.',
                  f'Maximum allocated CUDA memory: BF16 {summary["bf16_max_allocated_gib"]:.2f} GiB; '
                  f'INT4 {summary["int4_max_allocated_gib"]:.2f} GiB.', '',
                  'Raw records: comparison.csv, bf16/records.jsonl, int4/records.jsonl.',
                  'Visual notes belong in qualitative.md and are written after inspecting the images.', ''])
    (out / 'report.md').write_text('\n'.join(lines) + '\n', encoding='utf-8')
    (out / 'comparison.html').write_text('\n'.join(gallery), encoding='utf-8')
    print(json.dumps(summary, indent=2))


if __name__ == '__main__':
    main()