#!/usr/bin/env python3
"""Typeset consistent, beautifully-sized equation SVGs, usable offline without font files.
Every line uses the same STIX math font at the same physical point size.
Long expressions are typeset as semantic multiline equations, never scaled smaller.
"""
from pathlib import Path
from matplotlib.mathtext import MathTextParser
from matplotlib.font_manager import FontProperties
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

ROOT=Path(__file__).resolve().parents[1]/'assets'/'equations'
ROOT.mkdir(parents=True,exist_ok=True)
# Same latex definitions as V1. These are the public mathematical expressions (not source-code identities).
EQUATIONS={
'q':r'\mathbf{q}=\mathbf{k}_{H}-\mathbf{k}_{0},\qquad |\mathbf{q}|=\frac{4\pi\sin\theta}{\lambda}',
'bragg':r'2d_{hkl}\sin\theta_B=m\lambda',
'lattice':r'A(\mathbf{q})=\sum_n e^{i\mathbf{q}\cdot\mathbf{R}_n}\sum_j f_j(\mathbf{q})e^{i\mathbf{q}\cdot\mathbf{r}_j}',
'film':r'A_f(\Delta q_z)\propto\frac{e^{i\Delta q_z t}-1}{i\Delta q_z}',
'fringes':r'I_f(\Delta q_z)\propto t^2\left[\frac{\sin(\Delta q_z t/2)}{\Delta q_z t/2}\right]^2',
'atomic':r'f_j(s,\lambda)=f_j^0(s)+f_j\prime(\lambda)+if_j\prime\prime(\lambda),\quad s=\sin\theta/\lambda',
'structure':r'F_H=\sum_j o_j f_j\exp[2\pi i(hx_j+ky_j+lz_j)]\,e^{-W_j(H)}',
'chi':r'\chi_H=-\frac{r_e\lambda^2}{\pi V_c}F_H,\quad \chi_0=-\frac{r_e\lambda^2}{\pi V_c}F_{000}',
'wave':r'\nabla^2\Psi+k^2[1+\chi(\mathbf{r})]\Psi=0,\qquad k=2\pi/\lambda',
'fourier':r'\chi(\mathbf{r})\simeq\chi_0+\chi_H e^{i\mathbf{H}\cdot(\mathbf{r}-\mathbf{u})}+\chi_{-H}e^{-i\mathbf{H}\cdot(\mathbf{r}-\mathbf{u})}',
'ansatz':r'\Psi(\mathbf{r})=D_0(\mathbf{r})e^{i\mathbf{k}_0\cdot\mathbf{r}}+D_H(\mathbf{r})e^{i(\mathbf{k}_0+\mathbf{H})\cdot\mathbf{r}}',
'tt1':r'\partial_{s_0}D_0=ia_0D_0+ib_0D_H e^{-i\mathbf{H}\cdot\mathbf{u}}',
'tt2':r'\partial_{s_H}D_H=ia_HD_H+ib_HD_0 e^{+i\mathbf{H}\cdot\mathbf{u}}',
'normal1':r'-i\frac{dU}{dT}=-V,\qquad -i\frac{dV}{dT}=U+2\eta V',
'quotient':r'X=\frac{V}{U},\qquad \frac{dX}{dT}=\frac{UV\prime-VU\prime}{U^2}',
'riccati':r'-i\frac{dX}{dT}=X^2+2\eta X+1',
'fixed':r'X_{\pm}=-\eta\pm\sqrt{\eta^2-1},\qquad \xi=\sqrt{\eta^2-1}',
'linear':r'X=i\,U\prime/U\quad\Rightarrow\quad U\prime\prime-2i\eta U\prime-U=0',
'closed':r'\frac{X(T)-X_+}{X(T)-X_-}=\frac{X(0)-X_+}{X(0)-X_-}\,e^{2i\xi T}',
'matrix':r'\frac{d\mathbf{Y}}{dT}=\mathbf{A}\mathbf{Y},\quad A_{11}=0,\ A_{12}=-i,\ A_{21}=i,\ A_{22}=2i\eta',
'mobius':r'X_{\rm out}=\frac{M_{21}+M_{22}X_{\rm in}}{M_{11}+M_{12}X_{\rm in}}',
'bartels1':r'S_1=(X_{\rm in}-\eta+\xi)e^{-iT\xi},\quad S_2=(X_{\rm in}-\eta-\xi)e^{iT\xi}',
'bartels2':r'X_{\rm out}=\eta+\xi\frac{S_1+S_2}{S_1-S_2}',
'T':r'T_j=\frac{\pi C_j\gamma_{M,j}\,t_j\sqrt{F_{H,j}F_{-H,j}}}{\lambda\sqrt{|\gamma_{0,j}\gamma_{H,j}|}}',
'sub':r'X_{\infty}=-\eta\pm\sqrt{\eta^2-1}\qquad(\mathrm{normalised\ convention})',
'intensity':r'I_H\propto\left|\frac{F_H}{F_{-H}}\right|\,|X_{\rm surface}|^2',
'extinction':r'\ell_{\rm dyn}\sim\frac{\lambda\sqrt{|\gamma_0\gamma_H|}}{\pi|C|\sqrt{|\chi_H\chi_{-H}|}}',
'strain':r'\frac{\Delta d_{hkl}}{d_{hkl}}\simeq\hat{\mathbf{n}}^{T}\boldsymbol{\varepsilon}\hat{\mathbf{n}}',
'shift':r'\Delta(2\theta_B)\simeq-2\tan\theta_B\,\frac{\Delta d_{hkl}}{d_{hkl}}\quad(\mathrm{radians})',
'elastic':r'\sigma_i=\sum_j C_{ij}\varepsilon_j,\qquad \sigma_{nn}=0\quad(\mathrm{free\ surface})',
'dw':r'\left\langle e^{i\mathbf{H}\cdot\mathbf{u}}\right\rangle=e^{-\frac{1}{2}\mathbf{H}^T\mathbf{U}\mathbf{H}}',
'implant':r'\varepsilon(z)=\sum_{k=1}^{K}A_k\exp\left[-\frac{(z-\mu_k)^2}{2\sigma_k(z)^2}\right]',
'asym':r'\sigma_k(z)=\sigma_{L,k}\ [z\leq\mu_k];\qquad \sigma_{R,k}\ [z>\mu_k]',
'convolve':r'I_{\rm measured}(x)=B(x)+A\int I_{\rm dyn}(u)R(x-u)\,du,\quad x=2\theta',
'pv':r'PV(x)=(1-\beta)G(x;w)+\beta L(x;w),\quad 0\leq\beta\leq1',
'gauss':r'G(x;w)=\frac{2\sqrt{\ln2}}{w\sqrt{\pi}}\exp\left(-4\ln2\frac{x^2}{w^2}\right)',
'lorentz':r'L(x;w)=\frac{2}{\pi w}\frac{1}{1+4x^2/w^2}',
'variable':r'\sigma(x)=\sigma_{\min}+(\sigma_{\max}-\sigma_{\min})\left(\frac{|x-x_B|}{D_{\max}}\right)^p',
'kernel':r'I_{\rm res}(x_i)=\frac{\sum_jI(x_j)\exp[-(x_j-x_i)^2/(2\sigma_i^2)]}{\sum_j\exp[-(x_j-x_i)^2/(2\sigma_i^2)]}',
'period':r'\Delta(2\theta)_{\rm fringe}\simeq\frac{\lambda}{t\cos\theta_B}\quad(\mathrm{radians})',
'error':r'E_{\Delta z}=\frac{\max_i|I_{\Delta z}(x_i)-I_{\Delta z/2}(x_i)|}{\max(\epsilon,\max_iI_{\Delta z/2}(x_i))}',
'fit':r'F=0.60L_{\rm Huber}+0.25(1-\rho_I)+0.15(1-\rho_{\Delta I})',
'huber':r'H_{\delta}(e)=e^2/2\ [|e|\leq\delta];\quad \delta(|e|-\delta/2)\ [|e|>\delta]',
'loglike':r'\log p(D|\boldsymbol{\vartheta},\sigma)=-N\log\sigma-\frac{1}{2\sigma^2}\sum_i r_i^2+\mathrm{constant}',
}

