#!/usr/bin/env python
#
# file: MLE vs Bayesian Estimation Comparison
#
# revision history:
#
# 20260120 (AA): initial version using Bernoulli/Beta conjugate example
#------------------------------------------------------------------------------

# import system modules
#
import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import beta

#------------------------------------------------------------------------------
#
# global variables are listed here
#
#------------------------------------------------------------------------------

# set the filename using basename
#
__FILE__ = os.path.basename(__file__)

# define problem parameters
# Case: A coin is flipped 3 times and lands on Heads every time.
#
DEF_DATA          = [1, 1, 1]  # 1 = Head, 0 = Tail
DEF_PRIOR_ALPHA   = 2.0        # Prior successes (fictional)
DEF_PRIOR_BETA    = 2.0        # Prior failures (fictional)
DEF_OUT_FILE_NAME = "mle_vs_bayes_1d.png"
DEF_PLOT_TITLE    = "MLE vs. Bayesian Estimation (3 Heads in a Row)"

#------------------------------------------------------------------------------
#
# functions are listed here
#
#------------------------------------------------------------------------------

# compute the maximum likelihood estimate
# formula: heads / total_flips
#
def compute_mle(data):
    
    n = len(data)
    k = np.sum(data)
    
    if n == 0:
        return 0.0
    
    theta_mle = float(k) / n
    return theta_mle

# compute the bayesian estimate (Posterior Mean)
# formula: (heads + prior_alpha) / (total_flips + prior_alpha + prior_beta)
#
def compute_bayes_mean(data, alpha_prior, beta_prior):
    
    n = len(data)
    k = np.sum(data)
    
    # posterior parameters
    #
    alpha_post = alpha_prior + k
    beta_post = beta_prior + (n - k)
    
    # expected value of Beta distribution (Mean)
    #
    theta_bayes = alpha_post / (alpha_post + beta_post)
    
    return theta_bayes, alpha_post, beta_post

# plot the distributions
#
def plot_comparison(mle_val, bayes_val, post_alpha, post_beta, 
                    prior_alpha, prior_beta, outfile):
    
    # x-axis range (probability of heads from 0 to 1)
    #
    x = np.linspace(0, 1, 500)
    
    plt.figure(figsize=(10, 6))

    # 1. Plot the Prior Distribution
    #
    y_prior = beta.pdf(x, prior_alpha, prior_beta)
    plt.plot(x, y_prior, 'g--', label=f'Prior (Beta {prior_alpha},{prior_beta})')
    plt.fill_between(x, y_prior, color='green', alpha=0.1)

    # 2. Plot the Posterior Distribution
    #
    y_post = beta.pdf(x, post_alpha, post_beta)
    plt.plot(x, y_post, 'b-', linewidth=2, 
             label=f'Posterior (Beta {post_alpha},{post_beta})')
    plt.fill_between(x, y_post, color='blue', alpha=0.1)

    # 3. Mark the Estimates
    #
    # MLE Estimate line
    plt.axvline(mle_val, color='red', linestyle='-', linewidth=2, 
                label=f'MLE Estimate = {mle_val:.2f}')
    
    # Bayesian Estimate line
    plt.axvline(bayes_val, color='blue', linestyle='--', linewidth=2, 
                label=f'Bayes Estimate = {bayes_val:.2f}')

    # labels
    #
    plt.title(DEF_PLOT_TITLE)
    plt.xlabel("Probability of Heads (theta)")
    plt.ylabel("Density")
    plt.legend(loc='upper left')
    plt.grid(True, alpha=0.3)
    
    # save
    #
    plt.savefig(outfile)
    print("Saved plot to: %s" % outfile)

# function: main
#
def main(argv):
    
    print("Running Estimation Comparison...")

    # 1. Define Data (Small non-Gaussian set)
    #
    data = np.array(DEF_DATA)
    print(f"  Data: {data} ({len(data)} flips, {np.sum(data)} Heads)")

    # 2. Compute MLE
    #
    theta_mle = compute_mle(data)
    print(f"  MLE Estimate: {theta_mle:.4f}")
    print("    -> Implies Tails is IMPOSSIBLE (Overfitting)")

    # 3. Compute Bayesian Estimate
    #
    # We use a Beta(2,2) prior which represents a weak belief that
    # the coin is fair (centered at 0.5).
    #
    theta_bayes, a_post, b_post = compute_bayes_mean(data, 
                                                     DEF_PRIOR_ALPHA, 
                                                     DEF_PRIOR_BETA)
    
    print(f"  Prior: Beta({DEF_PRIOR_ALPHA}, {DEF_PRIOR_BETA})")
    print(f"  Posterior: Beta({a_post}, {b_post})")
    print(f"  Bayes Estimate: {theta_bayes:.4f}")
    print("    -> Implies Tails is still possible (Regularization)")

    # 4. Visualization
    #
    plot_comparison(theta_mle, theta_bayes, a_post, b_post, 
                    DEF_PRIOR_ALPHA, DEF_PRIOR_BETA, DEF_OUT_FILE_NAME)

    # exit gracefully
    #
    return True

# begin gracefully
#
if __name__ == '__main__':
    main(sys.argv[0:])

#
# end of file