"""Plot Week 8's original access/use and offer-environment examples.

Run after the lab smoke, from the book directory:
    conda run -n research python assets/scripts/build-week8-voice.py

Exact exported CSVs are preserved in assets/data. No article estimates are used.
Local quit elasticities are derived from the plotted probabilities under the
specified logistic scale 10 and fixed arrival probabilities .20 and .04.
"""
from pathlib import Path
import csv
import math
import os
import sys

BOOK=Path(__file__).resolve().parents[2]
REPO=BOOK.parents[1]
LAB=BOOK/'labs/08-wage-setting-retention-outside-options-and-worker-voice'
MODE='SYNTHETIC_TEACHING_CONSTRUCTION_NOT_PUBLISHED_REPLICATION'
os.environ.setdefault('MPLCONFIGDIR',str(BOOK/'.mpl-cache'))
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.ticker import PercentFormatter
sys.path.insert(0,str(REPO))
from shared.figure_style import COLORS,apply_style,save_figure,style_axis


def read_and_preserve(name,target):
    source=LAB/'output/diagnose'/name
    path=BOOK/'assets/data'/target
    path.parent.mkdir(parents=True,exist_ok=True)
    path.write_bytes(source.read_bytes())
    with path.open(newline='') as stream:rows=list(csv.DictReader(stream))
    assert rows and all(r['provenance']==MODE for r in rows)
    return rows


def save(fig,stem):
    path=BOOK/'assets/figures'/stem
    save_figure(fig,path.with_suffix('.png'))
    fig.savefig(path.with_suffix('.pdf'),bbox_inches='tight',facecolor='white',
                metadata={'CreationDate':None,'ModDate':None})
    plt.close(fig)


