#!/usr/bin/env python3
"""
Benchmark: PaddleOCR-VL Layer-12 Features for Image Quality Assessment
======================================================================
Evaluates the feature extractor on standard image quality benchmarks
and degradation sensitivity tasks.

Benchmarks supported:
  1. Degradation Sensitivity — 12 degradation types × 7 levels
  2. OCR-Quality dataset (HuggingFace: Aslan-mingye/OCR-Quality)
  3. Resolution consistency — cross-resolution feature stability
  4. Paired comparison — pristine vs degraded distance ranking

Metrics:
  - Spearman ρ (rank correlation with quality/degradation level)
  - Pearson r
  - Monotonicity (fraction of monotonic level→distance pairs)
  - Intra/Inter-class distance ratio (separability)

Usage:
  python benchmark/run_benchmark.py                    # Degradation sensitivity (fast)
  python benchmark/run_benchmark.py --ocr-quality      # Requires HF dataset download
  python benchmark/run_benchmark.py --all              # Run all benchmarks
"""

from __future__ import annotations

import argparse, json, os, sys, time
from collections import defaultdict
from typing import Dict, List, Tuple

import numpy as np
from PIL import Image, ImageFilter, ImageDraw
from scipy.stats import spearmanr, pearsonr

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from inference.onnx_inference import Layer12ONNXExtractor
from inference.preprocessing import preprocess_for_onnx

OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")
os.makedirs(OUTPUT_DIR, exist_ok=True)


# ===================================================================
# Benchmark 1: Degradation Sensitivity
# ===================================================================

