"""Reproduce the preregistered IBDMDB audit; Python standard library only."""
import csv, gzip, hashlib, json, math, pathlib, random, statistics, sys
from collections import Counter, defaultdict
ROOT = pathlib.Path(__file__).resolve().parents[1]
DATA = pathlib.Path(sys.argv[1]) if len(sys.argv) > 1 else ROOT / 'data'
TARGETS = ['PWY-5676', 'P163-PWY', 'CENTFERM-PWY', 'PWY-5677']
IMM = 'Immunosuppressants (e.g. oral corticosteroids)'

def mean(xs): return statistics.mean(xs) if xs else None

def interval(xs):
    if len(xs) < 5: return None
    rng = random.Random(903)
    boot = sorted(mean(rng.choices(xs, k=len(xs))) for _ in range(5000))
    def q(p):
        i = (len(boot)-1)*p; lo = int(i); return boot[lo]+(i-lo)*(boot[min(lo+1,len(boot)-1)]-boot[lo])
    return [q(.025), q(.975)]

def read_table(name):
    targets = {p: {} for p in TARGETS}; normalizer = None; feature_count = 0
    with gzip.open(DATA/(name+'.tsv.gz'), 'rt') as f:
        rows = csv.reader(f, delimiter='\t'); header = next(rows)
        samples = [s.removesuffix('_pathabundance_cpm') for s in header[1:]]
        assert len(set(samples)) == len(samples)
        normalizer = [0.]*len(samples)
        for r in rows:
            assert len(r) == len(header)
            feature_count += 1
            base, sep, taxon = r[0].partition('|'); pathway = base.split(':',1)[0]
            if pathway not in TARGETS and (sep or pathway in ('UNMAPPED','UNINTEGRATED')): continue
            vals = [float(x) for x in r[1:]]
            assert all(math.isfinite(x) and x >= 0 for x in vals)
            if not sep and pathway not in ('UNMAPPED','UNINTEGRATED'):
                normalizer = [a+b for a,b in zip(normalizer, vals)]
            if pathway in TARGETS: targets[pathway][taxon if sep else 'total'] = vals
    return {'samples':samples, 'index':{s:i for i,s in enumerate(samples)},'targets':targets,
            'normalizer':normalizer,'feature_count':feature_count}

def cohort(dna,rna):
    rows = list(csv.DictReader((DATA/'metadata.csv').open()))
    meta = defaultdict(list); rmeta = defaultdict(list)
    for r in rows:
        if r['data_type'] == 'metagenomics': meta[r['External ID']].append(r)
        if r['data_type'] == 'metatranscriptomics': rmeta[r['External ID']].append(r)
    audit=Counter(); kept=[]
    common=sorted(set(dna['samples']) & set(rna['samples']));audit['shared_samples']=len(common)
    for s in common:
        if s.endswith('_TR'): audit['technical_replicates_excluded']+=1;continue
        ms=meta[s];rs=rmeta[s]
        if len(ms)!=1 or len(rs)!=1: audit['missing_or_duplicate_metadata']+=1;continue
        m=ms[0]
        if any(not m[k] or m[k]!=rs[0][k] for k in ['Participant ID','diagnosis','week_num']):
            audit['inconsistent_metadata']+=1;continue
        try: w=float(m['week_num'])
        except ValueError: audit['invalid_week']+=1;continue
        if not math.isfinite(w):audit['invalid_week']+=1;continue
        kept.append({'sample':s,'participant':m['Participant ID'],'diagnosis':m['diagnosis'],
                     'week':w,'antibiotics':m['Antibiotics'],'immunosupp':m[IMM],
                     'calprotectin':float(m['fecalcal']) if m['fecalcal'] else None})
    unique={}
    for r in kept:
        k=(r['participant'],r['week'])
        if k in unique:audit['same_week_excluded']+=1
        else:unique[k]=r
    out=sorted(unique.values(),key=lambda r:(r['participant'],r['week'],r['sample']))
    audit['retained_samples']=len(out)
    return out,dict(audit)

