from pathlib import Path
import math,csv,json,sys,hashlib
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
P=Path(__file__).resolve().parent
SEEDS=list(range(41001,41011))
N_TRAIN=N_TEST=5000
FEATURES=6
TRUE_BETA=np.array([-.25,.9,-.6,.45,.25,-.3,.2])
L2=.01
ECE_BINS=10
def sigmoid(z: np.ndarray) -> np.ndarray:
    z = np.clip(z, -30.0, 30.0)
    return 1.0 / (1.0 + np.exp(-z))

def make_data(rng: np.random.Generator, n: int) -> tuple[np.ndarray, np.ndarray]:
    x = rng.normal(size=(n, FEATURES))
    design = np.column_stack((np.ones(n), x))
    p = sigmoid(design @ TRUE_BETA)
    y = rng.binomial(1, p).astype(float)
    return design, y

def make_shifted_data(rng: np.random.Generator, n: int) -> tuple[np.ndarray, np.ndarray]:
    """Synthetic covariate/prevalence shift; deliberately uses no people data."""
    x = rng.normal(loc=0.75, scale=1.25, size=(n, FEATURES))
    design = np.column_stack((np.ones(n), x))
    shifted_beta = TRUE_BETA.copy()
    shifted_beta[0] -= 0.8
    p = sigmoid(design @ shifted_beta)
    y = rng.binomial(1, p).astype(float)
    return design, y

def introduce_missingness(x: np.ndarray, rng: np.random.Generator, rate: float = 0.15) -> np.ndarray:
    """MCAR feature loss with mean-imputation based only on training data."""
    out = x.copy()
    mask = rng.random(out[:, 1:].shape) < rate
    features = out[:, 1:]
    features[mask] = np.nan
    return out

