#!/usr/bin/env python3
"""
Download the RaSP proteome-wide ΔΔG dataset (Blaabjerg et al., eLife 2023)
and upload it as a partitioned HF dataset table (ddg/).

Source: https://sid.erda.dk/sharelink/fFPJWflLeE
  File: rasp_preds_alphafold_UP000005640_9606_HUMAN_v2_vaex_dataframe.zip
  ~230M variants, 23 391 human proteins, AlphaFold2 structures, kcal/mol

Strategy
--------
  1. Download the Vaex bulk file from ERDA (~8.6 GB zip)
  2. Convert Vaex (HDF5) -> flat Parquet using vaex
  3. DuckDB COPY partitions by protein_id (UniProt extracted from pdbid)
  4. upload_large_folder -> ddg/ on HF (same pattern as esm1b/)
  5. Backend JOIN: COALESCE(variants.pred_ddg, ddg.pred_ddg)
     -> ESM-IF1 values kept for 13 enriched proteins, RaSP for all others

Column mapping
--------------
  RaSP column  ->  MutVar column
  pdbid        ->  protein_id  (AF-{UNIPROT}-F1-... -> extract UniProt)
  variant      ->  mutation_code  (e.g. "V600E" — exact match)
  score_ml     ->  pred_ddg  (kcal/mol, positive = destabilising)

Threshold: pred_ddg > 1.5 kcal/mol -> destabilising (same as ESM-IF1 proxy)

Estimated runtime
-----------------
  Download       :  ~30-60 min  (8.6 GB from ERDA)
  Vaex convert   :  ~10-20 min  (HDF5 -> flat parquet)
  DuckDB COPY    :  ~10-20 min  (partition 230M rows)
  upload_large   :  ~1-3h       (chunked, resumable)
  Total          :  ~3-5h

Resume behaviour
----------------
  Output dir defaults to ~/mutvar_rasp_upload/
  Partition step is skipped if ddg/ subfolder already has parquet files.
  Upload is resumable (upload_large_folder uses .cache/ for state).

Usage
-----
  pip install h5py duckdb huggingface_hub pyarrow
  huggingface-cli login
  python scripts/join_rasp_ddg.py [--dry-run] [--output-dir PATH]
"""

import os, re, argparse, time, zipfile
from pathlib import Path
import duckdb

HF_DATASET      = os.environ.get("HF_DATASET", "edohollou/mutvar-variants")
ERDA_BASE       = "https://sid.erda.dk/share_redirect/fFPJWflLeE"
VAEX_FILENAME   = "rasp_preds_alphafold_UP000005640_9606_HUMAN_v2_vaex_dataframe.zip"
DDG_THRESHOLD   = 1.5   # kcal/mol — destabilising threshold (Blaabjerg et al.)

DEFAULT_OUT_DIR = Path.home() / "mutvar_rasp_upload"


# ── Step 1: Download ───────────────────────────────────────────────────────────

