#!/usr/bin/env python3
"""Reproducible exact-sufficient-statistics Monte Carlo for rate-parameter exponential MLE.
Each worker is partitioned into four quarters; sums are sampled from their exact Gamma law.
No data movement/communication/network benchmark is claimed.
"""
import csv, numpy as np
from pathlib import Path
import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt
ROOT=Path(__file__).parent
SEED=20261028
N=2**20
THETA=1.0
REPS=500
M_LIST=(256,64,16)
rng=np.random.default_rng(SEED)

def newton(x, mle):
    return 2.0*x-x*x/mle

rows=[]; summary=[]
for m in M_LIST:
    K=N//m
    obs={key: [] for key in ('raw','one_step_raw','two_step_raw','jack2','one_step_jack2','jack3','one_step_jack3')}
    for rep in range(REPS):
        # Gamma is a *sum* of m/4 exponential(rate=theta) observations.
        quarters=rng.gamma(shape=m//4, scale=1/THETA,size=(K,4))
        halves=quarters.reshape(K,2,2).sum(axis=2)
        totals=halves.sum(axis=1)
        full_mle=N/np.sum(totals)
        local_full=m/totals
        local_halves=(m/2)/halves
        local_quarts=(m/4)/quarters
        raw=float(local_full.mean())
        J2=float((2*local_full - local_halves.mean(axis=1)).mean())
        J3=float(((8/3)*local_full -2*local_halves.mean(axis=1) +(1/3)*local_quarts.mean(axis=1)).mean())
        values={'raw':raw,
          'one_step_raw':newton(raw,full_mle),
          'two_step_raw':newton(newton(raw,full_mle),full_mle),
          'jack2':J2,'one_step_jack2':newton(J2,full_mle),
          'jack3':J3,'one_step_jack3':newton(J3,full_mle)}
        for key,val in values.items():
            delta=np.sqrt(N)*(val-full_mle)
            obs[key].append(delta)
            rows.append((m,K,rep,key,delta))
    for name, data in obs.items():
        a=np.array(data)
        summary.append(dict(m=m,K=K,method=name,mean_scaled_error=float(a.mean()),
                mc_standard_error=float(a.std(ddof=1)/np.sqrt(REPS)),
                rmse_scaled_error=float(np.sqrt(np.mean(a*a))),
                p95_abs_scaled_error=float(np.quantile(np.abs(a),.95))))
    print('m',m,'K',K,[(k,round(np.mean(v),6)) for k,v in obs.items()],flush=True)
with open(ROOT/'simulation_replicates.csv','w',newline='') as f:
    w=csv.writer(f); w.writerow(['m','K','rep','method','scaled_error_relative_to_centralized_mle']); w.writerows(rows)
with open(ROOT/'simulation_summary.csv','w',newline='') as f:
    w=csv.DictWriter(f,fieldnames=list(summary[0]));w.writeheader();w.writerows(summary)

palette={'one_step_raw':'#A55040','two_step_raw':'#9B7C5A','one_step_jack2':'#176F86','one_step_jack3':'#407B58'}
labels={'one_step_raw':'Raw averaging + 1 Newton','two_step_raw':'Raw averaging + 2 Newton','one_step_jack2':'Half-block jackknife + 1 Newton','one_step_jack3':'Quarter-block jackknife + 1 Newton'}
fig,axs=plt.subplots(1,2,figsize=(11.5,4.0),layout='constrained')
ks=np.array([N//m for m in M_LIST])
for method in palette:
    vals=[next(r for r in summary if r['K']==int(k) and r['method']==method) for k in ks]
    axs[0].plot(ks,[abs(x['mean_scaled_error']) for x in vals],'-o',lw=2,ms=5,
                color=palette[method],label=labels[method])
    axs[1].plot(ks,[x['rmse_scaled_error'] for x in vals],'-o',lw=2,ms=5,color=palette[method])
for ax in axs:
    ax.set_xscale('log',base=2)
    ax.set_yscale('log')
    ax.set_xticks(ks)
    ax.set_xticklabels([f'{k:,}' for k in ks])
    ax.grid(alpha=.2,which='both')
    ax.set_xlabel('Number of partitions K')
axs[0].set_ylabel('Absolute mean scaled gap')
axs[1].set_ylabel('RMSE of scaled gap')
axs[0].set_title('Bias relative to centralized MLE')
axs[1].set_title('Error relative to centralized MLE')
fig.legend(loc='outside lower center',bbox_to_anchor=(.5,-.09),ncol=2,frameon=False,fontsize=9)
fig.savefig(ROOT/'figures'/'mc_comparison.pdf',bbox_inches='tight')
fig.savefig(ROOT/'figures'/'mc_comparison.png',dpi=170,bbox_inches='tight')
plt.close(fig)
print('wrote outputs')