"""J Exp Med 2025, doi 10.1084/jem.20250460 (GSE305300): human memory CD4+ T cells sorted CD4+CD69+CD40L+ after co-culture with autologous H37Rv-infected macrophages,
with or without added Mtb lysate (GEO design statement). 18 Cell Ranger filtered_contig_annotations files, one per donor and condition; donor and condition come from the file name
(`CL003-Lysate`: -Infected = infected macrophages, -Lysate = infected macrophages plus Mtb lysate; GEO treatment field). CL010-Control: the GEO treatment field says Mtb-infected macrophages, the meaning of "Control"
is not explained, so the raw label is kept in `context`. LTBI status per donor is copied from the GEO sample characteristics (CL002, CL010, Lkpk2 to Lkpk5 non-LTBI; the others LTBI).
Contigs: is_cell, high_confidence and productive only; the top-UMI TRA and TRB of each barcode pair a cell (barcodes are only unique within a file, so donor and condition are part of the cell key).
One record per donor, condition and TCR, `cells` = number of cells. Level "antigen": activation (AIM) after infected macrophages; there is no per-cell antigen flag and no mapped epitope, and
a clonotype activated by infected macrophages may recognise any Mtb or other antigen they present.
Supplemental Data 4 (7 cloned paired TCRs, GLIPH2 group named, no antigen in the table): the TCRs whose alpha and beta CDR3 are found paired in the contigs of the same donor are already loaded from the contigs and only
get "cloned and recombinantly expressed (Supplemental Data 4, GLIPH2 group ...)" appended to `specificity_evidence`; a cloned TCR with no matching cell is added once from the table (alpha and beta as printed, alleles removed).
Not loaded: Supplemental Data 1 and 2 (GLIPH2 groups, mostly beta only, bulk visits), 3 (CDR3 beta lists), 5 (peptides), the Dryad deposit (not retrievable)."""
import glob
import pandas as pd
from atlas import SRC, pair_contigs, clones

LTBI_NEG = {"CL002", "CL010", "Lkpk2", "Lkpk3", "Lkpk4", "Lkpk5"}
COND = {"Infected": "Mtb-infected macrophages", "Lysate": "Mtb-infected macrophages plus Mtb lysate", "Control": "Mtb-infected macrophages (file label Control)"}


def parse():
    out = []
    for f in sorted(glob.glob(str(SRC / "Stetsenko2025/*-filtered_contig_annotations.csv.gz"))):
        donor, cond = f.split("/")[-1].split("_", 1)[1].split("-filtered")[0].rsplit("-", 1)
        c = pd.read_csv(f)
        c = c[c.is_cell.astype(str).eq("True") & c.high_confidence.astype(str).eq("True") & c.productive.astype(str).eq("True")].assign(donor=donor, cond=cond)
        out.append(pair_contigs(c, ["donor", "cond", "barcode"]))
    d = pd.concat(out, ignore_index=True)
    d["context"] = "memory CD4+ CD69+CD40L+ sorted, " + d.donor.map(lambda x: "non-LTBI" if x in LTBI_NEG else "LTBI") + " donor, " + d.cond.map(COND)

    cl = pd.read_excel(SRC / "Stetsenko2025/jem_20250460_datas4.xlsx", header=2).rename(columns={"patient_id": "donor"})
    for k in ("TRAV", "TRAJ", "TRBV", "TRBJ"): cl[k] = cl[k].str.split("*").str[0]
    cloned = d.merge(cl[["donor", "CDR3a", "CDR3b", "GLIPH2 Group"]], on=["donor", "CDR3a", "CDR3b"], how="left")["GLIPH2 Group"]
    d["specificity_evidence"] = "CD69+CD40L+ (AIM) after Mtb-infected macrophages, TCR not tested"
    d.loc[cloned.notna().values, "specificity_evidence"] += "; cloned and recombinantly expressed (Supplemental Data 4, GLIPH2 group " + cloned.dropna() + ")"
    r = clones(d, ["donor", "context", "specificity_evidence"])

    seen = r[r.specificity_evidence.str.contains("Supplemental Data 4")]
    add = cl[~(cl.donor + cl.CDR3a + cl.CDR3b).isin(seen.donor + seen.CDR3a + seen.CDR3b)]
    add = add.assign(context="cloned TCR, " + add.donor.map(lambda x: "non-LTBI" if x in LTBI_NEG else "LTBI") + " donor", cells=None,
                     specificity_evidence="cloned and recombinantly expressed (Supplemental Data 4, GLIPH2 group " + add["GLIPH2 Group"] + "); the table gives no reactivity result or antigen; cell of origin not found among the loaded contigs")
    r = pd.concat([r, add[["donor", "context", "specificity_evidence", "TRAV", "TRAJ", "CDR3a", "TRBV", "TRBJ", "CDR3b", "cells"]]])
    return r.assign(source="Stetsenko 2025 (J Exp Med)", study="doi:10.1084/jem.20250460", species="human", cd4_evidence="sorted memory CD4+", antigen="Mycobacterium tuberculosis (infected macrophages)",
                    organism="Mycobacterium tuberculosis", annotation_level="antigen")