"""An exploratory, frozen aggregate-data audit. No participant records or network use."""
from copy import deepcopy
import json
from pathlib import Path
from legacy_analysis import effect, pool
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / 'results'

def save(name, value):
    (OUT / name).write_text(json.dumps(value, indent=2, allow_nan=False) + '\n')

def display(value):
    if isinstance(value, float):
        return round(value, 3)
    if isinstance(value, dict):
        return {k: display(v) for k, v in value.items()}
    if isinstance(value, list):
        return [display(v) for v in value]
    return value

def main():
    OUT.mkdir(exist_ok=True)
    rows = json.loads((ROOT / 'data/trials.json').read_text())
    old = json.loads((ROOT / 'data/previous-trials.json').read_text())
    extra = json.loads((ROOT / 'data/supplement-audit.json').read_text())
    primary = [r for r in rows if r['primary_pool']]
    by_id = {r['id']: r for r in rows}
    assert len(by_id) == len(rows)
    assert all(r['followup_days'] == 28 for r in primary)
    assert len({r['id'] for r in primary}) == len(primary)
    revised = pool(primary)
    previous = pool([r for r in old if r['primary_pool']])
    sensitivity = {'leave_out_' + r['id']: pool([s for s in primary if s['id'] != r['id']]) for r in primary}
    sensitivity['blinded_only'] = pool([r for r in primary if r['blinded']])
    sensitivity['replace_covidout_with_day14'] = pool([r for r in primary if r['id'] != 'covidout'] + [r for r in old if r['id'] == 'covidout'])
    for name, updates in [
        ('remove_accidental_death', {'events_t': 6}),
        ('highdose_missing_control_nonevents', {'n_c': 724}),
        ('highdose_missing_control_events', {'n_c': 724, 'events_c': 4}),
    ]:
        changed = deepcopy(primary)
        next(r for r in changed if r['id'] == 'activ600').update(updates)
        sensitivity[name] = pool(changed)
    invalid = deepcopy(by_id['principle'])
    invalid.update(invalid['other_control'])
    invalid = effect(invalid)

    tg = by_id['together']
    h, d = tg['hospital_only'], tg['death_only']
    amin, amax = max(h['events_t'], d['events_t']), min(tg['n_t'], h['events_t'] + d['events_t'])
    cmin, cmax = max(h['events_c'], d['events_c']), min(tg['n_c'], h['events_c'] + d['events_c'])
    scenarios = []
    for a in range(amin, amax + 1):
        for c in range(cmin, cmax + 1):
            hypothetical = {**tg, 'events_t': a, 'events_c': c, 'endpoint': 'hypothetical hospitalization/death union'}
            scenarios.append({'union_events_t': a, 'union_events_c': c, **pool(primary + [hypothetical])})
    envelope = dict(
        scenarios=len(scenarios), pooled_n=scenarios[0]['n'],
        intervals_including_null=sum(r['lower'] <= 1 <= r['upper'] for r in scenarios),
        pooled_rr_min=min(r['rr'] for r in scenarios), pooled_rr_max=max(r['rr'] for r in scenarios),
        lower_envelope=min(r['lower'] for r in scenarios), upper_envelope=max(r['upper'] for r in scenarios),
        events_t_min=amin, events_t_max=amax, events_c_min=cmin, events_c_max=cmax)
    save('together-union-scenarios.json', scenarios)

    # This is conditional denominator sensitivity, not an estimate of missing outcomes.
    co = by_id['covidout']
    nt, nc = co['analysis_n_t'], co['analysis_n_c']
    mt, mc = nt - co['n_t'], nc - co['n_c']
    assignments = []
    for at in range(mt + 1):
        for ac in range(mc + 1):
            assignments.append({'additional_events_t': at, 'additional_events_c': ac,
                                'rr': ((co['events_t'] + at) / nt) / ((co['events_c'] + ac) / nc)})
    missing = dict(
        assumption='Registry day-28 endpoint sample is nested in the published modified-intention-to-treat arms; all remaining records have unknown binary endpoints.',
        reference_n_t=nt, reference_n_c=nc, denominator_gap_t=mt, denominator_gap_c=mc,
        denominator_gap_total=mt + mc, denominator_gap_percent=100 * (mt + mc) / (nt + nc),
        scenarios=len(assignments), rr_min=min(r['rr'] for r in assignments), rr_max=max(r['rr'] for r in assignments),
        all_additional_nonevents_rr=assignments[0]['rr'],
        first_treatment_events_reversing_direction=min(r['additional_events_t'] for r in assignments if r['additional_events_c'] == 0 and r['rr'] > 1))
    save('covidout-denominator-scenarios.json', assignments)

    values = {r['id']: effect(r) for r in rows}
    values['ivercor_separate'] = effect(extra['ivercor'])
    values['together_covid_hospital_only'] = effect({**tg, **extra['together_supplement']['covid_hospitalization'], 'endpoint': 'COVID hospitalization alone'})
    full = dict(primary=revised, previous=previous, trials=values, sensitivities=sensitivity,
                principle_invalid=invalid, together=envelope, covidout_denominator=missing)
    save('full-results.json', full)
    summary = {**full, 'interval_percent': 95, 'source_comparisons': len(rows),
               'source_platforms': len({r['platform'] for r in rows}),
               'source_analysis_records': sum(r['analysis_n_t'] + r['analysis_n_c'] for r in rows),
               'added_endpoint_records': revised['n'] - previous['n'],
               'event_total': revised['events_t'] + revised['events_c'],
               'covidout_weight_percent': 100 * revised['weights']['covidout'],
               'ivercor_n': extra['ivercor']['n_t'] + extra['ivercor']['n_c'],
               'ivercor_mean_followup': extra['ivercor']['mean_end_of_study_days']}
    save('R1.json', display(summary))
    plot(primary, full)
    print(json.dumps({k: display(summary[k]) for k in ['primary', 'together', 'covidout_denominator']}, indent=2))

