#!/usr/bin/env python3
"""SOTA quality metrics — blur, exposure, contrast, resolution.

Her image için skor hesaplar, manifest'e yazar. Training'de filter.
Reference: Pech-Pacheco 2000, PreP-OCR ACL 2025, CHURRO-DS EMNLP 2025.
"""
import json, os, argparse
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor
import numpy as np
from PIL import Image
import cv2

Image.MAX_IMAGE_PIXELS = None
ROOT = Path("/arf/scratch/stakan/hitit-proje")
SOURCES = ROOT / "datasets" / "sources"

def compute_quality(item):
    """
    Returns: (rid, blur_score, exposure_mean, contrast_std, width, height, mode)
    blur_score: Laplacian variance (higher = sharper)
    exposure_mean: mean brightness on grayscale
    contrast_std: std dev of grayscale
    """
    rid, path = item
    try:
        with Image.open(path) as img:
            img_arr = np.array(img.convert('L'))
        if img_arr.size == 0:
            return (rid, None, None, None, 0, 0, None)
        h, w = img_arr.shape[:2]
        blur = float(cv2.Laplacian(img_arr, cv2.CV_64F).var())
        exposure = float(img_arr.mean())
        contrast = float(img_arr.std())
        with Image.open(path) as img:
            mode = img.mode
        return (rid, blur, exposure, contrast, w, h, mode)
    except Exception:
        return (rid, None, None, None, 0, 0, None)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--workers', type=int, default=200)
    ap.add_argument('--limit-per-source', type=int, default=None)
    args = ap.parse_args()
    
    # Tüm image path'leri topla (unique)
    all_items = []
    seen_paths = set()
    for d in sorted(SOURCES.iterdir()):
        if not d.is_dir(): continue
        mp = d / "manifest.jsonl"
        if not mp.exists(): continue
        src_items = []
        with open(mp) as f:
            for line in f:
                r = json.loads(line)
                p = r.get('path')
                if p and r.get('storage') == 'fs' and r.get('integrity_ok') is True:
                    if p in seen_paths: continue
                    seen_paths.add(p)
                    src_items.append((r['id'], p))
                    if args.limit_per_source and len(src_items) >= args.limit_per_source:
                        break
        all_items.extend(src_items)
    
    print(f"Quality metrics: {len(all_items):,} unique images")
    
    # Paralel compute
    results = {}
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        for r in ex.map(compute_quality, all_items, chunksize=500):
            results[r[0]] = r
    print(f"Hesaplanan: {len(results):,}")
    
    # Manifest'lere yaz
    for d in sorted(SOURCES.iterdir()):
        if not d.is_dir(): continue
        for mf in ['manifest.jsonl', 'manifest_classification.jsonl', 'manifest_detection.jsonl']:
            mp = d / mf
            if not mp.exists(): continue
            records = []
            with open(mp) as f:
                for line in f:
                    r = json.loads(line)
                    res = results.get(r.get('id'))
                    if res:
                        _, blur, expo, cont, w, h, mode = res
                        r['blur_score'] = round(blur, 2) if blur is not None else None
                        r['exposure_mean'] = round(expo, 2) if expo is not None else None
                        r['contrast_std'] = round(cont, 2) if cont is not None else None
                        # width/height zaten var, güvenli ekle
                        if w and not r.get('width'): r['width'] = w
                        if h and not r.get('height'): r['height'] = h
                    records.append(r)
            with open(mp, 'w') as f:
                for r in records:
                    f.write(json.dumps(r, ensure_ascii=False) + '\n')
    
    # Dataset stats
    blurs = [v[1] for v in results.values() if v[1] is not None]
    expos = [v[2] for v in results.values() if v[2] is not None]
    conts = [v[3] for v in results.values() if v[3] is not None]
    
    def percentiles(arr):
        arr = sorted(arr)
        n = len(arr)
        return {f"p{p}": round(arr[int(n*p/100)], 2) for p in [1, 5, 25, 50, 75, 95, 99]}
    
    summary = {
        "n_images_scored": len(results),
        "blur_score": percentiles(blurs) if blurs else {},
        "exposure_mean": percentiles(expos) if expos else {},
        "contrast_std": percentiles(conts) if conts else {},
        "thresholds_applied": {
            "blur_min": 100.0,
            "exposure_min": 20.0,
            "exposure_max": 235.0,
            "contrast_min": 15.0,
        },
        "n_failing_blur": sum(1 for b in blurs if b < 100),
        "n_failing_exposure": sum(1 for e in expos if e < 20 or e > 235),
        "n_failing_contrast": sum(1 for c in conts if c < 15),
    }
    
    with open(ROOT / "datasets" / "processed" / "quality_metrics_summary.json", 'w') as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)
    
    print(f"Blur: min={min(blurs):.1f}, median={np.median(blurs):.1f}, max={max(blurs):.1f}")
    print(f"Exposure: min={min(expos):.1f}, median={np.median(expos):.1f}, max={max(expos):.1f}")
    print(f"Contrast: min={min(conts):.1f}, median={np.median(conts):.1f}, max={max(conts):.1f}")
    print(f"Failing blur (<100): {summary['n_failing_blur']:,}")
    print(f"Failing exposure: {summary['n_failing_exposure']:,}")
    print(f"Failing contrast (<15): {summary['n_failing_contrast']:,}")

if __name__ == '__main__':
    main()