prepare_gallery.py

"""
Prepare gallery: download dataset, extract embeddings, build FAISS index.

Downloads real multi-modal satellite data (optical RGB, SAR 2ch, MS 13ch)
and builds a cross-modal retrieval gallery.

Usage:
    python prepare_gallery.py --samples 50
"""

import argparse
import json
import shutil
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from tqdm import tqdm
from torchvision import transforms

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
DATA_DIR = Path("data")
RAW_DIR = DATA_DIR / "raw"
PROCESSED_DIR = DATA_DIR / "processed"
GALLERY_DIR = DATA_DIR / "gallery"

CLASSES = [
    "AnnualCrop", "Forest", "HerbaceousVegetation", "Highway",
    "Industrial", "Pasture", "PermanentCrop", "Residential",
    "River", "SeaLake",
]

BATCH_SIZE = 64
EMBED_DIM = 768  # CLIP ViT-L/14 output dim


def _samples_exist(dir: Path, n_per_class: int) -> bool:
    """Check if directory has n_per_class images per class."""
    if not dir.exists():
        return False
    return all(len(list((dir / c).glob("*.*"))) >= n_per_class for c in CLASSES)


def download_optical(n_per_class: int = 50) -> Path:
    """Download optical RGB from HuggingFace."""
    from datasets import load_dataset
    out_dir = RAW_DIR / "eurosat"
    if _samples_exist(out_dir, n_per_class):
        print(f"Optical already at {out_dir}, skipping.")
        return out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    print("Downloading optical RGB from blanchon/EuroSAT_RGB ...")
    ds = load_dataset("blanchon/EuroSAT_RGB", split="train")
    selected, counts = [], {c: 0 for c in range(10)}
    for row in ds:
        lbl = row["label"]
        if counts[lbl] < n_per_class:
            selected.append(row)
            counts[lbl] += 1
        if all(v >= n_per_class for v in counts.values()):
            break
    for i, row in enumerate(tqdm(selected, desc="Saving optical")):
        cls_name = CLASSES[row["label"]]
        cls_dir = out_dir / cls_name
        cls_dir.mkdir(exist_ok=True)
        row["image"].save(cls_dir / f"{cls_name}_{i}.tif")
    print(f"Optical saved to {out_dir}")
    return out_dir


def download_sar(n_per_class: int = 50) -> Path:
    """Download real SAR (2ch VV/VH) from HuggingFace dataset or zip."""
    out_dir = RAW_DIR / "eurosat_sar"
    if _samples_exist(out_dir, n_per_class):
        print(f"SAR already at {out_dir}, skipping.")
        return out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    print("Downloading SAR from wangyi111/EuroSAT-SAR ...")
    try:
        from datasets import load_dataset
        ds = load_dataset("wangyi111/EuroSAT-SAR", split="train", streaming=True)
        selected, counts = [], {c: 0 for c in range(10)}
        for row in ds:
            # SAR labels are class names directly
            lbl_name = row.get("label", row.get("label_name", ""))
            if isinstance(lbl_name, int) and lbl_name < 10:
                cls_name = CLASSES[lbl_name]
            elif isinstance(lbl_name, str) and lbl_name in CLASSES:
                cls_name = lbl_name
            else:
                continue
            idx = CLASSES.index(cls_name)
            if counts[idx] < n_per_class:
                # Convert to 2-channel grayscale (VV, VH from RGBA)
                img = np.array(row["image"].convert("L"))
                selected.append((img, cls_name))
                counts[idx] += 1
            if all(v >= n_per_class for v in counts.values()):
                break
        for i, (img_arr, cls_name) in enumerate(tqdm(selected, desc="Saving SAR")):
            cls_dir = out_dir / cls_name
            cls_dir.mkdir(exist_ok=True)
            # Save as 2-channel TIFF (stack Luminance as VV/VH)
            two_ch = np.stack([img_arr, img_arr], axis=-1).astype(np.uint8)
            Image.fromarray(two_ch[:, :, 0], mode="L").save(
                cls_dir / f"{cls_name}_{i}.tif")
        print(f"SAR saved to {out_dir}")
    except Exception as e:
        print(f"SAR download failed ({e}), using local fallback.")
        # Fallback: convert optical to 2-channel SAR-like
        optical_dir = RAW_DIR / "eurosat"
        if optical_dir.exists():
            for cls_name in CLASSES:
                cls_in = optical_dir / cls_name
                cls_out = out_dir / cls_name
                cls_out.mkdir(parents=True, exist_ok=True)
                for path in list(cls_in.glob("*.*"))[:n_per_class]:
                    arr = np.array(Image.open(path).convert("L"))
                    noise = np.random.rayleigh(1.0, arr.shape).astype(np.float32)
                    sar = np.clip(arr * noise, 0, 255).astype(np.uint8)
                    Image.fromarray(sar, mode="L").save(
                        cls_out / f"{cls_name}_{path.stem}.tif")
            print(f"SAR fallback saved to {out_dir}")
    return out_dir


