"""Reproduce the SR-MMP assessment offline from the bundled frozen public inputs.
Run: python reproduce_analysis.py --output tox21_reproduced
No package installation, network access, model download, or application modification.
"""
import argparse, sys, json, hashlib, datetime, importlib
from pathlib import Path
import numpy as np
import pandas as pd

BASE=Path(__file__).resolve().parent
parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--output',default='tox21_reproduced',help='New or empty output directory')
args=parser.parse_args()
OUTPUT_DIR=Path(args.output)
if OUTPUT_DIR.exists() and any(OUTPUT_DIR.iterdir()):
    raise RuntimeError('Output directory must be empty; original results will not be overwritten.')
OUTPUT_DIR.mkdir(parents=True,exist_ok=True)
source_dir=BASE/'sources'
source_records=json.loads((BASE/'source_provenance.json').read_text(encoding='utf-8'))
for record in source_records:
    path=source_dir/record['filename']
    assert hashlib.sha256(path.read_bytes()).hexdigest()==record['sha256'], f"Source checksum mismatch: {record['filename']}"
raw=pd.read_csv(source_dir/'tox21.csv.gz')
versions={'python':sys.version}
for name in ['rdkit','sklearn','pandas','numpy','scipy','matplotlib','requests']:
    versions[name]=getattr(importlib.import_module(name),'__version__','available')
print('Replay software versions:',json.dumps(versions))


import numpy as np
from rdkit import Chem, DataStructs, RDLogger
from rdkit.Chem import rdMolDescriptors, Descriptors, Crippen, Lipinski, rdFingerprintGenerator
from rdkit.Chem.MolStandardize import rdMolStandardize
from rdkit.Chem.Scaffolds import MurckoScaffold
from sklearn.model_selection import StratifiedGroupKFold
from collections import Counter
RDLogger.DisableLog('rdApp.warning')
RDLogger.DisableLog('rdApp.error')
SEED = 20260928
out = OUTPUT_DIR
table_dir = out/'tables'
table_dir.mkdir(parents=True, exist_ok=True)
uncharger = rdMolStandardize.Uncharger()
tautomers = rdMolStandardize.TautomerEnumerator()
allowed_atoms = {1,5,6,7,8,9,14,15,16,17,34,35,53}
audit_records = []
for row_id, row in raw.iterrows():
    rec = {'source_row':int(row_id), 'mol_id':str(row['mol_id']), 'original_smiles':row['smiles'],
           'label':row['SR-MMP'], 'structure_status':'', 'standardized_smiles':'',
           'identity':'', 'scaffold_group':'', 'removed_fragments':0, 'error':''}
    try:
        mol = Chem.MolFromSmiles(str(row['smiles']))
        if mol is None:
            rec['structure_status'] = 'invalid_smiles'
        else:
            clean = rdMolStandardize.Cleanup(mol)
            fragments = Chem.GetMolFrags(clean, asMols=True, sanitizeFrags=True)
            carbon_fragments = [m for m in fragments if any(a.GetAtomicNum()==6 for a in m.GetAtoms())]
            if len(carbon_fragments)==0:
                rec['structure_status'] = 'no_carbon_parent'
            elif len(carbon_fragments)>1:
                rec['structure_status'] = 'multiple_carbon_fragments'
            else:
                parent = carbon_fragments[0]
                rec['removed_fragments'] = len(fragments)-1
                if any(a.GetAtomicNum() not in allowed_atoms for a in parent.GetAtoms()):
                    rec['structure_status'] = 'unsupported_parent_element'
                else:
                    parent = uncharger.uncharge(parent)
                    parent = tautomers.Canonicalize(parent)
                    Chem.SanitizeMol(parent)
                    rec['standardized_smiles'] = Chem.MolToSmiles(parent, isomericSmiles=True)
                    Chem.RemoveStereochemistry(parent)
                    rec['identity'] = Chem.MolToSmiles(parent, isomericSmiles=False)
                    scaffold = MurckoScaffold.MurckoScaffoldSmiles(mol=parent, includeChirality=False)
                    rec['scaffold_group'] = scaffold if scaffold else '__ACYCLIC__'
                    rec['structure_status'] = 'valid'
    except Exception as exc:
        rec['structure_status'] = 'standardization_error'
        rec['error'] = type(exc).__name__+': '+str(exc)
    if pd.isna(rec['label']):
        rec['disposition'] = 'missing_label'
    elif rec['label'] not in [0,1]:
        rec['disposition'] = 'unexpected_label'
    elif rec['structure_status'] != 'valid':
        rec['disposition'] = rec['structure_status']
    else:
        rec['disposition'] = 'eligible'
    audit_records.append(rec)
audit = pd.DataFrame(audit_records)
eligible = audit[audit.disposition=='eligible']
conflicting = set(eligible.groupby('identity').label.nunique().loc[lambda x:x>1].index)
audit.loc[(audit.disposition=='eligible') & audit.identity.isin(conflicting), 'disposition'] = 'conflicting_identity_labels'
eligible = audit[audit.disposition=='eligible']
representatives = eligible.sort_values('source_row').drop_duplicates('identity', keep='first').index
audit.loc[(audit.disposition=='eligible') & ~audit.index.isin(representatives), 'disposition'] = 'consistent_duplicate_collapsed'
audit.loc[representatives, 'disposition'] = 'included_representative'
clean = audit.loc[representatives].copy().sort_values('source_row').reset_index(drop=True)
clean['y'] = clean.label.astype(int)
clean['compound_id'] = clean.identity.map(lambda s:'M_'+hashlib.sha256(s.encode()).hexdigest()[:16])
source_groups = eligible.groupby('identity').agg(source_rows=('source_row',lambda x:';'.join(map(str,x))), mol_ids=('mol_id',lambda x:';'.join(x)),source_record_count=('source_row','size'))
clean = clean.join(source_groups,on='identity')
assert clean.compound_id.is_unique and clean.identity.is_unique
splitter = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=SEED)
clean['outer_fold'] = -1
for fold, (_, indices) in enumerate(splitter.split(np.zeros(len(clean)),clean.y,clean.scaffold_group)):
    clean.loc[indices,'outer_fold'] = fold
