#!/usr/bin/env python3
"""Reproduce original, mathematical/theoretical figures used in the MROX website.
Figures are demonstration calculations, not the MROX experimental or production output.
Python 3 with numpy, scipy and matplotlib. Matplotlib SVG output is fully offline.
"""
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, FancyArrowPatch, Arc
from scipy.ndimage import gaussian_filter1d

ROOT=Path(__file__).resolve().parents[1]/'assets'/'figures'; ROOT.mkdir(parents=True,exist_ok=True)
BLUE='#156ec3'; CYAN='#10a6cd'; ORANGE='#e38839'; PURPLE='#7867c7'; DARK='#173356'; MUTED='#597189'; GRID='#e7eef6'; BG='#ffffff'
plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'axes.titlesize':13,'axes.titleweight':'bold','axes.labelsize':10,'axes.edgecolor':'#cedbe9','axes.labelcolor':DARK,'xtick.color':MUTED,'ytick.color':MUTED,'text.color':DARK,'svg.fonttype':'none','figure.facecolor':BG,'axes.facecolor':BG,'axes.spines.top':False,'axes.spines.right':False,'axes.grid':True,'grid.color':GRID,'grid.linewidth':.8,'savefig.facecolor':BG})
def done(name,fig):
 fig.savefig(ROOT/name,bbox_inches='tight',pad_inches=.23,format='svg'); plt.close(fig)

def base(size=(8.2,3.5)):
 f,a=plt.subplots(figsize=size,dpi=130);a.grid(True,alpha=.95);return f,a

# 1 Bragg geometry and multilayer concept (actual physical schematic, not article figures)
f,ax=plt.subplots(figsize=(9.4,3.7),dpi=120);ax.set_xlim(0,10);ax.set_ylim(0,5);ax.axis('off')
for k,(y,c,label) in enumerate([(2.25,'#cbeafb','Film 1'),(1.68,'#a7d0ed','Film 2'),(1.1,'#7cafe0','Film 3'),(.36,'#b0c3d8','Semi-infinite substrate')]):
 h=.56 if k!=3 else .73; ax.add_patch(Rectangle((1.1,y),7.6,h,facecolor=c,edgecolor='#7897b7',lw=.85)); ax.text(8.46,y+h/2,label,ha='right',va='center',fontsize=9,color='#173c65')
ax.plot([1.1,8.7],[2.84,2.84],color=DARK,lw=2)
ax.annotate('',(4.1,2.84),(2.1,4.55),arrowprops=dict(arrowstyle='-|>',color=ORANGE,lw=2.5));ax.annotate('',(6.1,4.55),(4.1,2.84),arrowprops=dict(arrowstyle='-|>',color=CYAN,lw=2.5))
ax.annotate('',(5.65,.65),(4.1,2.84),arrowprops=dict(arrowstyle='-|>',color=BLUE,lw=1.8,linestyle='dashed'))
ax.plot([4.1,4.1],[2.84,4.4],ls=':',color=MUTED,lw=1.3)
ax.text(2.05,4.67,'Incident $D_0$',color=ORANGE,fontsize=11,weight='bold');ax.text(6.1,4.69,'Diffracted $D_H$',color=CYAN,fontsize=11,weight='bold')
ax.text(5.62,1.05,'Transmission',color=BLUE,fontsize=10); ax.text(.75,2,'$z\\downarrow$',fontsize=13,color=DARK)
ax.text(3.22,2.46,'surface',color=DARK,fontsize=9)
done('01-bragg-multilayers.svg',f)

# 2 thin-film kinematic fringe family; x = deltaq * t_ref/(2pi)
f,a=base(); x=np.linspace(-4.3,4.3,2000)
for thick,color in [(1,BLUE),(2,CYAN),(3,ORANGE)]:
 y=np.sinc(thick*x)**2
 a.plot(x,y,color=color,lw=2.4,label=f'$t / t_0 = {thick}$')
a.set(xlabel=r'Detuning $\Delta q_z t_0 / (2\pi)$',ylabel='Normalised intensity',ylim=(-.04,1.07),xlim=(-4.2,4.2));a.set_title('Thickness fringes: kinematic reference');a.legend(frameon=False,ncol=3,loc='upper right')
done('02-kinematic-fringes.svg',f)

# 3 ideal Darwin plateau from X=-eta+/-sqrt(eta^2-1)
f,a=base();eta=np.linspace(-3.6,3.6,1101);X=np.empty_like(eta,dtype=complex)
for i,z in enumerate(eta):
 if z>=1: X[i]=-z+np.sqrt(z*z-1)
 elif z<=-1: X[i]=-z-np.sqrt(z*z-1)
 else: X[i]=-z+1j*np.sqrt(1-z*z)