def series(table,pathway,sample,renorm=False):
    i=table['index'][sample]; multiplier=1e6/table['normalizer'][i] if renorm else 1
    return {k:v[i]*multiplier for k,v in table['targets'][pathway].items()}

def named(x):return sorted(k for k in x if k not in ('unclassified','total'))

def make_events(rows,dna,rna,pathway,dx='CD',cutoff=.5,maxgap=8,minimum=1,
                no_antibiotics=False,stable_immunosupp=False,rna_leader=False,
                nonoverlap=False,renorm=False):
    by=defaultdict(list)
    for r in rows:
        if r['diagnosis']==dx:by[r['participant']].append(r)
    events=[];audit=Counter()
    dn=named(dna['targets'][pathway]);rn=named(rna['targets'][pathway]); taxa=sorted(set(dn)|set(rn))
    for participant,visits in sorted(by.items()):
        last_used=-1
        for i in range(len(visits)-2):
            triple=visits[i:i+3];audit['candidate_triples']+=1
            if nonoverlap and i<=last_used:continue
            if not all(1<=triple[j+1]['week']-triple[j]['week']<=maxgap for j in (0,1)):continue
            audit['time_eligible']+=1
            if no_antibiotics and not all(r['antibiotics']=='No' for r in triple):continue
            if stable_immunosupp and not (len({r['immunosupp'] for r in triple})==1 and triple[0]['immunosupp'] in ('Yes','No')):continue
            if renorm and any(t['normalizer'][t['index'][r['sample']]] <= 0 for t in [dna,rna] for r in triple):
                audit['zero_normalizer_excluded']+=1;continue
            ds=[series(dna,pathway,r['sample'],renorm) for r in triple]
            rs=[series(rna,pathway,r['sample'],renorm) for r in triple]
            if not dn:continue
            leader=sorted(dn,key=lambda t:(-ds[0][t],t))[0]
            leader0=ds[0][leader];others=[t for t in taxa if t!=leader]
            dalt=[sum(d.get(t,0) for t in others) for d in ds]
            ralt=[sum(r.get(t,0) for t in others) for r in rs]
            total=[r.get('total',0) for r in rs]
            if not (leader0>=minimum and dalt[0]>=minimum and leader0>=.5*sum(ds[0][t] for t in dn) and min(total[:2])>=1):continue
            if rna_leader and leader not in rn:continue
            audit['baseline_eligible']+=1
            ratios=[ds[j].get(leader,0)/leader0 for j in (1,2)]
            kind='decline' if all(x<=cutoff for x in ratios) else 'stable' if all(.8<=x<=1.25 for x in ratios) else 'other'
            if kind=='other':continue
            last_used=i+2
            increase=ralt[2]-ralt[1]>=1 and ralt[2]>=1.25*ralt[1]
            maintain=total[2]>=.8*total[0]
            concurrent=ralt[1]-ralt[0]>=1 and ralt[1]>=1.25*ralt[0] and total[1]>=.8*total[0]
            ev={'participant':participant,'samples':[r['sample'] for r in triple],'weeks':[r['week'] for r in triple],
                'kind':kind,'leader':leader,'leader_dna':[d.get(leader,0) for d in ds],
                'alternative_dna':dalt,'alternative_rna':ralt,'community_rna':total,
                'signature':increase and maintain,'alternative_increase':increase,'rna_maintained':maintain,
                'contemporaneous_signature':concurrent,'calprotectin_t2':triple[2]['calprotectin']}
            ev['relative_expression_log2']={str(pc):math.log2((ralt[2]+pc)/(dalt[2]+pc))-math.log2((ralt[1]+pc)/(dalt[1]+pc)) for pc in [.01,.1,1]}
            later=[r for r in visits[i+3:] if 1<=r['week']-triple[2]['week']<=12 and r['calprotectin'] is not None]
            ev['calprotectin_future']=later[0]['calprotectin'] if later else None
            events.append(ev)
    return events,dict(audit)