clean['split'] = clean.outer_fold.map({0:'train',1:'train',2:'train',3:'calibration',4:'test'})
identity_to_split = clean.set_index('identity')['split']
identity_to_id = clean.set_index('identity')['compound_id']
audit['retained_identity_id'] = audit.identity.map(identity_to_id)
audit['identity_split'] = audit.identity.map(identity_to_split)
split_summary = clean.groupby('split',sort=False).agg(n=('y','size'),active=('y','sum'),scaffolds=('scaffold_group','nunique'))
split_summary['inactive'] = split_summary.n-split_summary.active
split_summary['prevalence'] = split_summary.active/split_summary.n
split_summary['fraction'] = split_summary.n/len(clean)
for a,b in [('train','calibration'),('train','test'),('calibration','test')]:
    assert set(clean.loc[clean.split==a,'identity']).isdisjoint(set(clean.loc[clean.split==b,'identity']))
    assert set(clean.loc[clean.split==a,'scaffold_group']).isdisjoint(set(clean.loc[clean.split==b,'scaffold_group']))
assert (split_summary.active>=20).all() and (split_summary.inactive>=20).all()
assert set(audit.disposition) <= {'missing_label','invalid_smiles','no_carbon_parent','multiple_carbon_fragments','unsupported_parent_element','standardization_error','conflicting_identity_labels','consistent_duplicate_collapsed','included_representative'}
audit.to_csv(table_dir/'row_audit.csv',index=False)
clean.to_csv(table_dir/'split_assignments.csv',index=False)
split_summary.to_csv(table_dir/'split_summary.csv')
audit.groupby(['disposition','structure_status'],dropna=False).size().rename('n').reset_index().to_csv(table_dir/'exclusion_counts.csv',index=False)
audit[audit.identity.isin(conflicting)].to_csv(table_dir/'conflicting_labels.csv',index=False)
print('Disposition counts:',audit.disposition.value_counts().to_dict())
print('Unique identity conflicts:',len(conflicting))
print(split_summary.to_string())
print('Largest scaffold groups:',clean.groupby(['scaffold_group','split']).size().nlargest(6).to_dict())
print('Split checksum:',hashlib.sha256((table_dir/'split_assignments.csv').read_bytes()).hexdigest())


from scipy import sparse
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score, roc_auc_score, brier_score_loss, log_loss, precision_recall_curve
from sklearn.exceptions import ConvergenceWarning
import warnings, time
config = {
    'seed':SEED, 'endpoint':'SR-MMP', 'positive_label':1,
    'fingerprint':{'kind':'Morgan bit vector','radius':2,'bits':2048,'includeChirality':False},
    'C_grid':[0.01,0.1,1.0,10.0], 'cv_folds':4,'selection_metric':'mean fold average precision',
    'class_weight':None,'solver':'liblinear','max_iter':2000,
    'calibration':'logistic sigmoid on independent calibration decision scores, C=1e6',
    'operating_threshold':'maximum calibration-set F1, ties use highest threshold',
    'outer_split':'StratifiedGroupKFold(5, shuffle=True, seed), folds 0-2 train, 3 calibration, 4 test',
    'similarity_bands':[0,0.3,0.5,0.7,1.00000001],
    'coverage_rule':'maximum training Tanimoto at least calibration 10th percentile AND within every training descriptor 1st-99th percentile',
    'bootstrap':{'unit':'scaffold group','replicates':1000,'confidence':0.95},
    'deadline_utc':'2026-09-28T03:25:33.897Z'
}
(out/'analysis_config.json').write_text(json.dumps(config,indent=2),encoding='utf-8')
mols = [Chem.MolFromSmiles(s) for s in clean.identity]
generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048, includeChirality=False)
fps = [generator.GetFingerprint(m) for m in mols]
X = sparse.csr_matrix(np.asarray([generator.GetFingerprintAsNumPy(m) for m in mols],dtype=np.float64))
y = clean.y.to_numpy()
idx_train = np.flatnonzero(clean.split.to_numpy()=='train')
idx_cal = np.flatnonzero(clean.split.to_numpy()=='calibration')
idx_test = np.flatnonzero(clean.split.to_numpy()=='test')
cv_splitter = StratifiedGroupKFold(n_splits=4,shuffle=True,random_state=SEED+1)
cv_folds = list(cv_splitter.split(X[idx_train],y[idx_train],clean.iloc[idx_train].scaffold_group))
cv_assign = clean.iloc[idx_train][['compound_id','scaffold_group','y']].copy()
cv_assign['cv_validation_fold']=-1
for k, (tr, va) in enumerate(cv_folds):
    assert set(clean.iloc[idx_train[tr]].scaffold_group).isdisjoint(set(clean.iloc[idx_train[va]].scaffold_group))
    assert len(np.unique(y[idx_train[tr]]))==2 and len(np.unique(y[idx_train[va]]))==2
    cv_assign.iloc[va,cv_assign.columns.get_loc('cv_validation_fold')]=k
