"""eLife 2020, doi 10.7554/eLife.57063 (GSE158848): bulk TCR alpha and TCR beta of FACS-sorted CD4 effector/memory subsets (Tfh, Th1, Th1-17, Th17, Th22, Th2a, Th2, Treg)
of 5 healthy blood donors, UMI-based (MiGEC). 86 clonesets (43 per chain). Donor, subset, chain and library come from the file name `GSM_Donor<d>_<n>[_2]_<TRA|TRB>_<subset>.txt.gz`.
TRA and TRB are separate libraries, so records are single chain (CDR3a or CDR3b empty, never paired). One record per sample, chain and (V, J, CDR3 amino acid): rows that differ only
in the nucleotide sequence are merged. `count` is a UMI count, not a cell count, so `cells` is left empty. Donor 4 has a second library for Treg, Th1-17 and Th17 (extra `_2` in the name):
the label "library 1" / "library 2" is kept in `context` for those subsets only; the manifest calls `_2` a replicate, which of the two is the original is not stated, so the label is the file-name order.
Gene names: the source writes TRAV14DV4 etc.; `/` is inserted before DV (TRAV14/DV4) to match IMGT naming. The naive CD4 subsets of the paper (Figshare) are not downloaded."""
import re
import pandas as pd
from atlas import SRC

NAME = re.compile(r"GSM\d+_Donor(\d+)_(\d+)(_2)?_(TR[AB])_(.+)\.txt\.gz$")


def parse():
    parts = []
    for f in sorted((SRC / "Kasatskaya2020/clonesets").glob("*.txt.gz")):
        donor, _, rep, chain, subset = NAME.match(f.name).groups()
        d = pd.read_csv(f, sep="\t", usecols=["count", "cdr3aa", "v", "j"])
        parts.append(d.assign(donor=donor, chain=chain, subset=subset, rep=rep or ""))
    d = pd.concat(parts, ignore_index=True)
    d["subset"] = d.subset.replace({"TREG": "Treg"})
    d = d.groupby(["donor", "chain", "subset", "rep", "v", "j", "cdr3aa"], as_index=False)["count"].sum()
    has_rep = d.assign(r=d.rep.eq("_2")).groupby(["donor", "subset"]).r.transform("any")
    lib = has_rep.map({True: ", library "}).fillna("") + d.rep.map({"": "1", "_2": "2"}).where(has_rep, "")
    d["context"] = "healthy blood, sorted CD4 subset " + d.subset + lib
    d["donor"] = "Donor" + d.donor
    v = d.v.str.replace(r"^(TRAV\d+(?:-\d+)?)DV(\d)", r"\1/DV\2", regex=True)
    a = d.chain.eq("TRA")
    out = pd.DataFrame({"TRAV": v.where(a), "TRAJ": d.j.where(a), "CDR3a": d.cdr3aa.where(a), "TRBV": v.where(~a), "TRBJ": d.j.where(~a), "CDR3b": d.cdr3aa.where(~a),
                        "donor": d.donor, "context": d.context})
    return out.assign(source="Kasatskaya 2020 (eLife)", study="GEO:GSE158848", species="human", cd4_evidence="sorted CD4+", specificity_evidence="none", annotation_level="none")