"""J Clin Invest 2018, doi 10.1172/JCI98819 (Sollid and Qiao labs): public gluten-specific TCRs of celiac disease patients, HLA-DQ2.5:gluten tetramer-sorted CD4+ T cells.
Source: Supplemental data PDF, text transcribed with `pdftotext -layout` into sources/Risnes2018/derived/ (printed values only):
- risnes2018_tableS3_paired.csv: Table S3, one row per sequenced TCR (203 rows, 39 distinct pairs, 16 patients). In the PDF the pair labels (AV.._CDR3a_AJ..,
  BV.._CDR3b_BJ..) sit in a side column, not on each row; each row was assigned to its pair by translating the printed nucleotides and matching the label CDR3s
  (188 rows match exactly). 15 rows (patients CD1339 and CD1283, pair AV41_AVEGGSNYKLT_AJ53) carry a label whose beta CDR3 (SVGAGGSGELF) differs by one amino acid from
  their own nucleotides (SVGAGGTGELF); they were assigned by their alpha chain, and the CDR3 used here is always the translation of the printed nucleotides.
- risnes2018_tableS2_public_chains.csv: Table S2, public single chains (111 alpha, 104 beta) with the number of patients and sequences; the per-patient counts are not loaded.
  Only chains whose CDR3 is not already in a Table S3 pair are loaded, as alpha-only or beta-only records without donor (`cells` = total number of sequences).
CDR3 is printed without the leading C and the trailing F/W (IMGT 105 to 117), so `cdr3a_core` and `cdr3b_core` are set to the printed CDR3. V and J names are printed without 'TR'; it is added.
Epitope: none per TCR in these tables (the tetramer set mixed several gluten peptides): annotation "antigen". Time point and tissue are in `context`: UCD = untreated celiac disease,
'<n> yr GFD' or 'wk GFD' = time on a gluten-free diet, Baseline / Day 6 / Day 14 / Day 28 = stage of gluten challenge (footnote b of the supplement)."""
import itertools

import pandas as pd
from atlas import SRC, clones

AA = "FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG"
CODON = {"".join(c): a for c, a in zip(itertools.product("tcag", repeat=3), AA)}
tr = lambda s: "".join(CODON[s[i:i + 3]] for i in range(0, len(s) - len(s) % 3, 3))


def label(s):
    """'AV4_LVGD..._AJ32' -> (TRAV4, CDR3, TRAJ32)"""
    v, c, j = s.split("_")
    return "TR" + v, c, "TR" + j


def parse():
    d = SRC / "Risnes2018/derived"
    p = pd.read_csv(d / "risnes2018_tableS3_paired.csv")
    a, b = p.alpha_label.map(label), p.beta_label.map(label)
    p["TRAV"], p["TRAJ"], p["TRBV"], p["TRBJ"] = a.str[0], a.str[2], b.str[0], b.str[2]
    p["CDR3a"], p["CDR3b"] = p.cdr3a_nt.map(tr), p.cdr3b_nt.map(tr)
    p["context"] = "celiac disease, " + p.source + ", " + p.time_point
    r = clones(p, ["donor", "context"])

    s = pd.read_csv(d / "risnes2018_tableS2_public_chains.csv")
    s = s.join(s.label.map(label).apply(pd.Series).set_axis(["V", "CDR3", "J"], axis=1))
    s = s[~s.CDR3.isin(set(r.CDR3a.dropna()) | set(r.CDR3b.dropna()))]
    al = s.label.str.startswith("AV")
    one = pd.DataFrame({"TRAV": s.V.where(al), "TRAJ": s.J.where(al), "CDR3a": s.CDR3.where(al), "TRBV": s.V.where(~al), "TRBJ": s.J.where(~al), "CDR3b": s.CDR3.where(~al),
                        "context": "celiac disease, public single chain (Table S2), found in " + s.n_patients.astype(str) + " patients", "cells": s.n_sequences})
    r = pd.concat([r, one], ignore_index=True)
    r["cdr3a_core"], r["cdr3b_core"] = r.CDR3a, r.CDR3b
    return r.assign(source="Risnes 2018 (J Clin Invest)", study="doi:10.1172/JCI98819", species="human", cd4_evidence="HLA-DQ2.5:gluten tetramer-sorted CD4+ T cells",
                    restriction="HLA-DQ2.5", antigen="gluten (gliadin)", organism="Triticum aestivum (wheat)", specificity_evidence="HLA-DQ2.5:gluten tetramer sort", annotation_level="antigen")