def download_multispectral(n_per_class: int = 50) -> Path:
    """Download real multispectral (13ch) from HuggingFace."""
    out_dir = RAW_DIR / "eurosat_ms"
    if _samples_exist(out_dir, n_per_class):
        print(f"Multispectral already at {out_dir}, skipping.")
        return out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    print("Downloading MS from giswqs/EuroSAT_MS ...")
    try:
        from datasets import load_dataset
        ds = load_dataset("giswqs/EuroSAT_MS", split="train", streaming=True)
        selected, counts = [], {c: 0 for c in range(10)}
        for row in ds:
            lbl = row["label"]
            if counts[lbl] < n_per_class:
                # Image is a list of 13 arrays (one per band)
                bands = [np.array(b) for b in row["image"]]
                selected.append((bands, lbl))
                counts[lbl] += 1
            if all(v >= n_per_class for v in counts.values()):
                break
        for bands, lbl in tqdm(selected, desc="Saving MS"):
            cls_name = CLASSES[lbl]
            cls_dir = out_dir / cls_name
            cls_dir.mkdir(exist_ok=True)
            # Stack bands into single array, save as multi-channel TIFF
            arr = np.stack(bands, axis=0)  # (13, 64, 64)
            import tifffile
            tifffile.imwrite(
                cls_dir / f"{cls_name}_{len(list(cls_dir.glob('*.*')))}.tif",
                arr.astype(np.uint16))
        print(f"MS saved to {out_dir}")
    except Exception as e:
        print(f"MS download failed ({e}), using local fallback.")
        optical_dir = RAW_DIR / "eurosat"
        if optical_dir.exists():
            for cls_name in CLASSES:
                cls_in = optical_dir / cls_name
                cls_out = out_dir / cls_name
                cls_out.mkdir(parents=True, exist_ok=True)
                for path in list(cls_in.glob("*.*"))[:n_per_class]:
                    shutil.copy2(path, cls_out / path.name)
            print(f"MS fallback saved to {out_dir}")
    return out_dir


def load_satclip():
    """Load SatCLIP encoder."""
    from src.features.satclip_encoder import SatCLIPEncoder
    print("Loading SatCLIP encoder...")
    encoder = SatCLIPEncoder()
    print(f"SatCLIP loaded on {encoder.device}")
    return encoder


