"""MIT licensed. Reproduce a source-coded endpoint map and exact table reanalysis."""
from fractions import Fraction
from math import comb
from pathlib import Path
import itertools
import json

ROOT = Path(__file__).resolve().parents[1]

def tail(n, x, p, upper):
    values = range(x, n + 1) if upper else range(x + 1)
    return sum(comb(n, k) * p**k * (1-p)**(n-k) for k in values)

def cp_interval(x, n, alpha):
    """Equal-tailed Clopper-Pearson interval by monotone binomial inversion."""
    def invert(x, upper):
        lo, hi = 0.0, 1.0
        for _ in range(100):
            mid = (lo + hi) / 2
            probability = tail(n, x, mid, upper)
            if (probability < alpha/2) == upper:
                lo = mid
            else:
                hi = mid
        return (lo + hi) / 2
    return [0.0 if x == 0 else invert(x, True),
            1.0 if x == n else invert(x, False)]

def fisher(a, n1, c, n0):
    """Two-sided, probability-ordering exact Fisher test using rational arithmetic."""
    total_success = a + c
    denom = comb(n1 + n0, n1)
    probabilities = {k: Fraction(comb(total_success, k) * comb(n1+n0-total_success, n1-k), denom)
                     for k in range(max(0, n1-(n1+n0-total_success)), min(n1, total_success)+1)}
    observed = probabilities[a]
    return sum((p for p in probabilities.values() if p <= observed), Fraction())

def holm(pvalues):
    ordered = sorted(pvalues, key=pvalues.get)
    result, running = {}, Fraction()
    for i, key in enumerate(ordered):
        running = max(running, min(Fraction(1), (len(ordered)-i)*pvalues[key]))
        result[key] = running
    return result

def validate_fisher():
    # Independent enumeration of binary assignments conditional on total successes.
    cases = 0
    for n1 in range(1,5):
        for n0 in range(1,5):
            for successes in range(n1+n0+1):
                hist = {}
                for positions in itertools.combinations(range(n1+n0), successes):
                    k = sum(i < n1 for i in positions)
                    hist[k] = hist.get(k, 0) + 1
                for a, freq in hist.items():
                    expected = Fraction(sum(v for v in hist.values() if v <= freq), sum(hist.values()))
                    assert fisher(a, n1, successes-a, n0) == expected
                    cases += 1
    return cases

def main():
    studies = json.loads((ROOT/'data/studies.json').read_text())
    source = json.loads((ROOT/'data/tumor_counts.json').read_text())
    assert len({s['study_id'] for s in studies}) == len(studies)
    assert len({s['doi'] for s in studies}) == len(studies)
    arms = {a['id']: a for a in source['arms']}
    assert len(arms) == len(source['arms'])
    for a in arms.values():
        assert isinstance(a['tumors'], int) and 0 <= a['tumors'] <= a['total']
        a['risk'] = a['tumors']/a['total']
        # Bonferroni over all three risks gives a simultaneous >=95% rectangle.
        a['simultaneous_risk_interval'] = cp_interval(a['tumors'], a['total'], 0.05/len(arms))
    tests, comparisons = {}, {}
    for spec in source['comparisons']:
        t, c = arms[spec['treated']], arms[spec['control']]
        name = spec['treated']+'_vs_'+spec['control']
        p = fisher(t['tumors'], t['total'], c['tumors'], c['total'])
        tests[name] = p
        lt, ut = t['simultaneous_risk_interval']; lc, uc = c['simultaneous_risk_interval']
        comparisons[name] = {
            'treated':t['id'], 'control':c['id'],
            'risk_ratio':t['risk']/c['risk'],
            'risk_difference':t['risk']-c['risk'],
            'risk_ratio_interval':[lt/uc, ut/lc],
            'risk_difference_interval':[lt-uc, ut-lc],
            'fisher_p':float(p), 'fisher_p_fraction':str(p)}
    for name, p in holm(tests).items():
        comparisons[name]['holm_p'] = float(p)
        comparisons[name]['holm_p_fraction'] = str(p)
    classes = sorted({s['endpoint_class'] for s in studies})
    data = {'map':{'papers':len(studies),'endpoint_classes':len(classes),
                   'clinical_incident_cancer_papers':sum(s['clinical_incident_cancer_endpoint'] for s in studies),
                   'primary_cell_reprogramming_papers':sum(s['normal_primary_cells_at_start'] for s in studies),
                   'classes':{k:[s['study_id'] for s in studies if s['endpoint_class']==k] for k in classes}},
            'arms':arms, 'comparisons':comparisons,
            'validation':{'fisher_exhaustive_tables':validate_fisher(),
                          'interval_method':'Three Bonferroni-adjusted Clopper-Pearson risk intervals transformed monotonically; at least 95% joint coverage under independent binomial arms.',
                          'interval_coverage_percent':95,
                          'p_method':'Two-sided probability-ordering Fisher exact; Holm adjustment across the two prespecified-for-this-reanalysis comparisons.'}}
    # Boundary and symmetry checks test interval inversion independently of study outcomes.
    for n in [1,2,5,16,30]:
        for x in range(n+1):
            low, high = cp_interval(x,n,0.05)
            other = cp_interval(n-x,n,0.05)
            assert abs(low-(1-other[1])) < 1e-12
            assert abs(high-(1-other[0])) < 1e-12
            if x: assert abs(tail(n,x,low,True)-0.025) < 1e-12
            if x<n: assert abs(tail(n,x,high,False)-0.025) < 1e-12
    result = ROOT/'results'; result.mkdir(exist_ok=True)
    (result/'R1.json').write_text(json.dumps(data,indent=2,sort_keys=True)+'\n')
    (result/'R2.json').write_text(json.dumps({name:{k:round(v,4) if isinstance(v,float) else [round(x,4) for x in v] if isinstance(v,list) else v for k,v in item.items()} for name,item in comparisons.items()},indent=2,sort_keys=True)+'\n')
    print(json.dumps(data,indent=2))

if __name__ == '__main__':
    main()