#!/usr/bin/env python3
"""Collect comprehensive statistics about the Terrain subset of UAV-GeoLoc."""

import os
import json
from collections import defaultdict
from PIL import Image

ROOT = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Terrain"

def get_image_size_safe(path):
    try:
        with Image.open(path) as im:
            return im.size
    except Exception:
        return None

def parse_db_position(path):
    """Parse db_postion.txt -> list of (lon, lat)."""
    coords = []
    if not os.path.isfile(path):
        return coords
    with open(path) as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) >= 3:
                try:
                    lon, lat = float(parts[1]), float(parts[2])
                    coords.append((lon, lat))
                except ValueError:
                    pass
    return coords

def count_files_in_dir(d, exts=None):
    if not os.path.isdir(d):
        return 0
    if exts is None:
        return len(os.listdir(d))
    return sum(1 for f in os.listdir(d) if os.path.splitext(f)[1].lower() in exts)

def analyze_scene(scene_path):
    info = {}
    # DB crops
    db_img_dir = os.path.join(scene_path, "DB", "img")
    info["db_crops"] = count_files_in_dir(db_img_dir, {".png", ".jpg", ".jpeg", ".tif"})

    # DB crop size (sample first image)
    info["crop_size"] = None
    if os.path.isdir(db_img_dir):
        for f in sorted(os.listdir(db_img_dir)):
            sz = get_image_size_safe(os.path.join(db_img_dir, f))
            if sz:
                info["crop_size"] = sz
                break

    # merge.tif size
    merge_path = os.path.join(scene_path, "DB", "merge.tif")
    info["merge_size"] = get_image_size_safe(merge_path) if os.path.isfile(merge_path) else None

    # Query variants
    query_dir = os.path.join(scene_path, "query")
    variants = []
    frames_per_variant = {}
    if os.path.isdir(query_dir):
        for v in sorted(os.listdir(query_dir)):
            vpath = os.path.join(query_dir, v)
            if os.path.isdir(vpath):
                footage_dir = os.path.join(vpath, "footage")
                n = count_files_in_dir(footage_dir, {".png", ".jpg", ".jpeg"})
                variants.append(v)
                frames_per_variant[v] = n
    info["variants"] = variants
    info["num_variants"] = len(variants)
    info["frames_per_variant"] = frames_per_variant
    # Use first variant's frame count as representative
    info["frames_per_variant_sample"] = list(frames_per_variant.values())[0] if frames_per_variant else 0
    info["total_query_frames"] = sum(frames_per_variant.values())

    # db_postion.txt
    db_pos_path = os.path.join(scene_path, "DB", "db_postion.txt")
    coords = parse_db_position(db_pos_path)
    if coords:
        lons = [c[0] for c in coords]
        lats = [c[1] for c in coords]
        info["gps"] = {
            "lon_min": min(lons), "lon_max": max(lons),
            "lat_min": min(lats), "lat_max": max(lats),
            "num_entries": len(coords)
        }
    else:
        info["gps"] = None

    # positive.json
    pos_path = os.path.join(scene_path, "positive.json")
    if os.path.isfile(pos_path):
        with open(pos_path) as f:
            pos = json.load(f)
        counts = [len(v) if isinstance(v, list) else 1 for v in pos.values()]
        info["positive"] = {
            "num_frames": len(pos),
            "total_positives": sum(counts),
            "avg_per_frame": sum(counts) / len(counts) if counts else 0,
            "min_per_frame": min(counts) if counts else 0,
            "max_per_frame": max(counts) if counts else 0,
        }
    else:
        info["positive"] = None

    # semi_positive.json
    sp_path = os.path.join(scene_path, "semi_positive.json")
    if os.path.isfile(sp_path):
        with open(sp_path) as f:
            sp = json.load(f)
        counts = [len(v) if isinstance(v, list) else 1 for v in sp.values()]
        info["semi_positive"] = {
            "num_frames": len(sp),
            "total_semi_positives": sum(counts),
            "avg_per_frame": sum(counts) / len(counts) if counts else 0,
            "min_per_frame": min(counts) if counts else 0,
            "max_per_frame": max(counts) if counts else 0,
        }
    else:
        info["semi_positive"] = None

    return info