R=abs(X)**2
a.axvspan(-1,1,facecolor='#e5f6fb',alpha=.95,label='Darwin region')
a.plot(eta,R,color=BLUE,lw=2.8,label=r'$|X_\infty(\eta)|^2$')
a.axvline(-1,color=CYAN,ls='--',lw=1);a.axvline(1,color=CYAN,ls='--',lw=1)
a.set(xlabel=r'Normalised deviation $\eta$',ylabel='Ideal reflectivity',ylim=(-.025,1.08),xlim=(-3.6,3.6));a.set_title('Semi-infinite, non-absorbing crystal');a.legend(frameon=False,loc='upper right')
done('03-darwin-plateau.svg',f)

# finite homogeneous slab from normalized two-wave matrix (ratio X = M21 / M11)
def finite_reflectivity(eta,T):
 eta=np.asarray(eta,dtype=complex); s=np.sqrt(1-eta*eta)
 with np.errstate(divide='ignore',invalid='ignore',over='ignore'):
  h=np.where(np.abs(s)>1e-9,np.sinh(s*T)/s,T+0*s)
  c=np.cosh(s*T)
  return np.abs((1j*h)/(c-1j*eta*h))**2
# 4 dynamical thickness family
f,a=base(); eta=np.linspace(-3.7,3.7,1600)
for T,color in [(.35,ORANGE),(1.0,CYAN),(2.5,BLUE),(4.2,PURPLE)]:
 a.plot(eta,finite_reflectivity(eta,T),lw=2.15,color=color,label=f'T={T:g}')
a.axvspan(-1,1,color='#e8f6fe',alpha=.7,zorder=-5)
a.set(xlabel=r'Normalised deviation $\eta$',ylabel=r'$|X|^2$',ylim=(-.035,1.08),xlim=(-3.7,3.7));a.set_title('Finite perfect slab: exact normalised two-wave solution');a.legend(frameon=False,ncol=4,loc='upper right')
done('04-finite-slab-dynamics.svg',f)

# 5 2D heatmap
f,a=plt.subplots(figsize=(8.5,4.15),dpi=130);xx=np.linspace(-3.2,3.2,330);yy=np.linspace(.08,4.7,210);XX,YY=np.meshgrid(xx,yy);ZZ=finite_reflectivity(XX,YY)
im=a.imshow(ZZ,origin='lower',extent=[xx.min(),xx.max(),yy.min(),yy.max()],aspect='auto',vmin=0,vmax=1,cmap='YlGnBu',interpolation='bilinear')
a.axvline(-1,color='white',ls='--',alpha=.75);a.axvline(1,color='white',ls='--',alpha=.75)
a.set(xlabel=r'Normalised deviation $\eta$',ylabel='Optical thickness $T$',title='Two-wave reflectivity landscape (ideal finite slab)');f.colorbar(im,ax=a,label=r'$|X|^2$',fraction=.046,pad=.03)
done('05-dynamical-map.svg',f)

# 6 strain derivative dtheta
f,a=base();thetaB=np.linspace(12,75,450);eps=np.array([.001,.003,.005]);colors=[CYAN,BLUE,ORANGE]
for e,c in zip(eps,colors):
 delta2=-2*np.tan(np.deg2rad(thetaB))*e*180/np.pi
 a.plot(2*thetaB,delta2,lw=2.3,color=c,label=f'$\\epsilon_\\perp$={e*100:.1f}%')
a.set(xlabel=r'Bragg angle $2\theta_B$ (degrees)',ylabel=r'$\Delta(2\theta_B)$ (degrees)',xlim=(24,150));a.set_title('Peak displacement from normal strain (first order)');a.legend(frameon=False)
done('06-strain-peak-shift.svg',f)

# 7 depth implant profile
f,a=base((8.2,3.7));z=np.linspace(0,270,700)
def gauss_asym(z,A,mu,sL,sR): return A*np.exp(-.5*((z-mu)/np.where(z<=mu,sL,sR))**2)
g1=gauss_asym(z,.55,84,28,52);g2=gauss_asym(z,.20,163,18,32);total=g1+g2
a.plot(z,total,color=BLUE,lw=2.7,label='Total strain profile');a.plot(z,g1,color=CYAN,ls='--',lw=1.9,label='Asymmetric component 1');a.plot(z,g2,color=ORANGE,ls='--',lw=1.9,label='Component 2')
a.fill_between(z,total,color=BLUE,alpha=.07);a.set(xlabel='Depth z (nm)',ylabel='Strain (%)',xlim=(0,260),ylim=(0,.7));a.set_title('Illustrative implantation strain profile');a.legend(frameon=False,fontsize=8,ncol=2)
done('07-implantation-strain.svg',f)