# Equations that would otherwise be compressed in a narrow column are split
# into mathematically meaningful rows. Their copyable LaTeX remains unchanged.
LAYOUT={
 'q': [r'\mathbf{q}=\mathbf{k}_{H}-\mathbf{k}_{0}',r'|\mathbf{q}|=\frac{4\pi\sin\theta}{\lambda}'],
 'atomic':[r'f_j(s,\lambda)=f_j^0(s)+f_j\prime(\lambda)+if_j\prime\prime(\lambda)',r's=\sin\theta/\lambda'],
 'chi':[r'\chi_H=-\frac{r_e\lambda^2}{\pi V_c}F_H',r'\chi_0=-\frac{r_e\lambda^2}{\pi V_c}F_{000}'],
 'normal1':[r'-i\frac{dU}{dT}=-V',r'-i\frac{dV}{dT}=U+2\eta V'],
 'matrix':[r'\frac{d\mathbf{Y}}{dT}=\mathbf{A}\mathbf{Y}',r'A_{11}=0,\quad A_{12}=-i,\quad A_{21}=i,\quad A_{22}=2i\eta'],
 'bartels1':[r'S_1=(X_{\rm in}-\eta+\xi)e^{-iT\xi}',r'S_2=(X_{\rm in}-\eta-\xi)e^{iT\xi}'],
 'convolve':[r'I_{\rm measured}(x)=B(x)+A\int I_{\rm dyn}(u)R(x-u)\,du',r'x=2\theta'],
 'pv':[r'PV(x)=(1-\beta)G(x;w)+\beta L(x;w)',r'0\leq\beta\leq1'],
 'elastic':[r'\sigma_i=\sum_j C_{ij}\varepsilon_j',r'\sigma_{nn}=0\qquad(\mathrm{free\ surface})'],
 'fixed':[r'X_{\pm}=-\eta\pm\sqrt{\eta^2-1}',r'\xi=\sqrt{\eta^2-1}'],
 'sub':[r'X_{\infty}=-\eta\pm\sqrt{\eta^2-1}',r'(\mathrm{normalised\ convention})'],
 'wave':[r'\nabla^2\Psi+k^2[1+\chi(\mathbf{r})]\Psi=0',r'k=2\pi/\lambda'],
 'period':[r'\Delta(2\theta)_{\rm fringe}\simeq\frac{\lambda}{t\cos\theta_B}',r'(\mathrm{radians})'],
}


