"""Map pedestal energy in the stated additive rectangular-pulse model."""
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt

duration = np.geomspace(1, 1e9, 401)
contrast = np.geomspace(10, 1e9, 401)
x, y = np.meshgrid(duration, contrast)
fraction = x / (y - 1 + x)
assert np.all((fraction > 0) & (fraction < 1))
assert abs(10/(1001-1+10)-1/101) < 1e-15
assert 1000/(1001-1+1000) == .5
for c in [10, 1001, 1e6]:
    assert abs((c-1)/((c-1)+(c-1))-.5) < 1e-15

# Verify the ghost weights by exact discrete cross-correlation of disjoint pulses.
epsilon = .01
g = np.array([0, 1, 1, 1, 0], dtype=float)
i = np.zeros(50)
i[5:10] = g
i[30:35] = epsilon*g
correlation = np.correlate(i, i*i, mode='full')
mid = len(i)-1
np.testing.assert_allclose(correlation[[mid-25, mid, mid+25]],
                           [3*epsilon**2, 3*(1+epsilon**3), 3*epsilon], rtol=1e-12)

fig, ax = plt.subplots(figsize=(9, 4.5), dpi=100, layout='constrained')
fig.set_facecolor('white')
mesh = ax.pcolormesh(x, y, fraction, shading='auto', cmap='viridis', vmin=0, vmax=1,
                     rasterized=True)
contours = ax.contour(x, y, fraction, levels=[.1, .5, .9], colors='white', linewidths=1)
labels = ax.clabel(contours, fmt={.1: '10%', .5: '50%', .9: '90%'}, inline=True, fontsize=9,
                  manual=[(1e4, 9e4+1), (1e6, 1e6+1), (1e7, 1e7/9+1)])
for label in labels:
    label.set_bbox({'facecolor': '#202020', 'edgecolor': 'none', 'alpha': .7, 'pad': 1})
ax.set(xscale='log', yscale='log', xlabel='Pedestal duration / main duration (T / τ)',
       ylabel='Peak-to-pedestal power contrast C',
       title='Weak light can carry substantial energy if it lasts long enough')
for t, label, labelpos in [(10, '0.99% energy', (80, 2e4)),
                           (1000, '50% energy', (1.6e4, 100))]:
    ax.plot(t, 1001, 'o', color='white', markeredgecolor='black', ms=6)
    ax.annotate(label, (t, 1001), labelpos, color='white', fontsize=9,
                arrowprops={'arrowstyle': '->', 'color': 'white'},
                bbox={'facecolor': '#202020', 'alpha': .8, 'edgecolor': 'none', 'pad': 3})
fig.colorbar(mesh, ax=ax, label='Pedestal energy / total energy')
fig.savefig(Path.cwd()/'optical-pedestal-energy-and-contrast.png', facecolor='white')
plt.close(fig)