def _load_multichannel_image(path: Path, modality: str) -> torch.Tensor:
    """
    Load an image handling different channel counts.

    Returns tensor of shape (C, H, W) normalized to [0, 1].
    """
    # Try tifffile first for multi-channel TIFFs
    try:
        import tifffile
        arr = tifffile.imread(str(path))
        if arr.ndim == 3 and arr.shape[-1] in [2, 3, 4, 13]:
            # Channels-last format
            arr = np.transpose(arr, (2, 0, 1))
        elif arr.ndim == 2:
            arr = arr[np.newaxis, :, :]
        # Normalize uint to [0, 1]
        if arr.dtype in [np.uint8, np.uint16]:
            arr = arr.astype(np.float32) / np.float32(np.iinfo(arr.dtype).max)
        else:
            arr = arr.astype(np.float32)
            arr = (arr - arr.min()) / (arr.max() - arr.min() + 1e-8)
        tensor = torch.from_numpy(arr).float()
    except Exception:
        # Fallback to PIL (handles RGB)
        img = Image.open(path).convert("RGB")
        tensor = transforms.ToTensor()(img)

    # Resize to 224x224
    if tensor.shape[1] != 224 or tensor.shape[2] != 224:
        tensor = F.interpolate(tensor.unsqueeze(0), size=(224, 224),
                               mode="bilinear", align_corners=False).squeeze(0)
                               
    # Enforce strict channel counts per modality to prevent torch.stack failures
    c = tensor.shape[0]
    if modality == "optical":
        if c == 1:
            tensor = tensor.repeat(3, 1, 1)
        elif c > 3:
            tensor = tensor[:3]
        elif c == 2:
            tensor = torch.cat([tensor, tensor[:1]], dim=0)
    elif modality == "sar":
        if c == 1:
            tensor = tensor.repeat(2, 1, 1)
        elif c > 2:
            tensor = tensor[:2]
    elif modality == "multispectral":
        if c < 13:
            pad = torch.zeros(13 - c, tensor.shape[1], tensor.shape[2])
            tensor = torch.cat([tensor, pad], dim=0)
        elif c > 13:
            tensor = tensor[:13]
            
    return tensor


def _pad_to_13ch(tensor: torch.Tensor, modality: str) -> torch.Tensor:
    """Pad tensor to 13 channels for SatCLIP."""
    n_channels = tensor.shape[1]
    if n_channels >= 13:
        return tensor[:, :13]
    # Repeat channels if single-channel (SAR fallback)
    if n_channels == 1:
        tensor = tensor.repeat(1, 3, 1, 1)
        n_channels = 3
    pad_channels = 13 - n_channels
    padding = torch.zeros(
        tensor.shape[0], pad_channels, tensor.shape[2], tensor.shape[3])
    return torch.cat([tensor, padding], dim=1)


def _make_rgb_preview(path: Path, modality: str, size=(128, 128)) -> Image.Image:
    """Create an RGB preview image from any modality file."""
    try:
        import tifffile
        arr = tifffile.imread(str(path))
        if arr.ndim == 3 and arr.shape[-1] >= 3:
            preview = arr[:, :, :3]
        elif arr.ndim == 3 and arr.shape[0] >= 3:
            preview = np.transpose(arr[:3], (1, 2, 0))
        elif arr.ndim == 2:
            preview = np.stack([arr] * 3, axis=-1)
        else:
            preview = np.stack([arr[:, :, 0]] * 3, axis=-1)

        # Normalize for display
        if preview.dtype == np.uint16:
            preview = (preview / 65535.0 * 255).astype(np.uint8)
        elif preview.dtype == np.uint8:
            pass
        else:
            preview = (np.clip(preview, 0, 1) * 255).astype(np.uint8)

        # Special handling for SAR: grayscale with colormap feel
        if modality == "sar":
            preview = preview  # Keep as is

        return Image.fromarray(preview).resize(size, Image.LANCZOS)
    except Exception:
        # Fallback to PIL
        return Image.open(path).convert("RGB").resize(size, Image.LANCZOS)


@torch.no_grad()
def extract_embeddings_satclip(images, encoder, modality="optical"):
    """Extract L2-normalized embeddings from image tensors using SatCLIP."""
    all_feats = []
    for i in range(0, len(images), BATCH_SIZE):
        batch = images[i: i + BATCH_SIZE]
        tensors = torch.stack(batch)
        # Pad to 13 channels if needed
        if tensors.shape[1] < 13:
            tensors = _pad_to_13ch(tensors, modality)
        feats = encoder.encode(tensors, normalize=True)
        all_feats.append(feats.cpu())
    return torch.cat(all_feats, dim=0)


