"""Online Resource 1. Simulation code for
"Small numbers, big bubbles: Gambler's fallacy and hot-hand traders in a heterogeneous-agent market".
Reproduces Figs. 1-4 and Tables 2-3. Requires numpy, scipy, matplotlib.
Run: python ESM_1_simulation_code.py
"""
import numpy as np
from scipy import stats
def simulate(Tn=50, T=7, gamma=2., theta=0.07, tau=0.9, M=100., N=10000., s2=2.5,
             F0=10., F1=60., liq=None, kappa=0.0):
    """Price path of the LSN bubble model.
    tau = share of momentum traders among LSN traders.
    Momentum belief: lam_m(t) = (M + dP_{t-1})/(2M), dP_{t-1}=P_{t-1}-P_{t-2}.
    Contrarian belief: lam_c(t) = lam_c(t-1) - dP_{t-1}/(2N) + kappa*(0.5-lam_c(t-1)).
    """
    F=np.array([F0 if t<T else F1 for t in range(Tn)],float)
    XL=np.zeros(Tn)
    if liq: 
        for t,v in liq.items(): XL[t]=v
    P=np.zeros(Tn); lm=np.full(Tn,0.5); lc=np.full(Tn,0.5)
    Xf=np.zeros(Tn); Xm=np.zeros(Tn); Xc=np.zeros(Tn)
    d=lambda l:(2*l-1)/(4*gamma*l*(1-l))
    for t in range(Tn):
        if t>=2:
            dp=P[t-1]-P[t-2]
            lm[t]=(M+dp)/(2*M)
            lc[t]=lc[t-1]-dp/(2*N)+kappa*(0.5-lc[t-1])
        elif t==1:
            lc[t]=lc[0]; lm[t]=0.5
        Xm[t]=d(lm[t]); Xc[t]=d(lc[t])
        P[t]=F[t]+gamma*s2/theta*(XL[t]+(1-theta)*(tau*Xm[t]+(1-tau)*Xc[t]))
        Xf[t]=(F[t]-P[t])/(gamma*s2)
    V=0.5*(theta*abs(Xf)+(1-theta)*tau*abs(Xm)+(1-theta)*(1-tau)*abs(Xc)+abs(XL))
    return dict(P=P,F=F,B=P-F,V=V,lm=lm,lc=lc,Xm=Xm,Xc=Xc,Xf=Xf)
def a_coef(theta=0.07,tau=0.9,M=100.,s2=2.5): return s2*(1-theta)*tau/(theta*M)

def stoch(Tn=20000, gamma=2., theta=0.07, tau=0.9, M=100., N=10000., s2=2.5, sig_nu=1.0, sig_eps=0.02, kappa=0.05, seed=1):
    rng=np.random.default_rng(seed)
    F=10+np.cumsum(rng.normal(0,sig_nu,Tn)); XL=rng.normal(0,sig_eps,Tn)
    P=np.zeros(Tn); lc=0.5; V=np.zeros(Tn)
    d=lambda l:(2*l-1)/(4*gamma*l*(1-l))
    P[0]=F[0]; P[1]=F[1]
    for t in range(2,Tn):
        dp=P[t-1]-P[t-2]; dp=np.clip(dp,-0.99*M,0.99*M)
        lm=(M+dp)/(2*M); lc=lc-dp/(2*N)+kappa*(0.5-lc)
        Xm=d(lm); Xc=d(lc)
        P[t]=F[t]+gamma*s2/theta*(XL[t]+(1-theta)*(tau*Xm+(1-tau)*Xc))
        Xf=(F[t]-P[t])/(gamma*s2)
        V[t]=0.5*(theta*abs(Xf)+(1-theta)*tau*abs(Xm)+(1-theta)*(1-tau)*abs(Xc)+abs(XL[t]))
    r=np.diff(P)[100:]; dF=np.diff(F)[100:]; B=(P-F)[100:]
    ac=lambda x,k: np.corrcoef(x[:-k],x[k:])[0,1]
    return dict(ac=[round(ac(r,k),3) for k in (1,2,3,4)], acF=[round(ac(dF,k),3) for k in (1,2)],
        sdr=round(r.std(),3), sdF=round(dF.std(),3), kurt=round(stats.kurtosis(r),2),
        corrVB=round(np.corrcoef(V[101:],abs(B[1:]))[0,1],3), corrVr=round(np.corrcoef(V[101:],abs(r))[0,1],3),
        acabs=[round(ac(abs(r),k),3) for k in (1,5)])

import json, matplotlib
matplotlib.use('Agg'); import matplotlib.pyplot as plt
plt.rcParams.update({'font.size':9,'axes.spines.top':False,'axes.spines.right':False})
K=0.05
out={}
def two_panel(r, fname, title_note):
    t=np.arange(len(r['P']))
    fig,ax=plt.subplots(2,1,figsize=(6.85,4.6),sharex=True,gridspec_kw={'height_ratios':[2,1]})
    ax[0].plot(t,r['F'],color='0.35',lw=1.4,ls='--',label='Fundamental value $F_t$')
    ax[0].plot(t,r['P'],color='#1f4e9c',lw=1.6,marker='o',ms=2.8,label='Price $P_t$')
    ax[0].set_ylabel('Value'); ax[0].legend(frameon=False,loc='upper right')
    ax[1].bar(t,r['V'],color='#d9541e',width=0.7)
    ax[1].set_ylabel('Volume $V_t$'); ax[1].set_xlabel('Period $t$')
    ax[0].set_xlim(0,40)
    fig.tight_layout(); fig.savefig(f'{fname}.tif',dpi=600,pil_kwargs={'compression':'tiff_lzw'}); fig.savefig(f'{fname}.png',dpi=300); plt.close(fig)
