from pyscript import document
from pyscript import display
import matplotlib.pyplot as plt
import numpy as np
from scipy.special import beta as beta_func, comb
from scipy.stats import binom

def beta_binomial_pmf(k, n, alpha, beta_param):
    return comb(n, k) * beta_func(k + alpha, n - k + beta_param) / beta_func(alpha, beta_param)

def plot_beta_binomial(p, n, korr):
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # Berechnung der Alpha- und Beta-Parameter für die Beta-Binomialverteilung
    alph = p * (1 - korr) / korr
    bet = (1 - p) * (1 - korr) / korr
    
    # Werte für die x-Achse (Anzahl der Ausfälle)
    k_values = np.arange(0, int(max(23, 1.3 * n * p)))
    
    # Beta-Binomialverteilung berechnen und plotten
    probabilities = beta_binomial_pmf(k_values, n, alph, bet)
    ax.plot(k_values, probabilities, color="#00794f", marker='o', label='Beta-Binomial')
    
    # Binomialverteilung berechnen und als Balkendiagramm darstellen
    probabilities_binom = binom.pmf(k_values, n, p)
    ax.bar(k_values, probabilities_binom, color="#00794f", alpha=0.5, label='Binomial')
    
    # Diagrammeinstellungen
    ax.set_title(f'Beta-Binomialverteilung $\\rho$={korr}, PD={p}, N={n}')
    ax.set_xlabel('Ausfälle')
    ax.set_ylabel('Wahrscheinlichkeit')
    ax.legend()
    ax.grid(True)
    
    return fig  # Rückgabe der Figur

def BetaBin_Auf(event):
    input_text1 = document.querySelector("#pd")
    input_text2 = document.querySelector("#defaultcorr")
    input_text3 = document.querySelector("#anzahl")
    document.getElementById("BinPlot").innerHTML = ""
    p = float(input_text1.value)
    korr = float(input_text2.value)
    n = int(input_text3.value)
    
    fig = plot_beta_binomial(p, n, korr)
    display(fig, target="BinPlot")