#!/usr/bin/env python3
"""Dataset-specific mean/std hesapla (ImageNet yerine).
Reference: DINOv3 satellite example, Meta AI 2025.

Per-source ve global stats; hem RGB hem grayscale.
"""
import json, os, argparse
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor
import numpy as np
from PIL import Image

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

def accumulate(item):
    """Tek image için channel-wise sum, sum_sq, pixel_count."""
    rid, path = item
    try:
        with Image.open(path) as img:
            arr = np.array(img.convert('RGB'), dtype=np.float64) / 255.0
        h, w, c = arr.shape
        n = h * w
        # Flatten HW
        flat = arr.reshape(-1, c)
        return (n, flat.sum(axis=0), (flat**2).sum(axis=0))
    except Exception:
        return None

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--sample-per-source', type=int, default=1000)
    ap.add_argument('--workers', type=int, default=200)
    args = ap.parse_args()
    
    import random
    random.seed(42)
    
    all_global_n = 0
    all_global_sum = np.zeros(3)
    all_global_sq = np.zeros(3)
    per_source_stats = {}
    
    for d in sorted(SOURCES.iterdir()):
        if not d.is_dir(): continue
        mp = d / "manifest.jsonl"
        if not mp.exists(): continue
        paths = []
        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:
                    paths.append((r['id'], p))
        if not paths: continue
        
        sample = random.sample(paths, min(args.sample_per_source, len(paths)))
        
        n_total = 0
        sum_ = np.zeros(3)
        sq_ = np.zeros(3)
        with ProcessPoolExecutor(max_workers=args.workers) as ex:
            for res in ex.map(accumulate, sample, chunksize=50):
                if res is None: continue
                n, s, sq = res
                n_total += n
                sum_ += s
                sq_ += sq
        
        if n_total:
            mean = sum_ / n_total
            var = (sq_ / n_total) - mean**2
            std = np.sqrt(np.maximum(var, 0))
            per_source_stats[d.name] = {
                "n_pixels": int(n_total),
                "n_images": len(sample),
                "mean": [round(float(x), 4) for x in mean],
                "std": [round(float(x), 4) for x in std],
            }
            print(f"  {d.name}: mean={per_source_stats[d.name]['mean']}, std={per_source_stats[d.name]['std']}")
            all_global_n += n_total
            all_global_sum += sum_
            all_global_sq += sq_
    
    global_mean = (all_global_sum / all_global_n).tolist() if all_global_n else [0,0,0]
    global_var = (all_global_sq / all_global_n) - (all_global_sum / all_global_n)**2 if all_global_n else [0,0,0]
    global_std = np.sqrt(np.maximum(global_var, 0)).tolist()
    
    out = {
        "strategy": "DINOv3 dataset-specific (not ImageNet)",
        "reference": "DINOv3 satellite example, Meta AI 2025",
        "sample_per_source": args.sample_per_source,
        "seed": 42,
        "global_rgb": {
            "mean": [round(float(x), 4) for x in global_mean],
            "std": [round(float(x), 4) for x in global_std],
            "n_pixels_total": int(all_global_n),
        },
        "per_source": per_source_stats,
        "imagenet_reference": {
            "mean": [0.485, 0.456, 0.406],
            "std": [0.229, 0.224, 0.225],
        },
    }
    with open(ROOT / "datasets" / "processed" / "normalization_stats.json", 'w') as f:
        json.dump(out, f, indent=2, ensure_ascii=False)
    
    print(f"\nGLOBAL mean: {out['global_rgb']['mean']}")
    print(f"GLOBAL std: {out['global_rgb']['std']}")

if __name__ == '__main__':
    main()