scripts/download_met_v2.py

"""
The Met Museum Open Access downloader (v2 — rate-limit safe).

Improvements over v1:
  - 1-second delay between API calls (avoids 403 rate-limiting)
  - Relaxed filters — downloads ALL Indian paintings, not just ones with "mughal" in title
  - Better search queries per style
  - Retries failed downloads once
  - Logs progress more clearly

Usage:
    python scripts/download_met_v2.py --all --count 40
    python scripts/download_met_v2.py --style mughal --count 60
"""
from __future__ import annotations

import argparse
import logging
import re
import time
from pathlib import Path
from typing import List

import requests
from PIL import Image
from io import BytesIO

log = logging.getLogger(__name__)

MET_API = "https://collectionapi.metmuseum.org/public/collection/v1"

HEADERS = {
    "User-Agent": "IndicHeritageStudio/2.0 (educational hackathon project; contact: ai_dev_contests@amd.com)",
    "Accept": "application/json",
}

# Broad search queries per style — get many candidates, filter later
STYLE_QUERIES = {
    "madhubani": ["madhubani", "mithila"],
    "warli": ["warli", "tribal india painting"],
    "pattachitra": ["pattachitra", "orissa painting", "odisha painting"],
    "mughal": [
        "mughal", "mughal painting", "mughal miniature",
        "imperial mughal", "mughal album page", "hamzanama",
        "akbar painting", "jahangir painting", "shah jahan painting",
        "india miniature", "indian miniature painting",
    ],
    "tanjore": ["tanjore", "thanjavur", "south india painting", "tamil painting"],
}


def search_met(query: str, retry: int = 2) -> List[int]:
    """Search The Met API for object IDs matching the query."""
    for attempt in range(retry):
        try:
            resp = requests.get(
                f"{MET_API}/search",
                params={"q": query, "hasImages": "true"},
                headers=HEADERS,
                timeout=30,
            )
            if resp.status_code == 403:
                log.warning(f"  Rate-limited, waiting 5s... (attempt {attempt+1}/{retry})")
                time.sleep(5)
                continue
            resp.raise_for_status()
            return resp.json().get("objectIDs", []) or []
        except Exception as exc:
            log.warning(f"  Search error for '{query}': {exc}")
            time.sleep(2)
    return []


def get_object(object_id: int, retry: int = 2) -> dict:
    """Fetch full metadata for a Met object."""
    for attempt in range(retry):
        try:
            resp = requests.get(
                f"{MET_API}/objects/{object_id}",
                headers=HEADERS,
                timeout=30,
            )
            if resp.status_code == 403:
                time.sleep(3)
                continue
            resp.raise_for_status()
            return resp.json()
        except Exception:
            time.sleep(2)
    return {}


def is_relevant_for_style(obj: dict, style_id: str) -> bool:
    """Check if a Met object is relevant for our style (relaxed filter)."""
    if not obj:
        return False

    title = (obj.get("title", "") or "").lower()
    country = (obj.get("country", "") or "").lower()
    culture = (obj.get("culture", "") or "").lower()
    department = obj.get("department", "")
    classification = (obj.get("classification", "") or "").lower()
    object_name = (obj.get("objectName", "") or "").lower()

    # Must have an image
    if not (obj.get("primaryImage") or obj.get("primaryImageSmall")):
        return False

    # Must be in Asian Art or Paintings department
    if department and department not in ("Asian Art", "Paintings"):
        return False

    # Combine all text for style matching
    full_text = f"{title} {country} {culture} {classification} {object_name}"

    # Style-specific relevance check (relaxed)
    if style_id == "mughal":
        # Mughal is broad — accept any Mughal/Indian miniature painting
        return any(kw in full_text for kw in [
            "mughal", "india", "indian", "akbar", "jahangir", "shah jahan",
            "hamzanama", "rajput", "rajasthani", "pahari",
        ])
    elif style_id == "tanjore":
        return any(kw in full_text for kw in [
            "tanjore", "thanjavur", "tamil", "south india",
            "vijayanagara", "madras", "madurai",
        ])
    elif style_id == "madhubani":
        return any(kw in full_text for kw in [
            "madhubani", "mithila", "bihar",
        ])
    elif style_id == "warli":
        return any(kw in full_text for kw in [
            "warli", "tribal", "maharashtra", "adivasi",
        ])
    elif style_id == "pattachitra":
        return any(kw in full_text for kw in [
            "pattachitra", "orissa", "odisha", "bengal scroll",
        ])
    return False


