#!/usr/bin/env python3
"""
Precompute ML mechanistic label predictions for all 216M variants.

Reads variants from HF dataset (or local cache), runs the lite XGBoost model,
writes results as hive-partitioned parquet:
    ~/mutvar_ml_upload/ml_preds/protein_id=X/data_0.parquet

Columns:  mutation_code | ml_mechLabel | ml_confidence

Then uploads to edohollou/mutvar-variants dataset.

Usage:
    python scripts/precompute_ml_predictions.py [--dry-run] [--output-dir PATH]
"""

import argparse
import io
import json
import os
import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import requests
from huggingface_hub import HfApi, HfFileSystem, list_repo_files, snapshot_download

HF_DATASET   = os.environ.get("HF_DATASET", "edohollou/mutvar-variants")
HF_SPACE_ID  = "edohollou/mutvar"
DEFAULT_OUT  = Path.home() / "mutvar_ml_upload"

# ── AA properties (same as training) ──────────────────────────────────────────
AA_PROPS = {
    'A': {'hydrophobicity': 1.8,  'charge': 0,   'size': 89,  'polarity': 0},
    'R': {'hydrophobicity': -4.5, 'charge': 1,   'size': 174, 'polarity': 1},
    'N': {'hydrophobicity': -3.5, 'charge': 0,   'size': 132, 'polarity': 1},
    'D': {'hydrophobicity': -3.5, 'charge': -1,  'size': 133, 'polarity': 1},
    'C': {'hydrophobicity': 2.5,  'charge': 0,   'size': 121, 'polarity': 0},
    'Q': {'hydrophobicity': -3.5, 'charge': 0,   'size': 146, 'polarity': 1},
    'E': {'hydrophobicity': -3.5, 'charge': -1,  'size': 147, 'polarity': 1},
    'G': {'hydrophobicity': -0.4, 'charge': 0,   'size': 75,  'polarity': 0},
    'H': {'hydrophobicity': -3.2, 'charge': 0.5, 'size': 155, 'polarity': 1},
    'I': {'hydrophobicity': 4.5,  'charge': 0,   'size': 131, 'polarity': 0},
    'L': {'hydrophobicity': 3.8,  'charge': 0,   'size': 131, 'polarity': 0},
    'K': {'hydrophobicity': -3.9, 'charge': 1,   'size': 146, 'polarity': 1},
    'M': {'hydrophobicity': 1.9,  'charge': 0,   'size': 149, 'polarity': 0},
    'F': {'hydrophobicity': 2.8,  'charge': 0,   'size': 165, 'polarity': 0},
    'P': {'hydrophobicity': -1.6, 'charge': 0,   'size': 115, 'polarity': 0},
    'S': {'hydrophobicity': -0.8, 'charge': 0,   'size': 105, 'polarity': 1},
    'T': {'hydrophobicity': -0.7, 'charge': 0,   'size': 119, 'polarity': 1},
    'W': {'hydrophobicity': -0.9, 'charge': 0,   'size': 204, 'polarity': 1},
    'Y': {'hydrophobicity': -1.3, 'charge': 0,   'size': 181, 'polarity': 1},
    'V': {'hydrophobicity': 4.2,  'charge': 0,   'size': 117, 'polarity': 0},
}

