# Figures of the Hertz check from hertz_cylinder.jl: the contact pressure of the finite
# element model against Hertz's ellipse, and the von Mises stress under the contact with
# the depth of its maximum. The model is the quarter x >= 0; both figures mirror it so the
# whole contact shows.
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.tri as mtri
import numpy as np

here = Path(__file__).parent
out = here/'out'
graphics = here.parent/'graphics'
BLUE, RED = '#4472c4', '#c00000'
plt.rcParams['svg.fonttype'] = 'none'

S = dict(np.genfromtxt(here/'hertz_summary.csv', delimiter=',', dtype=None, encoding=None))
S = {k: float(v) if k != 'plane' else v for k, v in S.items()}
a, pmax = S['a_hertz_mm'], S['pmax_hertz_MPa']
pr = np.genfromtxt(here/'hertz_pressure.csv', delimiter=',', names=True)


def end_ticks(ax, xt, yt, xfmt='{:g}', yfmt='{:g}'):
    """Ticks at the given places, which include both ends of each axis."""
    ax.set_xticks(xt); ax.set_xticklabels([xfmt.format(t) for t in xt])
    ax.set_yticks(yt); ax.set_yticklabels([yfmt.format(t) for t in yt])


# Pressure: FE at the quadrature points of the arc, both penalties, against Hertz
fig, ax = plt.subplots(figsize=(6.4, 3.4))
xx = np.linspace(-a, a, 401)
ax.plot(xx*1e3, pmax*np.sqrt(1 - (xx/a)**2), color='k', lw=1.4, label='Hertz')
x, p1, p2 = pr['x_mm'], pr['p_fe_kappa1_MPa'], pr['p_fe_kappa2_MPa']
sel = (x < 1.4*a) & (np.arange(len(x)) % 2 == 0)     # every second quadrature point
# the stiffer penalty on the left half, the softer on the right: the contact is symmetric
ax.plot(x[sel]*1e3, p1[sel], 'o', ms=3, color=BLUE, label=r'FE, $\kappa = 10E/h$')
ax.plot(-x[sel]*1e3, p2[sel], 'o', ms=3, color=RED, label=r'FE, $\kappa = 100E/h$')
for sgn in (1, -1):
    ax.axvline(sgn*a*1e3, color='0.4', ls='--', lw=0.8)
ax.annotate('', xy=(a*1e3, 120), xytext=(0, 120),
            arrowprops=dict(arrowstyle='->', color='0.2', lw=0.8))
ax.text(a*1e3/2, 150, f'$a = {a*1e3:.1f}$ µm', ha='center', fontsize=9)
xm = 1.4*a*1e3
ax.set_xlim(-xm, xm); ax.set_ylim(0, 1100)
xt = [-round(xm), -round(a*1e3, 1), 0, round(a*1e3, 1), round(xm)]
ax.set_xlim(xt[0], xt[-1])
end_ticks(ax, xt, [0, 250, 500, 750, round(pmax), 1100])
ax.set_xlabel(r'position $x$ across the contact (µm)')
ax.set_ylabel('contact pressure $p$ (MPa)')
for side in ('top', 'right'):
    ax.spines[side].set_visible(False)
ax.legend(loc='upper right', frameon=False, fontsize=8)
fig.tight_layout()
fig.savefig(graphics/'GridapHertz_pressure.svg')
plt.close(fig)

# Von Mises stress near the contact, plane strain, quarter mirrored about the axis
X = np.loadtxt(out/'hertz_nodes.csv', delimiter=',')
T = np.loadtxt(out/'hertz_tris.csv', delimiter=',', dtype=int)
vm = np.loadtxt(out/'hertz_vm.csv')
n = len(X)
Xm = np.r_[X, X*[-1, 1]]
Tm = np.r_[T, T[:, ::-1] + n]
vmm = np.r_[vm, vm]
tri = mtri.Triangulation(Xm[:, 0]*1e3, Xm[:, 1]*1e3, Tm)

L = 2.5*a*1e3
vmax = 50*np.ceil(S['vm_max_fe_MPa']/50)
fig = plt.figure(figsize=(7.2, 4.0), facecolor='white')
ax = fig.add_axes([0.08, 0.13, 0.78, 0.82])
im = ax.tripcolor(tri, vmm, shading='gouraud', cmap='jet', vmin=0, vmax=vmax, rasterized=True)
ax.tricontour(tri, vmm, levels=np.arange(100, vmax, 100), colors='k', linewidths=0.4, alpha=0.5)
ax.axhspan(-0.3*L, 0, color='0.75', lw=0)          # the rigid plane
ax.plot([-L, L], [0, 0], color='k', lw=1.0)
# contact width and the depth of the largest von Mises stress
ax.plot([-a*1e3, a*1e3], [0, 0], color='w', lw=3, solid_capstyle='butt')
zf, zH = S['vm_depth_fe_mm']*1e3, S['vm_depth_hertz_mm']*1e3
ax.plot(0, zf, 'o', ms=6, mfc='white', mec='k', mew=1.2)
ax.annotate(f"max {S['vm_max_fe_MPa']:.0f} MPa at $z = {zf:.0f}$ µm $= {zf/(a*1e3):.2f}\\,a$\n"
            f"Hertz: {S['vm_max_hertz_MPa']:.0f} MPa at {zH:.0f} µm",
            xy=(0, zf), xytext=(-0.97*L, 0.95*L), fontsize=9, color='w', va='top',
            arrowprops=dict(arrowstyle='-', color='w', lw=0.8))
ax.text(0, -0.15*L, f'contact width $2a = {2*a*1e3:.0f}$ µm', ha='center', va='center', fontsize=9)
ax.set_xlim(-L, L); ax.set_ylim(-0.3*L, L); ax.set_aspect('equal')
end_ticks(ax, [-round(L), -round(a*1e3), 0, round(a*1e3), round(L)],
          [0, round(zf), round(a*1e3), round(L)])
ax.set_xlabel('$x$ (µm)'); ax.set_ylabel('height $z$ above the plane (µm)')
cax = fig.add_axes([0.88, 0.13, 0.025, 0.82])
cb = fig.colorbar(im, cax=cax)
cb.set_ticks(np.linspace(0, vmax, 7))
cb.set_label('von Mises stress, plane strain (MPa)')
fig.savefig(graphics/'GridapHertz_vonmises.png', dpi=200, facecolor='white')
plt.close(fig)