"""Illustrate original binary lever-rule sensitivity and covariance examples."""
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt

def phase_fraction(y, a, b, weight):
    d = b - a
    return d @ weight @ (y - a) / (d @ weight @ d)

# Independent finite-difference check of all residual-aware GLS derivatives.
a = np.array([0.60, 0.30])
b = np.array([0.20, 0.40])
y = np.array([0.470, 0.325])
weight = np.diag([1 / 0.010**2, 1 / 0.005**2])
d = b - a
q = d @ weight @ d
lam = phase_fraction(y, a, b, weight)
r = y - a - lam * d
g = d @ weight / q
h = r @ weight / q
expected = np.concatenate([g, -(1 - lam) * g - h, -lam * g + h])
x = np.concatenate([y, a, b])
numerical = []
for i in range(6):
    perturb = np.zeros(6)
    perturb[i] = 1e-6
    lo, hi = x - perturb, x + perturb
    numerical.append((phase_fraction(hi[:2], hi[2:4], hi[4:], weight)
                      - phase_fraction(lo[:2], lo[2:4], lo[4:], weight)) / 2e-6)
np.testing.assert_allclose(numerical, expected, rtol=1e-8, atol=1e-8)

offset = np.linspace(-0.006, 0.006, 401)
rho = np.linspace(-1, 1, 401)
uncertainty = 0.01 * np.sqrt(0.52 + 0.48 * rho)
np.testing.assert_allclose(uncertainty[[0, 200, -1]], [0.002, np.sqrt(0.52)*0.01, 0.01])
for m, c in [(1.0, 0.006), (1.03, -0.01), (0.9, 0.05)]:
    assert abs(((m*.44+c)-(m*.2+c))/((m*.8+c)-(m*.2+c))-.4) < 1e-14

fig, axes = plt.subplots(1, 2, figsize=(9, 4), dpi=100, layout='constrained')
fig.set_facecolor('white')
for sep, color in [(0.60, '#1769aa'), (0.06, '#b24a20')]:
    shifted = (0.4*sep - offset)/(sep - offset)
    axes[0].plot(offset*1000, shifted, color=color, label=f'Endpoint separation {sep:.2f}')
axes[0].axhline(0.4, color='#777777', ls=':', lw=1)
axes[0].set(title='Closer endpoints amplify errors', xlabel='Error in endpoint a (×0.001 mole fraction)',
            ylabel='Inferred phase fraction', ylim=(0.325, 0.465))
axes[0].legend(fontsize=8, loc='upper right')
axes[1].plot(rho, uncertainty, color='#1769aa')
axes[1].plot([-1, 0, 1], uncertainty[[0, 200, -1]], 'o', color='#1769aa', ms=4)
axes[1].set(title='Endpoint correlation matters', xlabel='Endpoint correlation coefficient',
            ylabel='Standard uncertainty of phase fraction', ylim=(0, 0.012))
axes[1].text(0.04, 0.94, 'Overall composition exact\nEndpoint standard uncertainties: 0.006',
             transform=axes[1].transAxes, va='top', fontsize=8)
for ax in axes:
    ax.grid(alpha=0.2)
fig.savefig(Path.cwd()/'tie-line-sensitivity-and-covariance.png', facecolor='white')
plt.close(fig)