base=simulate(kappa=K); seq=simulate(s2=3.65,kappa=K); liq=simulate(F1=10.,liq={9:0.25},kappa=K)
two_panel(base,'Fig1','base'); two_panel(seq,'Fig2','seq'); two_panel(liq,'Fig3','liq')
for k,r in (('base',base),('seq',seq),('liq',liq)):
    out[k]={'B':[round(x,2) for x in r['B'][5:25]],'V':[round(x,3) for x in r['V'][5:25]],'lm':[round(x,3) for x in r['lm'][5:15]],'lc':[round(x,5) for x in r['lc'][5:15]]}
# comparative statics
fig,ax=plt.subplots(1,2,figsize=(6.85,2.9))
th=np.linspace(0.04,0.4,200)
for tau,ls in ((0.5,':'),(0.7,'--'),(0.9,'-')):
    pk=[simulate(theta=x,tau=tau,kappa=K)['B'].max() for x in th]
    ax[0].plot(th,pk,ls=ls,color='#1f4e9c',label=f'$\\tau$ = {tau}')
ax[0].set_xlabel('Share of fundamentalists $\\theta$'); ax[0].set_ylabel('Peak deviation $P_t-F_t$'); ax[0].legend(frameon=False)
s2g=np.linspace(0.5,12,300); thg=np.linspace(0.02,0.25,300)
S,TH=np.meshgrid(s2g,thg); A=S*(1-TH)*0.9/(TH*100)
cs=ax[1].contourf(S,TH,A<1,levels=[-0.5,0.5,1.5],colors=['#f2c4b1','#cfe0f5'])
ax[1].contour(S,TH,A,levels=[1],colors='k',linewidths=1)
ax[1].plot([2.5],[0.07],'o',color='#1f4e9c'); ax[1].annotate('baseline',(2.5,0.07),xytext=(1.0,0.13),fontsize=8,arrowprops=dict(arrowstyle='-',lw=0.6))
ax[1].plot([3.65],[0.07],'s',color='#1f4e9c'); ax[1].annotate('sequential bubbles',(3.65,0.07),xytext=(3.4,0.17),fontsize=8,arrowprops=dict(arrowstyle='-',lw=0.6))
ax[1].text(9.0,0.2,'stable ($a<1$)',ha='center',fontsize=8); ax[1].text(10.0,0.035,'unstable ($a>1$)',ha='center',fontsize=8)
ax[1].set_xlabel('Fundamental variance $\\sigma_\\nu^2$'); ax[1].set_ylabel('Share of fundamentalists $\\theta$')
fig.tight_layout(); fig.savefig('Fig4.tif',dpi=600,pil_kwargs={'compression':'tiff_lzw'}); fig.savefig('Fig4.png',dpi=300); plt.close(fig)
# comparative statics table
rows=[]
for lab,kw in [('Baseline',{}),('theta=0.14',{'theta':0.14}),('tau=0.5',{'tau':0.5}),('s2=3.65',{'s2':3.65}),('M=200',{'M':200.}),('dF=25',{'F1':35.})]:
    r=simulate(kappa=K,**kw); B=r['B']; a=a_coef(theta=kw.get('theta',0.07),tau=kw.get('tau',0.9),M=kw.get('M',100.),s2=kw.get('s2',2.5))
    neg=B.min(); 
    # periods until |B|<0.5 for good
    idx=[t for t in range(len(B)) if abs(B[t])>=0.5]; settle=(max(idx)-7) if idx else 0
    rows.append(dict(case=lab,a=round(a,3),peak=round(B.max(),2),trough=round(neg,2),settle=settle,peakV=round(r['V'].max(),3),
        hl=round(np.log(0.5)/np.log(np.sqrt(a)),2), period=round(2*np.pi/np.arccos(np.sqrt(a)/2),2)))
out['cs']=rows
# stochastic
st={}
for lab,kw in [('Baseline',{}),('s2=3.65',{'s2':3.65}),('tau=0.5',{'tau':0.5}),('theta=0.2',{'theta':0.2})]:
    s2=kw.get('s2',2.5); st[lab]=stoch(sig_nu=np.sqrt(s2),kappa=K,**kw)
out['stoch']=st
json.dump(out,open('results.json','w'),indent=1,default=float)
print(json.dumps(out['cs'],indent=0)); print(json.dumps(st,default=float))
print('base B',out['base']['B'][:12]); print('seq B',out['seq']['B'][:12]); print('liq B',out['liq']['B'][:12])
print('base V',out['base']['V'][:8]); print('lm',out['base']['lm'], 'lc',out['base']['lc'])