def main():
    rows=read_and_preserve('access_option_curve.csv','08-access-option-curve.csv')
    quit_rows=read_and_preserve('quit_curves.csv','08-quit-curves.csv')
    shares=[float(r['potential_user_share']) for r in rows]
    itt=[float(r['access_ITT']) for r in rows]
    wald=[float(r['invalid_Wald_target']) for r in rows]
    assert shares==[.05,.1,.2,.4,.8]
    for r,p,a,v in zip(rows,shares,itt,wald):
        assert float(r['direct_access_effect'])==.6 and float(r['use_effect'])==.8
        assert math.isclose(a,.6+.8*p,abs_tol=1e-9)
        assert math.isclose(v,.8+.6/p,abs_tol=1e-9)
    focal=shares.index(.2)
    assert math.isclose(itt[focal],.76) and math.isclose(wald[focal],3.8)
    wages=[float(r['wage']) for r in quit_rows]
    assert wages==[10+k/2 for k in range(41)]
    probabilities={k:[float(r['quit_'+k]) for r in quit_rows] for k in ['A','B']}
    arrivals={'A':.20,'B':.04}
    locations={'A':20-10*math.log(9),'B':20}
    elasticities={}
    for env in ['A','B']:
        for w,q in zip(wages,probabilities[env]):
            cdf=1/(1+math.exp(-(w-locations[env])/10))
            assert math.isclose(q,arrivals[env]*(1-cdf),abs_tol=1e-9)
        elasticities[env]=[-w*(1-q/arrivals[env])/10 for w,q in zip(wages,probabilities[env])]
    center=wages.index(20)
    assert all(math.isclose(probabilities[e][center],.02,abs_tol=1e-9) for e in ['A','B'])
    assert math.isclose(elasticities['A'][center],-1.8,abs_tol=1e-8)
    assert math.isclose(elasticities['B'][center],-1,abs_tol=1e-8)

    apply_style()
    fig,axes=plt.subplots(1,2,figsize=(10.5,5.6))
    for ax,values,title,ylabel,color,ylim in [
        (axes[0],itt,'Access effect in the population','Access ITT (score units per worker)',COLORS['navy'],(0,1.45)),
        (axes[1],wald,'Invalid interpretation as a use effect','Wald ratio (score units per unit use)',COLORS['rust'],(0,15))]:
        # CSV points anchor the curve; interpolate with the specified model,
        # not straight segments that distort its reciprocal term.
        grid=[shares[0]+k*(shares[-1]-shares[0])/300 for k in range(301)]
        direct=float(rows[0]['direct_access_effect']);use=float(rows[0]['use_effect'])
        fitted=[direct+use*p if ax is axes[0] else use+direct/p for p in grid]
        ax.plot(grid,fitted,color=color,linewidth=2.2)
        ax.scatter(shares,values,color=color,s=25,zorder=3)
        ax.scatter([.2],[values[focal]],color=COLORS['teal'],s=65,zorder=4)
        ax.axvline(.2,color=COLORS['muted'],linestyle=':',linewidth=1)
        ax.set(title=title,xlabel='Possible-user share / take-up under access',ylabel=ylabel,
               xlim=(0,.85),ylim=ylim,xticks=[.05,.2,.4,.6,.8])
        ax.xaxis.set_major_formatter(PercentFormatter(1,decimals=0))
        style_axis(ax)
    axes[0].axhline(.6,color=COLORS['muted'],linestyle='--',linewidth=1.2)
    axes[0].text(.43,.49,'Direct access component = 0.6',fontsize=9,color=COLORS['muted'])
    axes[0].annotate('20%: ITT = 0.76',(.2,.76),xytext=(.31,1.1),fontsize=10,
                     arrowprops={'arrowstyle':'->','color':COLORS['teal']})
    axes[0].text(.06,1.3,r'$ITT(p)=0.6+0.8p$',fontsize=11)
    axes[1].axhline(.8,color=COLORS['teal'],linestyle='--',linewidth=1.4)
    axes[1].text(.34,1.22,'Controlled-use parameter = 0.8',fontsize=9,color=COLORS['teal'])
    axes[1].annotate('20%: ratio = 3.8',(.2,3.8),xytext=(.32,7.2),fontsize=10,
                     arrowprops={'arrowstyle':'->','color':COLORS['teal']})
    axes[1].text(.34,12.9,r'$Wald(p)=0.8+0.6/p$',fontsize=11)
    fig.text(.5,.086,'Direct access violates exclusion. Separate vertical scales; the controlled-use parameter is not identified by assignment.',
             ha='center',fontsize=9)
    fig.text(.5,.027,'Original synthetic teaching example • population targets, not estimated mediation',ha='center',fontsize=10,color=COLORS['muted'])
    fig.tight_layout(rect=(0,.14,1,1),w_pad=3)
    save(fig,'08-access-and-use')

    fig,axes=plt.subplots(1,2,figsize=(10.5,5.6))
    for env,color,style in [('A',COLORS['navy'],'-'),('B',COLORS['teal'],'--')]:
        label=f"{env}: arrival {arrivals[env]:.0%}"
        axes[0].plot(wages,probabilities[env],color=color,linestyle=style,linewidth=2.3,label=label)
        axes[1].plot(wages,elasticities[env],color=color,linestyle=style,linewidth=2.3,label=label)
        axes[1].scatter([20],[elasticities[env][center]],color=color,s=40,zorder=4)
        axes[1].text(20.65,elasticities[env][center]+.045,f"{env}: {elasticities[env][center]:.1f}",fontsize=10,color=color)
    axes[0].scatter([20],[.02],color=COLORS['rust'],s=55,zorder=4)
    axes[0].annotate('Both: 2% at wage 20',(20,.02),xytext=(15,.037),fontsize=10,
                     arrowprops={'arrowstyle':'->','color':COLORS['rust']})
    axes[0].set(title='Equal quit levels at the starting wage',ylabel='One-period voluntary quit probability',ylim=(0,.056))
    axes[0].yaxis.set_major_formatter(PercentFormatter(1,decimals=0))
    axes[1].set(title='Different local wage responses',ylabel='Local wage elasticity of quitting',ylim=(-3.15,0))
    for ax in axes:
        ax.axvline(20,color=COLORS['muted'],linestyle=':',linewidth=1)
        ax.set(xlabel='Current hourly wage (invented currency units)',xlim=(10,30),xticks=[10,15,20,25,30])
        ax.legend(loc='lower left',frameon=False,fontsize=9)
        style_axis(ax)
    fig.text(.5,.1,'Offer arrival and logistic distributions stay fixed; scale = 10. At wage 20: F(A) = 0.90, F(B) = 0.50.',ha='center',fontsize=9.5)
    fig.text(.5,.055,'Quit elasticities are not employer labor-supply elasticities.',ha='center',fontsize=9.5)
    fig.text(.5,.009,'Original synthetic teaching example • no published estimates',ha='center',fontsize=10,color=COLORS['muted'])
    fig.tight_layout(rect=(0,.16,1,1),w_pad=3)
    save(fig,'08-quit-environments')
    print('PASS: access/use targets and quit environments checked; exact two CSVs and two PNG/PDF figures saved.')


if __name__=='__main__':
    main()