BLOSUM62 = {
    ('A','R'):-1,('A','N'):-2,('A','D'):-2,('A','C'):0,('A','Q'):-1,('A','E'):-1,('A','G'):0,
    ('A','H'):-2,('A','I'):-1,('A','L'):-1,('A','K'):-1,('A','M'):-1,('A','F'):-2,('A','P'):-1,
    ('A','S'):1,('A','T'):0,('A','W'):-3,('A','Y'):-2,('A','V'):0,('R','N'):-1,('R','D'):-2,
    ('R','C'):-3,('R','Q'):1,('R','E'):0,('R','G'):-2,('R','H'):0,('R','I'):-3,('R','L'):-2,
    ('R','K'):2,('R','M'):-1,('R','F'):-3,('R','P'):-2,('R','S'):-1,('R','T'):-1,('R','W'):-3,
    ('R','Y'):-2,('R','V'):-3,('N','D'):1,('N','C'):-3,('N','Q'):0,('N','E'):0,('N','G'):0,
    ('N','H'):1,('N','I'):-3,('N','L'):-3,('N','K'):0,('N','M'):-2,('N','F'):-3,('N','P'):-2,
    ('N','S'):1,('N','T'):0,('N','W'):-4,('N','Y'):-2,('N','V'):-3,('D','C'):-3,('D','Q'):0,
    ('D','E'):2,('D','G'):-1,('D','H'):-1,('D','I'):-3,('D','L'):-4,('D','K'):-1,('D','M'):-3,
    ('D','F'):-3,('D','P'):-1,('D','S'):0,('D','T'):-1,('D','W'):-4,('D','Y'):-3,('D','V'):-3,
    ('C','Q'):-3,('C','E'):-4,('C','G'):-3,('C','H'):-3,('C','I'):-1,('C','L'):-1,('C','K'):-3,
    ('C','M'):-1,('C','F'):-2,('C','P'):-3,('C','S'):-1,('C','T'):-1,('C','W'):-2,('C','Y'):-2,
    ('C','V'):-1,('Q','E'):2,('Q','G'):-2,('Q','H'):0,('Q','I'):-3,('Q','L'):-2,('Q','K'):1,
    ('Q','M'):0,('Q','F'):-3,('Q','P'):-1,('Q','S'):0,('Q','T'):-1,('Q','W'):-2,('Q','Y'):-1,
    ('Q','V'):-2,('E','G'):-2,('E','H'):0,('E','I'):-3,('E','L'):-3,('E','K'):1,('E','M'):-2,
    ('E','F'):-3,('E','P'):-1,('E','S'):0,('E','T'):-1,('E','W'):-3,('E','Y'):-2,('E','V'):-2,
    ('G','H'):-2,('G','I'):-4,('G','L'):-4,('G','K'):-2,('G','M'):-3,('G','F'):-3,('G','P'):-2,
    ('G','S'):0,('G','T'):-2,('G','W'):-2,('G','Y'):-3,('G','V'):-3,('H','I'):-3,('H','L'):-3,
    ('H','K'):-1,('H','M'):-2,('H','F'):-1,('H','P'):-2,('H','S'):-1,('H','T'):-2,('H','W'):-2,
    ('H','Y'):2,('H','V'):-3,('I','L'):2,('I','K'):-1,('I','M'):1,('I','F'):0,('I','P'):-3,
    ('I','S'):-2,('I','T'):-1,('I','W'):-3,('I','Y'):-1,('I','V'):3,('L','K'):-2,('L','M'):2,
    ('L','F'):0,('L','P'):-3,('L','S'):-2,('L','T'):-1,('L','W'):-2,('L','Y'):-1,('L','V'):1,
    ('K','M'):-1,('K','F'):-3,('K','P'):-1,('K','S'):0,('K','T'):-1,('K','W'):-3,('K','Y'):-2,
    ('K','V'):-2,('M','F'):0,('M','P'):-2,('M','S'):-1,('M','T'):-1,('M','W'):-1,('M','Y'):-1,
    ('M','V'):1,('F','P'):-4,('F','S'):-2,('F','T'):-2,('F','W'):1,('F','Y'):3,('F','V'):-1,
    ('P','S'):-1,('P','T'):-1,('P','W'):-4,('P','Y'):-3,('P','V'):-2,('S','T'):1,('S','W'):-3,
    ('S','Y'):-2,('S','V'):-2,('T','W'):-2,('T','Y'):-2,('T','V'):0,('W','Y'):2,('W','V'):-3,
    ('Y','V'):-1,
}
BLOSUM62.update({(b, a): v for (a, b), v in list(BLOSUM62.items())})

