#!/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, all_transcripts=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
    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:
            hits = [t for t in gtf_transcripts if t["chrom"] == chrom and t["start"] <= pos1 <= t["end"]]
            if hits:
                if all_transcripts:
                    for t in hits:
                        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
                # pick best transcript (MANE > appris1 > canonical > basic > protein_coding)
                def _score(t):
                    return (t["is_mane"], t["appris"] == 1, t["appris"] == 2, t["is_canonical"], t["is_basic"], t["is_ccds"], t["biotype"] == "protein_coding")
                best = max(hits, key=_score)
                hgvs_c = f"{best['tx']}:c.{pos1}{ref}>{alt}"
                hgvs_p = "p.(?)"
                effect = best["biotype"] or "transcript_variant"
                rows.append({
                    "_orig_idx": r.name,
                    "Ген": best["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
        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:
                    if all_transcripts:
                        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
                    def _score_tc(tc):
                        return (tc.get("mane_select") is not None, tc.get("canonical") == 1, tc.get("biotype") == "protein_coding", tc.get("impact") == "HIGH")
                    tc = max(tcs, key=_score_tc)
                    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, clinvar_path=None):
    """ACMG via ClinVar variant_summary if available, else VUS placeholder."""
    clinvar = {}
    if clinvar_path:
        cp = Path(clinvar_path).expanduser()
        if cp.is_file():
            print(f"  loading ClinVar {cp} ...", flush=True)
            clinvar = _load_clinvar(cp)
            print(f"  ClinVar loaded: {len(clinvar)//2} variants", flush=True)
    out = []
    for _, r in df.iterrows():
        chrom = str(r["chrom"]); pos = int(r["position"]) + 1
        ref = str(r["ref"]); alt = str(r["alt"])
        key = (chrom, pos, ref, alt)
        key2 = (chrom.replace("chr",""), pos, ref, alt)
        if key in clinvar:
            out.append(clinvar[key])
        elif key2 in clinvar:
            out.append(clinvar[key2])
        else:
            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
    import re
    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]
                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:
                    tags = set(re.findall(r'tag "([^"]+)"', attr))
                    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 "",
                        "is_mane": "MANE_Select" in tags,
                        "is_canonical": "Ensembl_canonical" in tags,
                        "is_basic": "basic" in tags,
                        "is_ccds": "CCDS" in tags,
                        "appris": next((int(t.split("_")[-1]) for t in tags if t.startswith("appris_principal_")), 99),
                    })
    except Exception as e:
        print(f"[gtf] failed to load {gtf_path}: {e}", file=sys.stderr)
        return []
    return transcripts


def _load_clinvar(clinvar_path):
    import gzip
    m = {}
    try:
        opener = gzip.open if str(clinvar_path).endswith(".gz") else open
        with opener(clinvar_path, "rt") as fh:
            header = None
            for line in fh:
                if line.startswith("#"):
                    continue
                if header is None:
                    header = line.rstrip("\n").split("\t")
                    # Find column indices
                    try:
                        ci_chr = header.index("Chromosome")
                        ci_start = header.index("Start")
                        ci_ref = header.index("ReferenceAllele")
                        ci_alt = header.index("AlternateAllele")
                        ci_sig = header.index("ClinicalSignificance")
                    except ValueError:
                        # Fallback for older format: try different names
                        continue
                    continue
                parts = line.rstrip("\n").split("\t")
                if len(parts) <= max(ci_chr, ci_start, ci_ref, ci_alt, ci_sig):
                    continue
                chrom = parts[ci_chr]
                # ClinVar Chromosome is 1,2.. not chr1
                chrom_norm = f"chr{chrom}" if not chrom.startswith("chr") else chrom
                try:
                    pos = int(parts[ci_start])
                except ValueError:
                    continue
                ref = parts[ci_ref]; alt = parts[ci_alt]
                sig = parts[ci_sig]
                # Normalize significance to ACMG
                # ClinVar: Pathogenic, Likely pathogenic, Uncertain significance, Likely benign, Benign, etc.
                # Map to our 5-tier
                sig_lower = sig.lower()
                if "pathogenic" in sig_lower and "likely" not in sig_lower:
                    acmg = "Pathogenic"
                elif "likely pathogenic" in sig_lower:
                    acmg = "Likely pathogenic"
                elif "benign" in sig_lower and "likely" not in sig_lower:
                    acmg = "Benign"
                elif "likely benign" in sig_lower:
                    acmg = "Likely benign"
                else:
                    acmg = "Uncertain significance"
                key = (chrom_norm, pos, ref, alt)
                # Also add without chr prefix for matching
                key2 = (chrom, pos, ref, alt)
                m[key] = acmg
                m[key2] = acmg
    except Exception as e:
        print(f"[clinvar] failed to load {clinvar_path}: {e}", file=__import__("sys").stderr)
        return {}
    return m

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("--all-transcripts", action="store_true", help="expand all transcripts (default: pick best MANE/canonical per variant)")
    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 expansion (pick best when not --all-transcripts)")
    ap.add_argument("--clinvar", default="~/Projects/ffpe_damage/references/clinvar_variant_summary.txt.gz",
                    help="ClinVar variant_summary.txt.gz for ACMG P/L (auto if exists)")
    args = ap.parse_args()

    # VAF>20% default for online (pan-cancer), leave --min-vaf configurable, offline keeps None
    if args.min_vaf is None and not args.offline:
        args.min_vaf = 0.20
        print(f"[info] online default --min-vaf 0.20 (use --min-vaf 0.05 to keep more)", flush=True)

    clean_path = Path(args.clean)
    df_clean = pd.read_csv(clean_path)
    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, all_transcripts=args.all_transcripts)

    # 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, clinvar_path=args.clinvar)
    # 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()