#!/usr/bin/env python3
"""Pixel comparison of two images (reference vs result), optionally per region.

Works for any pair: live Kodi screenshot vs design reference, live vs offline
render, before vs after. Prints metrics and writes side_by_side.png (reference
| result | diff heatmap), diff.png, and a zoomed stack for small regions.

Usage:
    python3 compare_images.py REF.png TEST.png [--region x0,y0,x1,y1]
        [--ref-region x0,y0,x1,y1] [--out-dir DIR] [--zoom N] [--threshold 32]

Regions are in each image's own pixel coordinates (x1/y1 exclusive).
--region applies to both images unless --ref-region is given for the reference.
If the two crops differ in size, the reference crop is resized to the test crop
and the report says so (not pixel-exact in that case).
"""
import argparse
import json
import os
import sys

import numpy as np
from PIL import Image, ImageDraw

HERE = os.path.dirname(os.path.abspath(__file__))


def parse_box(text):
    if not text:
        return None
    parts = [int(round(float(v))) for v in text.split(',')]
    if len(parts) != 4 or parts[2] <= parts[0] or parts[3] <= parts[1]:
        raise SystemExit('region must be x0,y0,x1,y1 with x1>x0, y1>y0: {}'.format(text))
    return tuple(parts)


def load_rgb(path):
    img = Image.open(path)
    if img.mode == 'RGBA':
        flat = Image.new('RGBA', img.size, (0, 0, 0, 255))
        flat.alpha_composite(img)
        img = flat
    return img.convert('RGB')


def compare(ref, test, threshold=32):
    """Return metrics dict for two same-size RGB images."""
    a = np.asarray(ref).astype(np.int16)
    b = np.asarray(test).astype(np.int16)
    diff = np.abs(a - b)
    per_px = diff.max(axis=2)
    total = per_px.size
    over = per_px > threshold
    mse = float((diff.astype(np.float64) ** 2).mean())
    metrics = {
        'size': list(ref.size),
        'mean_abs_diff': round(float(diff.mean()), 3),
        'max_diff': int(per_px.max()),
        'identical_pct': round(100.0 * float((per_px == 0).sum()) / total, 3),
        'over_8_pct': round(100.0 * float((per_px > 8).sum()) / total, 3),
        'over_threshold_pct': round(100.0 * float(over.sum()) / total, 3),
        'threshold': threshold,
        'psnr_db': round(10 * np.log10(255.0 ** 2 / mse), 2) if mse > 0 else None,
        'diff_bbox': None,
    }
    if over.any():
        ys, xs = np.where(over)
        metrics['diff_bbox'] = [int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1]
    return metrics, diff.astype(np.uint8), over


def heatmap(diff, over):
    amp = np.clip(diff.astype(np.int32) * 4, 0, 255).astype(np.uint8)
    img = Image.fromarray(amp, 'RGB')
    red = np.zeros(diff.shape, dtype=np.uint8)
    red[over] = (255, 0, 64)
    return Image.blend(img, Image.fromarray(red, 'RGB'), 0.35) if over.any() else img


def side_by_side(ref, test, heat, labels=('reference', 'result', 'diff x4')):
    gap, bar = 12, 28
    w, h = ref.size
    out = Image.new('RGB', (w * 3 + gap * 2, h + bar), (24, 24, 28))
    draw = ImageDraw.Draw(out)
    for index, (img, label) in enumerate(zip((ref, test, heat), labels)):
        x = index * (w + gap)
        out.paste(img, (x, bar))
        draw.text((x + 4, 6), label, fill=(230, 230, 230))
    return out


def zoom_stack(ref, test, factor):
    w, h = ref.size[0] * factor, ref.size[1] * factor
    out = Image.new('RGB', (w, h * 2 + 6), (255, 0, 0))
    out.paste(ref.resize((w, h), Image.NEAREST), (0, 0))
    out.paste(test.resize((w, h), Image.NEAREST), (0, h + 6))
    draw = ImageDraw.Draw(out)
    step = 10 * factor
    for gx in range(0, w, step):
        draw.line([(gx, 0), (gx, out.height)], fill=(255, 0, 0), width=1)
    for gy in range(0, h, step):
        draw.line([(0, gy), (w, gy)], fill=(0, 200, 0), width=1)
        draw.line([(0, h + 6 + gy), (w, h + 6 + gy)], fill=(0, 200, 0), width=1)
    return out


