"""Selection on a noisy denominator: exact Gaussian model and ratio bounds."""
from pathlib import Path
import math

import matplotlib.pyplot as plt
import numpy as np


def main():
    mean, sigma, threshold, side = 5., 1., 6., .05
    z = (threshold-mean)/sigma
    phi = math.exp(-z*z/2)/math.sqrt(2*math.pi)
    accepted = math.erfc(z/math.sqrt(2))/2
    selected_mean = mean + sigma*phi/accepted
    lower, upper = side/selected_mean, side/threshold
    assert np.isclose(accepted, .15865525393145707)
    assert np.isclose(selected_mean, 6.525135276160981)
    nodes, weights = np.polynomial.legendre.leggauss(128)
    # The omitted Gaussian tail above 15 is smaller than 8e-24 in probability.
    c = (15-threshold)/2*nodes+(15+threshold)/2
    pdf = np.exp(-.5*((c-mean)/sigma)**2)/(sigma*np.sqrt(2*np.pi))
    ratio_mean = (15-threshold)/2*np.sum(weights*side/c*pdf)/accepted
    assert lower < ratio_mean < upper < side/mean
    fig, (density, ratio) = plt.subplots(1, 2, figsize=(9.4, 4.2), dpi=100,
                                        layout="constrained")
    fig.set_facecolor("white")
    x = np.linspace(0, 9, 901)
    f = np.exp(-.5*((x-mean)/sigma)**2)/(sigma*np.sqrt(2*np.pi))
    density.plot(x, f, color="#1565a5", lw=2)
    keep = x >= threshold
    density.fill_between(x[keep], 0, f[keep], color="#cc741c", alpha=.4)
    density.axvline(mean, color="#555555", ls="--", lw=1.2, label="True central signal = 5")
    density.axvline(threshold, color="#ad4f00", lw=1.3, label="Acceptance cutoff = 6")
    density.text(6.35, .18, f"Accepted\n{accepted:.2%}", color="#8d4300", fontsize=10)
    density.set(xlabel="Measured central signal C", ylabel="Probability density",
                title="A cutoff retains the upward-noise tail", xlim=(0, 9), ylim=(0, .44))
    density.legend(fontsize=8, loc="upper left")
    measured = np.linspace(threshold, 9, 301)
    ratio.plot(measured, side/measured, color="#1565a5", lw=2,
               label="Individual accepted ratio: 0.05/C")
    ratio.axhspan(lower, upper, color="#e5c76e", alpha=.4,
                  label="Bounds on conditional mean ratio")
    ratio.axhline(side/mean, color="#a33434", ls="--", lw=1.3,
                  label="True-signal ratio = 0.01")
    ratio.set(xlabel="Accepted central signal C", ylabel="Normalized side signal",
              title="Every accepted ratio is below 0.01", xlim=(6, 9), ylim=(.0045, .0107))
    ratio.legend(fontsize=8, loc="lower left")
    for ax in [density, ratio]:
        ax.grid(alpha=.16)
    fig.savefig(Path(__file__).with_suffix(".png").name, facecolor="white")
    plt.close(fig)


if __name__ == "__main__":
    main()