_MUT_RE = re.compile(r'^([A-Z])(\d+)([A-Z])$')

SCHEMA = pa.schema([
    ("mutation_code",  pa.string()),
    ("ml_mechLabel",   pa.string()),
    ("ml_confidence",  pa.float32()),
])


# ── Feature engineering ────────────────────────────────────────────────────────

def build_features(df: pd.DataFrame) -> np.ndarray:
    """Vectorised feature extraction — matches the lite model training exactly."""
    mc = df["mutation_code"].str.extract(r'^([A-Z])(\d+)([A-Z])$')
    mc.columns = ["aa_from", "pos_str", "aa_to"]
    mc["position"] = pd.to_numeric(mc["pos_str"], errors="coerce").fillna(0).astype(int)

    prot_len = max(mc["position"].max(), 1)

    def prop(aa_series, p):
        return aa_series.map(lambda a: AA_PROPS.get(a, {}).get(p, 0)).astype(float)

    h_from = prop(mc["aa_from"], "hydrophobicity")
    h_to   = prop(mc["aa_to"],   "hydrophobicity")
    c_from = prop(mc["aa_from"], "charge")
    c_to   = prop(mc["aa_to"],   "charge")
    s_from = prop(mc["aa_from"], "size")
    s_to   = prop(mc["aa_to"],   "size")
    p_from = prop(mc["aa_from"], "polarity")
    p_to   = prop(mc["aa_to"],   "polarity")

    bl62 = mc.apply(lambda r: float(BLOSUM62.get((r["aa_from"], r["aa_to"]), 0)), axis=1)

    am  = pd.to_numeric(df.get("am_pathogenicity",  0), errors="coerce").fillna(0)
    esm = pd.to_numeric(df.get("esm1b_llr",         0), errors="coerce").fillna(0)
    ddg = pd.to_numeric(df.get("pred_ddg",          0), errors="coerce").fillna(0)

    cons = (
        (am > 0.564).astype(int) +
        (esm < -4.0).astype(int) +
        (ddg > 1.5).astype(int)
    )

    X = np.column_stack([
        am,
        esm,
        ddg,
        mc["position"],
        mc["position"] / prot_len,
        h_to - h_from,
        c_to - c_from,
        s_to - s_from,
        (p_from != p_to).astype(int),
        (np.sign(c_from) != np.sign(c_to)).astype(int),
        bl62,
        cons,
        am * np.clip(ddg,  -10, 10),
        am * np.clip(esm,  -30,  5),
    ])
    return X.astype(np.float32)


# ── Model loading ──────────────────────────────────────────────────────────────

def load_model():
    """Load lite classifier from local models/ or HF Space."""
    local = Path(__file__).parent.parent / "models"
    for p in [local, Path("models")]:
        lite = p / "mechanistic_classifier_lite.pkl"
        enc  = p / "label_encoder.json"
        if lite.exists() and enc.exists():
            clf = joblib.load(lite)
            classes = json.load(open(enc))["classes"]
            print(f"Loaded model from {p}")
            return clf, classes

    print("Downloading model from HF Space...")
    from huggingface_hub import hf_hub_download
    clf = joblib.load(hf_hub_download(HF_SPACE_ID, "models/mechanistic_classifier_lite.pkl", repo_type="space"))
    classes = json.load(open(hf_hub_download(HF_SPACE_ID, "models/label_encoder.json", repo_type="space")))["classes"]
    return clf, classes


# ── Main pipeline ──────────────────────────────────────────────────────────────

