#!/usr/bin/env python3
"""Experiment 3: PR histograms for 1/t, 8/t, and t^(-rho) mean estimation.

Theory is printed and saved before simulation. All methods share each
replication's sample stream and average pre-update SGD iterates.
"""
import argparse
import json
import time
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from pr_histograms import optimum, summarize


def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--reps',type=int,default=10000)
    p.add_argument('--horizons',type=int,nargs='+',default=[50000,200000,500000,1000000])
    p.add_argument('--rho',type=float,default=.6)
    p.add_argument('--alpha',type=float,default=.65)
    p.add_argument('--step-scale',type=float,default=.5)
    p.add_argument('--initial',type=float,default=2.)
    p.add_argument('--mean',type=float,default=0.)
    p.add_argument('--seed',type=int,default=20261018)
    p.add_argument('--output',type=Path,default=Path(__file__).resolve().parent/'results'/'experiment_3')
    args=p.parse_args()
    args.horizons=sorted(set(args.horizons))
    if args.reps<2 or min(args.horizons)<1 or not .5<args.alpha<1 or not max(.5,2*args.alpha-1)<args.rho<1 or args.step_scale<=0:
        p.error('Require positive horizons/step-scale, reps >= 2, 0.5 < alpha < 1, and max(0.5,2*alpha-1) < rho < 1.')
    args.output.mkdir(parents=True,exist_ok=True)
    star=optimum(args.mean)
    h=float(1+1/np.cosh(star)**2)
    names=['Running average','8/t recursive average',f't^(-{args.rho:g}) recursive average']
    variances=[2/h**2,(16/15)/h**2,1/h**2]
    config={k:str(v) if isinstance(v,Path) else v for k,v in vars(args).items()}
    theory={'H':h,'theta_star':star,'methods':names,'pr_variances':variances,
            'single_step_mse_rates':['t^-1','t^-1',f't^-{args.rho:g}']}
    (args.output/'theory.json').write_text(json.dumps({'config':config,'theory':theory},indent=2)+'\n')
    print('Theoretical PR variances (before simulation):',flush=True)
    for name,v in zip(names,variances): print(f'  {name}: {v:.8f}',flush=True)
    rng=np.random.default_rng(args.seed)
    theta=np.full((3,args.reps),args.initial)
    mhat=np.zeros_like(theta)
    total=np.zeros_like(theta)
    draws={}
    rows=[]
    start=time.perf_counter()
    for t in range(1,max(args.horizons)+1):
        total+=theta
        x=rng.normal(args.mean,1.,args.reps)
        gamma=np.array([1/t,8/t,t**(-args.rho)])[:,None]
        mhat+=gamma*(x[None,:]-mhat)
        if t in args.horizons:
            z=np.sqrt(t)*(total/t-star)
            draws[f'scaled_pr_T{t}']=z.copy()
            for i,name in enumerate(names):
                row={'T':t,'method':name,'theory_variance':variances[i],
                     **summarize(z[i]),'single_step_mse':float(np.mean((mhat[i]-args.mean)**2))}
                rows.append(row)
                print(f'T={t:,}: {name}, empirical variance={row["variance"]:.6f}',flush=True)
        theta-=args.step_scale*t**(-args.alpha)*(theta+np.tanh(theta)-mhat)
        if t%50000==0: print(f'Progress {t:,}/{max(args.horizons):,}, elapsed {time.perf_counter()-start:.1f}s',flush=True)
    report={'config':config,'theory':theory,'results':rows,'elapsed_seconds':time.perf_counter()-start,
            'numpy_version':np.__version__,'matplotlib_version':matplotlib.__version__}
    (args.output/'summary.json').write_text(json.dumps(report,indent=2)+'\n')
    np.savez_compressed(args.output/'scaled_pr_draws.npz',**draws)
    plt.rcParams.update({'font.family':'serif','font.size':10,'pdf.fonttype':42,
                         'axes.spines.top':False,'axes.spines.right':False})
    last=draws[f'scaled_pr_T{max(args.horizons)}']
    extent=max(float(np.abs(last).max()),4.3*np.sqrt(max(variances)))*1.03
    bins=np.linspace(-extent,extent,66)
    grid=np.linspace(-extent,extent,1000)
    fig,axes=plt.subplots(1,3,figsize=(12.6,3.8),sharex=True,sharey=True)
    titles=[r'Running average: $1/t$',r'Recursive average: $8/t$',rf'Recursive average: $t^{{-{args.rho:g}}}$']
    for ax,z,v,title,color in zip(axes,last,variances,titles,['#3569a8','#c46b35','#35846c']):
        ax.hist(z,bins=bins,density=True,color=color,alpha=.57,edgecolor='white',linewidth=.3,label='Monte Carlo')
        ax.plot(grid,np.exp(-grid**2/(2*v))/np.sqrt(2*np.pi*v),color=color,linewidth=2,label='Asymptotic normal')
        ax.text(.97,.96,f'Empirical variance: {z.var(ddof=1):.4f}\nAsymptotic variance: {v:.4f}',
                transform=ax.transAxes,ha='right',va='top',fontsize=8,
                bbox={'facecolor':'white','alpha':.85,'edgecolor':'none'})
        ax.set_title(title)
        ax.set_xlabel(r'$\sqrt{T}(\bar\theta_T-\theta^*)$')
        ax.set_xlim(-extent,extent)
        ax.legend(loc='upper left',fontsize=7,frameon=False)
    axes[0].set_ylabel('Density')
    fig.suptitle(f'Attaining the downstream variance bound | T = {max(args.horizons):,} | {args.reps:,} replications',fontsize=12)
    fig.tight_layout()
    for ext in ['png','pdf']: fig.savefig(args.output/f'attainment_histograms.{ext}',dpi=220,bbox_inches='tight')
    plt.close(fig)
    print(f'Output: {args.output}',flush=True)


if __name__=='__main__': main()