"""Exact finite-table sensitivity analysis. No clinical or missingness model is fitted."""
import json
from fractions import Fraction as F
from pathlib import Path
ROOT=Path(__file__).resolve().parents[1]
def validate_arm(a):
    e,n,N=(a[k] for k in ['recovered','observed','cohort'])
    if any(type(x) is not int for x in [e,n,N]) or not 0<=e<=n<=N or n==0:
        raise ValueError('Require integer 0 <= recovered <= observed <= cohort, observed > 0')
    return e,n,N,N-n

def calculate(row):
    a,n,A,m=validate_arm(row['treatment']);b,k,B,l=validate_arm(row['control'])
    rd=lambda q,r:F(a,A)+F(m,A)*q-F(b,B)-F(l,B)*r
    lo,hi=rd(0,1),rd(1,0)
    common=sorted([rd(0,0),rd(1,1)])
    control_anchor=F(b,k)
    if m:
        treatment_tip=(F(A)*(F(b,B)+F(l,B)*control_anchor)-a)/m
    else:treatment_tip=None
    boundaries=[]
    for x in range(m+1):
        threshold=F(B*(a+x),A)-b
        minimum=-(-threshold.numerator//threshold.denominator)
        feasible=max(0,minimum) if minimum<=l else None
        boundaries.append({'treatment_missing_recovered':x,'control_minimum_for_nonpositive':feasible})
    exact={'respondent_rd':F(a,n)-F(b,k),'lower':lo,'upper':hi,
           'common_lower':common[0],'common_upper':common[1],
           'control_anchor':control_anchor,'treatment_tipping':treatment_tip,
           'anchor_difference':None if treatment_tip is None else control_anchor-treatment_tip}
    res={'cohort':row['cohort'],'month':row['month'],'n_t':A,'n_c':B,
         'observed_t':n,'observed_c':k,'recovered_t':a,'recovered_c':b,
         'unavailable_t':m,'unavailable_c':l,'availability_t':n/A,'availability_c':k/B,
         'respondent_recovery_t':a/n,'respondent_recovery_c':b/k,
         'contains_zero':lo<=0<=hi,'common_all_positive':common[0]>0,
         'anchor_feasible':treatment_tip is not None and 0<=treatment_tip<=1,
         'exact':{key:None if v is None else str(v) for key,v in exact.items()},
         **{key:None if v is None else float(v) for key,v in exact.items()},
         'null_frontier':None if l==0 else {'intercept':float(F(B*a-A*b,A*l)),'slope':float(F(B*m,A*l)),
             'q_t_min':float(max(F(0),F(A*b-B*a,B*m))) if m else None,
             'q_t_max':float(min(F(1),F(A*(b+l)-B*a,B*m))) if m else None},
         'completion_pairs':(m+1)*(l+1)}
    return res,boundaries

def main():
    invalid=[{'recovered':-1,'observed':1,'cohort':2}, {'recovered':2,'observed':1,'cohort':2}, {'recovered':0,'observed':3,'cohort':2}, {'recovered':0,'observed':0,'cohort':2}, {'recovered':0.5,'observed':1,'cohort':2}, {'recovered':True,'observed':1,'cohort':2}]
    for arm in invalid:
        try: validate_arm(arm)
        except ValueError: pass
        else: raise AssertionError('invalid input accepted')
    rows=json.loads((ROOT/'data/extraction.json').read_text())
    results={};boundaries={};display={}
    for row in rows:
        v,b=calculate(row);results[row['id']]=v;boundaries[row['id']]=b
        display[row['id']]={key:round(v[key]*100,3) for key in ['respondent_rd','lower','upper','common_lower','common_upper','control_anchor','treatment_tipping','anchor_difference','availability_t','availability_c']}
    r={'rows':results,'row_count':len(results),'input_guard_checks':len(invalid),'all_bounds_include_zero':all(v['contains_zero'] for v in results.values()),'all_common_positive':all(v['common_all_positive'] for v in results.values()),'total_completion_pairs':sum(v['completion_pairs'] for v in results.values())}
    for name,data in [('R1.json',r),('R2.json',display),('integer-frontiers.json',boundaries)]:
        (ROOT/'results'/name).write_text(json.dumps(data,indent=2)+'\n')
    print(json.dumps(display,indent=2))
if __name__=='__main__':main()