import numpy as np
from matplotlib import pyplot as plt
from scipy.stats import binom
from matplotlib.ticker import LogFormatterMathtext, LogLocator

plt.rcParams['text.usetex'] = True
plt.rcParams['text.latex.preamble'] = r'\usepackage{amsfonts}'


nRange = [100, 2000]

nVals = np.arange(start=nRange[0], stop=nRange[1], step=5)
estCounts = nVals * 0.4

P = np.array([0.5, 0.5])
P_r = np.array([0.4, 0.6])

targKLD = (P_r[0] * np.log2(P_r[0] / P[0])) + (P_r[1] * np.log2(P_r[1] / P[1]))

typeUpperBounds = np.zeros(shape=nVals.shape)
typeLowerBounds = np.zeros(shape=nVals.shape)
exactVals = np.zeros(shape=nVals.shape)

for i in range(len(nVals)):
    targNVal = nVals[i]
    targCount = estCounts[i]

    kldExponent = 2**(-1* targNVal * targKLD)

    exactVals[i] = binom.pmf(k=targCount, n=targNVal, p=0.5)
    typeLowerBounds[i] = kldExponent / ((targNVal+1)**(2))
    typeUpperBounds[i] = kldExponent

style = {
    "font.family": "DejaVu Sans",
    "font.size": 11,
    "axes.labelsize": 12,
    "axes.titlesize": 12,
    "legend.fontsize": 9.5,
    "xtick.labelsize": 10,
    "ytick.labelsize": 10,
    "mathtext.fontset": "dejavusans",
    "svg.fonttype": "none",  # Preserve text as editable SVG text.
}

with plt.rc_context(style):
    fig, ax = plt.subplots(
        figsize=(7.2, 4.7),
        constrained_layout=True,
    )

    ax.plot(
        nVals,
        exactVals,
        linewidth=2.2,
        label=r"$\mathbb{P}\left[B(n, 0.5) = \frac{2n}{5}\right]$",
    )
    ax.plot(
        nVals,
        typeUpperBounds,
        linestyle="--",
        linewidth=1.9,
        label=r"Upper bound: $2^{-nD(P_r\Vert P)}$",
    )
    ax.plot(
        nVals,
        typeLowerBounds,
        linestyle=":",
        linewidth=2.2,
        label=r"Lower bound: $(n+1)^{-2}2^{-nD(P_r\Vert P)}$",
    )

    # Base-2 logarithmic probability scale.
    ax.set_yscale("log", base=2)
    ax.yaxis.set_major_locator(LogLocator(base=2, numticks=9))
    ax.yaxis.set_major_formatter(LogFormatterMathtext(base=2))

    ax.set_xlim(nVals[0], nVals[-1])
    ax.set_xticks([100, 500, 1000, 1500, 2000])

    ax.set_xlabel(r"Sample size $n$")
    ax.set_ylabel("Probability")

    ax.grid(
        axis="y",
        which="major",
        linewidth=0.6,
        alpha=0.35,
    )

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

    ax.legend(frameon=True, loc="upper right")

    metadata = {
        "Title": "Binary type probability and method-of-types bounds",
        "Description": (
            "The exact probability that a Bernoulli(1/2) sample has "
            "type (2/5, 3/5), together with the standard "
            "method-of-types bounds."
        ),
    }

    fig.savefig(
        "./MethodOfTypesBound.svg",
        bbox_inches="tight",
        metadata=metadata,
    )

print("This is here to catch a break point")