cv_assign.to_csv(table_dir/'training_cv_assignments.csv',index=False)
cv_rows=[]
cv_oof=[]
fit_warnings=[]
started=time.perf_counter()
for C in config['C_grid']:
    for fold,(tr,va) in enumerate(cv_folds):
        estimator=LogisticRegression(C=C,solver='liblinear',max_iter=2000,random_state=SEED)
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter('always')
            estimator.fit(X[idx_train[tr]],y[idx_train[tr]])
        fit_warnings.extend([str(w.message) for w in caught])
        assert not any(issubclass(w.category,ConvergenceWarning) for w in caught)
        pv=estimator.predict_proba(X[idx_train[va]])[:,1]
        cv_rows.append({'C':C,'fold':fold,'n_train':len(tr),'n_validation':len(va),'validation_active':int(y[idx_train[va]].sum()),
                        'average_precision':average_precision_score(y[idx_train[va]],pv),
                        'roc_auc':roc_auc_score(y[idx_train[va]],pv),'brier':brier_score_loss(y[idx_train[va]],pv),
                        'n_iter':int(estimator.n_iter_.max())})
        cv_oof.extend([{'C':C,'fold':fold,'compound_id':clean.iloc[int(i)].compound_id,'y':int(y[i]),'p_raw':float(p)} for i,p in zip(idx_train[va],pv)])
cv_results=pd.DataFrame(cv_rows)
cv_summary=cv_results.groupby('C').agg(mean_AP=('average_precision','mean'),sd_AP=('average_precision','std'),mean_ROC_AUC=('roc_auc','mean'),mean_Brier=('brier','mean')).reset_index()
best_C=float(cv_summary.sort_values(['mean_AP','C'],ascending=[False,True]).iloc[0].C)
model=LogisticRegression(C=best_C,solver='liblinear',max_iter=2000,random_state=SEED)
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter('always')
    model.fit(X[idx_train],y[idx_train])
assert not any(issubclass(w.category,ConvergenceWarning) for w in caught)
fit_warnings.extend([str(w.message) for w in caught])
cal_scores=model.decision_function(X[idx_cal])
calibrator=LogisticRegression(C=1e6,solver='lbfgs',max_iter=2000,random_state=SEED)
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter('always')
    calibrator.fit(cal_scores.reshape(-1,1),y[idx_cal])
assert not any(issubclass(w.category,ConvergenceWarning) for w in caught)
fit_warnings.extend([str(w.message) for w in caught])
cal_prob=calibrator.predict_proba(cal_scores.reshape(-1,1))[:,1]
pr,rc,th=precision_recall_curve(y[idx_cal],cal_prob)
f1=2*pr[:-1]*rc[:-1]/np.maximum(pr[:-1]+rc[:-1],1e-15)
best_f1=float(np.max(f1))
selected_threshold=float(np.max(th[np.isclose(f1,best_f1,rtol=0,atol=1e-12)]))
assert calibrator.coef_[0,0]>0, 'Non-monotone reversal: investigate before testing.'
model_selection={'best_C':best_C,'calibration_intercept':float(calibrator.intercept_[0]),'calibration_slope':float(calibrator.coef_[0,0]),
                 'selected_threshold':selected_threshold,'calibration_F1_at_selected_threshold':best_f1,
                 'training_prevalence':float(y[idx_train].mean()),'fit_seconds':time.perf_counter()-started,
                 'model_n_iter':int(model.n_iter_.max()),'warnings':fit_warnings}
cv_results.to_csv(table_dir/'model_selection_folds.csv',index=False)
cv_summary.to_csv(table_dir/'model_selection_summary.csv',index=False)
pd.DataFrame(cv_oof).to_csv(table_dir/'training_oof_predictions.csv',index=False)
cal_predictions=clean.iloc[idx_cal][['compound_id','identity','scaffold_group','y']].copy()
cal_predictions['decision_score']=cal_scores
cal_predictions['p_raw']=model.predict_proba(X[idx_cal])[:,1]
cal_predictions['p_calibrated']=cal_prob
cal_predictions.to_csv(table_dir/'calibration_predictions.csv',index=False)
pd.DataFrame({'threshold':th,'precision':pr[:-1],'recall':rc[:-1],'F1':f1}).to_csv(table_dir/'calibration_threshold_selection.csv',index=False)
(out/'model_selection.json').write_text(json.dumps(model_selection,indent=2),encoding='utf-8')
print(cv_summary.to_string(index=False))
print(json.dumps(model_selection,indent=2))
print('Test probabilities and test performance have not yet been computed.')


from sklearn.metrics import roc_curve, confusion_matrix
from scipy.special import logit
test_predictions=clean.iloc[idx_test][['compound_id','source_rows','mol_ids','identity','scaffold_group','y']].copy().reset_index(drop=True)
test_score=model.decision_function(X[idx_test])
test_predictions['decision_score']=test_score
test_predictions['p_raw']=model.predict_proba(X[idx_test])[:,1]
test_predictions['p_calibrated']=calibrator.predict_proba(test_score.reshape(-1,1))[:,1]
test_predictions['p_prevalence']=float(y[idx_train].mean())
test_predictions['predicted_at_0_5']=(test_predictions.p_calibrated>=0.5).astype(int)
test_predictions['predicted_at_development_threshold']=(test_predictions.p_calibrated>=selected_threshold).astype(int)
yt=test_predictions.y.to_numpy()
probabilities={k:test_predictions[c].to_numpy() for k,c in [('prevalence','p_prevalence'),('raw','p_raw'),('calibrated','p_calibrated')]}
BIN_EDGES=np.linspace(0,1,11)
def metrics(yy,pp):
    pp=np.clip(pp,1e-12,1-1e-12)
    both=len(np.unique(yy))==2
    ids=np.minimum(np.searchsorted(BIN_EDGES,pp,side='right')-1,9)
    ece=sum(np.sum(ids==j)*abs(yy[ids==j].mean()-pp[ids==j].mean()) for j in range(10) if np.any(ids==j))/len(yy)
    return {'average_precision':float(average_precision_score(yy,pp)) if both else np.nan,
            'roc_auc':float(roc_auc_score(yy,pp)) if both else np.nan,
            'brier':float(np.mean((yy-pp)**2)),
            'log_loss':float(-np.mean(yy*np.log(pp)+(1-yy)*np.log1p(-pp))),
            'ECE_10_fixed_bins':float(ece),
            'mean_probability':float(pp.mean()),
            'observed_prevalence':float(yy.mean())}