# Additional editorial line breaks: each row remains at exactly the same
# mathematical type size. This particularly improves phones and narrow tablets.
LAYOUT.update({
 'lattice':[
  r'A(\mathbf{q})=\sum_n e^{i\mathbf{q}\cdot\mathbf{R}_n}',
  r'\times\sum_j f_j(\mathbf{q})e^{i\mathbf{q}\cdot\mathbf{r}_j}'],
 'structure':[
  r'F_H=\sum_j o_j f_j',
  r'\times\exp[2\pi i(hx_j+ky_j+lz_j)]\,e^{-W_j(H)}\quad(\mathrm{inside\ the\ sum})'],
 'fourier':[
  r'\chi(\mathbf{r})\simeq\chi_0+\chi_H e^{i\mathbf{H}\cdot(\mathbf{r}-\mathbf{u})}',
  r'+\chi_{-H}e^{-i\mathbf{H}\cdot(\mathbf{r}-\mathbf{u})}'],
 'ansatz':[
  r'\Psi(\mathbf{r})=D_0(\mathbf{r})e^{i\mathbf{k}_0\cdot\mathbf{r}}',
  r'+D_H(\mathbf{r})e^{i(\mathbf{k}_0+\mathbf{H})\cdot\mathbf{r}}'],
 'tt1':[
  r'\partial_{s_0}D_0=ia_0D_0',
  r'+ib_0D_H e^{-i\mathbf{H}\cdot\mathbf{u}}'],
 'tt2':[
  r'\partial_{s_H}D_H=ia_HD_H',
  r'+ib_HD_0 e^{+i\mathbf{H}\cdot\mathbf{u}}'],
 'quotient':[
  r'X=\frac{V}{U}',
  r'\frac{dX}{dT}=\frac{UV\prime-VU\prime}{U^2}'],
 'linear':[
  r'X=i\,U\prime/U',
  r'\Rightarrow\quad U\prime\prime-2i\eta U\prime-U=0'],
 'shift':[
  r'\Delta(2\theta_B)\simeq-2\tan\theta_B\,\frac{\Delta d_{hkl}}{d_{hkl}}',
  r'(\mathrm{radians})'],
 'asym':[
  r'\sigma_k(z)=\sigma_{L,k}\qquad (z\leq\mu_k)',
  r'\sigma_k(z)=\sigma_{R,k}\qquad (z>\mu_k)'],
 'gauss':[
  r'G(x;w)=\frac{2\sqrt{\ln2}}{w\sqrt{\pi}}',
  r'\times\exp\left(-4\ln2\frac{x^2}{w^2}\right)'],
 'variable':[
  r'\sigma(x)=\sigma_{\min}+',
  r'(\sigma_{\max}-\sigma_{\min})\left(\frac{|x-x_B|}{D_{\max}}\right)^p'],
 'loglike':[
  r'\log p(D|\boldsymbol{\vartheta},\sigma)=-N\log\sigma',
  r'-\frac{1}{2\sigma^2}\sum_i r_i^2+\mathrm{constant}'],
 'fit':[
  r'F=0.60L_{\rm Huber}+0.25(1-\rho_I)',
  r'+0.15(1-\rho_{\Delta I})'],
 'huber':[
  r'H_{\delta}(e)=e^2/2\qquad (|e|\leq\delta)',
  r'H_{\delta}(e)=\delta(|e|-\delta/2)\qquad (|e|>\delta)'],
 'pv':[
  r'PV(x)=(1-\beta)G(x;w)+\beta L(x;w)',
  r'0\leq\beta\leq1'],
 'convolve':[
  r'I_{\rm measured}(x)=B(x)',
  r'+A\int I_{\rm dyn}(u)R(x-u)\,du',
  r'x=2\theta'],
 'matrix':[
  r'\frac{d\mathbf{Y}}{dT}=\mathbf{A}\mathbf{Y}',
  r'A_{11}=0,\qquad A_{12}=-i',
  r'A_{21}=i,\qquad A_{22}=2i\eta'],
})