def plot(primary, result):
    labels = [r['label'] for r in primary] + ['Pooled, modified Hartung-Knapp']
    effects = [effect(r) for r in primary] + [result['primary']]
    fig, ax = plt.subplots(figsize=(12, 5.8))
    fig.subplots_adjust(left=.35, right=.79, top=.77, bottom=.18)
    for i, r in enumerate(effects):
        ax.errorbar(r['rr'], i, xerr=[[r['rr'] - r['lower']], [r['upper'] - r['rr']]],
                    color='#146977', marker='D' if i == len(primary) else 'o', capsize=4, linestyle='none')
        ax.text(1.03, i, f"{r['rr']:.2f} [{r['lower']:.2f}, {r['upper']:.2f}]", transform=ax.get_yaxis_transform(), va='center', fontsize=10)
    ax.set_yticks(range(len(labels)), labels)
    ax.invert_yaxis()
    ax.set_xscale('log')
    ax.set_xlim(.2, 22)
    ax.set_xticks([.25, .5, 1, 2, 5, 10], ['0.25', '0.5', '1', '2', '5', '10'])
    ax.axvline(1, color='#657780', linestyle='--', linewidth=1)
    ax.set_xlabel('Risk ratio with 95% confidence interval; below 1 favors ivermectin')
    ax.grid(axis='x', alpha=.16)
    ax.spines[['top', 'right']].set_visible(False)
    fig.text(.05, .94, 'Adding the public COVID-OUT day-28 result', fontsize=18, weight='bold')
    fig.text(.05, .865, f"{result['primary']['n']:,} analysis records; {result['primary']['events_t'] + result['primary']['events_c']} hospital-or-death events", fontsize=12)
    fig.text(.05, .06, 'Targeted aggregate audit. Different cause definitions remain; this is not an equivalence test.', fontsize=10, color='#40515b')
    fig.savefig(OUT / 'trial-audit.png', dpi=160)
    plt.close(fig)

if __name__ == '__main__':
    main()