def group_summary(events):
    by=defaultdict(list)
    for e in events:by[e['participant']].append(e)
    out={'events':len(events),'participants':len(by),'population_inference_adequate':len(by)>=10,
         'leaders':dict(sorted(Counter(e['leader'] for e in events).items()))}
    for field in ['signature','alternative_increase','rna_maintained','contemporaneous_signature']:
        ps=[mean([int(e[field]) for e in es]) for _,es in sorted(by.items())]
        out[field]={'count':sum(e[field] for e in events),'event_fraction':mean([int(e[field]) for e in events]),
                    'participant_mean':mean(ps),'participant_bootstrap_95':interval(ps)}
    out['relative_expression_log2']={pc:{'participant_mean':mean(ps:=[mean([e['relative_expression_log2'][pc] for e in es]) for _,es in sorted(by.items())]),'participant_bootstrap_95':interval(ps)} for pc in ['0.01','0.1','1']}
    return out

def summarize(events,audit):
    out={'audit':audit,**{k:group_summary([e for e in events if e['kind']==k]) for k in ['decline','stable']}}
    by=defaultdict(lambda:defaultdict(list))
    for e in events:by[e['participant']][e['kind']].append(e['signature'])
    diffs=[mean([int(x) for x in v['decline']])-mean([int(x) for x in v['stable']]) for _,v in sorted(by.items()) if v['decline'] and v['stable']]
    out['within_participant_decline_minus_stable']={'participants':len(diffs),'mean':mean(diffs),'bootstrap_95':interval(diffs)}
    clinical={}
    for flag in [False,True]:
        es=[e for e in events if e['kind']=='decline' and e['signature']==flag and e['calprotectin_t2'] is not None and e['calprotectin_future'] is not None]
        byp=defaultdict(list)
        for e in es:byp[e['participant']].append(math.log2((e['calprotectin_future']+1)/(e['calprotectin_t2']+1)))
        clinical[str(flag).lower()]={'events':len(es),'participants':len(byp),'_values':{p:mean(v) for p,v in sorted(byp.items())}}
    enough=all(c['participants']>=10 for c in clinical.values())
    if enough:
        a=list(clinical['true']['_values'].values());b=list(clinical['false']['_values'].values());rng=random.Random(903)
        # A participant can occur in both groups. Resample union of participants, retaining both groups.
        participants=sorted(set(clinical['true']['_values'])|set(clinical['false']['_values']));boot=[]
        for _ in range(5000):
            picked=rng.choices(participants,k=len(participants))
            v=[[clinical[k]['_values'][p] for p in picked if p in clinical[k]['_values']] for k in ['true','false']]
            if all(v):boot.append(mean(v[0])-mean(v[1]))
        boot.sort();clinical['contrast']={'mean_log2_difference':mean(a)-mean(b),'bootstrap_95':[boot[int(.025*(len(boot)-1))],boot[int(.975*(len(boot)-1))]]}
    for k in ['true','false']:del clinical[k]['_values']
    clinical['adequate']=enough;out['inflammation_feasibility']=clinical
    return out

def coverage(table,pathway,rows):
    x=table['targets'][pathway];ns=named(x)
    vals=[];namedpositive=twonamed=totalpositive=unclassifiedpositive=0
    for r in rows:
        if r['sample'] not in table['index']:continue
        v=series(table,pathway,r['sample']);total=v.get('total',0);s=sum(v.get(t,0) for t in ns)
        assert not (s>0 and total==0), 'Named signal without community total'
        totalpositive+=total>0;namedpositive+=s>0;twonamed+=sum(v.get(t,0)>0 for t in ns)>=2;unclassifiedpositive+=v.get('unclassified',0)>0
        if total>0:vals.append(s/total)
    return {'named_strata_count':len(ns),'named_strata':ns,'samples_with_positive_total':totalpositive,
            'samples_with_named_signal':namedpositive,'samples_with_two_named_strata':twonamed,
            'samples_with_unclassified_signal':unclassifiedpositive,
            'median_named_sum_over_total':statistics.median(vals) if vals else None,
            'note':'Ratio is descriptive; HUMAnN pathway totals are not additive species sums.'}