def process_protein(pid: str, parquet_path: Path, clf, classes, out_dir: Path) -> int:
    """Run prediction for one protein, write parquet. Returns variant count."""
    df = pd.read_parquet(parquet_path, columns=[
        "mutation_code", "am_pathogenicity", "esm1b_llr", "pred_ddg"
    ])
    if df.empty:
        return 0

    X = build_features(df)
    labels  = [classes[i] for i in clf.predict(X)]
    probas  = clf.predict_proba(X).max(axis=1).astype(np.float32)

    table = pa.table({
        "mutation_code": pa.array(df["mutation_code"].tolist(), type=pa.string()),
        "ml_mechLabel":  pa.array(labels,  type=pa.string()),
        "ml_confidence": pa.array(probas,  type=pa.float32()),
    }, schema=SCHEMA)

    out = out_dir / f"protein_id={pid}" / "data_0.parquet"
    out.parent.mkdir(parents=True, exist_ok=True)
    pq.write_table(table, str(out), compression="snappy")
    return len(df)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true", help="Process 5 proteins only")
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--variants-dir", type=Path, default=None,
                        help="Local variants/ dir (default: download from HF)")
    args = parser.parse_args()

    ml_dir = args.output_dir / "ml_preds"
    ml_dir.mkdir(parents=True, exist_ok=True)

    # ── Load model ─────────────────────────────────────────────────────────────
    clf, classes = load_model()
    print(f"Classes: {classes}")

    # ── Find variants parquets ─────────────────────────────────────────────────
    variants_dir = args.variants_dir
    if variants_dir is None:
        local_dl = args.output_dir / "variants_cache"
        local_dl.mkdir(parents=True, exist_ok=True)
        # Variants are at repo root as protein_id=X/data_0.parquet
        # Use HfFileSystem to enumerate and download (snapshot_download allow_patterns
        # doesn't handle the hive-partitioned root-level layout reliably).
        # Only count proteins where the parquet file actually exists (not just the dir)
        existing_pids = {
            p.parent.name.replace("protein_id=", "")
            for p in local_dl.glob("protein_id=*/data_0.parquet")
        }
        if existing_pids:
            print(f"Using cached variants at {local_dl} ({len(existing_pids):,} proteins already present)")

        # List all variants proteins via list_repo_files.
        # The batch-uploaded 21K are double-nested: variants/variants/protein_id=X/
        # A handful of singles are at variants/protein_id=X/
        print(f"Listing variants proteins in {HF_DATASET} ...")
        all_files = list(list_repo_files(HF_DATASET, repo_type="dataset"))
        all_pids_double = sorted({
            f.split("variants/variants/protein_id=")[1].split("/")[0]
            for f in all_files
            if f.startswith("variants/variants/protein_id=") and f.endswith(".parquet")
        })
        all_pids_single = sorted({
            f.split("variants/protein_id=")[1].split("/")[0]
            for f in all_files
            if f.startswith("variants/protein_id=") and f.endswith(".parquet")
        })
        # Resolve URL per protein: prefer double-nested (21K batch), fall back to single
        pid_url_map = {}
        for pid in all_pids_single:
            pid_url_map[pid] = f"https://huggingface.co/datasets/{HF_DATASET}/resolve/main/variants/protein_id={pid}/data_0.parquet"
        for pid in all_pids_double:
            pid_url_map[pid] = f"https://huggingface.co/datasets/{HF_DATASET}/resolve/main/variants/variants/protein_id={pid}/data_0.parquet"
        all_pids = sorted(pid_url_map.keys())
        missing = [pid for pid in all_pids if pid not in existing_pids]
        print(f"Found {len(all_pids):,} proteins total ({len(all_pids_double):,} batch + {len(all_pids_single):,} single), {len(missing):,} to download")
        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        headers = {"Authorization": f"Bearer {token}"} if token else {}
        session = requests.Session()
        session.headers.update(headers)

        def _dl_one(pid):
            url = pid_url_map[pid]
            dst = local_dl / f"protein_id={pid}" / "data_0.parquet"
            dst.parent.mkdir(parents=True, exist_ok=True)
            for attempt in range(4):
                r = session.get(url, timeout=60)
                if r.status_code == 429:
                    time.sleep(2 ** attempt * 5)  # 5s, 10s, 20s, 40s
                    continue
                r.raise_for_status()
                dst.write_bytes(r.content)
                return
            r.raise_for_status()  # raise on final attempt

        if missing:
            workers = min(4, len(missing))
            done = n_fail = 0
            t0_dl = time.time()
            with ThreadPoolExecutor(max_workers=workers) as pool:
                futs = {pool.submit(_dl_one, pid): pid for pid in missing}
                for fut in as_completed(futs):
                    pid = futs[fut]
                    try:
                        fut.result()
                        done += 1
                    except Exception as e:
                        n_fail += 1
                        print(f"  Warning: {pid} download failed ({e})")
                    if (done + n_fail) % 1000 == 0 or (done + n_fail) == len(missing):
                        elapsed = time.time() - t0_dl
                        rate = done / elapsed if elapsed > 0 else 0
                        remain = (len(missing) - done - n_fail) / rate if rate else 0
                        print(f"  Downloaded {done:,}/{len(missing):,}  "
                              f"[{elapsed/60:.0f}min, ~{remain/60:.0f}min left]", flush=True)
            print(f"Download complete: {local_dl}  ({n_fail} failures)")
        variants_dir = local_dl

    protein_dirs = sorted([
        p for p in variants_dir.rglob("protein_id=*")
        if p.is_dir()
    ])
    print(f"Found {len(protein_dirs):,} proteins in {variants_dir}")

    # Resume support
    done = {p.parent.name.replace("protein_id=", "") for p in ml_dir.rglob("data_0.parquet")}
    if done:
        print(f"Resuming — {len(done):,} already done")

    max_prot = 5 if args.dry_run else None
    t0 = time.time()
    n_written = n_skipped = n_variants = 0

    for pdir in protein_dirs:
        if max_prot and n_written >= max_prot:
            break

        pid = pdir.name.replace("protein_id=", "")
        if pid in done:
            n_skipped += 1
            continue

        parquets = list(pdir.glob("*.parquet"))
        if not parquets:
            continue

        try:
            count = process_protein(pid, parquets[0], clf, classes, ml_dir)
            n_variants += count
            n_written  += 1
        except Exception as e:
            print(f"  Warning: {pid} failed ({e})")
            continue

        if args.dry_run:
            sample = pd.read_parquet(ml_dir / f"protein_id={pid}" / "data_0.parquet").head(3)
            print(f"  {pid}: {count:,} variants")
            print(sample[["mutation_code", "ml_mechLabel", "ml_confidence"]].to_string(index=False))
        elif n_written % 1000 == 0:
            elapsed = time.time() - t0
            rate = n_written / elapsed
            remain = (len(protein_dirs) - n_written - n_skipped) / rate if rate else 0
            print(f"  {n_written:,} proteins / {n_variants:,} variants  "
                  f"[{elapsed/60:.0f}min, ~{remain/60:.0f}min left]", flush=True)

    elapsed = time.time() - t0
    print(f"\nDone: {n_written:,} proteins, {n_variants:,} variants in {elapsed/60:.1f} min")
    print(f"Output: {ml_dir}")

    if args.dry_run:
        print("[dry-run] Skipping upload.")
        return

    # ── Upload ─────────────────────────────────────────────────────────────────
    n_files = len(list(ml_dir.rglob("*.parquet")))
    print(f"\nUploading {n_files:,} parquets to {HF_DATASET}/ml_preds/ ...")
    api = HfApi()
    api.upload_large_folder(
        folder_path=str(args.output_dir),
        repo_id=HF_DATASET,
        repo_type="dataset",
    )
    print("Upload done.")
    print(f"\nNext: bump CACHEBUST in Dockerfile.backend to pick up ml_preds/")


if __name__ == "__main__":
    main()