import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import jax
import tensorflow as tf
import seaborn as sns
from io import BytesIO
import tensorflow_probability.substrates.jax as tfp
tfd = tfp.distributions

# plt.rcParams['image.composite_image_limit'] = 200000000
st.set_page_config(layout="wide")
st.set_option('deprecation.showPyplotGlobalUse', False)  # Disable a deprecated warning
# st.set_option('plotly.use_container_width', True)  

def plot_distributions(alpha, beta, left_count, right_count, num_people):
    # Calculate the prior distribution
    prng_key = jax.random.PRNGKey(0)
    prior = tfd.Beta(alpha, beta)
    x = np.linspace(0, 1, 1000)

    posterior_left = tfd.Beta(alpha + left_count, beta + right_count)
    likelihood = tfd.Beta(left_count+1, right_count+1)
    if(left_count+right_count == 0):
        maximum_likelihood_estimate = 0
        maximum_aposteriori_estimate = 0
    else: 
        maximum_likelihood_estimate = (left_count)/(left_count+right_count)
        maximum_aposteriori_estimate = (left_count+alpha)/(left_count+right_count+alpha+beta)
    
    total_count = left_count + right_count

    ## For MLE based prediction
    if(total_count<=num_people):
        samples_posterior_mle = tfd.Binomial(num_people - total_count, probs = maximum_likelihood_estimate).prob(np.arange(0,num_people-total_count+1))
        samples_prior_pred = tfd.BetaBinomial(num_people - total_count, alpha, beta).prob(np.arange(0,num_people-total_count+1))
        samples_posterior_pred = tfd.BetaBinomial(num_people - total_count, alpha + left_count, beta + right_count).prob(np.arange(0,num_people-total_count+1))
    else:
        samples_posterior_mle = tfd.Binomial(num_people, probs = maximum_likelihood_estimate).prob(np.arange(0,num_people+1))
        samples_prior_pred = tfd.BetaBinomial(num_people, alpha, beta).prob(np.arange(0,num_people+1))
        samples_posterior_pred = tfd.BetaBinomial(num_people, alpha + left_count, beta + right_count).prob(np.arange(0,num_people+1))

    # ## For prior predictive distribution
    # dist_pripre = tfd.BetaBinomial(num_people, alpha, beta)

    # ## For posterior predictive distribution
    # dist_postpre = tfd.BetaBinomial(num_people, alpha + left_count, beta + right_count)

    fig, ax = plt.subplots(2, 3, figsize=(15, 12))

    if(left_count+right_count != 0):
        ax[0,0].plot(x, likelihood.prob(x), label='Likelihood')
        ax[0,0].axvline(x=maximum_likelihood_estimate, color='black', label='MLE', linestyle='--')
    ax[0,0].set_ylabel('Probability Density')
    ax[0,0].set_xlabel('Left Handedness Probability')
    ax[0,0].set_title('Likelihood (Left)')
    ax[0,0].legend()

    ax[0,1].plot(x, prior.prob(x), label='Prior')
    ax[0,1].set_xlabel('Handedness')
    ax[0,1].set_ylabel('Probability Density')
    ax[0,1].set_title('Prior for Left Handedness')

    ax[0,2].plot(x, posterior_left.prob(x), label='Posterior')
    ax[0,2].set_ylabel('Probability Density')
    if(left_count+right_count != 0):
        ax[0,2].axvline(x=posterior_left.mean(), color='black', label='Posterior mean', linestyle='--')
        # print(posterior_left.mean())
    ax[0,2].set_xlabel('Left Handedness Probability')
    ax[0,2].set_title('Closed-Form Posterior (Left)')
    ax[0,2].legend()

    if(left_count+right_count != 0):
        if(total_count<=num_people):
            ax[1,0].bar(np.arange(left_count,num_people-total_count+left_count+1), samples_posterior_mle)
            ax[1, 0].set_xlim(-0.5, num_people+0.5)
        else:
            ax[1,0].bar(np.arange(0,num_people+1), samples_posterior_mle)
            ax[1,0].set_xlim(-0.5,num_people+0.5)

    ax[1,0].set_xlabel('Number of Left Handed Student')
    ax[1,0].set_title('Predictive Distribution given MLE')

## Change here
    if(total_count<=num_people):
        ax[1,1].bar(np.arange(left_count,num_people-total_count+left_count+1), samples_prior_pred)
        ax[1,1].set_xlim(-0.5,num_people+0.5)
    else:
        ax[1,1].bar(np.arange(0,num_people+1), samples_prior_pred)
        ax[1,1].set_xlim(-0.5,num_people+0.5)

    ax[1,1].set_xlabel('Number of Left Handed Student')
    ax[1,1].set_title('Prior Predictive Distribution')

## Change here
    if(total_count<=num_people):
        ax[1,2].bar(np.arange(left_count,num_people-total_count+left_count+1), samples_posterior_pred)
        ax[1,2].set_xlim(-0.5,num_people+0.5)
    else:
        ax[1,2].bar(np.arange(0,num_people+1), samples_posterior_pred)
        ax[1,2].set_xlim(-0.5,num_people+0.5)
        
    ax[1,2].set_xlabel('Number of Left Handed Student')
    ax[1,2].set_title('Posterior Predictive Distribution')

    st.pyplot(fig)


def main():
    # Title and description
    st.markdown(
    """
    <div style="text-align:center">
        <h1>Handedness Analysis</h1>
    </div>
    """,
    unsafe_allow_html=True
)
    # st.write('This application calculates the posterior distribution of handedness based on a beta-binomial model and count information.')
    st.sidebar.title("Parameter selection Window")
    # Beta prior parameters
    alpha = st.sidebar.number_input('Alpha:', min_value=0.0, step=0.1, value=1.0)
    beta = st.sidebar.number_input('Beta:', min_value=0.0, step=0.1, value=1.0)

    # Left-handed counter
    left_count = st.sidebar.number_input('Left-handed count:', min_value=0, step=1)

    # Right-handed counter
    right_count = st.sidebar.number_input('Right-handed count:', min_value=0, step=1)

    # Number of people for predictive distribution
    num_people = st.sidebar.number_input('Total Population size:', min_value=1, step=1, value=4)

    # Calculate and plot the distributions
    plot_distributions(alpha, beta, left_count, right_count,num_people)

if __name__ == '__main__':
    main()