def operating_metrics(yy,pp,threshold):
    pred=pp>=threshold
    tn,fp,fn,tp=confusion_matrix(yy,pred,labels=[0,1]).ravel()
    return {'tn':int(tn),'fp':int(fp),'fn':int(fn),'tp':int(tp),
            'precision':float(tp/(tp+fp)) if tp+fp else np.nan,
            'recall':float(tp/(tp+fn)) if tp+fn else np.nan,
            'specificity':float(tn/(tn+fp)) if tn+fp else np.nan,
            'F1':float(2*tp/(2*tp+fp+fn)) if 2*tp+fp+fn else np.nan,
            'balanced_accuracy':float((tp/(tp+fn)+tn/(tn+fp))/2) if (tp+fn)*(tn+fp) else np.nan,
            'predicted_positive_fraction':float(pred.mean())}
point={name:metrics(yt,p) for name,p in probabilities.items()}
groups=test_predictions.scaffold_group.to_numpy()
group_names,group_codes=np.unique(groups,return_inverse=True)
group_indices=[np.flatnonzero(group_codes==i) for i in range(len(group_names))]
rng=np.random.default_rng(SEED+2)
bootstrap_indices=[np.concatenate([group_indices[j] for j in rng.integers(0,len(group_names),size=len(group_names))]) for _ in range(config['bootstrap']['replicates'])]
bootstrap_rows=[]
operating_boot=[]
bin_boot=[]
for b,ii in enumerate(bootstrap_indices):
    for name,p in probabilities.items():
        bootstrap_rows.append({'replicate':b,'model':name,'n':len(ii),'active':int(yt[ii].sum()),**metrics(yt[ii],p[ii])})
        bi=np.minimum(np.searchsorted(BIN_EDGES,p[ii],side='right')-1,9)
        for j in range(10):
            within=bi==j
            if within.any():
                bin_boot.append({'replicate':b,'model':name,'bin':j,'observed_fraction':float(yt[ii][within].mean())})
    for name,threshold_label,threshold in [('raw','0.5',0.5),('calibrated','0.5',0.5),('calibrated','development_F1',selected_threshold)]:
        operating_boot.append({'replicate':b,'model':name,'threshold_label':threshold_label,**operating_metrics(yt[ii],probabilities[name][ii],threshold)})
boot=pd.DataFrame(bootstrap_rows)
boot_operating=pd.DataFrame(operating_boot)
boot_bins=pd.DataFrame(bin_boot)
metric_rows=[]
for name,vals in point.items():
    for metric,value in vals.items():
        samples=boot.loc[boot.model==name,metric].dropna()
        lo,hi=np.quantile(samples,[.025,.975])
        metric_rows.append({'model':name,'metric':metric,'estimate':value,'ci_low':lo,'ci_high':hi,
                            'n_test':len(yt),'n_scaffolds':len(group_names),'bootstrap_valid':len(samples),
                            'bootstrap_invalid':len(bootstrap_indices)-len(samples),'CI_method':'95% percentile scaffold-cluster bootstrap; fixed model'})
test_metrics=pd.DataFrame(metric_rows)
delta_rows=[]
for metric in ['brier','log_loss','ECE_10_fixed_bins']:
    paired=boot.pivot(index='replicate',columns='model',values=metric)
    for comparator in ['raw','prevalence']:
        dif=paired.calibrated-paired[comparator]
        lo,hi=np.quantile(dif,[.025,.975])
        delta_rows.append({'metric':metric,'comparison':'calibrated minus '+comparator,'estimate':point['calibrated'][metric]-point[comparator][metric],
                           'ci_low':lo,'ci_high':hi,'bootstrap_valid':len(dif)})
calibration_differences=pd.DataFrame(delta_rows)
operating_rows=[]
for name,label,threshold in [('raw','0.5',.5),('calibrated','0.5',.5),('calibrated','development_F1',selected_threshold)]:
    op=operating_metrics(yt,probabilities[name],threshold)
    for metric,value in op.items():
        samples=boot_operating.loc[(boot_operating.model==name)&(boot_operating.threshold_label==label),metric].dropna()
        lo,hi=np.quantile(samples,[.025,.975]) if len(samples) else (np.nan,np.nan)
        operating_rows.append({'model':name,'threshold_label':label,'threshold':threshold,'metric':metric,'estimate':value,'ci_low':lo,'ci_high':hi,'bootstrap_valid':len(samples)})
operating_table=pd.DataFrame(operating_rows)
calibration_rows=[]
calibration_diagnostics=[]
for name,p in probabilities.items():
    ids=np.minimum(np.searchsorted(BIN_EDGES,p,side='right')-1,9)
    for j in range(10):
        inside=ids==j
        samples=boot_bins.loc[(boot_bins.model==name)&(boot_bins.bin==j),'observed_fraction']
        lo,hi=np.quantile(samples,[.025,.975]) if len(samples) else (np.nan,np.nan)
        calibration_rows.append({'model':name,'bin':j,'lower':BIN_EDGES[j],'upper':BIN_EDGES[j+1],
                                 'n':int(inside.sum()),'active':int(yt[inside].sum()),
                                 'n_scaffolds':int(len(np.unique(groups[inside]))),
                                 'mean_probability':float(p[inside].mean()) if inside.any() else np.nan,
                                 'observed_fraction':float(yt[inside].mean()) if inside.any() else np.nan,
                                 'ci_low':lo,'ci_high':hi,'bootstrap_valid':len(samples)})
    if name!='prevalence':
        diagnostic=LogisticRegression(C=1e6,solver='lbfgs',max_iter=2000).fit(logit(np.clip(p,1e-12,1-1e-12)).reshape(-1,1),yt)
        calibration_diagnostics.append({'model':name,'intercept_joint':float(diagnostic.intercept_[0]),'slope_joint':float(diagnostic.coef_[0,0]),
                                        'note':'Test-set diagnostic only, not used to change predictions; ideal intercept0 slope1.'})