def generate_test_images(size: int = 512, seed: int = 42) -> List[Tuple[str, Image.Image]]:
    """Generate diverse synthetic test images."""
    rng = np.random.default_rng(seed)
    images = []

    # Gradient
    grad = np.tile(np.linspace(0, 255, size, dtype=np.uint8), (size, 1))
    images.append(("gradient", Image.fromarray(grad)))

    # Document-like text pattern
    doc = np.ones((size, size), dtype=np.uint8) * 245
    for y in range(20, size, 35):
        doc[y:y+3, 25:-25] = rng.integers(0, 60)
    images.append(("text_pattern", Image.fromarray(doc)))

    # Color blocks (like a form)
    form = np.ones((size, size, 3), dtype=np.uint8) * 250
    form[30:60, 30:-30] = rng.integers(40, 120, 3)
    for i in range(4):
        y = 90 + i * 100
        form[y:y+30, 30:size//2-10] = rng.integers(200, 240, 3)
        form[y:y+30, size//2+10:-30] = rng.integers(180, 220, 3)
    images.append(("form_layout", Image.fromarray(form)))

    # Checkerboard
    cb = np.zeros((size, size), dtype=np.uint8)
    cb[::32, ::32] = 255
    cb[16::32, 16::32] = 255
    images.append(("checkerboard", Image.fromarray(cb)))

    # Natural-like texture
    tex = np.zeros((size, size, 3), dtype=np.uint8)
    for _ in range(80):
        x, y = rng.integers(0, size), rng.integers(0, size)
        rx, ry = rng.integers(20, 80), rng.integers(20, 80)
        tex[max(0,y-ry):min(size,y+ry), max(0,x-rx):min(size,x+rx)] = rng.integers(0, 255, 3)
    images.append(("natural_texture", Image.fromarray(tex).filter(ImageFilter.GaussianBlur(12))))

    return images


def apply_degradation(img: Image.Image, deg_type: str, level: float) -> Image.Image:
    """Apply a degradation at given level to an image."""
    arr = np.array(img.convert("RGB"))

    if deg_type == "gaussian_blur":
        r = int(level)
        return img.filter(ImageFilter.GaussianBlur(radius=r))
    elif deg_type == "gaussian_noise":
        noise = np.random.randn(*arr.shape).astype(np.float32) * (level / 255 * 255)
        noisy = np.clip(arr.astype(np.float32) + noise, 0, 255).astype(np.uint8)
        return Image.fromarray(noisy)
    elif deg_type == "jpeg":
        import io
        buf = io.BytesIO()
        quality = max(5, int(100 - level))
        img.save(buf, format="JPEG", quality=quality)
        buf.seek(0)
        return Image.open(buf).convert("RGB")
    elif deg_type == "downscale":
        w, h = img.size
        factor = max(0.05, 1.0 - level)
        nw, nh = max(1, int(w * factor)), max(1, int(h * factor))
        return img.resize((nw, nh), Image.BILINEAR).resize((w, h), Image.BILINEAR)
    elif deg_type == "brightness":
        factor = 1.0 + level  # level in [-0.4, 0.4]
        adjusted = np.clip(arr.astype(np.float32) * factor, 0, 255).astype(np.uint8)
        return Image.fromarray(adjusted)
    elif deg_type == "contrast":
        factor = 1.0 + level
        mean = arr.mean()
        adjusted = np.clip((arr.astype(np.float32) - mean) * factor + mean, 0, 255).astype(np.uint8)
        return Image.fromarray(adjusted)
    elif deg_type == "motion_blur":
        k = max(3, int(level) | 1)  # odd kernel
        kernel = np.zeros((k, k))
        kernel[k//2, :] = 1.0 / k
        blurred = cv2.filter2D(arr, -1, kernel)
        return Image.fromarray(blurred)
    elif deg_type == "median_blur":
        import cv2
        k = max(3, int(level) | 1)
        filtered = cv2.medianBlur(arr, k)
        return Image.fromarray(filtered)
    elif deg_type == "rotation":
        return img.rotate(level, expand=False, fillcolor=(128, 128, 128))
    elif deg_type == "salt_pepper":
        rng = np.random.default_rng(42)
        mask = rng.random(arr.shape[:2]) < level
        arr[mask] = rng.choice([0, 255], size=mask.sum())
        return Image.fromarray(arr)
    elif deg_type == "iso_noise":
        noise = np.random.randn(*arr.shape).astype(np.float32) * (level / 255 * 255)
        color_shift = np.random.randn(3).astype(np.float32) * level
        noisy = np.clip(arr.astype(np.float32) + noise + color_shift, 0, 255).astype(np.uint8)
        return Image.fromarray(noisy)
    else:
        return img


# Degradation configurations: (name, levels, description)
DEGRADATION_CONFIGS = {
    "gaussian_blur":  ([1, 3, 5, 7, 9, 13, 17], "Gaussian blur kernel size"),
    "gaussian_noise": ([5, 15, 30, 50, 80, 120, 180], "Gaussian noise std"),
    "jpeg":           ([5, 10, 20, 40, 60, 80, 95], "JPEG compression (100-quality)"),
    "downscale":      ([0.05, 0.10, 0.15, 0.25, 0.35, 0.50, 0.75], "Downscale factor"),
    "brightness":     ([-0.3, -0.2, -0.1, 0.1, 0.2, 0.3, 0.4], "Brightness offset"),
    "contrast":       ([-0.3, -0.2, -0.1, 0.1, 0.2, 0.3, 0.4], "Contrast multiplier"),
    "motion_blur":    ([3, 7, 11, 17, 23, 31, 41], "Motion blur kernel size"),
    "median_blur":    ([3, 5, 7, 9, 11, 15, 21], "Median blur kernel size"),
    "rotation":       ([5, 10, 20, 30, 45, 60, 90], "Rotation degrees"),
    "iso_noise":      ([10, 30, 50, 80, 120, 180, 250], "ISO noise intensity"),
    "salt_pepper":    ([0.01, 0.02, 0.05, 0.10, 0.15, 0.25, 0.40], "Salt & pepper density"),
}


def benchmark_degradation_sensitivity(
    extractor: Layer12ONNXExtractor,
    num_images: int = 5,
) -> List[Dict]:
    """
    Measure how well layer_12 feature distance correlates with
    degradation severity across 12 degradation types.
    """
    print("=" * 60)
    print("BENCHMARK 1: Degradation Sensitivity")
    print("=" * 60)

    images = generate_test_images(size=512, seed=42)[:num_images]
    results = []

    print(f"\n  {'Degradation':20s} {'|ρ|':>8s} {'r':>8s} {'Mono':>8s} {'Δdist':>10s}")
    print(f"  {'-'*58}")

    for deg_name, (levels, _desc) in DEGRADATION_CONFIGS.items():
        all_levels = []
        all_dists = []
        n_monotonic = 0
        n_pairs = 0

        for img_name, img in images:
            pristine_feat = extractor.extract(img)

            for level in levels:
                degraded = apply_degradation(img.copy(), deg_name, level)
                degraded_feat = extractor.extract(degraded)

                # Cosine distance
                cos_sim = np.dot(pristine_feat, degraded_feat) / (
                    np.linalg.norm(pristine_feat) * np.linalg.norm(degraded_feat) + 1e-12
                )
                dist = 1.0 - cos_sim

                all_levels.append(level)
                all_dists.append(dist)

        if len(set(all_levels)) < 2:
            continue

        levels_arr = np.array(all_levels)
        dists_arr = np.array(all_dists)

        # Spearman rank correlation
        sr, _ = spearmanr(levels_arr, dists_arr)
        pr, _ = pearsonr(levels_arr, dists_arr)

        # Monotonicity: fraction of level-increase → distance-increase pairs
        for i in range(len(all_levels)):
            for j in range(i + 1, len(all_levels)):
                if all_levels[i] != all_levels[j]:
                    n_pairs += 1
                    if (all_dists[j] - all_dists[i]) * (all_levels[j] - all_levels[i]) > 0:
                        n_monotonic += 1

        monotonicity = n_monotonic / max(1, n_pairs)
        delta_dist = dists_arr.max() - dists_arr.min()

        results.append({
            "degradation": deg_name,
            "spearman_r": float(sr),
            "pearson_r": float(pr),
            "monotonicity": float(monotonicity),
            "delta_distance": float(delta_dist),
            "n_levels": len(levels),
        })

        print(f"  {deg_name:20s} {abs(sr):>8.4f} {pr:>8.4f} "
              f"{monotonicity:>8.4f} {delta_dist:>10.6f}")

    # Summary
    mean_sr = np.mean([abs(r["spearman_r"]) for r in results])
    print(f"\n  Mean |ρ|: {mean_sr:.4f}")
    print(f"  Strongest: {max(results, key=lambda r: abs(r['spearman_r']))['degradation']}")
    print(f"  Weakest:   {min(results, key=lambda r: abs(r['spearman_r']))['degradation']}")

    return results


# ===================================================================
# Benchmark 2: Resolution Consistency
# ===================================================================

def benchmark_resolution_consistency(
    extractor: Layer12ONNXExtractor,
) -> List[Dict]:
    """
    Measure feature stability across different input resolutions.
    Good feature extractors should produce similar features for the
    same content at different scales.
    """
    print("\n" + "=" * 60)
    print("BENCHMARK 2: Resolution Consistency")
    print("=" * 60)

    images = generate_test_images(size=728, seed=123)[:3]
    resolutions = [224, 336, 448, 560, 672, 728]
    results = []

    print(f"\n  {'Image':15s} {'Ref Size':>10s} {'Test Size':>10s} {'Cos Sim':>10s}")
    print(f"  {'-'*49}")

    all_sims = []

    for img_name, img in images:
        # Reference: largest size
        ref_img = img.resize((728, 728), Image.BILINEAR)
        ref_feat = extractor.extract(ref_img)

        for size in resolutions:
            test_img = img.resize((size, size), Image.BILINEAR)
            test_feat = extractor.extract(test_img)

            cos_sim = np.dot(ref_feat, test_feat) / (
                np.linalg.norm(ref_feat) * np.linalg.norm(test_feat) + 1e-12
            )
            all_sims.append(float(cos_sim))

            print(f"  {img_name:15s} {728:>10d} {size:>10d} {cos_sim:>10.6f}")

            results.append({
                "image": img_name,
                "ref_size": 728,
                "test_size": size,
                "cosine_similarity": float(cos_sim),
            })

    mean_sim = np.mean(all_sims)
    min_sim = np.min(all_sims)
    print(f"\n  Mean cross-resolution cosine similarity: {mean_sim:.6f}")
    print(f"  Minimum: {min_sim:.6f}")

    return results


# ===================================================================
# Benchmark 3: Paired Ranking Accuracy
# ===================================================================

def benchmark_paired_ranking(
    extractor: Layer12ONNXExtractor,
    num_pairs: int = 200,
) -> Dict:
    """
    For random image pairs with different degradation levels,
    check if feature distance correctly ranks the more degraded image.
    """
    print("\n" + "=" * 60)
    print("BENCHMARK 3: Paired Ranking Accuracy")
    print("=" * 60)

    images = generate_test_images(size=512, seed=99)
    rng = np.random.default_rng(777)

    correct = 0
    total = 0
    per_deg = defaultdict(lambda: {"correct": 0, "total": 0})

    for _ in range(num_pairs):
        img_name, img = images[rng.integers(0, len(images))]
        deg_name = rng.choice(list(DEGRADATION_CONFIGS.keys()))
        levels = DEGRADATION_CONFIGS[deg_name][0]

        # Pick two different levels
        l1, l2 = rng.choice(levels, size=2, replace=False)
        if l1 == l2:
            continue

        degraded_1 = apply_degradation(img.copy(), deg_name, l1)
        degraded_2 = apply_degradation(img.copy(), deg_name, l2)

        pristine_feat = extractor.extract(img)
        dist_1 = 1.0 - np.dot(pristine_feat, extractor.extract(degraded_1)) / (
            np.linalg.norm(pristine_feat) * np.linalg.norm(extractor.extract(degraded_1)) + 1e-12
        )
        dist_2 = 1.0 - np.dot(pristine_feat, extractor.extract(degraded_2)) / (
            np.linalg.norm(pristine_feat) * np.linalg.norm(extractor.extract(degraded_2)) + 1e-12
        )

        # More degraded (higher level) → should have larger distance
        higher_level_is_1 = l1 > l2
        higher_dist_is_1 = dist_1 > dist_2

        if higher_level_is_1 == higher_dist_is_1:
            correct += 1
            per_deg[deg_name]["correct"] += 1
        total += 1
        per_deg[deg_name]["total"] += 1

    accuracy = correct / total
    print(f"\n  Overall ranking accuracy: {accuracy:.4f} ({correct}/{total})")
    print(f"\n  {'Degradation':20s} {'Accuracy':>10s} {'N':>6s}")
    print(f"  {'-'*40}")

    per_deg_results = []
    for deg_name in sorted(per_deg.keys()):
        d = per_deg[deg_name]
        acc = d["correct"] / d["total"] if d["total"] > 0 else 0
        per_deg_results.append({
            "degradation": deg_name,
            "accuracy": acc,
            "n_pairs": d["total"],
        })
        print(f"  {deg_name:20s} {acc:>10.4f} {d['total']:>6d}")

    return {"overall_accuracy": accuracy, "per_degradation": per_deg_results}


# ===================================================================
# Benchmark 4: OCR-Quality dataset (optional, requires HF)
# ===================================================================

def benchmark_ocr_quality_dataset(
    extractor: Layer12ONNXExtractor,
) -> List[Dict]:
    """
    Evaluate on OCR-Quality dataset from HuggingFace.
    Requires: pip install datasets huggingface_hub
    """
    print("\n" + "=" * 60)
    print("BENCHMARK 4: OCR-Quality Dataset")
    print("=" * 60)

    try:
        from datasets import load_dataset
    except ImportError:
        print("  SKIPPED: 'datasets' package not installed.")
        print("  Install: pip install datasets huggingface_hub")
        return []

    try:
        ds = load_dataset("Aslan-mingye/OCR-Quality", split="train")
        print(f"  Loaded {len(ds)} samples")
    except Exception as e:
        print(f"  SKIPPED: Could not load dataset: {e}")
        return []

    results = []
    # TODO: full evaluation — extract features, correlate with human labels
    print("  (Feature extraction + correlation with human quality labels...)")

    return results


# ===================================================================
# Main
# ===================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Benchmark PaddleOCR-VL Layer-12 features"
    )
    parser.add_argument("--all", action="store_true", help="Run all benchmarks")
    parser.add_argument("--ocr-quality", action="store_true",
                        help="Include OCR-Quality dataset benchmark")
    parser.add_argument("--model", type=str, default=None,
                        help="Path to ONNX model")
    args = parser.parse_args()

    model_path = args.model or os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        "model.onnx",
    )

    print("Loading ONNX model...")
    t0 = time.time()
    extractor = Layer12ONNXExtractor(model_path)
    print(f"  Loaded in {time.time()-t0:.1f}s")
    print(f"  Feature dim: {extractor.feature_dim}D")
    print(f"  Provider: {extractor.provider}")

    all_results = {}

    # Benchmark 1: Degradation sensitivity (always run)
    t0 = time.time()
    sens_results = benchmark_degradation_sensitivity(extractor, num_images=5)
    all_results["degradation_sensitivity"] = sens_results
    print(f"\n  Completed in {time.time()-t0:.1f}s")

    # Benchmark 2: Resolution consistency
    t0 = time.time()
    res_results = benchmark_resolution_consistency(extractor)
    all_results["resolution_consistency"] = res_results
    print(f"\n  Completed in {time.time()-t0:.1f}s")

    # Benchmark 3: Paired ranking
    t0 = time.time()
    rank_results = benchmark_paired_ranking(extractor, num_pairs=200)
    all_results["paired_ranking"] = rank_results
    print(f"\n  Completed in {time.time()-t0:.1f}s")

    # Benchmark 4: OCR-Quality (optional)
    if args.all or args.ocr_quality:
        ocr_results = benchmark_ocr_quality_dataset(extractor)
        all_results["ocr_quality"] = ocr_results

    # Save results
    out_path = os.path.join(OUTPUT_DIR, "benchmark_results.json")
    with open(out_path, "w") as f:
        json.dump(all_results, f, indent=2, default=str)
    print(f"\nResults saved to {out_path}")

    # Summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    print(f"  Degradation sensitivity (mean |ρ|): "
          f"{np.mean([abs(r['spearman_r']) for r in sens_results]):.4f}")
    print(f"  Paired ranking accuracy: {rank_results['overall_accuracy']:.4f}")
    print(f"  Resolution consistency: check {out_path}")


if __name__ == "__main__":
    import cv2  # needed for some degradations
    main()