def run(ref_path, test_path, region=None, ref_region=None, out_dir=None,
        zoom=0, threshold=32, quiet=False):
    ref = load_rgb(ref_path)
    test = load_rgb(test_path)
    ref_box = ref_region or region
    notes = []
    if ref_box:
        ref = ref.crop(ref_box)
    if region:
        test = test.crop(region)
    if ref.size != test.size:
        notes.append('reference {}x{} resized to {}x{} (not pixel-exact)'.format(
            ref.size[0], ref.size[1], test.size[0], test.size[1]))
        ref = ref.resize(test.size, Image.LANCZOS)
    metrics, diff, over = compare(ref, test, threshold)
    metrics['notes'] = notes
    metrics['region'] = list(region) if region else None
    metrics['ref_region'] = list(ref_box) if ref_box else None

    out_dir = out_dir or os.path.join(HERE, 'out', 'compare')
    os.makedirs(out_dir, exist_ok=True)
    heat = heatmap(diff, over)
    labels = ('ref: ' + os.path.basename(ref_path), 'test: ' + os.path.basename(test_path), 'diff x4 (red > threshold)')
    side_by_side(ref, test, heat, labels).save(os.path.join(out_dir, 'side_by_side.png'))
    heat.save(os.path.join(out_dir, 'diff.png'))
    outputs = ['side_by_side.png', 'diff.png']
    if zoom == 0 and max(test.size) <= 480:
        zoom = max(2, min(8, 960 // max(test.size)))
    if zoom:
        zoom_stack(ref, test, zoom).save(os.path.join(out_dir, 'zoom.png'))
        outputs.append('zoom.png')
    metrics['outputs'] = [os.path.join(out_dir, name) for name in outputs]
    with open(os.path.join(out_dir, 'metrics.json'), 'w') as handle:
        json.dump(metrics, handle, indent=2)

    if not quiet:
        print_report(metrics, ref_path, test_path)
    return metrics


def print_report(m, ref_path, test_path):
    print('=' * 72)
    print('PIXEL COMPARE  ref={}  test={}'.format(os.path.basename(ref_path), os.path.basename(test_path)))
    if m.get('region') or m.get('ref_region'):
        print('region test={}  ref={}'.format(m.get('region'), m.get('ref_region')))
    print('-' * 72)
    print('size              {} x {}'.format(*m['size']))
    print('identical pixels  {} %'.format(m['identical_pct']))
    print('mean abs diff     {} (0-255)'.format(m['mean_abs_diff']))
    print('max diff          {}'.format(m['max_diff']))
    print('pixels > 8        {} %'.format(m['over_8_pct']))
    print('pixels > {:<8} {} %'.format(m['threshold'], m['over_threshold_pct']))
    print('PSNR              {} dB'.format(m['psnr_db'] if m['psnr_db'] is not None else 'inf (identical)'))
    print('diff bbox (>{})   {}'.format(m['threshold'], m['diff_bbox']))
    for note in m.get('notes', []):
        print('NOTE: ' + note)
    print('-' * 72)
    print('outputs: ' + '  '.join(m['outputs']))


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('ref')
    parser.add_argument('test')
    parser.add_argument('--region')
    parser.add_argument('--ref-region')
    parser.add_argument('--out-dir')
    parser.add_argument('--zoom', type=int, default=0)
    parser.add_argument('--threshold', type=int, default=32)
    args = parser.parse_args()
    run(args.ref, args.test, parse_box(args.region), parse_box(args.ref_region),
        args.out_dir, args.zoom, args.threshold)


if __name__ == '__main__':
    sys.exit(main())