calibration_bins=pd.DataFrame(calibration_rows)
test_predictions.to_csv(table_dir/'held_out_predictions.csv',index=False)
test_metrics.to_csv(table_dir/'test_metrics.csv',index=False)
calibration_differences.to_csv(table_dir/'calibration_paired_differences.csv',index=False)
operating_table.to_csv(table_dir/'operating_points.csv',index=False)
calibration_bins.to_csv(table_dir/'calibration_bins.csv',index=False)
pd.DataFrame(calibration_diagnostics).to_csv(table_dir/'calibration_diagnostics.csv',index=False)
boot.to_csv(table_dir/'bootstrap_metrics.csv',index=False)
boot_operating.to_csv(table_dir/'bootstrap_operating_points.csv',index=False)
for name,p in probabilities.items():
    fpr,tpr,threshold=roc_curve(yt,p)
    pd.DataFrame({'false_positive_rate':fpr,'true_positive_rate':tpr,'threshold':threshold}).to_csv(table_dir/f'roc_curve_{name}.csv',index=False)
    prec,rec,threshold=precision_recall_curve(yt,p)
    pd.DataFrame({'precision':prec,'recall':rec,'threshold':np.r_[threshold,np.nan]}).to_csv(table_dir/f'pr_curve_{name}.csv',index=False)
saved_pred=pd.read_csv(table_dir/'held_out_predictions.csv')
assert len(saved_pred)==len(idx_test) and saved_pred.compound_id.is_unique
for name,column in [('raw','p_raw'),('calibrated','p_calibrated'),('prevalence','p_prevalence')]:
    recomputed=metrics(saved_pred.y.to_numpy(),saved_pred[column].to_numpy())
    for key,value in recomputed.items(): assert np.isclose(value,point[name][key],rtol=1e-10,atol=1e-12)
assert np.isclose(point['raw']['roc_auc'],point['calibrated']['roc_auc'])
assert np.isclose(point['raw']['average_precision'],point['calibrated']['average_precision'])
print(test_metrics[test_metrics.metric.isin(['average_precision','roc_auc','brier','log_loss'])][['model','metric','estimate','ci_low','ci_high']].to_string(index=False))
print('Paired differences (negative favors calibrated):')
print(calibration_differences.to_string(index=False))
print('Operating points:')
print(operating_table[operating_table.metric.isin(['tp','fp','tn','fn','precision','recall','specificity'])][['model','threshold_label','metric','estimate','ci_low','ci_high']].to_string(index=False))
print('Calibration diagnostics:',calibration_diagnostics)


descriptor_names=['molecular_weight','cLogP','TPSA','HBD','HBA','rotatable_bonds','ring_count','heavy_atoms']
descriptor_values=np.array([[Descriptors.MolWt(m),Crippen.MolLogP(m),rdMolDescriptors.CalcTPSA(m),
                             Lipinski.NumHDonors(m),Lipinski.NumHAcceptors(m),Lipinski.NumRotatableBonds(m),
                             rdMolDescriptors.CalcNumRings(m),m.GetNumHeavyAtoms()] for m in mols],dtype=float)
descriptor_table=pd.DataFrame(descriptor_values,columns=descriptor_names)
descriptor_table.insert(0,'compound_id',clean.compound_id.to_numpy())
descriptor_table.insert(1,'split',clean.split.to_numpy())
descriptor_table.insert(2,'y',y)
bounds=np.quantile(descriptor_values[idx_train],[.01,.99],axis=0)
coverage_rows=[]
for index in np.r_[idx_cal,idx_test]:
    similarities=np.asarray(DataStructs.BulkTanimotoSimilarity(fps[index],[fps[j] for j in idx_train]))
    nearest=int(np.argmax(similarities))
    within=(descriptor_values[index]>=bounds[0])&(descriptor_values[index]<=bounds[1])
    coverage_rows.append({'compound_id':clean.iloc[index].compound_id,'split':clean.iloc[index].split,
                          'max_training_tanimoto':float(similarities[nearest]),
                          'nearest_training_compound_id':clean.iloc[idx_train[nearest]].compound_id,
                          'nearest_training_identity':clean.iloc[idx_train[nearest]].identity,
                          'nearest_training_label':int(y[idx_train[nearest]]),
                          'training_neighbors_at_least_0_5':int(np.sum(similarities>=.5)),
                          'within_all_training_descriptor_bounds':bool(within.all()),
                          'outside_descriptors':';'.join(np.array(descriptor_names)[~within])})
coverage=pd.DataFrame(coverage_rows)
coverage_cutoff=float(np.quantile(coverage.loc[coverage.split=='calibration','max_training_tanimoto'],.10))
coverage['supported_by_screen']=((coverage.max_training_tanimoto>=coverage_cutoff)&coverage.within_all_training_descriptor_bounds)
band_labels=['[0,0.3)','[0.3,0.5)','[0.5,0.7)','[0.7,1]']
coverage['similarity_band']=pd.cut(coverage.max_training_tanimoto,bins=config['similarity_bands'],labels=band_labels,right=False,include_lowest=True).astype(str)
train_id_set=set(clean.iloc[idx_train].compound_id)
assert set(coverage.nearest_training_compound_id)<=train_id_set
assert not set(coverage.compound_id)&train_id_set
assert coverage.max_training_tanimoto.between(0,1).all()
test_annotated=test_predictions.merge(coverage.drop(columns='split'),on='compound_id',validate='one_to_one')
test_annotated=test_annotated.merge(descriptor_table.drop(columns=['split','y']),on='compound_id',validate='one_to_one')
assert test_annotated.compound_id.tolist()==test_predictions.compound_id.tolist()
scope_summary=[]
scope_boot=[]
scope_masks=[]
for band in band_labels:
    scope_masks.append(('similarity',band,test_annotated.similarity_band.eq(band).to_numpy()))
