# Book figures from gears_hertz.jl: the load share of one tooth pair along its path of
# contact, from the tooth stiffness, against the 50/50 split; and the mesh stiffness, the
# force per unit face width over the approach of the gears, against the same position.
from pathlib import Path

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

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

summary = dict(np.genfromtxt(here/'summary.csv', delimiter=',', dtype=None, encoding=None))
s_a, s_r, pb = summary['s_a_mm'], summary['s_r_mm'], summary['base_pitch_mm']
fr = np.genfromtxt(here/'frames2.csv', delimiter=',', names=True)
s, share, k = fr['s0'], fr['share'], fr['k_mesh_N_per_mm_um']
xt = [round(s_a, 1), round(s_r - pb, 1), 0, round(s_a + pb, 1), round(s_r, 1)]


def zones(ax):
    """Shade the two stretches of double contact of the tracked pair."""
    for lo, hi in ((s_a, s_r - pb), (s_a + pb, s_r)):
        ax.axvspan(lo, hi, color='0.9', lw=0)


def runs(mask):
    """Index runs over which the number of pairs stays the same, so steps are not joined."""
    idx = np.flatnonzero(mask)
    cut = np.flatnonzero((np.diff(idx) > 1) | (np.diff(fr['npairs'][idx]) != 0)) + 1
    return np.split(idx, cut)


def tidy(ax):
    for side in ('top', 'right'):
        ax.spines[side].set_visible(False)


# Load share of the tracked pair, only while it is in contact
fig, ax = plt.subplots(figsize=(7, 3.4))
zones(ax)
on = share > 0
for i, run in enumerate(runs(on)):
    ax.plot(s[run], share[run], color=BLUE, lw=1.8, label='from the tooth stiffness' if i == 0 else None)
ax.plot([s_a, s_r - pb, s_r - pb, s_a + pb, s_a + pb, s_r], [0.5, 0.5, 1, 1, 0.5, 0.5],
        color='k', ls='--', lw=1, label='equal split')
dbl = on & (fr['npairs'] == 2)
lo = share[dbl].min()
ax.set_xlim(s_a, s_r); ax.set_ylim(0, 1.1)
ax.set_xticks(xt)
ax.set_yticks([0, round(lo, 2), 0.5, 1, 1.1])
ax.set_xlabel('contact position $s$ of the pair on the line of action [mm]')
ax.set_ylabel('load share $f/F$')
ax.text(0.5*(s_a + s_r - pb), 0.15, 'double', ha='center', fontsize=9)
ax.text(0.5*(s_r - pb + s_a + pb), 0.15, 'single', ha='center', fontsize=9)
ax.text(0.5*(s_a + pb + s_r), 0.15, 'double', ha='center', fontsize=9)
ax.legend(loc='upper left', frameon=False, fontsize=9)
tidy(ax)
fig.tight_layout()
fig.savefig(graphics/'GridapGears_share.svg')
plt.close(fig)

# Mesh stiffness over the same stretch. Every frame is a position of the whole mesh, so the
# stiffness repeats every base pitch; it is shown against the tracked pair's position.
fig, ax = plt.subplots(figsize=(7, 3.2))
zones(ax)
inside = (s >= s_a) & (s <= s_r)
for run in runs(inside):
    ax.plot(s[run], k[run], color=BLUE, lw=1.8)
k1 = k[inside & (fr['npairs'] == 1)].mean()
k2 = k[inside & (fr['npairs'] == 2)].mean()
kmax = 2*np.ceil(k.max()/2)
ax.set_xlim(s_a, s_r); ax.set_ylim(0, kmax)
ax.set_xticks(xt)
ax.set_yticks([0, 5, round(k1, 1), round(k2, 1), kmax])
ax.set_xlabel('contact position $s$ of the pair on the line of action [mm]')
ax.set_ylabel(r'mesh stiffness [N/(mm$\cdot$µm)]')
tidy(ax)
fig.tight_layout()
fig.savefig(graphics/'GridapGears_stiffness.svg')
plt.close(fig)