# SCRIPT: BH-FIGS-V8-4
# Figure f18 (a potential that stiffens toward coplanarity), from the formulas and sampling of BH-QUTRIT-18.
# (a) the square lifted at fixed moment of inertia, in the signed coordinate sigma3 = 2h: gravity and like charges,
#     exact curves against the wall parabolas (2 sqrt2 - 1) sigma3^2 and -(2 sqrt2 - 1) sigma3^2; both curves are even,
#     and the ends sigma3 = +-1/sqrt3 are the two mirror-image regular tetrahedra (same G, fidelity 1).
# (b) dU/dw3 at the wall over 400 random planar configurations (seeded), gravity and like charges.
import os as _os
_HERE=_os.path.dirname(_os.path.abspath(__file__))
_FIG=_os.path.join(_HERE,'..','paper','fig') if _os.path.isdir(_os.path.join(_HERE,'..','paper','fig')) else _os.path.join(_HERE,'..','fig')
import numpy as np, matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
OUT=_FIG
BG='#0f1115'; INK='#e6e8ec'; MUT='#9aa3ae'; RULE='#2a313c'
C1='#3987e5'; C2='#d95926'; C3='#199e70'; C4='#c98500'
plt.rcParams.update({'figure.facecolor':BG,'axes.facecolor':BG,'savefig.facecolor':BG,'text.color':INK,
  'axes.labelcolor':INK,'xtick.color':MUT,'ytick.color':MUT,'axes.edgecolor':RULE,'font.family':'serif',
  'font.serif':['TeX Gyre Pagella','URW Palladio L','Palatino','DejaVu Serif'],'mathtext.fontset':'cm','font.size':11})
def clean(ax):
    for s in ['top','right']: ax.spines[s].set_visible(False)
k=2*np.sqrt(2)-1
def pair_sum(h):
    a=np.sqrt(0.25-h*h); q=[np.array([a,0,h]),np.array([-a,0,h]),np.array([0,a,-h]),np.array([0,-a,-h])]
    return sum(1/np.linalg.norm(q[i]-q[j]) for i in range(4) for j in range(i+1,4))
s=np.linspace(-1/np.sqrt(3),1/np.sqrt(3),801)          # sigma3 = 2h (w3 = 4h^2 while 4h^2 <= 2a^2)
S0=pair_sum(0.0)
Ug=np.array([-(pair_sum(x/2)-S0) for x in s]); Uc=-Ug
# check: evenness and the wall slope, printed
print('max |U(s)-U(-s)| =', np.abs(Ug-Ug[::-1]).max())
print('U/sigma3^2 near the wall =', Ug[400+1]/s[400+1]**2, 'vs', k)
print('tetrahedron rise =', Ug[-1], '(exact 2+4sqrt2-6sqrt(3/2) =', 2+4*np.sqrt(2)-6*np.sqrt(1.5), ')')
# (b) the BH-QUTRIT-18 sampler, same seed sequence position is not needed: a fresh seeded sample of 400
rng=np.random.default_rng(1804)
def bodies(X):
    r1,r2,r3=X[:,0],X[:,1],X[:,2]
    return [-r1/np.sqrt(2)-r2/np.sqrt(6)-r3/np.sqrt(12), r1/np.sqrt(2)-r2/np.sqrt(6)-r3/np.sqrt(12),
            2*r2/np.sqrt(6)-r3/np.sqrt(12), 3*r3/np.sqrt(12)]
def S(X):
    q=bodies(X); return sum(1/np.linalg.norm(q[i]-q[j]) for i in range(4) for j in range(i+1,4))
dg=[]
for _ in range(400):
    P=rng.normal(size=(3,3)); P[2,:]=0; P/=np.sqrt(np.trace(P.T@P))
    u,sv,vt=np.linalg.svd(P); n=u[:,2]; h=1e-4
    Xp=P+h*np.outer(n,vt[2]); Xp/=np.sqrt(np.trace(Xp.T@Xp)); w3=np.linalg.eigvalsh(Xp.T@Xp)[0]
    dg.append(-(S(Xp)-S(P))/w3)
dg=np.array(dg); print('gravity dU/dw3: min', dg.min(), 'fraction > 0', np.mean(dg>0))
fig=plt.figure(figsize=(11.6,4.6)); ax=fig.add_axes([0.065,0.14,0.41,0.74]); bx=fig.add_axes([0.575,0.14,0.40,0.74])
for a in (ax,bx): clean(a)
ax.plot(s,Ug,color=C1,lw=2.4,label='gravity')
ax.plot(s,Uc,color=C2,lw=2.4,label='like charges')
ax.plot(s,k*s**2,color=C1,lw=1.0,ls=(0,(4,3))); ax.plot(s,-k*s**2,color=C2,lw=1.0,ls=(0,(4,3)))
ax.axvline(0,color=RULE,lw=1.2,ls=(0,(4,3))); ax.axhline(0,color=RULE,lw=0.8)
for x in (-1/np.sqrt(3),1/np.sqrt(3)): ax.plot(x,Ug[-1],'o',color=C1,ms=5)
ax.text(0.02,0.335,'wall: the flat square',color=INK,fontsize=9.5)
ax.text(-0.575,Ug[-1]+0.035,'tetrahedron',color=MUT,fontsize=9); ax.text(0.40,Ug[-1]+0.035,'its mirror',color=MUT,fontsize=9)
ax.text(-0.24,0.385,r'dashed: $\pm(2\sqrt{2}-1)\,\sigma_3^2$',color=MUT,fontsize=9.5)
ax.set_ylim(-0.42,0.42)
ax.set_xlabel(r'signed distance from the wall, $\sigma_3$'); ax.set_ylabel(r'$U-U_{\rm square}$ at fixed $I$')
ax.legend(loc='center left',frameon=False,fontsize=9.5,labelcolor=INK)
ax.set_title(r'a  the square lifted: even in $\sigma_3$, $\omega=\sqrt{2(2\sqrt{2}-1)}$',color=INK,fontsize=10.5,loc='left')
bins=np.logspace(0,5,46)
bx.hist(dg,bins=bins,color=C1,alpha=0.85,label='gravity: all 400 positive')
bx.set_xscale('log'); bx.axvline(k,color=C4,lw=1.6,ls=(0,(4,3)))
bx.set_ylim(0,46); bx.text(k*1.12,42.5,r'square: $2\sqrt{2}-1$',color=C4,fontsize=9.5)
bx.text(0.40,0.70,'like charges: the mirror image,\nevery value negative',transform=bx.transAxes,color=C2,fontsize=9.5)
bx.set_xlabel(r'stiffness $\partial U/\partial w_3$ at the wall (log scale)'); bx.set_ylabel('planar configurations')
bx.legend(loc='upper right',frameon=False,fontsize=9.5,labelcolor=INK)
bx.set_title('b  400 random flat configurations: gravity always stiffens',color=INK,fontsize=10.5,loc='left')
fig.savefig(f'{OUT}/f18_stiffen.png',dpi=260,bbox_inches='tight',pad_inches=0.12)