for value in [True,False]:
    scope_masks.append(('coverage_screen','supported' if value else 'flagged',test_annotated.supported_by_screen.eq(value).to_numpy()))
for scheme,label,mask in scope_masks:
    n=int(mask.sum()); active=int(yt[mask].sum()); n_group=len(np.unique(groups[mask]))
    enough=n>=50 and active>=10 and (n-active)>=10 and n_group>=10
    vals=metrics(yt[mask],probabilities['calibrated'][mask]) if n else {k:np.nan for k in point['calibrated']}
    # Ranking estimates are suppressed when class/group counts are too small.
    if not enough:
        vals['average_precision']=np.nan
        vals['roc_auc']=np.nan
    scope_summary.append({'scheme':scheme,'stratum':label,'n':n,'active':active,'inactive':n-active,'n_scaffolds':n_group,
                          'test_fraction':n/len(yt),'ranking_estimates_reportable':enough,**vals})
    for b,ii in enumerate(bootstrap_indices):
        jj=ii[mask[ii]]
        if len(jj):
            bvals=metrics(yt[jj],probabilities['calibrated'][jj])
            if not enough:
                bvals['average_precision']=np.nan; bvals['roc_auc']=np.nan
            scope_boot.append({'replicate':b,'scheme':scheme,'stratum':label,**bvals})
scope_boot=pd.DataFrame(scope_boot)
scope_summary=pd.DataFrame(scope_summary)
scope_metric_rows=[]
for row in scope_summary.itertuples(index=False):
    for key in point['calibrated']:
        samples=scope_boot.loc[(scope_boot.scheme==row.scheme)&(scope_boot.stratum==row.stratum),key].dropna()
        lo,hi=np.quantile(samples,[.025,.975]) if len(samples) else (np.nan,np.nan)
        scope_metric_rows.append({'scheme':row.scheme,'stratum':row.stratum,'metric':key,'estimate':getattr(row,key),
                                 'ci_low':lo,'ci_high':hi,'n':row.n,'active':row.active,'n_scaffolds':row.n_scaffolds,
                                 'bootstrap_valid':len(samples),'bootstrap_invalid':len(bootstrap_indices)-len(samples)})
scope_metrics=pd.DataFrame(scope_metric_rows)
descriptor_bounds=pd.DataFrame({'descriptor':descriptor_names,'training_q01':bounds[0],'training_q99':bounds[1],
                                 'training_min':descriptor_values[idx_train].min(axis=0),'training_max':descriptor_values[idx_train].max(axis=0)})
descriptor_summary=descriptor_table.groupby('split')[descriptor_names].describe(percentiles=[.01,.1,.5,.9,.99])
descriptor_table.to_csv(table_dir/'molecular_descriptors.csv',index=False)
descriptor_bounds.to_csv(table_dir/'training_descriptor_bounds.csv',index=False)
descriptor_summary.to_csv(table_dir/'descriptor_summary.csv')
coverage.to_csv(table_dir/'coverage_annotations.csv',index=False)
test_annotated.to_csv(table_dir/'held_out_predictions_annotated.csv',index=False)
scope_summary.to_csv(table_dir/'applicability_summary.csv',index=False)
scope_metrics.to_csv(table_dir/'applicability_metrics.csv',index=False)
scope_boot.to_csv(table_dir/'bootstrap_applicability.csv',index=False)
rule={'training_reference_n':len(idx_train),'calibration_similarity_q10':coverage_cutoff,
      'descriptor_bounds_quantiles':[.01,.99],
      'definition':'Supported by heuristic screen iff similarity >= calibration q10 AND all descriptors within training q01-q99.',
      'limitation':'Heuristic flags, not validated guarantees. Low-coverage predictions remain in all primary metrics.',
      'ranking_minimums':{'n':50,'active':10,'inactive':10,'scaffolds':10}}
(out/'coverage_rule.json').write_text(json.dumps(rule,indent=2),encoding='utf-8')
test_annotated['error_category']=np.where(yt==1,np.where(test_annotated.predicted_at_development_threshold==1,'true_positive','false_negative'),np.where(test_annotated.predicted_at_development_threshold==1,'false_positive','true_negative'))
# Deterministic descriptive examples: most confident wrong/right predictions within each category.
examples=pd.concat([frame.sort_values('p_calibrated',ascending=(category in ['false_negative','true_negative'])).head(3) for category,frame in test_annotated.groupby('error_category')],ignore_index=True)
examples.to_csv(table_dir/'representative_predictions.csv',index=False)
excluded_label_counts=audit.groupby(['disposition','label'],dropna=False).size().rename('n').reset_index()
excluded_label_counts.to_csv(table_dir/'exclusion_counts_by_label.csv',index=False)
print('Coverage heuristic:',json.dumps(rule))
print(scope_summary[['scheme','stratum','n','active','n_scaffolds','average_precision','roc_auc','brier','log_loss']].to_string(index=False))
print('Training descriptor envelopes:')
print(descriptor_bounds.to_string(index=False))
print('Nearest-training similarity quantiles:',test_annotated.max_training_tanimoto.quantile([0,.1,.25,.5,.75,.9,1]).to_dict())
print('Exact fingerprint similarities of 1:',int((test_annotated.max_training_tanimoto>=1).sum()))


exact_matches=test_annotated[test_annotated.max_training_tanimoto>=1-1e-12].copy()
exact_matches.to_csv(table_dir/'cross_split_identical_fingerprints.csv',index=False)
sensitivity_rows=[]
for label,mask in [('full_test',np.ones(len(yt),dtype=bool)),('exclude_identical_training_fingerprints',test_annotated.max_training_tanimoto.to_numpy()<1-1e-12)]:
    for name,p in probabilities.items():
        sensitivity_rows.append({'analysis':label,'model':name,'n':int(mask.sum()),'active':int(yt[mask].sum()),**metrics(yt[mask],p[mask])})