def main():
    terrain_types = sorted([d for d in os.listdir(ROOT) if os.path.isdir(os.path.join(ROOT, d))])

    all_data = {}  # terrain_type -> {scene_name -> info}
    grand_total_db = 0
    grand_total_query = 0
    grand_total_scenes = 0

    for tt in terrain_types:
        tt_path = os.path.join(ROOT, tt)
        scenes = sorted([d for d in os.listdir(tt_path)
                         if os.path.isdir(os.path.join(tt_path, d))])
        all_data[tt] = {}
        for sc in scenes:
            sc_path = os.path.join(tt_path, sc)
            # Check it's actually a scene (has DB dir)
            if not os.path.isdir(os.path.join(sc_path, "DB")):
                continue
            info = analyze_scene(sc_path)
            all_data[tt][sc] = info
            grand_total_db += info["db_crops"]
            grand_total_query += info["total_query_frames"]
            grand_total_scenes += 1

    # ===================== PRINT RESULTS =====================

    print("=" * 120)
    print("UAV-GeoLoc TERRAIN SUBSET - COMPREHENSIVE STATISTICS")
    print("=" * 120)

    # 1. Hierarchy
    print("\n" + "=" * 120)
    print("TABLE 1: COMPLETE HIERARCHY (TerrainType -> Scenes)")
    print("=" * 120)
    print(f"{'TerrainType':<25} {'#Scenes':>7}  Scenes")
    print("-" * 120)
    for tt in terrain_types:
        scenes = list(all_data.get(tt, {}).keys())
        if not scenes:
            print(f"{tt:<25} {'0':>7}  (no valid scenes)")
            continue
        print(f"{tt:<25} {len(scenes):>7}  {', '.join(scenes)}")
    print(f"\n{'TOTAL TERRAIN TYPES:':<25} {len(terrain_types)}")
    print(f"{'TOTAL SCENES:':<25} {grand_total_scenes}")

    # 2. Per-scene counts
    print("\n" + "=" * 120)
    print("TABLE 2: PER-SCENE IMAGE COUNTS")
    print("=" * 120)
    print(f"{'TerrainType':<20} {'Scene':<40} {'DB Crops':>9} {'#Variants':>10} {'Frames/Var':>11} {'Total QFrames':>14}")
    print("-" * 120)
    for tt in terrain_types:
        for sc, info in sorted(all_data.get(tt, {}).items()):
            print(f"{tt:<20} {sc:<40} {info['db_crops']:>9} {info['num_variants']:>10} {info['frames_per_variant_sample']:>11} {info['total_query_frames']:>14}")

    print(f"\n{'GRAND TOTAL DB CROPS:':<50} {grand_total_db:>14}")
    print(f"{'GRAND TOTAL QUERY FRAMES:':<50} {grand_total_query:>14}")
    print(f"{'GRAND TOTAL ALL IMAGES:':<50} {grand_total_db + grand_total_query:>14}")

    # 3. Crop & merge sizes
    print("\n" + "=" * 120)
    print("TABLE 3: CROP AND MERGE.TIF SIZES (pixels)")
    print("=" * 120)
    print(f"{'TerrainType':<20} {'Scene':<40} {'Crop WxH':>12} {'Merge WxH':>14}")
    print("-" * 120)
    for tt in terrain_types:
        for sc, info in sorted(all_data.get(tt, {}).items()):
            cs = f"{info['crop_size'][0]}x{info['crop_size'][1]}" if info['crop_size'] else "N/A"
            ms = f"{info['merge_size'][0]}x{info['merge_size'][1]}" if info['merge_size'] else "N/A"
            print(f"{tt:<20} {sc:<40} {cs:>12} {ms:>14}")

    # Summary of unique sizes
    crop_sizes = defaultdict(int)
    merge_sizes = defaultdict(int)
    for tt in terrain_types:
        for sc, info in all_data.get(tt, {}).items():
            if info['crop_size']:
                crop_sizes[info['crop_size']] += 1
            if info['merge_size']:
                merge_sizes[info['merge_size']] += 1
    print(f"\nUnique crop sizes: {dict(crop_sizes)}")
    print(f"Unique merge.tif sizes: {dict(merge_sizes)}")

    # 4. GPS coordinate ranges
    print("\n" + "=" * 120)
    print("TABLE 4: GPS COORDINATE RANGES (from db_postion.txt)")
    print("=" * 120)
    print(f"{'TerrainType':<20} {'Scene':<35} {'#Entries':>8} {'Lon Min':>12} {'Lon Max':>12} {'Lat Min':>12} {'Lat Max':>12}")
    print("-" * 120)
    for tt in terrain_types:
        for sc, info in sorted(all_data.get(tt, {}).items()):
            g = info["gps"]
            if g:
                print(f"{tt:<20} {sc:<35} {g['num_entries']:>8} {g['lon_min']:>12.6f} {g['lon_max']:>12.6f} {g['lat_min']:>12.6f} {g['lat_max']:>12.6f}")
            else:
                print(f"{tt:<20} {sc:<35} {'N/A':>8} {'N/A':>12} {'N/A':>12} {'N/A':>12} {'N/A':>12}")

    # 5. positive.json stats
    print("\n" + "=" * 120)
    print("TABLE 5: positive.json STATS")
    print("=" * 120)
    print(f"{'TerrainType':<20} {'Scene':<35} {'#Frames':>8} {'TotalPos':>9} {'AvgPos':>8} {'MinPos':>7} {'MaxPos':>7}")
    print("-" * 120)
    all_pos_avg = []
    for tt in terrain_types:
        for sc, info in sorted(all_data.get(tt, {}).items()):
            p = info["positive"]
            if p:
                print(f"{tt:<20} {sc:<35} {p['num_frames']:>8} {p['total_positives']:>9} {p['avg_per_frame']:>8.2f} {p['min_per_frame']:>7} {p['max_per_frame']:>7}")
                all_pos_avg.append(p['avg_per_frame'])
            else:
                print(f"{tt:<20} {sc:<35} {'N/A':>8} {'N/A':>9} {'N/A':>8} {'N/A':>7} {'N/A':>7}")
    if all_pos_avg:
        print(f"\nOverall avg positives per frame across all scenes: {sum(all_pos_avg)/len(all_pos_avg):.2f}")

    # 6. semi_positive.json stats
    print("\n" + "=" * 120)
    print("TABLE 6: semi_positive.json STATS")
    print("=" * 120)
    print(f"{'TerrainType':<20} {'Scene':<35} {'#Frames':>8} {'TotalSP':>9} {'AvgSP':>8} {'MinSP':>7} {'MaxSP':>7}")
    print("-" * 120)
    all_sp_avg = []
    for tt in terrain_types:
        for sc, info in sorted(all_data.get(tt, {}).items()):
            sp = info["semi_positive"]
            if sp:
                print(f"{tt:<20} {sc:<35} {sp['num_frames']:>8} {sp['total_semi_positives']:>9} {sp['avg_per_frame']:>8.2f} {sp['min_per_frame']:>7} {sp['max_per_frame']:>7}")
                all_sp_avg.append(sp['avg_per_frame'])
            else:
                print(f"{tt:<20} {sc:<35} {'N/A':>8} {'N/A':>9} {'N/A':>8} {'N/A':>7} {'N/A':>7}")
    if all_sp_avg:
        print(f"\nOverall avg semi-positives per frame across all scenes: {sum(all_sp_avg)/len(all_sp_avg):.2f}")

    # 7. Variant breakdown (unique variant names across dataset)
    print("\n" + "=" * 120)
    print("TABLE 7: QUERY VARIANT NAMES (height/rot combinations)")
    print("=" * 120)
    all_variants = set()
    for tt in terrain_types:
        for sc, info in all_data.get(tt, {}).items():
            all_variants.update(info["variants"])
    for v in sorted(all_variants):
        print(f"  {v}")
    print(f"\nTotal unique variant names: {len(all_variants)}")

    print("\n" + "=" * 120)
    print("END OF REPORT")
    print("=" * 120)


if __name__ == "__main__":
    main()