LAYOUT.pop("structure",None)  # Keep the exact summand grouping intact; scroll if required.

plt.rcParams.update({
    'mathtext.fontset':'stix',
    'svg.fonttype':'path',
    'text.color':'#123c69',
    'figure.facecolor':'white',
    'savefig.transparent':True,
})
FONT_PT=19.2
DPI=96
PADDING_X=22
parser=MathTextParser('path')
font=FontProperties(size=FONT_PT)
for key,raw in EQUATIONS.items():
    pieces=LAYOUT.get(key,[raw])
    m=[parser.parse('$'+piece+'$',dpi=DPI,prop=font) for piece in pieces]
    canvas_width=max(e.width for e in m)+2*PADDING_X
    # 74px for one line; for two lines, an extra 53px. Every line has the same font size.
    line_separation=54
    canvas_height=74+line_separation*(len(pieces)-1)
    fig=plt.figure(figsize=(canvas_width/DPI,canvas_height/DPI),dpi=DPI)
    for idx,piece in enumerate(pieces):
        # Center each line in both axes with a constant physical font size.
        y=(canvas_height/2+(len(pieces)-1)/2*line_separation-idx*line_separation)/canvas_height
        fig.text(.5,y,'$'+piece+'$',ha='center',va='center',fontsize=FONT_PT,color='#143b68')
    fig.savefig(ROOT/f'{key}.svg',format='svg',transparent=True)
    plt.close(fig)
print('Generated',len(EQUATIONS),'equations, uniform',FONT_PT,'pt STIX math font, stable physical scale.')