pd.DataFrame(sensitivity_rows).to_csv(table_dir/'fingerprint_overlap_sensitivity.csv',index=False)
print('Identical-fingerprint pairs:')
print(exact_matches[['compound_id','identity','nearest_training_identity','y','nearest_training_label','p_calibrated']].to_string(index=False))
print('Sensitivity:')
print(pd.DataFrame(sensitivity_rows)[['analysis','model','n','active','average_precision','roc_auc','brier']].to_string(index=False))
print('Calibration bin counts:',calibration_bins[calibration_bins.model=='calibrated'][['bin','n','active','n_scaffolds','mean_probability','observed_fraction']].to_string(index=False))
print('Coverage uncertainty:')
print(scope_metrics[scope_metrics.metric.isin(['average_precision','roc_auc','brier'])][['scheme','stratum','metric','estimate','ci_low','ci_high']].to_string(index=False))


alias_rows=[]
id_to_index={cid:i for i,cid in enumerate(clean.compound_id)}
for row in exact_matches.itertuples():
    a=id_to_index[row.compound_id]; b=id_to_index[row.nearest_training_compound_id]
    fa=generator.GetSparseCountFingerprint(mols[a]).GetNonzeroElements()
    fb=generator.GetSparseCountFingerprint(mols[b]).GetNonzeroElements()
    alias_rows.append({'compound_id':row.compound_id,'nearest_training_compound_id':row.nearest_training_compound_id,
                      'test_smiles':row.identity,'training_smiles':row.nearest_training_identity,
                      'unfolded_feature_presence_equal':set(fa)==set(fb),'unfolded_feature_counts_equal':fa==fb,
                      'test_heavy_atoms':mols[a].GetNumHeavyAtoms(),'training_heavy_atoms':mols[b].GetNumHeavyAtoms()})
pd.DataFrame(alias_rows).to_csv(table_dir/'fingerprint_alias_diagnostics.csv',index=False)
print(pd.DataFrame(alias_rows).drop(columns=['test_smiles','training_smiles']).to_string(index=False))


import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from rdkit.Chem import Draw
figure_dir=out/'figures'
figure_dir.mkdir(parents=True,exist_ok=True)
BLUE='#0072B2'; ORANGE='#D55E00'; GRAY='#555555'
style={'font.size':10,'axes.titlesize':12,'axes.labelsize':11,'legend.fontsize':9,
       'axes.spines.top':False,'axes.spines.right':False,'pdf.fonttype':42,'ps.fonttype':42,
       'savefig.facecolor':'white'}
def export_figure(fig,name):
    fig.savefig(figure_dir/(name+'.png'),dpi=300,facecolor='white')
    fig.savefig(figure_dir/(name+'.pdf'),facecolor='white')
    plt.close(fig)
with plt.rc_context(style):
    fig,axes=plt.subplots(1,2,figsize=(10.5,4.6),layout='constrained')
    prtab=pd.read_csv(table_dir/'pr_curve_calibrated.csv')
    roctab=pd.read_csv(table_dir/'roc_curve_calibrated.csv')
    axes[0].step(prtab.recall,prtab.precision,where='post',color=BLUE,lw=1.6,label=f'Logistic model: AP = {point["calibrated"]["average_precision"]:.3f}')
    axes[0].axhline(yt.mean(),color=GRAY,ls='--',label=f'Test prevalence = {yt.mean():.3f}')
    axes[0].set(xlabel='Recall (sensitivity)',ylabel='Precision',xlim=(0,1),ylim=(0,1.02),title='A  Precision–recall')
    axes[0].legend(loc='upper right')
    axes[1].plot(roctab.false_positive_rate,roctab.true_positive_rate,color=BLUE,lw=1.8,label=f'ROC AUC = {point["calibrated"]["roc_auc"]:.3f}')
    axes[1].plot([0,1],[0,1],color=GRAY,ls='--',label='Random ranking')
    axes[1].set(xlabel='False-positive rate',ylabel='True-positive rate',xlim=(0,1),ylim=(0,1.02),title='B  Receiver operating characteristic')
    axes[1].legend(loc='lower right')
    fig.suptitle(f'SR-MMP: scaffold-held-out test (n = {len(yt)}, active = {yt.sum()})\nRaw and calibrated probabilities have identical rankings',fontsize=12)
    export_figure(fig,'discrimination')
with plt.rc_context(style):
    fig,axes=plt.subplots(2,2,figsize=(10.5,7.0),gridspec_kw={'height_ratios':[3,1]},layout='constrained')
    for j,(name,color,marker) in enumerate([('raw',ORANGE,'s'),('calibrated',BLUE,'o')]):
        bins=calibration_bins[calibration_bins.model==name]
        nonempty=bins.n>0
        reliable=(bins.n>=20)&(bins.n_scaffolds>=10)
        sparsebin=nonempty&~reliable
        ax=axes[0,j]
        ax.plot([0,1],[0,1],color=GRAY,ls='--',lw=1,label='Ideal agreement')
        # NaN rows leave explicit gaps at empty probability bins.
        ax.plot(bins.mean_probability,bins.observed_fraction,color=color,lw=1)
        z=bins[reliable]
        ax.errorbar(z.mean_probability,z.observed_fraction,
                    yerr=np.vstack([z.observed_fraction-z.ci_low,z.ci_high-z.observed_fraction]).clip(0),
                    fmt=marker,color=color,capsize=3,ms=5,label='Bin rate; 95% scaffold CI')
        z=bins[sparsebin]
        ax.plot(z.mean_probability,z.observed_fraction,marker='^',ls='none',mec=color,mfc='white',ms=7,label='Sparse bin; interval omitted')
        ax.set(xlabel='Mean predicted probability',ylabel='Observed active fraction',xlim=(0,1),ylim=(0,1.02),
               title=('A  Raw model' if name=='raw' else 'B  After sigmoid calibration'))
        ax.legend(loc='upper left',fontsize=8)
        axes[1,j].bar((BIN_EDGES[:-1]+BIN_EDGES[1:])/2,bins.n.to_numpy(),width=.085,color=color,alpha=.8,edgecolor='black',linewidth=.4)
        axes[1,j].set(xlabel='Predicted probability bin',ylabel='Compounds',xlim=(0,1))
        for k,n in enumerate(bins.n):
            if n: axes[1,j].text((BIN_EDGES[k]+BIN_EDGES[k+1])/2,n+3,str(n),ha='center',va='bottom',fontsize=7)
        axes[1,j].set_ylim(0,max(bins.n)*1.22)
    fig.suptitle('SR-MMP probability reliability on the untouched test set\nFixed-width bins; sparse = fewer than 20 compounds or 10 scaffolds',fontsize=12)
    export_figure(fig,'calibration')
