# SCRIPT: BH-FIGS-V8-5
# 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) the rhombus family: the stiffness changes sign at diagonal ratio k* (BH-STIFF-19).
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
import os
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 rhombus family (diagonal ratio k), closed form of BH-STIFF-19: dU/dw3 = (1/2)[4 sqrt2 - sqrt((1+k^2)/2)(1+1/k)]
kk=np.linspace(0.06,1,600); Sk=0.5*(4*np.sqrt(2)-np.sqrt((1+kk**2)/2)*(1+1/kk))
kst=[x.real for x in np.roots([1,2,-62,2,1]) if abs(x.imag)<1e-12 and 0<x.real<1][0]
print('rhombus root k* =',kst,' value at k=1:',Sk[-1])
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')
bx.plot(kk,Sk,color=C1,lw=2.4,label='gravity'); bx.plot(kk,-Sk,color=C2,lw=2.4,label='like charges')
bx.axhline(0,color=RULE,lw=1); bx.axvline(kst,color=C4,lw=1.4,ls=(0,(4,3)))
bx.plot(1,k,'o',color=C1,ms=5); bx.text(0.80,k+0.32,r'square: $2\sqrt{2}-1$',color=C1,fontsize=9.5)
bx.text(kst+0.015,2.9,r'$k_*=0.1446$',color=C4,fontsize=10)
def rh(x0,y0,kv,sc=0.045):
    a=9.07   # y-units per x-unit at this axes size, so the glyph keeps its true shape
    pts=np.array([[1,0],[0,kv],[-1,0],[0,-kv],[1,0]])*sc
    bx.plot(x0+pts[:,0],y0+pts[:,1]*a,color=MUT,lw=1)
    bx.plot(x0+pts[:4,0],y0+pts[:4,1]*a,'o',color=INK,ms=2.5)
    bx.text(x0,y0-0.62,f'k = {kv:g}',color=MUT,fontsize=8.8,ha='center')
rh(0.30,0.15,0.1); rh(0.62,0.15,0.4); rh(0.92,0.15,1.0)
bx.set_xlim(0,1.05); bx.set_ylim(-3.5,3.5)
bx.set_xlabel(r'diagonal ratio $k$ (short / long)'); bx.set_ylabel(r'stiffness $\partial U/\partial w_3$ at the wall')
bx.legend(loc='upper center',bbox_to_anchor=(0.55,1.0),frameon=False,fontsize=9.5,labelcolor=INK)
bx.set_title(r'b  the rhombus: gravity stiffens only above $k_*$',color=INK,fontsize=10.5,loc='left')
fig.savefig(f'{OUT}/f18b_stiffen.png',dpi=260,bbox_inches='tight',pad_inches=0.12)