#!/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):
    """Try VEP if installed and cache exists, else Ensembl REST, else fallback."""
    import shutil
    vep_bin = shutil.which("vep")
    if vep_bin and gtf_path and Path(gtf_path).is_file():
        pass
    rows = []
    for _, r in df.iterrows():
        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}"
        gene = "intergenic"
        hgvs_c = f"c.{pos1}{ref}>{alt}"
        hgvs_p = f"p.(?)"
        effect = "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel"
        is_synthetic = chrom.startswith("chr") and chrom[3:].isdigit() and int(chrom[3:]) <= 3 and pos1 < 6000
        if is_synthetic:
            gene = f"SYNTH_{chrom}"
        else:
            # Try Ensembl REST for real hg38 variants (one transcript, best)
            try:
                # Ensembl REST: GET /vep/homo_sapiens/hgvs/{hgvs_g}
                # Use chr without prefix for Ensembl: 7:g.140453136A>T
                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"}, timeout=5)
                if resp.status_code == 200:
                    j = resp.json()
                    if j and isinstance(j, list) and j[0].get("transcript_consequences"):
                        tc = j[0]["transcript_consequences"][0]
                        gene = tc.get("gene_symbol") or gene
                        hgvs_c = tc.get("hgvsc") or hgvs_c
                        hgvs_p = tc.get("hgvsp") or hgvs_p
                        cons = tc.get("consequence_terms", [])
                        effect = ", ".join(cons) if cons else effect
                        # Extract c. part from hgvsc like ENST00000288602.11:c.1799T>A -> c.1799T>A
                        if ":" in hgvs_c:
                            hgvs_c = hgvs_c.split(":")[-1]
                        if ":" in hgvs_p:
                            hgvs_p = hgvs_p.split(":")[-1]
            except Exception:
                pass
        # For all-transcripts mode, we currently emit one row per variant (best transcript).
        # With VEP cache, this would expand to N rows per variant.
        rows.append({
            "_orig_idx": r.name,
            "Ген": gene,
            "HGVS_c": hgvs_c,
            "HGVS_p": hgvs_p,
            "HGVS": f"{hgvs_c} {hgvs_p} ({hgvs_g})",
            "Тип варианта и эффект": effect,
        })
    return pd.DataFrame(rows)

def fetch_oncokb(hgvs_g_list, token, tumor_type="All Solid Tumors"):
    """Batch query OncoKB byHgvsVariant. Returns dict hgvs_g -> {amp, oncogenic}."""
    if not token:
        return {}
    headers = {"Authorization": f"Bearer {token}"}
    out = {}
    for hgvs in hgvs_g_list:
        try:
            # OncoKB expects e.g. 7:g.140453136A>T
            params = {"hgvsg": hgvs, "tumorType": tumor_type}
            resp = requests.get(ONCOKB_URL, headers=headers, params=params, timeout=10)
            if resp.status_code == 200:
                j = resp.json()
                # j contains levelOfEvidence, oncogenic, etc.
                amp = j.get("highestSensitiveLevel") or j.get("levelAssociated") or ""
                # Map OncoKB levels to AMP Tier
                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.2)
    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 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")
    args = ap.parse_args()

    clean_path = Path(args.clean)
    df_clean = pd.read_csv(clean_path)
    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")
        # Still create empty Excel with headers
        df_empty = pd.DataFrame(columns=["Ген","HGVS_c","HGVS_p","HGVS","Тип варианта и эффект","_orig_idx"])
        build_excel(df_clean, df_empty, out)
        return

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

    # 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) 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()