def main():
    sources=json.loads((ROOT/'data/sources.json').read_text())
    for s in sources:
        fn='metadata.csv' if s['assay']=='metadata' else s['assay']+'.tsv.gz';f=DATA/fn
        assert f.stat().st_size==s['bytes'],fn
        assert 'sha256:'+hashlib.sha256(f.read_bytes()).hexdigest()==s['sha256'],fn
    tables={n:read_table(n) for n in ['MGX','MTX','MGX3','MTX3']};dna=tables['MGX3'];rna=tables['MTX3']
    rows,attrition=cohort(dna,rna);cohorts={}
    for dx in ['CD','UC','nonIBD']:
        rs=[r for r in rows if r['diagnosis']==dx];cnt=Counter(r['participant'] for r in rs)
        cohorts[dx]={'samples':len(rs),'participants':len(cnt),'participants_with_three_samples':sum(v>=3 for v in cnt.values()),'calprotectin_observed_samples':sum(r['calprotectin'] is not None for r in rs)}
    cd=[r for r in rows if r['diagnosis']=='CD']
    result={'attrition':attrition,'cohorts':cohorts,'coverage':{p:{n:coverage(t,p,cd) for n,t in tables.items()} for p in TARGETS},'analyses':{}}
    primary_events=None
    for p in ['PWY-5676','P163-PWY']:
        result['analyses'][p]={}
        for dx in ['CD','UC','nonIBD']:
            ev,audit=make_events(rows,dna,rna,p,dx=dx);result['analyses'][p][dx]=summarize(ev,audit)
            if p=='PWY-5676' and dx=='CD':primary_events=ev
    specs={'decline_033':{'cutoff':.33},'decline_067':{'cutoff':.67},'gap_4':{'maxgap':4},'gap_12':{'maxgap':12},'min_05':{'minimum':.5},'min_2':{'minimum':2},'no_antibiotics':{'no_antibiotics':True},'stable_immunosupp':{'stable_immunosupp':True},'rna_resolved_leader':{'rna_leader':True},'nonoverlapping':{'nonoverlap':True},'annotated_pathway_scaling':{'renorm':True}}
    result['sensitivities']={}
    for name,kwargs in specs.items():
        ev,audit=make_events(rows,dna,rna,'PWY-5676',**kwargs);result['sensitivities'][name]=summarize(ev,audit)
    # The normalization sensitivity repeats all main calculations for both features and cohorts.
    result['normalized_main']={}
    for p in ['PWY-5676','P163-PWY']:
        result['normalized_main'][p]={}
        for dx in ['CD','UC','nonIBD']:
            ev,audit=make_events(rows,dna,rna,p,dx=dx,renorm=True);result['normalized_main'][p][dx]=summarize(ev,audit)
    out=ROOT/'results';out.mkdir(exist_ok=True)
    (out/'R1.json').write_text(json.dumps(result,indent=2,sort_keys=True,allow_nan=False)+'\n')
    # Event identifiers enable source checks; public deidentified source IDs only, no clinical values.
    audit_rows=[{k:e[k] for k in ['participant','samples','weeks','kind','leader','leader_dna','alternative_dna','alternative_rna','community_rna','signature']} for e in primary_events]
    (out/'R2.json').write_text(json.dumps({'primary_event_audit':audit_rows},indent=2,sort_keys=True,allow_nan=False)+'\n')
    print(json.dumps({'cohorts':cohorts,'primary':result['analyses']['PWY-5676']['CD'],'secondary':result['analyses']['P163-PWY']['CD']},indent=2))
if __name__=='__main__':main()