import numpy as np
import matplotlib.pyplot as plt
import scipy.stats


def heaviside(x):
    return x >= 0


def main():
    num_samples = 100000
    num_bins = 100
    start = -10
    stop = 10
    num_zs = 1000
    theta = 0.3

    # z
    u = np.random.uniform(size=(num_samples,))
    z = np.log(theta / (1 - theta)) + np.log(u / (1 - u))
    z_kde = scipy.stats.gaussian_kde(z)

    # z_density
    zs = np.linspace(start, stop, num=num_zs)
    z_density = (theta / (1 - theta)) * np.exp(-zs) / (1 + (theta / (1 - theta)) * np.exp(-zs))**2

    # z_tilde_b0
    v = np.random.uniform(size=(num_samples,))
    b = 0
    z_tilde_b0 = np.log(v / ((1 - v) * (1 - theta)) + 1) if b == 1 else -np.log(v / ((1 - v) * theta) + 1)
    z_tilde_b0_kde = scipy.stats.gaussian_kde(z_tilde_b0)

    # z_tilde_b1
    v = np.random.uniform(size=(num_samples,))
    b = 1
    z_tilde_b1 = np.log(v / ((1 - v) * (1 - theta)) + 1) if b == 1 else -np.log(v / ((1 - v) * theta) + 1)
    z_tilde_b1_kde = scipy.stats.gaussian_kde(z_tilde_b1)

    # z_tilde_b0_density
    z_tilde_b0_density = z_density * (1 / (1 - theta)) * (1 - heaviside(zs))

    # z_tilde_b1_density
    z_tilde_b1_density = z_density * (1 / theta) * heaviside(zs)

    # plotting
    fig, ax = plt.subplots(1, 1)
    fig.set_size_inches(4, 3)
    # ax.plot(zs, z_kde.evaluate(zs), label='$p(z)$', color='black', linestyle='dashed')
    ax.hist(z, bins=num_bins, density=True, alpha=0.8, color='black')
    ax.plot(zs, z_density, label='$p_{\\theta}^z(z)$', color='black', linestyle='solid')
    # ax.plot(zs, z_tilde_b0_kde.evaluate(zs), label='$p(z | b = 0)$', color='gray', linestyle='dashed')
    ax.hist(z_tilde_b0, bins=num_bins, density=True, alpha=0.8, color='gray')
    ax.plot(zs, z_tilde_b0_density, label='$p_{\\theta}^{z | b}(z | b = 0)$', color='gray', linestyle='solid')
    # ax.plot(zs, z_tilde_b1_kde.evaluate(zs), label='$p(z | b = 1)$', color='lightgray', linestyle='dashed')
    ax.hist(z_tilde_b1, bins=num_bins, density=True, alpha=0.8, color='lightgray')
    ax.plot(zs, z_tilde_b1_density, label='$p_{\\theta}^{z | b}(z | b = 1)$', color='lightgray', linestyle='solid')
    ax.set_xlim(-10, 10)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    # ax.set_yticks([])
    ax.get_yaxis().set_visible(False)
    ax.legend()
    fig.tight_layout()
    filenames = ['rebar_relax_dists.png', 'rebar_relax_dists.pdf']
    for filename in filenames:
        fig.savefig(filename, bbox_inches='tight', dpi=200)
        print('Saved to {}'.format(filename))


if __name__ == '__main__':
    main()