def build_gallery(n_per_class: int = 50):
    """Full pipeline: download, embed, build index."""
    t0 = time.time()

    # 1. Download data for all three modalities
    optical_dir = download_optical(n_per_class)
    sar_dir = download_sar(n_per_class)
    ms_dir = download_multispectral(n_per_class)

    # 2. Collect all images with proper multi-channel loading
    modalities = {
        "optical": optical_dir,
        "sar": sar_dir,
        "multispectral": ms_dir,
    }
    all_images = []  # list of (image_tensor, modality, class_name, path)
    for mod, base_dir in modalities.items():
        for cls_dir in sorted(base_dir.iterdir()):
            if not cls_dir.is_dir():
                continue
            paths = sorted(cls_dir.glob("*.*"))[:n_per_class]
            for img_path in paths:
                tensor = _load_multichannel_image(img_path, mod)
                all_images.append((tensor, mod, cls_dir.name, img_path))

    print(f"\nTotal gallery images: {len(all_images)}")
    for mod in modalities:
        count = sum(1 for _, m, _, _ in all_images if m == mod)
        print(f"  {mod}: {count}")

    # 3. Build gallery preview images and extract embeddings
    print("\nBuilding gallery previews ...")
    GALLERY_DIR.mkdir(parents=True, exist_ok=True)
    for i, (_, mod, cls, path) in enumerate(tqdm(all_images, desc="Previews")):
        preview = _make_rgb_preview(path, mod)
        preview.save(GALLERY_DIR / f"{i:05d}_{mod}_{cls}.png")

    # 4. Extract SatCLIP embeddings
    print("\nExtracting SatCLIP embeddings ...")
    encoder = load_satclip()
    embeddings_by_mod = {}
    for mod in ["optical", "sar", "multispectral"]:
        mod_tensors = [img for img, m, _, _ in all_images if m == mod]
        if mod_tensors:
            print(f"  Extracting {mod} ({len(mod_tensors)} images)...")
            embeddings_by_mod[mod] = extract_embeddings_satclip(
                mod_tensors, encoder, mod)
    embeddings = torch.cat(list(embeddings_by_mod.values()), dim=0)
    print(f"Embeddings shape: {embeddings.shape}")

    # 5. Build FAISS index
    print("\nBuilding FAISS index ...")
    import faiss
    embed_dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(embed_dim)
    index.add(embeddings.numpy().astype(np.float32))
    print(f"FAISS index size: {index.ntotal}")

    # 6. Save everything
    PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
    faiss.write_index(index, str(PROCESSED_DIR / "gallery.index"))
    torch.save(embeddings, PROCESSED_DIR / "gallery_embeddings.pt")

    metadata = []
    for i, (_, mod, cls, path) in enumerate(all_images):
        metadata.append({
            "index": i,
            "modality": mod,
            "class": cls,
            "gallery_path": str(GALLERY_DIR / f"{i:05d}_{mod}_{cls}.png"),
            "original_path": str(path),
        })
    with open(PROCESSED_DIR / "gallery_metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    elapsed = time.time() - t0
    print(f"\nDone in {elapsed:.1f}s")
    print(f"Index:     {PROCESSED_DIR / 'gallery.index'}")
    print(f"Embeddings:{PROCESSED_DIR / 'gallery_embeddings.pt'}")
    print(f"Metadata:  {PROCESSED_DIR / 'gallery_metadata.json'}")
    print(f"Gallery:   {GALLERY_DIR}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Build multi-modal satellite image gallery")
    parser.add_argument("--samples", type=int, default=50,
                       help="Images per class per modality")
    args = parser.parse_args()
    build_gallery(n_per_class=args.samples)