"""
Download CrisisLandMark dataset from HuggingFace.

Dataset: DarthReca/crisislandmark
Size: 647K paired Sentinel-1 (SAR) and Sentinel-2 (optical) images
Labels: Land-cover annotations for retrieval evaluation
"""

import os
from pathlib import Path
from datasets import load_dataset


DATA_DIR = Path("data/raw/crisislandmark")


def download_dataset(subset: str = "train", cache_dir: str = None) -> None:
    """
    Download CrisisLandMark dataset.
    
    Args:
        subset: Dataset split to download ("train", "validation", "test")
        cache_dir: Custom cache directory
    """
    print(f"Downloading CrisisLandMark dataset ({subset} split)...")
    
    # ponytail: using HuggingFace datasets library for clean download
    dataset = load_dataset(
        "DarthReca/crisislandmark",
        split=subset,
        cache_dir=cache_dir or str(DATA_DIR / "cache"),
        trust_remote_code=True
    )
    
    print(f"Downloaded {len(dataset)} samples")
    print(f"Columns: {dataset.column_names}")
    
    # Save to disk for faster loading later
    output_dir = DATA_DIR / subset
    output_dir.mkdir(parents=True, exist_ok=True)
    
    dataset.save_to_disk(str(output_dir))
    print(f"Saved to {output_dir}")
    
    return dataset


def verify_dataset(subset: str = "train") -> dict:
    """
    Verify downloaded dataset structure.
    
    Returns:
        Dictionary with dataset statistics
    """
    dataset = load_dataset(
        "DarthReca/crisislandmark",
        split=subset,
        cache_dir=str(DATA_DIR / "cache"),
        trust_remote_code=True
    )
    
    stats = {
        "total_samples": len(dataset),
        "columns": dataset.column_names,
        "features": {col: str(dataset.features[col]) for col in dataset.column_names}
    }
    
    print(f"Dataset stats: {stats}")
    return stats


if __name__ == "__main__":
    # Download train split
    download_dataset("train")
    
    # Verify
    stats = verify_dataset("train")
    print(f"\nVerification complete:")
    print(f"  Total samples: {stats['total_samples']}")
    print(f"  Columns: {stats['columns']}")