def download_image(url: str, out_path: Path, target_size: int = 1024) -> bool:
    """Download and resize an image."""
    try:
        resp = requests.get(url, headers=HEADERS, timeout=60, allow_redirects=True)
        if resp.status_code != 200 or len(resp.content) < 1000:
            return False
        img = Image.open(BytesIO(resp.content)).convert("RGB")
        w, h = img.size
        if w >= h:
            new_w, new_h = target_size, int(h * target_size / w)
        else:
            new_w, new_h = int(w * target_size / h), target_size
        img = img.resize((new_w, new_h), Image.LANCZOS)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        img.save(out_path, "JPEG", quality=95)
        return True
    except Exception as exc:
        log.debug(f"Download failed {url}: {exc}")
        return False


def download_style(style_id: str, total_count: int = 40) -> int:
    """Download `total_count` images for one style."""
    queries = STYLE_QUERIES.get(style_id, [])
    out_dir = Path(f"assets/datasets/raw/{style_id}")
    out_dir.mkdir(parents=True, exist_ok=True)

    downloaded = 0
    seen_urls = set()
    seen_ids = set()

    for query in queries:
        if downloaded >= total_count:
            break
        log.info(f"[{style_id}] Searching Met Museum: '{query}'")
        object_ids = search_met(query)
        log.info(f"[{style_id}] Found {len(object_ids)} candidate objects")
        time.sleep(1)  # be polite

        for oid in object_ids:
            if downloaded >= total_count:
                break
            if oid in seen_ids:
                continue
            seen_ids.add(oid)

            obj = get_object(oid)
            time.sleep(0.5)  # be polite

            if not is_relevant_for_style(obj, style_id):
                continue

            img_url = obj.get("primaryImage") or obj.get("primaryImageSmall")
            if not img_url or img_url in seen_urls:
                continue
            seen_urls.add(img_url)

            title = (obj.get("title", "") or "untitled")[:80]
            safe_title = re.sub(r"[^a-zA-Z0-9_-]", "_", title)[:50]
            out_path = out_dir / f"{style_id}_{downloaded:03d}_{safe_title}.jpg"

            if download_image(img_url, out_path):
                downloaded += 1
                culture = obj.get("culture", "")[:30]
                log.info(f"  ✓ [{downloaded}/{total_count}] {out_path.name} (culture={culture})")
                time.sleep(1)  # be polite — avoids 403 storms

    log.info(f"[{style_id}] Downloaded {downloaded}/{total_count} from Met Museum")
    return downloaded


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")

    p = argparse.ArgumentParser(description="Download heritage art from The Met Museum Open Access (v2)")
    p.add_argument("--style", choices=list(STYLE_QUERIES.keys()))
    p.add_argument("--all", action="store_true")
    p.add_argument("--count", type=int, default=40)
    args = p.parse_args()

    styles = [args.style] if args.style else list(STYLE_QUERIES.keys())

    log.info("=== The Met Museum Open Access Downloader v2 (rate-limit safe) ===")
    log.info(f"Styles: {styles}")
    log.info(f"Per style: {args.count}")
    log.info(f"Estimated time: ~{len(styles) * args.count * 1.5 / 60:.1f} min (1.5s per image)")

    total = 0
    for style in styles:
        log.info(f"\n--- {style.upper()} ---")
        n = download_style(style, total_count=args.count)
        total += n

    log.info(f"\n=== Done. Total: {total} images ===")


if __name__ == "__main__":
    main()