"""Quick test: run 12-class segmentation on 10 diverse images and save results."""

from __future__ import annotations

import sys
import shutil
from pathlib import Path

import torch
import numpy as np
from PIL import Image

# Ensure project root is on path.
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "src" / "nn" / "segearth_ov3"))

from src.conf.config_loader import load_all_configs
from src.augmentor.models import load_segmentation_model
from src.augmentor.inference import infer_segmentation_batch
from src.augmentor.io_utils import save_segmentation

INPUT_ROOT = Path("/mnt/data1tb/cvgl_datasets/UAV-GeoLoc")

# 10 diverse test images: satellite (DB) — urban + terrain
TEST_IMAGES = [
    # Urban (satellite)
    "Country/German/Munich/Ludwigs/DB/img/crop_12_4.png",
    "Country/French/Paris/LeMarais/DB/img/crop_12_4.png",
    "Country/Italy/Venice/SanMarco/DB/img/crop_12_4.png",
    "Country/USA/NewYork/Manhattan/DB/img/crop_12_4.png",
    "Country/Australia/Sydney/SydneyCBD/DB/img/crop_12_4.png",
    "Country/Korea/Busan/Gwangalli/DB/img/crop_12_4.png",
    # Terrain (satellite)
    "Terrain/Desert/GobiDesert/DB/img/crop_50_33.png",
    "Terrain/Glacier/AthabascaGlacier/DB/img/crop_12_4.png",
    "Terrain/Volcano/KilaueaVolcano/DB/img/crop_12_4.png",
    "Terrain/Danxia/GrandCanyon/DB/img/crop_50_33.png",
]

OUTPUT_DIR = ROOT / "test_seg_output"


def main():
    # Load configs (reads all .gin files)
    configs = load_all_configs(str(ROOT / "in" / "config_files") + "/")
    pipeline_conf = configs["pipeline"]
    hw_conf = configs["hardware"]
    models_conf = configs["models"]
    input_conf = configs["input"]
    seg_conf = configs["seg"]

    print(f"Prompts ({len(seg_conf.prompts)}): {seg_conf.prompts}")
    print(f"Threshold: {seg_conf.threshold}")
    print(f"Resolution: {seg_conf.default_resolution}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # Load model
    print("\nLoading SegEarth-OV3...")
    model, seg_config = load_segmentation_model(models_conf, hw_conf, seg_conf, device)
    num_classes = seg_config.get("num_classes", 150)
    print(f"Model loaded. Type: {seg_config['type']}, num_classes: {num_classes}")

    # Prepare output dir
    OUTPUT_DIR.mkdir(exist_ok=True)

    # Process each image
    for i, rel_path in enumerate(TEST_IMAGES):
        img_path = INPUT_ROOT / rel_path
        if not img_path.exists():
            print(f"[{i+1}/10] SKIP (not found): {rel_path}")
            continue

        print(f"\n[{i+1}/10] Processing: {rel_path}")

        # Load and preprocess
        pil_img = Image.open(img_path).convert("RGB")
        img_np = np.array(pil_img).astype(np.float32) / 255.0
        img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0)  # [1, 3, H, W]

        # Infer
        seg_ids = infer_segmentation_batch(model, seg_config, img_tensor, device)
        # seg_ids: [1, 1, H, W] uint8

        # Save original + segmentation side by side
        stem = img_path.stem
        tag = rel_path.split("/")[0]  # Country or Terrain
        if tag == "Country":
            tag = rel_path.split("/")[2]  # city name
        elif tag == "Terrain":
            tag = rel_path.split("/")[1]  # terrain type

        out_stem = f"{i:02d}_{tag}_{stem}"

        # Save original for reference
        pil_img.save(OUTPUT_DIR / f"{out_stem}_orig.png")

        # Save segmentation vis
        save_segmentation(
            seg_ids[0], OUTPUT_DIR, out_stem,
            save_npy=False, save_vis=True, num_classes=num_classes,
        )

        unique_classes = torch.unique(seg_ids[0]).tolist()
        class_names = [seg_conf.prompts[c] for c in unique_classes if c < len(seg_conf.prompts)]
        print(f"  -> Classes found: {class_names}")

    print(f"\nDone! Results saved to: {OUTPUT_DIR}")


if __name__ == "__main__":
    main()