"""Download REAL DrivAer CFD surface point clouds for RegDGCNN training.

The DrivAerNet paper trains RegDGCNN on surface point clouds sampled from
the industrial car CFD runs.  The full mesh+CFD release is huge (Globus /
Dataverse only), but a compact, self-labelled subset of real DrivAer
point clouds is published on Hugging Face:

    https://huggingface.co/datasets/Jrhoss/Drivaerml_point_clouds

Each CFD run ships a surface point cloud (point_cloud_<N>.npz, key
'points') and its CFD-computed coefficients (force_mom_<N>.csv with
header Cd,Cl,Clf,Clr,Cs), partitioned into train/val/test with a
top-level splits.json.  That is exactly the (point-cloud -> Cd)
supervision RegDGCNN needs — real geometry, real drag.

This script mirrors that repo into:

    DrivAerNet/PointClouds/
      splits.json
      train/run_<N>/point_cloud_<N>.npz + force_mom_<N>.csv
      val/...   test/...

so drivaernet_pointclouds.DrivAerNetPointCloudDataset finds it
automatically (it auto-detects this DrivAerML layout).

Uses `huggingface_hub`.  --limit N grabs only the first N runs for a
quick CPU smoke test.

Usage:
    pip install huggingface_hub
    python fetch_drivaernet_pointclouds.py                 # full subset (~175 runs)
    python fetch_drivaernet_pointclouds.py --limit 30      # smoke-test subset
    python fetch_drivaernet_pointclouds.py --repo <other/repo>
"""
from __future__ import annotations

import argparse
import os
import re
import shutil
import sys
from pathlib import Path

os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")

DATA_DIR = Path(__file__).parent / "DrivAerNet"
PC_DIR = DATA_DIR / "PointClouds"

DEFAULT_REPO = "Jrhoss/Drivaerml_point_clouds"
_RUN_RE = re.compile(r"run_(\d+)", re.IGNORECASE)


def _have_hf() -> bool:
    try:
        import huggingface_hub  # noqa: F401
        return True
    except Exception:
        return False


def _list_repo_files(repo: str):
    from huggingface_hub import HfApi
    return HfApi().list_repo_files(repo_id=repo, repo_type="dataset")


def _read_split_run_order(repo: str, files) -> list[str]:
    """Return run numbers (as strings) in splits.json order (train+val+test)."""
    if "splits.json" not in [f.rsplit("/", 1)[-1] for f in files]:
        return []
    from huggingface_hub import hf_hub_download
    import json
    sp_name = next(f for f in files if f.rsplit("/", 1)[-1] == "splits.json")
    try:
        local = hf_hub_download(repo_id=repo, filename=sp_name, repo_type="dataset")
        sd = json.loads(Path(local).read_text()).get("splits", {})
    except Exception:
        return []
    order: list[str] = []
    for key in ("train", "val", "validation", "test"):
        for rid in sd.get(key, []):
            m = _RUN_RE.search(str(rid))
            if m:
                order.append(m.group(1))
    return order


def _grab(repo: str, fname: str, dest_root: Path) -> bool:
    """Download one repo file, preserving its relative path under dest_root."""
    from huggingface_hub import hf_hub_download
    try:
        local = hf_hub_download(repo_id=repo, filename=fname, repo_type="dataset")
    except Exception as e:
        print(f"[fetch-pc]   skip {fname}: {e}")
        return False
    target = dest_root / fname
    target.parent.mkdir(parents=True, exist_ok=True)
    if not target.exists():
        shutil.copy(local, target)
    return True


def fetch(repo: str = DEFAULT_REPO, limit: int | None = None):
    PC_DIR.mkdir(parents=True, exist_ok=True)

    if not _have_hf():
        print("[fetch-pc] ERROR: huggingface_hub is not installed.\n"
              "  pip install huggingface_hub", file=sys.stderr)
        return 1

    print(f"[fetch-pc] listing files in dataset repo '{repo}' ...")
    try:
        files = _list_repo_files(repo)
    except Exception as e:
        print(f"[fetch-pc] ERROR listing repo '{repo}': {e}", file=sys.stderr)
        return 1

    # Group point-cloud + force_mom files by run.
    runs: dict[str, dict] = {}
    aux: list[str] = []
    for f in files:
        base = f.rsplit("/", 1)[-1]
        if base == "splits.json":
            aux.append(f)
            continue
        m = _RUN_RE.search(f)
        if not m:
            continue
        run = m.group(1)
        d = runs.setdefault(run, {})
        if base.startswith("point_cloud"):
            d["pc"] = f
        elif base.startswith("force_mom"):
            d["fm"] = f

    # Keep only complete runs (point cloud + label).
    complete = {r: d for r, d in runs.items() if "pc" in d and "fm" in d}

    # Prefer runs that the authors' splits.json actually partitions, so the
    # downloaded subset honours the real train/val/test split.  We read
    # splits.json straight from the repo file list (grab it first below).
    split_runs = _read_split_run_order(repo, files)
    in_splits = [r for r in split_runs if r in complete]
    rest = [r for r in sorted(complete, key=lambda r: int(r)) if r not in set(in_splits)]
    order = in_splits + rest
    if limit:
        order = order[:limit]

    print(f"[fetch-pc] {len(complete)} complete runs available; "
          f"downloading {len(order)}{' (limited)' if limit else ''} ...")

    # Grab splits.json first.
    for a in aux:
        _grab(repo, a, PC_DIR)

    n_ok = 0
    for i, run in enumerate(order, 1):
        d = complete[run]
        ok_pc = _grab(repo, d["pc"], PC_DIR)
        ok_fm = _grab(repo, d["fm"], PC_DIR)
        if ok_pc and ok_fm:
            n_ok += 1
        if i % 20 == 0 or i == len(order):
            print(f"[fetch-pc]   {i}/{len(order)} runs")

    print(f"[fetch-pc] done. {n_ok} runs downloaded into {PC_DIR}")

    # Cross-check via the dataset loader.
    try:
        import drivaernet_pointclouds as PC
        s = PC.summary(str(DATA_DIR))
        print(f"[fetch-pc] loader sees layout='{s['layout']}', "
              f"{s['n_matched']} trainable point clouds")
        if s.get("n_matched"):
            print(f"[fetch-pc]   Cd range [{s.get('cd_min'):.3f},"
                  f"{s.get('cd_max'):.3f}] mean {s.get('cd_mean'):.3f}")
        if "split_counts" in s:
            print(f"[fetch-pc]   splits.json counts: {s['split_counts']}")
    except Exception as e:
        print(f"[fetch-pc] note: loader cross-check skipped ({e})")

    return 0 if n_ok > 0 else 1


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default=DEFAULT_REPO,
                    help="Hugging Face dataset repo id")
    ap.add_argument("--limit", type=int, default=None,
                    help="download only the first N runs (smoke test)")
    args = ap.parse_args()
    try:
        rc = fetch(args.repo, args.limit)
    except KeyboardInterrupt:
        print("\n[fetch-pc] interrupted", file=sys.stderr)
        rc = 130
    sys.exit(rc)