def download_vaex(out_dir: Path) -> Path:
    """Download the Vaex zip from ERDA if not already present."""
    zip_path = out_dir / VAEX_FILENAME
    if zip_path.exists():
        print(f"[1/4] Vaex zip already downloaded: {zip_path}")
        return zip_path

    import urllib.request
    url = f"{ERDA_BASE}/{VAEX_FILENAME}"
    out_dir.mkdir(parents=True, exist_ok=True)
    print(f"[1/4] Downloading RaSP Vaex file (~8.6 GB)...")
    print(f"      {url}")
    print(f"      -> {zip_path}")

    t0 = time.time()

    def _progress(count, block, total):
        pct  = min(100, count * block * 100 // total)
        done = count * block / 1e9
        tot  = total / 1e9
        print(f"\r      {done:.2f} / {tot:.2f} GB  ({pct}%)", end="", flush=True)

    urllib.request.urlretrieve(url, zip_path, reporthook=_progress)
    print(f"\n      Done in {(time.time()-t0)/60:.1f} min")
    return zip_path


# ── Step 2: Vaex -> flat Parquet ───────────────────────────────────────────────

def _hdf5_is_valid(path: Path) -> bool:
    """Quick integrity check — opens the HDF5 and reads the row count."""
    try:
        import h5py
        with h5py.File(str(path), "r") as f:
            for col_path in [
                "table/columns/pdbid/data",
                "columns/pdbid/data",
                "pdbid",
            ]:
                if col_path in f:
                    _ = len(f[col_path])
                    return True
        return False
    except Exception:
        return False


def vaex_to_parquet(zip_path: Path, flat_parquet: Path) -> Path:
    """
    Unzip the Vaex HDF5 and convert to a flat Parquet using h5py.
    Reads in chunks of 5M rows to avoid OOM.

    Space management:
      - Skips extraction if a valid HDF5 already exists
      - Deletes the zip immediately after successful extraction
      - Deletes the HDF5 immediately after the parquet is written
    """
    if flat_parquet.exists():
        size_gb = flat_parquet.stat().st_size / 1e9
        print(f"[2/4] Flat parquet already exists ({size_gb:.2f} GB): {flat_parquet}")
        return flat_parquet

    import h5py
    import pyarrow as pa
    import pyarrow.parquet as pq

    # ── Find or extract HDF5 ──────────────────────────────────────────────────
    # Determine expected HDF5 path from zip contents without extracting
    with zipfile.ZipFile(zip_path, "r") as zf:
        names = zf.namelist()
        hdf5_files = [n for n in names if ".hdf5" in n or n.endswith(".h5")]
        if not hdf5_files:
            raise RuntimeError(f"No HDF5 in zip. Contents: {names[:20]}")
        hdf5_name = hdf5_files[0]

    hdf5_path = flat_parquet.parent / Path(hdf5_name).name

    if hdf5_path.exists() and _hdf5_is_valid(hdf5_path):
        print(f"[2/4] HDF5 already extracted and valid — skipping unzip.")
        print(f"      {hdf5_path}")
    else:
        # Remove any partial/corrupt file before extracting
        if hdf5_path.exists():
            print(f"      Removing corrupt/partial HDF5 ({hdf5_path.stat().st_size/1e9:.1f} GB)...")
            hdf5_path.unlink()

        print("[2/4] Extracting HDF5 from zip...")
        t0 = time.time()
        with zipfile.ZipFile(zip_path, "r") as zf:
            zf.extract(hdf5_name, flat_parquet.parent)
        print(f"      Extracted in {(time.time()-t0)/60:.1f} min  "
              f"({hdf5_path.stat().st_size/1e9:.1f} GB)")

        # Free the zip immediately — no longer needed
        try:
            zip_path.unlink()
            print(f"      Deleted zip to reclaim {zip_path.stat().st_size/1e9:.1f} GB")
        except Exception:
            pass

    # ── Convert HDF5 -> flat Parquet (chunked) ────────────────────────────────
    print("[2/4] Converting HDF5 -> flat Parquet (chunked, h5py)...")
    t0 = time.time()

    _uid_re = re.compile(r"AF-([A-Z0-9]+)-F\d")
    CHUNK = 5_000_000

    with h5py.File(str(hdf5_path), "r") as f:
        print(f"      HDF5 top-level keys: {list(f.keys())}")

        def _open_col(name: str):
            for path in [
                f"table/columns/{name}/data",
                f"columns/{name}/data",
                name,
            ]:
                if path in f:
                    return f[path]
            raise KeyError(f"Column '{name}' not found. Available: {list(f.keys())}")

        col_pdbid   = _open_col("pdbid")
        col_variant = _open_col("variant")
        col_score   = _open_col("score_ml")
        n_total     = len(col_pdbid)
        print(f"      {n_total:,} rows total")

        writer = None
        for start in range(0, n_total, CHUNK):
            end = min(start + CHUNK, n_total)

            pdbids   = col_pdbid[start:end].astype(str)
            variants = col_variant[start:end].astype(str)
            scores   = col_score[start:end].astype("float32")

            protein_ids = [
                (m.group(1) if (m := _uid_re.search(p)) else p)
                for p in pdbids
            ]

            chunk = pa.table({
                "protein_id":    pa.array(protein_ids,         type=pa.string()),
                "mutation_code": pa.array(variants.tolist(),   type=pa.string()),
                "pred_ddg":      pa.array(scores.tolist(),     type=pa.float32()),
                "pred_ddg_label": pa.array(
                    (scores > DDG_THRESHOLD).tolist(), type=pa.bool_()
                ),
            })

            if writer is None:
                writer = pq.ParquetWriter(str(flat_parquet), chunk.schema)
            writer.write_table(chunk)

            pct = end * 100 // n_total
            print(f"\r      {end:,} / {n_total:,} rows ({pct}%)", end="", flush=True)

    if writer:
        writer.close()
    print()

    # Delete HDF5 immediately — frees ~105 GB before DuckDB partitioning
    try:
        size = hdf5_path.stat().st_size / 1e9
        hdf5_path.unlink()
        print(f"      Deleted HDF5 ({size:.1f} GB freed)")
    except Exception:
        pass

    size_gb = flat_parquet.stat().st_size / 1e9
    print(f"      Flat parquet: {size_gb:.2f} GB in {(time.time()-t0)/60:.1f} min")
    return flat_parquet


# ── Step 3: DuckDB partition ───────────────────────────────────────────────────

def partition_locally(flat_parquet: Path, ddg_dir: Path, dry_run: bool) -> int:
    """
    DuckDB reads the flat parquet and writes per-protein partitions.
    Skipped if ddg_dir already contains parquet files.
    """
    existing = list(ddg_dir.rglob("*.parquet"))
    if existing and not dry_run:
        print(f"[3/4] Partition dir already exists ({len(existing):,} files) — skipping.")
        return len(existing)

    con = duckdb.connect()
    ddg_dir.mkdir(parents=True, exist_ok=True)

    if dry_run:
        print("[3/4] [dry-run] Partitioning 5-protein sample...")
        sample = con.execute(f"""
            SELECT DISTINCT protein_id FROM read_parquet('{flat_parquet}')
            ORDER BY protein_id LIMIT 5
        """).fetchall()
        sample_ids = [r[0] for r in sample]
        pid_list = ", ".join(f"'{p}'" for p in sample_ids)
        print(f"      Sample: {sample_ids}")

        con.execute(f"""
            COPY (
                SELECT protein_id, mutation_code,
                       ROUND(pred_ddg, 4) AS pred_ddg,
                       pred_ddg_label
                FROM read_parquet('{flat_parquet}')
                WHERE protein_id IN ({pid_list})
            )
            TO '{ddg_dir}'
            (FORMAT PARQUET, PARTITION_BY (protein_id), OVERWRITE_OR_IGNORE true)
        """)
    else:
        print("[3/4] Partitioning with DuckDB (~10-20 min)...")
        t0 = time.time()
        con.execute(f"""
            COPY (
                SELECT protein_id, mutation_code,
                       ROUND(pred_ddg, 4) AS pred_ddg,
                       pred_ddg_label
                FROM read_parquet('{flat_parquet}')
            )
            TO '{ddg_dir}'
            (FORMAT PARQUET, PARTITION_BY (protein_id), OVERWRITE_OR_IGNORE true)
        """)
        print(f"      Done in {(time.time()-t0)/60:.1f} min")

    n = len(list(ddg_dir.rglob("*.parquet")))
    print(f"      {n:,} parquet files in {ddg_dir}")
    con.close()
    return n


# ── Step 4: Upload ─────────────────────────────────────────────────────────────

def upload(upload_root: Path, dry_run: bool):
    if dry_run:
        print("[4/4] [dry-run] Skipping upload. Sample files:")
        import pandas as pd
        ddg_dir = upload_root / "ddg"
        for f in sorted(ddg_dir.rglob("*.parquet"))[:5]:
            df = pd.read_parquet(f)
            print(f"  {f.parent.name}/{f.name}: {len(df)} rows, "
                  f"pred_ddg sample={df['pred_ddg'].head(3).tolist()}")
        return

    from huggingface_hub import HfApi
    api = HfApi()
    print("[4/4] Uploading ddg/ to HuggingFace (upload_large_folder, resumable)...")
    print("      Safe to Ctrl-C and re-run — upload resumes from last checkpoint.")
    t0 = time.time()

    api.upload_large_folder(
        folder_path=str(upload_root),
        repo_id=HF_DATASET,
        repo_type="dataset",
    )

    print(f"      Done in {(time.time()-t0)/60:.1f} min")
    print()
    print("[done] Next steps:")
    print("  1. git push -> Space picks up backend/main.py JOIN for ddg/")
    print(f"  2. Delete local files: rm -rf {upload_root}")


# ── Main ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true",
                        help="Process 5 proteins only, skip upload")
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUT_DIR,
                        help=f"Persistent output dir (default: {DEFAULT_OUT_DIR})")
    args = parser.parse_args()

    upload_root  = args.output_dir
    ddg_dir      = upload_root / "ddg"
    flat_parquet = upload_root / "rasp_flat.parquet"

    print(f"      Upload root : {upload_root}")
    print(f"      Parquet dir : {ddg_dir}\n")

    zip_path = download_vaex(upload_root)
    vaex_to_parquet(zip_path, flat_parquet)
    n = partition_locally(flat_parquet, ddg_dir, dry_run=args.dry_run)
    if n > 0:
        upload(upload_root, dry_run=args.dry_run)


if __name__ == "__main__":
    main()