ffpe_damage/annotate_clean.py

#!/usr/bin/env python3
"""Annotate clean_variants.csv -> annotated Excel with 8 columns (pan-cancer, all transcripts).

Input: clean_variants.csv from ffpe_damage_v2.py (chrom 0-based pos)
Output: *.annotated.xlsx with columns:
  Ген | HGVS (c./p.) | Тип варианта и эффект | VAF | PAF | ACMG значимость | AMP уровень | Уровень онкогенности

- HGVS: c.6713C>T + p.(Pro2238Leu) style, all transcripts expanded (1 variant = N rows)
- PAF: empty (as requested)
- ACMG: ClinVar + InterVar placeholder (VUS if no ClinVar hit, requires manual curation for true Pathogenic)
- AMP/Oncogenicity: OncoKB (token ~/.config/oncokb/token, tumor_type pan-cancer) + CIViC fallback
- VAF from clean_variants.csv
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path

import pandas as pd
import requests
from openpyxl import Workbook
from openpyxl.styles import PatternFill, Font, Alignment, Border, Side
from openpyxl.utils import get_column_letter

ACMG_COLORS = {
    "Pathogenic": "FFC7CE",
    "Likely pathogenic": "FFEB9C",
    "Uncertain significance": "FFCC99",
    "VUS": "FFCC99",
    "Likely benign": "C6EFCE",
    "Benign": "C6EFCE",
}
AMP_COLORS = {
    "Tier I": "8B0000",
    "Tier II": "FF8C00",
    "Tier III": "FFD700",
    "Tier IV": "D3D3D3",
}
ONCO_COLORS = {
    "Oncogenic": "C6EFCE",
    "Likely Oncogenic": "FFEB9C",
    "Resistance": "FFC7CE",
}

ONCOKB_URL = "https://www.oncokb.org/api/v1/annotate/mutations/byHgvsVariant"
CIVIC_URL = "https://civicdb.org/api/variants?count=10000"

def load_token(path="~/.config/oncokb/token"):
    p = Path(path).expanduser()
    if p.is_file():
        return p.read_text().strip()
    return os.environ.get("ONCOKB_TOKEN", "").strip()

def vep_annotate(df, reference, gtf_path=None, offline=False):
    """Try VEP if installed and cache exists, else GTF offline, else Ensembl REST, else fallback."""
    import shutil
    vep_bin = shutil.which("vep")
    if vep_bin and gtf_path and Path(gtf_path).expanduser().is_file():
        pass
    # Load GTF for offline all-transcript expansion
    gtf_transcripts = None
    if offline and gtf_path:
        gtf_file = Path(gtf_path).expanduser()
        if gtf_file.is_file():
            print(f"  loading GTF {gtf_file} ...", flush=True)
            gtf_transcripts = _load_gtf(gtf_file)
            print(f"  GTF loaded: {len(gtf_transcripts)} transcripts", flush=True)
    rows = []
    total = len(df)
    for idx, (_, r) in enumerate(df.iterrows()):
        if idx % 500 == 0 and total > 500:
            print(f"  annotate {idx}/{total} ...", flush=True)
        chrom = str(r["chrom"])
        pos1 = int(r["position"]) + 1
        ref = str(r["ref"]); alt = str(r["alt"])
        hgvs_g = f"{chrom}:g.{pos1}{ref}>{alt}"
        is_synthetic = chrom.startswith("chr") and chrom[3:].isdigit() and int(chrom[3:]) <= 3 and pos1 < 6000
        if is_synthetic:
            rows.append({
                "_orig_idx": r.name,
                "Ген": f"SYNTH_{chrom}",
                "HGVS_c": f"c.{pos1}{ref}>{alt}",
                "HGVS_p": "p.(?)",
                "HGVS": f"c.{pos1}{ref}>{alt} p.(?) ({hgvs_g})",
                "Тип варианта и эффект": "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel",
            })
            continue
        if offline and gtf_transcripts is not None:
            # GTF offline: all transcripts overlapping pos
            hits = [t for t in gtf_transcripts if t["chrom"] == chrom and t["start"] <= pos1 <= t["end"]]
            if hits:
                for t in hits:
                    # Simple HGVS: use transcript ID + positional offset (placeholder, VEP would give exact c./p.)
                    hgvs_c = f"{t['tx']}:c.{pos1}{ref}>{alt}"
                    hgvs_p = "p.(?)"
                    effect = t["biotype"] or "transcript_variant"
                    rows.append({
                        "_orig_idx": r.name,
                        "Ген": t["gene"],
                        "HGVS_c": hgvs_c,
                        "HGVS_p": hgvs_p,
                        "HGVS": f"{hgvs_c} {hgvs_p} ({hgvs_g})",
                        "Тип варианта и эффект": effect,
                    })
                continue
            rows.append({
                "_orig_idx": r.name,
                "Ген": "intergenic",
                "HGVS_c": f"c.{pos1}{ref}>{alt}",
                "HGVS_p": "p.(?)",
                "HGVS": f"c.{pos1}{ref}>{alt} p.(?) ({hgvs_g})",
                "Тип варианта и эффект": "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel",
            })
            continue
        if offline:
            rows.append({
                "_orig_idx": r.name,
                "Ген": "intergenic",
                "HGVS_c": f"c.{pos1}{ref}>{alt}",
                "HGVS_p": "p.(?)",
                "HGVS": f"c.{pos1}{ref}>{alt} p.(?) ({hgvs_g})",
                "Тип варианта и эффект": "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel",
            })
            continue
        # Online: Ensembl REST, expand all transcripts
        try:
            hgvs_ens = f"{chrom.replace('chr','')}:g.{pos1}{ref}>{alt}"
            url = f"https://rest.ensembl.org/vep/homo_sapiens/hgvs/{hgvs_ens}?content-type=application/json"
            resp = requests.get(url, headers={"Content-Type": "application/json", "Accept": "application/json"}, timeout=(3, 5))
            if resp.status_code == 200:
                j = resp.json()
                tcs = j[0].get("transcript_consequences", []) if j and isinstance(j, list) and j[0] else []
                if tcs:
                    for tc in tcs:
                        gene = tc.get("gene_symbol") or "intergenic"
                        hgvs_c = tc.get("hgvsc") or f"c.{pos1}{ref}>{alt}"
                        hgvs_p = tc.get("hgvsp") or "p.(?)"
                        cons = tc.get("consequence_terms", [])
                        effect = ", ".join(cons) if cons else ("SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel")
                        if ":" in hgvs_c:
                            hgvs_c = hgvs_c.split(":")[-1]
                        if ":" in hgvs_p:
                            hgvs_p = hgvs_p.split(":")[-1]
                        rows.append({
                            "_orig_idx": r.name,
                            "Ген": gene,
                            "HGVS_c": hgvs_c,
                            "HGVS_p": hgvs_p,
                            "HGVS": f"{hgvs_c} {hgvs_p} ({hgvs_g})",
                            "Тип варианта и эффект": effect,
                        })
                    continue
        except Exception:
            pass
        rows.append({
            "_orig_idx": r.name,
            "Ген": "intergenic",
            "HGVS_c": f"c.{pos1}{ref}>{alt}",
            "HGVS_p": "p.(?)",
            "HGVS": f"c.{pos1}{ref}>{alt} p.(?) ({hgvs_g})",
            "Тип варианта и эффект": "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel",
        })
    return pd.DataFrame(rows)

def fetch_oncokb(hgvs_g_list, token, tumor_type="All Solid Tumors", offline=False):
    """Batch query OncoKB byHgvsVariant. Returns dict hgvs_g -> {amp, oncogenic}."""
    if not token or offline:
        return {}
    if len(hgvs_g_list) > 500:
        print(f"  [oncokb] {len(hgvs_g_list)} variants - skipping OncoKB (use --offline to silence, or annotate subset)", flush=True)
        return {}
    headers = {"Authorization": f"Bearer {token}"}
    out = {}
    for i, hgvs in enumerate(hgvs_g_list):
        if i % 100 == 0 and len(hgvs_g_list) > 100:
            print(f"  [oncokb] {i}/{len(hgvs_g_list)} ...", flush=True)
        try:
            params = {"hgvsg": hgvs, "tumorType": tumor_type}
            resp = requests.get(ONCOKB_URL, headers=headers, params=params, timeout=(3, 5))
            if resp.status_code == 200:
                j = resp.json()
                amp = j.get("highestSensitiveLevel") or j.get("levelAssociated") or ""
                level_map = {"LEVEL_1": "Tier I", "LEVEL_2": "Tier II", "LEVEL_3A": "Tier II", "LEVEL_3B": "Tier III", "LEVEL_4": "Tier IV", "LEVEL_R1": "Tier I", "LEVEL_R2": "Tier II"}
                amp_tier = level_map.get(amp, amp)
                oncogenic = j.get("oncogenic", "")
                out[hgvs] = {"AMP": amp_tier, "ONCO": oncogenic}
            else:
                out[hgvs] = {"AMP": "", "ONCO": ""}
        except Exception:
            out[hgvs] = {"AMP": "", "ONCO": ""}
        time.sleep(0.05)
    return out

def fetch_civic():
    try:
        resp = requests.get(CIVIC_URL, timeout=10)
        if resp.status_code == 200:
            return resp.json()
    except Exception:
        pass
    return None

def annotate_acmg(df):
    """Placeholder ACMG: ClinVar lookup would go here. For MVP, all VUS unless known pathogenic."""
    # Real implementation would join ClinVar variant_summary + InterVar
    # For synthetic data, mark as VUS
    out = []
    for _, r in df.iterrows():
        # Heuristic: if VAF low and not in ClinVar, VUS
        out.append("Uncertain significance")
    return out

def build_excel(df_clean, df_annot, out_path):
    wb = Workbook()
    ws = wb.active
    ws.title = "clean_variants"

    headers = ["Ген", "HGVS", "Тип варианта и эффект", "VAF", "PAF", "ACMG значимость", "AMP уровень", "Уровень онкогенности", "chrom", "position", "ref", "alt", "depth"]
    header_fill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid")
    header_font = Font(color="FFFFFF", bold=True, size=10)
    thin_border = Border(left=Side(style="thin", color="D0D7DE"), right=Side(style="thin", color="D0D7DE"),
                         top=Side(style="thin", color="D0D7DE"), bottom=Side(style="thin", color="D0D7DE"))

    ws.append(headers)
    for col in range(1, len(headers)+1):
        c = ws.cell(row=1, column=col)
        c.fill = header_fill
        c.font = header_font
        c.alignment = Alignment(horizontal="center", vertical="center", wrap_text=True)
        c.border = thin_border
    ws.freeze_panes = "A2"
    ws.auto_filter.ref = f"A1:{get_column_letter(len(headers))}1"

    # Merge clean + annot (all transcripts: currently 1 row per variant)
    # df_annot has _orig_idx linking to df_clean
    for idx, arow in df_annot.iterrows():
        orig = df_clean.loc[arow["_orig_idx"]]
        vaf = float(orig["VAF"])
        # ACMG/AMP/ONCO from annot or placeholder
        acmg = arow.get("ACMG", "Uncertain significance")
        amp = arow.get("AMP", "")
        onco = arow.get("ONCO", "")

        row = [
            arow["Ген"],
            arow["HGVS"],
            arow["Тип варианта и эффект"],
            f"{vaf*100:.2f}%" if vaf <=1 else str(vaf),
            "",  # PAF empty as requested
            acmg,
            amp,
            onco,
            str(orig["chrom"]),
            int(orig["position"])+1,
            str(orig["ref"]),
            str(orig["alt"]),
            int(orig["depth"]),
        ]
        ws.append(row)
        rnum = ws.max_row
        # Borders and alignment
        for col in range(1, len(headers)+1):
            c = ws.cell(row=rnum, column=col)
            c.border = thin_border
            c.alignment = Alignment(vertical="center", wrap_text=True)
            c.font = Font(size=9)
            # Color by ACMG
            if col == 6:
                colr = ACMG_COLORS.get(str(acmg), None)
                if colr:
                    c.fill = PatternFill(start_color=colr, end_color=colr, fill_type="solid")
            if col == 7 and amp:
                colr = AMP_COLORS.get(str(amp), None)
                if colr:
                    # AMP text white on dark
                    c.fill = PatternFill(start_color=colr, end_color=colr, fill_type="solid")
                    if amp == "Tier I":
                        c.font = Font(color="FFFFFF", size=9, bold=True)
            if col == 8 and onco:
                colr = ONCO_COLORS.get(str(onco), None)
                if colr:
                    c.fill = PatternFill(start_color=colr, end_color=colr, fill_type="solid")

    # Column widths
    widths = [12, 32, 28, 10, 8, 18, 12, 18, 8, 10, 6, 6, 8]
    for i, w in enumerate(widths, 1):
        ws.column_dimensions[get_column_letter(i)].width = w
    ws.row_dimensions[1].height = 28

    wb.save(out_path)
    print(f"written {out_path} ({ws.max_row-1} rows)")

def _load_gtf(gtf_path):
    import gzip
    transcripts = []
    try:
        opener = gzip.open if str(gtf_path).endswith(".gz") else open
        with opener(gtf_path, "rt") as fh:
            for line in fh:
                if line.startswith("#"):
                    continue
                parts = line.rstrip("\n").split("\t")
                if len(parts) < 9 or parts[2] != "transcript":
                    continue
                chrom = parts[0]
                start = int(parts[3]); end = int(parts[4])
                attr = parts[8]
                # parse gene_name, transcript_id
                import re
                m_gene = re.search(r'gene_name "([^"]+)"', attr)
                m_tx = re.search(r'transcript_id "([^"]+)"', attr)
                m_biotype = re.search(r'transcript_type "([^"]+)"', attr)
                if m_gene and m_tx:
                    # Normalize chr prefix: GTF uses chr1, our clean uses chr1 -> keep as is
                    transcripts.append({
                        "chrom": chrom,
                        "start": start,
                        "end": end,
                        "gene": m_gene.group(1),
                        "tx": m_tx.group(1),
                        "biotype": m_biotype.group(1) if m_biotype else "",
                    })
    except Exception as e:
        print(f"[gtf] failed to load {gtf_path}: {e}", file=sys.stderr)
        return []
    return transcripts


def main():
    ap = argparse.ArgumentParser(description="Annotate clean_variants.csv -> Excel")
    ap.add_argument("--clean", required=True, help="clean_variants.csv from ffpe_damage_v2.py")
    ap.add_argument("--reference", default=None, help="hg38.fa (optional, for VEP)")
    ap.add_argument("--out", default=None, help="output xlsx (default <clean>.annotated.xlsx)")
    ap.add_argument("--tumor-type", default="All Solid Tumors", help="OncoKB tumor type (pan-cancer default)")
    ap.add_argument("--token", default="~/.config/oncokb/token", help="OncoKB token file or env")
    ap.add_argument("--offline", action="store_true", help="skip Ensembl/OncoKB network calls (fast, offline)")
    ap.add_argument("--min-vaf", type=float, default=None, help="filter VAF > threshold (e.g. 0.10 for 10%%)")
    ap.add_argument("--min-depth", type=int, default=None, help="filter depth >= threshold")
    ap.add_argument("--gtf", default="~/Projects/ffpe_damage/references/gencode.v44.annotation.gtf.gz",
                    help="GENCODE GTF for offline all-transcript expansion")
    args = ap.parse_args()

    clean_path = Path(args.clean)
    df_clean = pd.read_csv(clean_path)
    # VAF/depth filter before annotation (VAF is alt_count/depth fraction)
    n_before = len(df_clean)
    if args.min_vaf is not None:
        df_clean = df_clean[df_clean["VAF"] > args.min_vaf]
    if args.min_depth is not None:
        df_clean = df_clean[df_clean["depth"] >= args.min_depth]
    if n_before != len(df_clean):
        print(f"filter VAF>{args.min_vaf} depth>={args.min_depth}: {n_before} -> {len(df_clean)} variants", flush=True)
        df_clean = df_clean.reset_index(drop=True)
    if df_clean.empty:
        print("clean_variants.csv is empty (or only header) - nothing to annotate")
        out = args.out or str(clean_path).replace(".csv", ".annotated.xlsx")
        df_empty = pd.DataFrame(columns=["Ген","HGVS_c","HGVS_p","HGVS","Тип варианта и эффект","_orig_idx"])
        build_excel(df_clean, df_empty, out)
        return

    if len(df_clean) > 500 and not args.offline:
        print(f"[warn] {len(df_clean)} variants - network annotation will be slow. Use --offline for fast placeholder.", flush=True)

    # HGVS / gene / effect
    df_annot = vep_annotate(df_clean, args.reference, gtf_path=args.gtf, offline=args.offline)

    # OncoKB
    token = load_token(args.token)
    hgvs_list = df_annot["HGVS"].tolist() if "HGVS" in df_annot else []
    # Extract g.HGVS for OncoKB: e.g. chr1:g.2505C>T -> 1:g.2505C>T (strip chr)
    hgvs_g_for_oncokb = []
    for _, r in df_clean.iterrows():
        hgvs_g_for_oncokb.append(f"{str(r['chrom']).replace('chr','')}:g.{int(r['position'])+1}{r['ref']}>{r['alt']}")
    oncokb_map = fetch_oncokb(hgvs_g_for_oncokb, token, args.tumor_type, offline=args.offline) if token else {}

    # Attach ACMG/AMP/ONCO to annot rows
    acmg_list = annotate_acmg(df_clean)
    # Map orig_idx -> acmg
    acmg_by_idx = {df_clean.index[i]: acmg_list[i] for i in range(len(acmg_list))}
    for idx, row in df_annot.iterrows():
        orig_idx = row["_orig_idx"]
        hgvs_g = hgvs_g_for_oncokb[df_clean.index.get_loc(orig_idx)] if orig_idx in df_clean.index else ""
        oc = oncokb_map.get(hgvs_g, {})
        df_annot.at[idx, "ACMG"] = acmg_by_idx.get(orig_idx, "Uncertain significance")
        df_annot.at[idx, "AMP"] = oc.get("AMP", "")
        df_annot.at[idx, "ONCO"] = oc.get("ONCO", "")

    out = args.out or str(clean_path).replace(".csv", ".annotated.xlsx")
    build_excel(df_clean, df_annot, out)

if __name__ == "__main__":
    main()