"""Independent check of the eligibility counts, written from the paper's Methods.
Usage: python3 -I hmp2_check.py <external dir> <out.json>"""
import csv, gzip, json, math, sys, collections, pathlib
ext = pathlib.Path(sys.argv[1])
F = {'meta': '656b7bd97660ddb875548805e30bede31f2d1208293f7170d2d5755e33862ec9',
     'dna': '9aa0f1ba96fd6d4ccd9e73a1bc0c9164336565effa936dbf682f05be162f8c60',
     'rna': '6110abb08823bb78a4ad63dcf98da1cae04fe758095e0b196aff1b2ed1e74331',
     'tax': 'd790ff15e46d61ca0cadc55d9f918de4e3415d7f97c992ac37610aaee02117ed'}
ECS = ['2.8.3.8', '2.7.2.7', '2.8.3.9', '2.3.1.19']
def num(v):
    try:
        x = float(v)
        return x if math.isfinite(x) else None
    except (TypeError, ValueError):
        return None
meta = {}
with open(ext / F['meta'], newline='') as f:
    for r in csv.DictReader(f):
        meta.setdefault(r['data_type'], {})[r['External ID']] = r
def read_table(path, keep):
    with gzip.open(path, 'rt') as f:
        head = f.readline().rstrip('\n').split('\t')[1:]
        rows = {}
        for line in f:
            k, rest = line.split('\t', 1)
            if keep(k):
                rows[k] = [float(v) for v in rest.rstrip('\n').split('\t')]
    return head, rows
def ec_of(k): return k.split('|')[0].split(':')[0]
dh, dr = read_table(ext / F['dna'], lambda k: ec_of(k) in ECS)
rh, rr = read_table(ext / F['rna'], lambda k: ec_of(k) in ECS)
th, tr = read_table(ext / F['tax'], lambda k: k.split('|')[-1].startswith('s__'))
def specimens(head, assay, suffix):
    md = meta[assay]
    recs = []
    for i, h in enumerate(head):
        assert h.endswith(suffix)
        x = h[:-len(suffix)]
        if x in md: recs.append((i, x, md[x]))
    nonpilot = {m['site_sub_coll'] for _, x, m in recs if not x.endswith('_P')}
    keep = collections.defaultdict(list)
    for i, x, m in recs:
        if '_TR' in x: continue
        if x.endswith('_P') and m['site_sub_coll'] in nonpilot: continue
        keep[m['site_sub_coll']].append((i, m))
    return {s: v[0] for s, v in keep.items() if s and len(v) == 1}
D = specimens(dh, 'metagenomics', '_level4ec')
R = specimens(rh, 'metatranscriptomics', '_level4ec')
T = specimens(th, 'metagenomics', '_profile')
sp = {k.split('|')[-1][3:]: v for k, v in tr.items()}
pairs = []
for s in sorted(D.keys() & R.keys()):
    (di, dm), (ri, rm) = D[s], R[s]
    if (dm['Participant ID'], dm['week_num'], dm['diagnosis']) != (rm['Participant ID'], rm['week_num'], rm['diagnosis']):
        continue
    if dm['diagnosis'] != 'CD' or num(dm['week_num']) is None: continue
    pairs.append(dict(s=s, p=dm['Participant ID'], w=num(dm['week_num']), di=di, ri=ri,
        ti=T[s][0] if s in T else None,
        reads=(num(dm['reads_filtered']), num(rm['reads_filtered'])),
        fc=(num(dm['fecalcal']), num(rm['fecalcal']))))
out = {'paired_CD_specimens': len(pairs)}
qc = [x for x in pairs if all(r is not None and r >= 1e6 for r in x['reads'])]
out['read_filtered_specimens'] = len(qc)
out['participants'] = len({x['p'] for x in qc})
def tax_ok(x):
    if x['ti'] is None: return False
    vals = [v[x['ti']] for v in sp.values()]
    return 99 <= sum(vals) <= 101 and sum(v > 0 for v in vals) > 1
weeks = collections.Counter((x['p'], x['w']) for x in qc)
lon = [x for x in qc if weeks[(x['p'], x['w'])] == 1 and tax_ok(x)]
byp = collections.defaultdict(list)
for x in lon: byp[x['p']].append(x)
edges = []
for xs in byp.values():
    xs.sort(key=lambda x: x['w'])
    edges += [(a, b) for a, b in zip(xs, xs[1:]) if 2 <= b['w'] - a['w'] <= 12]
out['intervals'] = len(edges)
def fc_ok(x): return x['fc'][0] is not None and x['fc'][0] >= 0 and x['fc'][0] == x['fc'][1]
ev, evp, evf = [], [], []
for ec in ECS:
    contrib = [k for k in dr if ec_of(k) == ec and '|g__' in k]
    n = nf = 0; people = set()
    for a, b in edges:
        base = {k: dr[k][a['di']] for k in contrib if dr[k][a['di']] > 0}
        if len(base) < 2: continue
        dom = max(base, key=lambda k: (base[k], -contrib.index(k)))
        if base[dom] / sum(base.values()) < 0.5: continue
        name = dom.split('s__')[-1]
        if name not in sp: continue
        t0, t1 = sp[name][a['ti']], sp[name][b['ti']]
        if t0 < 0.1 or t1 > 0.5 * t0: continue
        if rr.get(dom, [0] * len(rh))[a['ri']] <= 0: continue
        if not any(k != dom and k in rr and rr[k][a['ri']] > 0 and rr[k][b['ri']] > 0 for k in base): continue
        n += 1; people.add(a['p'])
        if fc_ok(a) and fc_ok(b): nf += 1
    ev.append(n); evp.append(len(people)); evf.append(nf)
out.update(eligible_events=ev, eligible_participants=evp, eligible_fecalcal=evf)
pathlib.Path(sys.argv[2]).write_text(json.dumps(out, indent=2) + '\n')
print(json.dumps(out))