"""Nat Commun 2026, doi 10.1038/s41467-026-72396-7: CD4 TCRs found with HLA-DRB1*01:01 class II single-chain-trimer (SCT) tetramers
(CEFT library, SARS-CoV-2 spike peptides, influenza HA clusters). Source: 41467_2026_72396_MOESM3_ESM.xlsx (Supplementary Data).
Loaded:
- S Data 4: 16 CEFT CD4 TCRs, one underscore-joined string TRAV_CDR3a_TRAJ_TRAC_TRBV_CDR3b_TRBJ_TRBC. The sheet gives no antigen per TCR (S Data 5 is only the
  key of the five antigens A to E), so these get level "antigen" (CEFT peptide library). Only TCR 3 gets an epitope: its beta (V, CDR3, J) is listed in S Data 13 with
  peptide YDVPDYASLRSLVASS ("TCRs validated in Fig. 1i"); the antigen name comes from S Data 5 (ID A).
- S Data 8: 10 SARS-CoV-2 CD4 TCRs with peptide and spike residue range as printed (TRAV_CDR3a_TRAJ_TRBV_CDR3b_TRBJ_TRBC, no TRAC); level "epitope".
- S Data 12: only the paper's own cells (Sample = barcode_pathogen_peptide, 29 distinct cells), beta V, J and CDR3 plus the alpha CDR3 (no TRAV or TRAJ in the sheet;
  alpha is empty for cells without alpha). Peptide as printed in the label; antigen name from S Data 5 when the peptide is listed there, otherwise the pathogen of the label. Level "epitope".
  The same cell is listed in several GLIPH groups: one cell is counted once; `cells` counts distinct barcodes.
Not loaded, and why: S Data 12 rows with Sample = VDJdb_n (public TCRs from VDJdb, not this study); S Data 13 (used only for TCR 3 above; its other rows are the S Data 12 cells again);
S Data 14 (the 5 HPV TCRs H1 to H5 with peptides, but the file holds no HPV TCR sequence anywhere, so no HPV record can be written; sequences are only in the GEO h5ad, not downloaded);
S Data 2 (previously published DR1 TCRs, other studies); S Data 1, 3, 5, 6, 7, 9, 10, 11, 15, 16, 17, 18 (primers, antigen libraries, participant metadata, expression, HLA, depth: no TCR).
Restriction is left empty: no sheet states an HLA per TCR (the screens use DRB1*01:01 SCT per the paper's methods, kept in `specificity_evidence`). No donor in the sheets (donor only in the GEO h5ad), so donor is empty.
Reporter-cell confirmation (Jurkat NFAT-GFP) is described by the paper but is not tabulated per TCR."""
import pandas as pd
from atlas import SRC

F = SRC / "Zhang2026/41467_2026_72396_MOESM3_ESM.xlsx"
SCT = "DRB1*01:01 class II single-chain-trimer tetramer"


def split(s):
    """TRAV_CDR3a_TRAJ[_TRAC]_TRBV_CDR3b_TRBJ[_TRBC] -> dict, each part located by its gene prefix."""
    p = s.split("_")
    a, b = next(i for i, x in enumerate(p) if x.startswith("TRAV")), next(i for i, x in enumerate(p) if x.startswith("TRBV"))
    assert p[a + 2].startswith("TRAJ") and p[b + 2].startswith("TRBJ") and p[a + 1][0] == "C" and p[b + 1][0] == "C", s
    return {"TRAV": p[a], "CDR3a": p[a + 1], "TRAJ": p[a + 2], "TRBV": p[b], "CDR3b": p[b + 1], "TRBJ": p[b + 2]}


def parse():
    sheet = lambda n: pd.read_excel(F, sheet_name=f"S Data {n}", header=1)
    key = sheet(5).assign(Antigen=lambda d: d.Antigen.str.strip()).set_index("Sequence").Antigen
    pathogen = sheet(3).assign(Peptide=lambda d: d.Peptide.str.strip()).set_index("Peptide").Pathogen

    s4 = sheet(4)
    s4 = pd.concat([s4, pd.DataFrame([split(x) for x in s4.Sequence])], axis=1)
    s13 = sheet(13).assign(beta=lambda d: d["TCR sequences"].str.split("_").str[:3].str.join("_")).drop_duplicates("beta").set_index("beta")
    pep = (s4.TRBV + "_" + s4.CDR3b + "_" + s4.TRBJ).map(s13["Epitope sequence"])
    s4 = s4.assign(epitope=pep, antigen=pep.map(key).fillna("CEFT peptide library (peptide not given per TCR)"), organism=pep.map(pathogen).fillna("CMV, EBV, influenza, tetanus (CEFT library)"),
                   context="Supplementary Data 4, CEFT screen", cells=None,
                   specificity_evidence=pep.map(lambda p: f"{SCT}; peptide from S Data 13, TCR listed as validated in Fig. 1i" if isinstance(p, str) else f"{SCT}, CEFT library pool; peptide not given per TCR"))

    s8 = sheet(8)
    s8 = pd.concat([s8, pd.DataFrame([split(x) for x in s8["TCR sequences"]])], axis=1).assign(
        epitope=lambda d: d["Peptide sequence"], antigen=lambda d: d.Antigen, organism="SARS-CoV-2", context="Supplementary Data 8, SARS-CoV-2 study (HLA-DR1+ participants)", cells=None,
        specificity_evidence=f"{SCT} with the listed spike peptide")

    s12 = sheet(12)
    s12 = s12[~s12.Sample.str.startswith("VDJdb_")].drop_duplicates("Sample")
    lab = s12.Sample.str.extract(r"^([ACGT]+-\d)_(.+)_([A-Z]+)$")
    s12 = pd.DataFrame({"barcode": lab[0], "pathogen": lab[1], "epitope": lab[2], "TRBV": s12.V, "TRBJ": s12.J, "CDR3b": s12.TcRb, "CDR3a": s12.TcRa})
    s12 = s12.assign(antigen=s12.epitope.map(key).fillna(s12.pathogen + " peptide (CEFT library)"), organism=s12.pathogen, context="Supplementary Data 12, single-cell barcode with peptide label",
                     specificity_evidence=f"{SCT} with the peptide in the barcode label; GLIPH clustering input")
    s12 = s12.groupby(["epitope", "antigen", "organism", "context", "specificity_evidence", "TRBV", "TRBJ", "CDR3b", "CDR3a"], dropna=False).size().rename("cells").reset_index()

    d = pd.concat([s4, s8, s12])
    d["annotation_level"] = d.epitope.notna().map({True: "epitope", False: "antigen"})
    return d.assign(source="Zhang 2026 (Nat Commun)", study="doi:10.1038/s41467-026-72396-7", species="human", cd4_evidence="MHC class II restriction (class II tetramer)")