import numpy as np
import matplotlib.pyplot as plt
def Plot_eta(eta):
    Theta_set = np.linspace(1 / 2 + 0.01, 0.99, 500)
    Prob = Theta_set * eta + (1 - Theta_set) * (1 - eta)
    KL_divergence = eta * np.log(eta / Prob) + (1 - eta) * np.log((1 - eta) / (1 - Prob))
    Lower_Bound = (1 - Theta_set) ** 2 * np.abs(2 * eta - 1) ** 2 * 2
    Denomi = np.array([max(eta * (1 - eta), i * (1 - i)) for i in Theta_set])
    Upper_Bound = (1 - Theta_set) ** 2 * np.abs(2 * eta - 1) ** 2 / Denomi
    # Upper_Bound = np.exp(eta * np.log(eta/Prob) + (1-eta) * np.log((1-eta)/(1-Prob)))-1
    Upper_Bound_1 = (1 - Theta_set) ** 2 * np.abs(2 * eta - 1) ** 2 / (eta*(1-eta))
    Upper_Bound_2 = (1 - Theta_set) ** 2 * np.abs(2 * eta - 1) ** 2 / (Theta_set*(1-Theta_set))

    epsi_set = np.log(Theta_set / (1 - Theta_set))

    plt.plot(epsi_set, Upper_Bound_1, '--', label='$U_1(\epsilon,x)$', color='green')
    plt.plot(epsi_set, Upper_Bound_2, '--',label='$U_2(\epsilon,x)$', color='yellow')
    plt.plot(epsi_set, KL_divergence, label='$D_{KL}(\mathbb{P}_{Y|X=x}| \mathbb{P}_{\widetilde{Y}|X=x})$', color='blue')
    plt.plot(epsi_set, Lower_Bound, label='Lower Bound: $\mathcal{L}(\epsilon,x)$', color='black')
    plt.plot(epsi_set, Upper_Bound, label='Upper Bound: $\mathcal{U}(\epsilon,x)$', color='red')
    plt.xlabel('$\epsilon$')
    plt.grid()
    plt.legend()

Plot_eta(0.65)
Plot_eta(0.75)
Plot_eta(0.85)
Plot_eta(0.95)