"""Front Immunol 2021, doi 10.3389/fimmu.2021.639672 (Sollid lab): gluten-specific CD4+ T-cell clones and clonotypes, celiac disease, HLA-DQ2.5 tetramer-sorted.
Table S3 (328 T-cell clone rows, 15 patients): one record per patient, tissue, status, epitope and TCR; `cells` is the number of T-cell clone (TCC) rows, not cells.
Chains: primary alpha/beta of each row. The table also has a second ('dual') beta; when the primary CDR3 beta does not start with A, S or T (out of frame, 12 rows) the dual one
is used if it does. Rows with an alpha or beta missing are completed from the other rows of the same Clonotype ID when those agree on a single chain.
Table S7 (public paired TCRs, 78): the 10 pairs already present in Table S3 are skipped, the other 68 are loaded. Donor is left empty (the TCR is shared by several patients); the patient list is in `context`.
Epitope is set only when the specificity is stated without parentheses; in S7 an epitope in parentheses means 'highly similar to a TCR with known epitope' and is kept in
`context` only (annotation "antigen": tetramer-sorted gluten-specific, epitope not established for that TCR); S7 rows without any epitope are annotation "antigen" as well.
CDR3 convention: the tables print CDR3 without the leading C and the trailing F/W (checked against the pattern of the sequences). The paper does not state this, so the C and F are NOT
added back: `cdr3a_core` and `cdr3b_core` are set to the printed CDR3.
Tables S4 to S6 (motif clusters) and the full 2,918-clonotype compilation (not published) are not loaded. Epitope peptide sequences are from Table S2 (gamma epitopes have none there)."""
import pandas as pd
from atlas import SRC, clones

F = SRC / "DahalKoirala2021/DataSheet_2.xlsx"


def pep(ep, seq):
    """'DQ2.5-glia-a2|DQ2.5-glia-w2' -> same with the Table S2 sequence in parentheses where there is one."""
    if not isinstance(ep, str): return None
    return " | ".join(f"{e} ({seq[e]})" if e in seq else e for e in (x.strip() for x in ep.split("|")))


def parse():
    t2 = pd.read_excel(F, sheet_name="Table S2", header=2).iloc[:, 1:3].dropna()
    seq = {a.strip(): b.strip() for a, b in zip(t2.iloc[:, 0], t2.iloc[:, 1]) if a.strip().startswith("DQ2.5-glia")}

    s = pd.read_excel(F, sheet_name="Table S3", header=1)
    s.columns = [c.strip().rstrip("*") for c in s.columns]
    s = s[s.Tissue.isin(["Blood", "Gut"])].copy()
    ok = s["CDR3β"].fillna("").str.match(r"^[AST]")
    dual = s["CDR3"].fillna("").str.match(r"^[AST]")
    use_dual = ~ok & dual
    for k, d in (("TRBV", "Dual V-gene"), ("TRBJ", "Dual J-gene"), ("CDR3β", "CDR3")): s[k] = s[k].where(~use_dual, s[d])
    s = s.rename(columns={"CDR3α": "CDR3a", "CDR3β": "CDR3b"})
    s.loc[~s.CDR3b.fillna("").str.match(r"^[AST]"), ["TRBV", "TRBJ", "CDR3b"]] = None   # out-of-frame primary beta without a usable dual
    for chain in (["TRAV", "TRAJ", "CDR3a"], ["TRBV", "TRBJ", "CDR3b"]):
        one = s.groupby("Clonotype ID")[chain].agg(lambda x: x.dropna().drop_duplicates().iloc[0] if len(x.dropna().drop_duplicates()) == 1 else None)
        miss = s[chain[2]].isna()
        s.loc[miss, chain] = s.loc[miss, ["Clonotype ID"]].join(one, on="Clonotype ID")[chain].values
    s["donor"] = s["Patient ID"].str.strip()
    s["epitope"] = s["Epitope specificity"].map(lambda e: pep(e, seq))
    s["antigen_ev"] = s["Epitope selection tetramer"].astype(str).str.strip()
    s["context"] = "celiac disease, " + s.Tissue + ", gluten status " + s["Gluten status"].str.strip()
    s["specificity_evidence"] = s.antigen_ev.map(lambda a: "culture-derived T-cell clone, epitope specificity listed in Table S3" if a == "Culture derived"
                                                 else f"HLA-DQ2.5:gluten tetramer sort ({a}), epitope specificity listed in Table S3")
    a = clones(s[s.CDR3a.notna() | s.CDR3b.notna()].copy(), ["donor", "context", "epitope", "specificity_evidence"])
    a["annotation_level"] = "epitope"

    p = pd.read_excel(F, sheet_name="Table S7", header=1)
    p.columns = [c.strip() for c in p.columns]
    p = p[p["CDR3β"].notna()].rename(columns={"CDR3α": "CDR3a", "CDR3β": "CDR3b"})
    done = set(zip(s.CDR3a, s.CDR3b))
    p = p[[(x, y) not in done for x, y in zip(p.CDR3a, p.CDR3b)]].copy()
    ep = p["Epitope specificity"].astype(str).str.strip()
    stated = ep.ne("nan") & ~ep.str.startswith("(")
    b = p.assign(donor=None, epitope=[pep(e, seq) if k else None for e, k in zip(ep, stated)],
                 context="celiac disease public paired TCR (patients " + p.Patients.str.replace(" ", "") + "), " + p.Tissue.str.strip() + p["Epitope specificity"].map(lambda e: f"; similar to a TCR with known epitope {e}" if isinstance(e, str) and e.strip().startswith("(") else ""),
                 specificity_evidence="HLA-DQ2.5:gluten tetramer sort (" + p["Epitope selection"].fillna("selection not stated in Table S7") + "), public paired TCR, Table S7",
                 annotation_level=stated.map({True: "epitope", False: "antigen"}), cells=None)
    b = b[["TRAV", "TRAJ", "CDR3a", "TRBV", "TRBJ", "CDR3b", "donor", "epitope", "context", "specificity_evidence", "annotation_level", "cells"]]
    r = pd.concat([a, b], ignore_index=True)
    r["cdr3a_core"], r["cdr3b_core"] = r.CDR3a, r.CDR3b
    return r.assign(source="Dahal-Koirala 2021 (Front Immunol)", study="doi:10.3389/fimmu.2021.639672", species="human", cd4_evidence="HLA-DQ2.5 tetramer-sorted CD4+ T cells and clones",
                    restriction="HLA-DQ2.5", antigen="gluten (gliadin)", organism="Triticum aestivum (wheat)")