#!/usr/bin/env python3
"""COMPLEXA C360: exact synthetic reconstruction, Python standard library only.

Not clinical data. Counts and potential outcomes are deterministic and prescribed in
chapter 360; no random seed is needed. Re-execution verifies the published tables.
Run: python reproduce.py --out output
"""
from __future__ import annotations
import argparse, collections, csv, json
from pathlib import Path

# Each tuple: baseline severity L; (Y0,Y1); population count; program share.
TYPES = [
    (0, 1, 1, 350, 1/5), (0, 0, 1, 50, 1/5), (0, 0, 0, 100, 1/5),
    (1, 1, 1, 100, 4/5), (1, 0, 1, 100, 4/5), (1, 0, 0, 300, 4/5),
]

def write_csv(path, rows, fields):
    with open(path, 'w', encoding='utf-8', newline='') as f:
        w=csv.DictWriter(f,fieldnames=fields,lineterminator='\n');w.writeheader();w.writerows(rows)

def generate():
    rows=[]
    for l,y0,y1,total,share in TYPES:
        count1=int(round(total*share))
        assert count1==total*share, 'Program assignment must have integer cell counts'
        # Subjects are enumerated within each type; assignment is exact-stratified,
        # rather than pretending to be a random realization.
        for i in range(total):
            a=int(i<count1); y=y1 if a else y0
            rows.append(dict(id=f'C360-{len(rows)+1:04d}', L=l,A=a,Y=y,Y0=y0,Y1=y1))
    assert len(rows)==1000
    return rows

def main(out):
    out.mkdir(parents=True,exist_ok=True)
    rows=generate()
    observed=[{k:r[k] for k in ('id','L','A','Y')} for r in rows]
    truepot=[{k:r[k] for k in ('id','L','Y0','Y1')} for r in rows]
    altpot=[dict(id=r['id'],L=r['L'],Y0=r['Y'],Y1=r['Y']) for r in rows]
    weights=[]
    for r in rows:
        e=0.2 if r['L']==0 else 0.8
        a=r['A']
        weights.append(dict(id=r['id'],L=r['L'],A=a,Y=r['Y'],
          propensity=f'{e:.6f}', weight_ate=f'{1/e if a else 1/(1-e):.6f}',
          weight_att=f'{1 if a else e/(1-e):.6f}',
          weight_atc=f'{(1-e)/e if a else 1:.6f}'))
    write_csv(out/'observados.csv',observed,['id','L','A','Y'])
    write_csv(out/'potenciales_mecanismo_efecto.csv',truepot,['id','L','Y0','Y1'])
    write_csv(out/'potenciales_mecanismo_nulo.csv',altpot,['id','L','Y0','Y1'])
    write_csv(out/'pesos_individuales.csv',weights,['id','L','A','Y','propensity','weight_ate','weight_att','weight_atc'])
    counts=collections.Counter((r['L'],r['A'],r['Y']) for r in rows)
    groups={}
    for l in (0,1):
        for a in (0,1):
            sub=[r for r in rows if r['L']==l and r['A']==a]
            groups[f'L{l}_A{a}']={'n':len(sub),'positive':sum(r['Y'] for r in sub),'prob':sum(r['Y'] for r in sub)/len(sub)}
    def mean(field,where=lambda r:True):
        xs=[r for r in rows if where(r)];return sum(r[field] for r in xs)/len(xs)
    aobs=mean('Y',lambda r:r['A']==1)
    cobs=mean('Y',lambda r:r['A']==0)
    true_ate=mean('Y1')-mean('Y0')
    att=mean('Y1',lambda r:r['A']==1)-mean('Y0',lambda r:r['A']==1)
    atc=mean('Y1',lambda r:r['A']==0)-mean('Y0',lambda r:r['A']==0)
    p={l:{a:groups[f'L{l}_A{a}']['prob'] for a in (0,1)} for l in (0,1)}
    std1=(p[0][1]+p[1][1])/2;std0=(p[0][0]+p[1][0])/2
    results={'observed_program':aobs,'observed_control':cobs,'observed_difference':aobs-cobs,
             'standardized_program':std1,'standardized_control':std0,'ate':std1-std0,
             'att':att,'atc':atc,'actual_effect_mean':true_ate,'null_mechanism_effect':0.0,
             'totals':{'n':len(rows),'potential_positive_program':sum(r['Y1'] for r in rows),
                       'potential_positive_control':sum(r['Y0'] for r in rows),
                       'observed_positive_program':sum(r['Y'] for r in rows if r['A']==1),
                       'observed_positive_control':sum(r['Y'] for r in rows if r['A']==0)},
             'groups':groups,'generation':'Deterministic exact counts; no random seed or Monte Carlo sampling'}
    for key,val in {'observed_difference':-.12,'ate':.15,'att':.18,'atc':.12,
                    'standardized_program':.60,'standardized_control':.45}.items():
        assert abs(results[key]-val)<1e-12,(key,results[key],val)
    assert all(r['Y']==(r['Y1'] if r['A'] else r['Y0']) for r in rows)
    assert all(r['Y']==p['Y0']==p['Y1'] for r,p in zip(rows,altpot))
    # Standardization from observed columns only, no access to truth in this estimator.
    weights_table=[]
    for target,weight in [('ATE','weight_ate'),('ATT','weight_att'),('ATC','weight_atc')]:
        t=[]
        for a in (0,1):
            select=[(r,w) for r,w in zip(rows,weights) if r['A']==a]
            weighted=sum(r['Y']*float(w[weight]) for r,w in select)/sum(float(w[weight]) for _,w in select)
            t.append(weighted)
        weights_table.append(dict(target=target, control=round(t[0],8), program=round(t[1],8),difference=round(t[1]-t[0],8)))
    assert all(abs(r['difference']-{'ATE':.15,'ATT':.18,'ATC':.12}[r['target']])<1e-10 for r in weights_table)
    write_csv(out/'estimandos.csv',weights_table,['target','control','program','difference'])
    (out/'comprobacion.json').write_text(json.dumps(results,ensure_ascii=False,indent=2,sort_keys=True)+'\n',encoding='utf8')
    print('PASS: C360 1000 synthetic rows; ATE +0.15, ATT +0.18, ATC +0.12, raw −0.12')
if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--out',type=Path,default=Path('output'));args=p.parse_args();main(args.out)