# Independent recount of the cohort and PWY-5676 RNA attribution, written apart from the bundle's code.
import csv, gzip, sys, os, collections
ext = sys.argv[1]
f = {'MGX3': 'dd983871b0e155255844b91ec10d50fb09230d2f4e915464ab680fa3a9c9ddb3', 'MTX3': 'd30419937c2fcfd6b0bfd026f529ba898cc19456b703af8fc64a5a91d88a5ffc', 'meta': '656b7bd97660ddb875548805e30bede31f2d1208293f7170d2d5755e33862ec9'}
path = lambda h: next(os.path.join(ext, n) for n in os.listdir(ext) if h in n)
def cols(h):
    with gzip.open(path(h), 'rt') as fh:
        head = fh.readline().rstrip('\n').split('\t')
    return [c[:-len('_pathabundance_cpm')] if c.endswith('_pathabundance_cpm') else c for c in head[1:]]
dna, rna = set(cols(f['MGX3'])), set(cols(f['MTX3']))
meta = collections.defaultdict(lambda: collections.defaultdict(list))
with open(path(f['meta']), newline='') as fh:
    for r in csv.DictReader(fh):
        meta[r['data_type']][r['External ID']].append(r)
kept = {}
for s in sorted(dna & rna):
    if s.endswith('_TR'): continue
    m, t = meta['metagenomics'].get(s, []), meta['metatranscriptomics'].get(s, [])
    if len(m) != 1 or len(t) != 1: continue
    m, t = m[0], t[0]
    if any(not m[k] or m[k] != t[k] for k in ('Participant ID', 'diagnosis', 'week_num')): continue
    key = (m['Participant ID'], float(m['week_num']))
    if key not in kept: kept[key] = (s, m['diagnosis'])
cd = [s for (p, w), (s, dx) in kept.items() if dx == 'CD']
print('CD samples', len(cd), 'participants', len({p for (p, w), (s, dx) in kept.items() if dx == 'CD'}))
with gzip.open(path(f['MTX3']), 'rt') as fh:
    rd = csv.reader(fh, delimiter='\t'); head = next(rd)
    idx = {c[:-len('_pathabundance_cpm')] if c.endswith('_pathabundance_cpm') else c: i for i, c in enumerate(head)}
    total = None; named = []
    for r in rd:
        name = r[0]
        if not name.startswith('PWY-5676:'): continue
        if '|' not in name: total = r
        elif not name.endswith('|unclassified'): named.append(r)
pos = sum(float(total[idx[s]]) > 0 for s in cd)
withnamed = sum(any(float(r[idx[s]]) > 0 for r in named) for s in cd)
print('PWY-5676 RNA: named strata', len(named), 'positive total', pos, 'with named signal', withnamed)