# 8 Debye-Waller physical amplitude vs r.m.s. disorder
f,a=base();u=np.linspace(0,.8,400)
for q,c in [(2,BLUE),(4,CYAN),(8,ORANGE)]:
 DW=np.exp(-.5*(q*u)**2)
 a.plot(u,DW,color=c,lw=2.5,label=f'|H| = {q} Å⁻¹')
a.set(xlabel='RMS projected displacement (Å)',ylabel='Coherent amplitude factor',ylim=(-.03,1.05),xlim=(0,.8));a.set_title('Physical Debye–Waller dependence on reciprocal vector');a.legend(frameon=False)
done('08-debye-waller.svg',f)

# 9 kernels and convolution
f,a=base();x=np.linspace(-.18,.18,1300);sigma=.025;gamma=.026
G=np.exp(-.5*(x/sigma)**2); L=1/(1+(x/gamma)**2);pv=.65*G+.35*L
for y,c,l in [(G,BLUE,'Gaussian'),(L,ORANGE,'Lorentzian'),(pv,CYAN,'Pseudo-Voigt 65/35')]: a.plot(x,y,color=c,lw=2.5,label=l)
a.set(xlabel='Angular offset (degrees)',ylabel='Peak-normalised response',ylim=(0,1.08));a.set_title('Instrumental resolution kernels');a.legend(frameon=False)
done('09-pseudo-voigt.svg',f)

# 10 resolution variable function
f,a=base();x=np.linspace(-.6,.6,601);smin=.0005;smax=.060;p=2.5;sigma=smin+(smax-smin)*(abs(x)/.6)**p
a.plot(x,sigma,color=BLUE,lw=2.8,label='Width σ(x)');a.axhline(smin,color=CYAN,ls='--',lw=1.5,label='σ min');a.axhline(smax,color=ORANGE,ls=':',lw=1.5,label='σ max')
a.set(xlabel=r'Offset from substrate $2\theta_B$ (degrees)',ylabel=r'Local $\sigma$ (degrees)',ylim=(0,.066));a.set_title('MROX documented empirical variable-resolution law');a.legend(frameon=False)
done('10-variable-resolution.svg',f)

# 11 ideal vs broadened synthetic fringes
f,a=base();x=np.linspace(-.25,.25,1401)
y=np.exp(-.5*(x/.023)**2)*(1+.35*np.cos(2*np.pi*x/.026))**2+ .2*np.exp(-.5*((x+.09)/.025)**2)
y/=y.max();yy=gaussian_filter1d(y,.017/(x[1]-x[0]),mode='nearest');
a.plot(x,y,color=BLUE,lw=1.8,label='Ideal mathematical signal');a.plot(x,yy,color=ORANGE,lw=2.6,label='After Gaussian broadening')
a.set(xlabel=r'Angular offset $\Delta(2\theta)$ (degrees)',ylabel='Normalised intensity',ylim=(0,1.05));a.set_title('Intensity convolution suppresses fine fringes');a.legend(frameon=False)
done('11-resolution-broadening.svg',f)

# 12 computational flow chart
f,ax=plt.subplots(figsize=(9.5,2.8),dpi=125);ax.axis('off'); ax.set_xlim(0,10);ax.set_ylim(0,3)
labels=['Crystal + CIF','Structure factors','Deviation η & T','Substrate root','Bottom → top','Intensity + resolution']
for i,label in enumerate(labels):
 x=.15+i*1.665
 ax.add_patch(Rectangle((x,1.15),1.46,.86,facecolor=['#e6f2ff','#dbeeff','#d2e8fb','#c5e4f6','#b6dbf3','#d4f4f2'][i],edgecolor='#99bfdc',lw=1.2))
 ax.text(x+.73,1.58,label,color='#123c66',ha='center',va='center',fontsize=9,weight='bold',wrap=True)
 if i<5: ax.annotate('',(x+1.63,1.57),(x+1.48,1.57),arrowprops={'arrowstyle':'->','color':ORANGE,'lw':2})
ax.text(5,2.57,'Forward model architecture · no GUI in the numerical kernel',ha='center',weight='bold',fontsize=12,color=DARK)
ax.text(5,.56,'Prepare once where possible · evaluate complex amplitudes · convolve intensities · preserve the measured angle grid',ha='center',fontsize=8.5,color=MUTED)
done('12-model-pipeline.svg',f)

print('Generated figures:',len(list(ROOT.glob('*.svg'))))