with plt.rc_context(style):
    fig,axes=plt.subplots(1,3,figsize=(14,4.9),layout='constrained')
    for split,color,ls in [('calibration',ORANGE,'--'),('test',BLUE,'-')]:
        vals=coverage.loc[coverage.split==split,'max_training_tanimoto']
        axes[0].hist(vals,bins=np.linspace(0,1,21),histtype='step',density=False,lw=1.7,color=color,ls=ls,label=f'{split.capitalize()} (n={len(vals)})')
    axes[0].axvline(coverage_cutoff,color=GRAY,ls=':',lw=1.5,label=f'Calibration 10th percentile\n({coverage_cutoff:.3f})')
    axes[0].set(xlabel='Maximum training Tanimoto similarity',ylabel='Compounds per 0.05 bin',xlim=(0,1),title='A  Similarity to training chemistry')
    axes[0].legend(fontsize=8)
    x=np.arange(len(band_labels))
    strata=scope_summary[scope_summary.scheme=='similarity'].set_index('stratum').loc[band_labels]
    for ax,metric,title in [(axes[1],'average_precision','B  Ranking within similarity bands'),(axes[2],'brier','C  Probability error within bands')]:
        tab=scope_metrics[(scope_metrics.scheme=='similarity')&(scope_metrics.metric==metric)].set_index('stratum').loc[band_labels]
        valid=tab.estimate.notna()
        ax.errorbar(x[valid],tab.estimate[valid],yerr=np.vstack([tab.estimate-tab.ci_low,tab.ci_high-tab.estimate])[:,valid].clip(0),
                    fmt='o',color=BLUE,capsize=4,ls='none',label='Calibrated model; 95% CI')
        if metric=='average_precision':
            ax.scatter(x,strata.observed_prevalence,marker='x',s=45,color=GRAY,label='Within-band prevalence')
            ax.set_ylim(0,1)
        else:
            ax.set_ylim(bottom=0)
        ax.set_xticks(x,[f'{label}\nn={int(row.n)}, a={int(row.active)}' for label,row in zip(band_labels,strata.itertuples())],fontsize=8)
        ax.set(xlabel='Maximum training similarity band (a = actives)',ylabel='Average precision' if metric=='average_precision' else 'Brier score (lower is better)',title=title)
        ax.legend(fontsize=8,loc='best')
    fig.suptitle('Chemical coverage on the scaffold-held-out test set\nIntervals use 1,000 scaffold bootstrap samples; similarity does not guarantee accuracy',fontsize=12)
    export_figure(fig,'chemical_coverage')
example_mols=[Chem.MolFromSmiles(s) for s in examples.identity]
example_legends=[f'{r.error_category.replace("_"," ")} | {r.compound_id}\ny={r.y}; p={r.p_calibrated:.2f}; similarity={r.max_training_tanimoto:.2f}' for r in examples.itertuples()]
example_image=Draw.MolsToGridImage(example_mols,molsPerRow=3,subImgSize=(450,330),legends=example_legends,returnPNG=False)
example_image.save(str(figure_dir/'illustrative_compounds.png'))
print('Exported figures:',[p.name for p in figure_dir.iterdir()])



# Export all fitted scalar parameters without executable model serialization.
pd.DataFrame({'fingerprint_bit':np.arange(X.shape[1]),'coefficient':model.coef_[0]}).to_csv(table_dir/'logistic_coefficients.csv',index=False)
linear_parameters={'intercept':float(model.intercept_[0]),'calibration_intercept':float(calibrator.intercept_[0]),'calibration_slope':float(calibrator.coef_[0,0])}
(out/'linear_parameters.json').write_text(json.dumps(linear_parameters,indent=2),encoding='utf-8')
# Verify the reproduced experiment against the supplied frozen outputs.

reference_split=pd.read_csv(BASE/'tables'/'split_assignments.csv')
pd.testing.assert_frame_equal(reference_split,pd.read_csv(table_dir/'split_assignments.csv'),check_dtype=False,rtol=1e-12,atol=1e-12)
reference_predictions=pd.read_csv(BASE/'tables'/'held_out_predictions_annotated.csv')
reproduced_predictions=pd.read_csv(table_dir/'held_out_predictions_annotated.csv')
pd.testing.assert_frame_equal(reference_predictions,reproduced_predictions,check_dtype=False,rtol=1e-9,atol=1e-10)
reference_metrics=pd.read_csv(BASE/'tables'/'test_metrics.csv')
pd.testing.assert_frame_equal(reference_metrics,pd.read_csv(table_dir/'test_metrics.csv'),check_dtype=False,rtol=1e-9,atol=1e-10)
checks={'source_checksums':'passed','split_assignments':'passed','held_out_predictions':'passed','test_metrics_and_intervals':'passed','software_versions':versions}
(OUTPUT_DIR/'replay_checks.json').write_text(json.dumps(checks,indent=2),encoding='utf-8')
print('REPLAY CHECKS PASSED')