def impute_from_train(train: np.ndarray, test: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    means = np.nanmean(train[:, 1:], axis=0)
    train_out = train.copy()
    test_out = test.copy()
    train_out[:, 1:] = np.where(np.isnan(train_out[:, 1:]), means, train_out[:, 1:])
    test_out[:, 1:] = np.where(np.isnan(test_out[:, 1:]), means, test_out[:, 1:])
    return train_out, test_out

def fit_logistic(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    beta = np.zeros(x.shape[1], dtype=float)
    penalty = np.diag([0.0] + [L2] * (x.shape[1] - 1))
    for _ in range(100):
        p = sigmoid(x @ beta)
        gradient = x.T @ (p - y) + penalty @ beta
        weights = np.maximum(p * (1.0 - p), 1e-8)
        hessian = (x.T * weights) @ x + penalty + np.eye(x.shape[1]) * 1e-10
        step = np.linalg.solve(hessian, gradient)
        beta -= step
        if np.max(np.abs(step)) < 1e-8:
            break
    return beta

def ece(y: np.ndarray, p: np.ndarray) -> float:
    bins = np.linspace(0.0, 1.0, ECE_BINS + 1)
    total = 0.0
    for i in range(ECE_BINS):
        mask = (p >= bins[i]) & ((p < bins[i + 1]) if i < ECE_BINS - 1 else (p <= bins[i + 1]))
        if np.any(mask):
            total += float(mask.mean()) * abs(float(y[mask].mean()) - float(p[mask].mean()))
    return total

def metrics(y: np.ndarray, p: np.ndarray) -> dict[str, float]:
    p = np.clip(p, 1e-9, 1.0 - 1e-9)
    return {
        "brier": float(np.mean((p - y) ** 2)),
        "log_loss": float(-np.mean(y * np.log(p) + (1.0 - y) * np.log(1.0 - p))),
        "ece_10_equal_width_bins": ece(y, p),
    }
baseline=json.loads((P/'numeric-baseline.json').read_text());rows=[];maxerror=0
for idx,seed in enumerate(SEEDS):
 rng=np.random.default_rng(seed);xt,yt=make_data(rng,N_TRAIN);x,y=make_data(rng,N_TEST);fit=fit_logistic(xt,yt);base=metrics(y,np.full(N_TEST,float(yt.mean())));candidate=metrics(y,sigmoid(x@fit));negative=metrics(y,sigmoid(x@fit_logistic(xt,rng.permutation(yt))))
 old=baseline['per_seed'][idx]
 for key,obj in [('baseline_train_prevalence',base),('logistic_candidate',candidate),('permuted_label_negative_control',negative)]:
  for metric,val in obj.items():maxerror=max(maxerror,abs(val-old[key][metric]))
 for label,obj in [('matched',candidate),('permuted_labels',negative)]:rows.append(dict(seed=seed,scenario=label,brier=obj['brier'],log_loss=obj['log_loss'],ece=obj['ece_10_equal_width_bins'],training_prevalence_baseline_brier=base['brier'],brier_skill=1-obj['brier']/base['brier']))
 rng=np.random.default_rng(seed+100000);xt,yt=make_data(rng,N_TRAIN);x,y=make_shifted_data(rng,N_TEST);candidate=metrics(y,sigmoid(x@fit_logistic(xt,yt)));base=metrics(y,np.full(N_TEST,float(yt.mean())));xtm=introduce_missingness(xt,rng);xm=introduce_missingness(x,rng);xtm,xm=impute_from_train(xtm,xm);missing=metrics(y,sigmoid(xm@fit_logistic(xtm,yt)));old=baseline['stress_test_per_seed'][idx]
 for key,obj in [('complete_model_on_shifted_test',candidate),('prevalence_baseline_on_shifted_test',base),('missing_model_on_shifted_test',missing)]:
  for metric,val in obj.items():maxerror=max(maxerror,abs(val-old[key][metric]))
 for label,obj in [('shifted',candidate),('shifted_missing',missing)]:rows.append(dict(seed=seed,scenario=label,brier=obj['brier'],log_loss=obj['log_loss'],ece=obj['ece_10_equal_width_bins'],training_prevalence_baseline_brier=base['brier'],brier_skill=1-obj['brier']/base['brier']))
assert maxerror<1e-12
summary=[]
labels=['matched','permuted_labels','shifted','shifted_missing']
for label in labels:
 rr=[r for r in rows if r['scenario']==label];summary.append(dict(scenario=label,median_brier_skill=float(np.median([r['brier_skill']for r in rr])),median_ece=float(np.median([r['ece']for r in rr]))))
for name,data in [('seed-metrics.csv',rows),('summary.csv',summary)]:
 with(P/name).open('w',newline='',encoding='utf8')as f:w=csv.DictWriter(f,fieldnames=list(data[0]));w.writeheader();w.writerows(data)
plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False,'savefig.dpi':180})
fig,axes=plt.subplots(1,2,figsize=(11,5),layout='constrained')
for ax,key,title in zip(axes,['brier_skill','ece'],['A  Brier improvement versus training prevalence','B  Ten-bin calibration discrepancy']):
 for i,label in enumerate(labels):
  vals=[r[key]for r in rows if r['scenario']==label];ax.scatter([i+(j-4.5)*.025 for j in range(10)],vals,s=20,alpha=.7);ax.plot([i-.18,i+.18],[np.median(vals)]*2,color='black',lw=2)
 ax.set(xticks=range(4),xticklabels=['Matched','Permuted\nlabels','Shifted','Shifted +\nmissing'],title=title,ylabel='Fraction'if key=='brier_skill'else 'Weighted absolute discrepancy');ax.grid(axis='y',alpha=.2)
axes[0].axhline(0,color='gray',ls='--');fig.suptitle('A better aggregate score need not mean calibrated probabilities\nSynthetic software rerun: ten seeds per scenario; black marks are medians')
for e in ('png','svg','pdf'):fig.savefig(P/f'calibration-comparison.{e}')
plt.close(fig)
(P/'validation.json').write_text(json.dumps(dict(all_pass=True,seed_rows=40,summary_rows=4,matched_original_metric_comparisons=180,max_original_metric_difference=maxerror,original_untouched=True),indent=2)+'\n')
(P/'study.json').write_text(json.dumps(dict(runtime=dict(python=sys.version,numpy=np.__version__,matplotlib=matplotlib.__version__),generator_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),baseline_sha256=hashlib.sha256((P/'numeric-baseline.json').read_bytes()).hexdigest(),original_result_sha256=baseline['original_result_sha256'],original_script_sha256=baseline['original_script_sha256'],seeds=SEEDS,train_rows=5000,test_rows=5000,features=6,synthetic_rerun=True,baseline='constant training-set prevalence, evaluated on each test set'),indent=2)+'\n')
print(json.dumps(dict(maxerror=maxerror,summary=summary)))