"""
//*************************************************************************
// Example program using OpenCV library
//					
// @file	bayes.py					
// @author Luis M. Jimenez
// @date 2025
//
// @brief Course: Computer Vision (1782)
// Dpo. of Systems Engineering and Automation
// Automation, Robotics and Computer Vision Lab (ARVC)
// http://arvc.umh.es
// University Miguel Hernandez
//
// @note Description: 
//	- Bayesian Classifier:  LDA, Naive Bayes, QDA 
//
//  Dependencies:
//     scipy, sklearn, matplotlib, numpy, argparse
//*************************************************************************
"""
import numpy as np
import matplotlib.pyplot as plt
from sklearn.naive_bayes import GaussianNB
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis
from sklearn.model_selection import train_test_split  


# Plot Data / Classification rules
def plotData(data, labels=None, val_data=None, val_labels=None, 
              model=None, fig=None, title="Data", show_ids=False):

    # Plot Data 
    plt.figure(fig, figsize=(9, 7))
    plt.clf()
    plt.title(title)
    plt.xlabel("x1")
    plt.ylabel("x2")
    plt.gcf().canvas.manager.set_window_title('Bayesian Classifier')

    # plot decision rule
    if (model is not None):
        all_data = data if val_data is None else np.concatenate((data, val_data), axis=0)
        x_min, x_max = all_data[:, 0].min() - 1, all_data[:, 0].max() + 1
        y_min, y_max = all_data[:, 1].min() - 1, all_data[:, 1].max() + 1
        xx, yy = np.meshgrid(np.linspace(x_min, x_max, 200),
                            np.linspace(y_min, y_max, 200))

        # predict the grid points
        Z = model.predict(np.c_[xx.ravel(), yy.ravel()])
        Z = Z.reshape(xx.shape)
        # plot decision rule regions
        plt.contourf(xx, yy, Z, alpha=0.2, cmap='coolwarm')

        # predict probabilities for the grid points
        Z_proba = model.predict_proba(np.c_[xx.ravel(), yy.ravel()])  # probabilities
        Z_diff = Z_proba[:, 1] - Z_proba[:, 0]  # Probability difference between classes
        Z_diff = Z_diff.reshape(xx.shape)
        # Plot decision rule limit (equal probability)
        plt.contour(xx, yy, Z_diff, levels=[0], colors='k', linewidths=1)

   # plot data points
    if labels is not None:
        plt.scatter(data[:,0], data[:,1], s=25, marker='o', c=labels, cmap='coolwarm', alpha=0.6, edgecolors='k')
    else:
        plt.scatter(data[:,0], data[:,1], s=25, marker='o', alpha=0.6, edgecolors='k')
 
    # Plot Validation Data (predict)
    if val_data is not None and val_labels is not None:
        plt.scatter(val_data[:,0],val_data[:,1], marker='x', s=40, c=val_labels, cmap='coolwarm', label="Validation Data")
        
        # Plot items ids
        if show_ids:
            for i, (x, y) in enumerate(zip(val_data[:, 0], val_data[:, 1])):
                plt.annotate(str(i), xy=(x, y), xytext=(-3, 3), textcoords='offset points', ha='right', va='bottom')
        plt.legend(loc='lower right')


    plt.draw()
    plt.pause(0.1)  # gives control to GUI event manager to show the plot
# end plotData

#########################################################

# Naive Bayes (Gaussian)
# Non correlated data generator (same variance)
np.random.seed(42)
mean_class0 = [-1, 0]  # mean class 0
mean_class1 = [2, 2]  # mean class 1
cov_matrix = [[1, 0], [0, 1]]  # both variables are correlated 

# Generate random samples for each class
m = 100  # number of samples
X0 = np.random.multivariate_normal(mean_class0, cov_matrix, m)
X1 = np.random.multivariate_normal(mean_class1, cov_matrix, m)

# Build dataset
X = np.vstack((X0, X1))
y = np.hstack((np.zeros(m), np.ones(m)))  # Labels {0  1}

# Split data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=42) 

#########################################
# Train LDA classificator
model_lda = LinearDiscriminantAnalysis()
model_lda.fit(X_train, y_train)

# Make predictions on the test data
y_pred = model_lda.predict(X_test)
accuracy = np.sum(y_pred == y_test)/len(y_pred)    # Calculate the accuracy (1 - error rate)

# obtain class probabilites
y_prob = model_lda.predict_proba(X_test)

print("\nLDA Linear Discriminant Analisys (same variance)")
print(f"Test Accuracy: {accuracy}")

print(f"Per classs priors:   (0):{model_lda.priors_[0]}  |  (1):{model_lda.priors_[1]}")
print(f"Per classs mean:     (0):{model_lda.means_[0]}  |  (1):{model_lda.means_[1]}")
print(f"Model Coefficients (w1, w2): {model_lda.coef_[0]}")
print(f"Intercept/Bias (w0): {model_lda.intercept_[0]}")
print(f"Covariance matrix:\n {model_lda.covariance_}")

plotData(X_train, y_train, val_data=X_test, val_labels=y_test, model=model_lda,
          fig=1, title="LDA Classifier (same variance)")

#########################################

# Train Naive Bayes classificator for LDA data (non correlated with the same variance)
model_nbayes = GaussianNB()
model_nbayes.fit(X_train, y_train)

# Make predictions on the test data
y_pred = model_nbayes.predict(X_test)
accuracy = np.sum(y_pred == y_test)/len(y_pred)    # Calculate the accuracy (1 - error rate)

# obtain class probabilites
y_prob = model_nbayes.predict_proba(X_test)

print("\nNaive Bayes Classifier (same variance, LDA)")
print(f"Test Accuracy: {accuracy}")
print(f"Per classs priors:   (0):{model_nbayes.class_prior_[0]}  |  (1):{model_nbayes.class_prior_[1]}")
print(f"Per classs mean:     (0):{model_nbayes.theta_[0]}  |  (1):{model_nbayes.theta_[1]}")
print(f"Per classs variance: (0):{model_nbayes.var_[0]}  |  (1):{model_nbayes.var_[1]}")

plotData(X_train, y_train, val_data=X_test, val_labels=y_test, model=model_nbayes, 
          fig=2, title="(LDA) Naive Bayes Classifier (same variance)")

#########################################################

# Naive Bayes (Gaussian)
# Non correlated data generator (independent variance)
np.random.seed(42)
mean_class0 = [-1, 0]  # mean class 0
mean_class1 = [2, 2]  # mean class 1
cov_matrix = [[1, 0], [0, 1.5]]  # both variables are independent with different variance  

# Generate random samples for each class
m = 100  # number of samples
X0 = np.random.multivariate_normal(mean_class0, cov_matrix, m)
X1 = np.random.multivariate_normal(mean_class1, cov_matrix, m)

# Build dataset
X = np.vstack((X0, X1))
y = np.hstack((np.zeros(m), np.ones(m)))  # Labels {0  1}

# Split data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) 

# Train Naive Bayes classificator
model_nbayes = GaussianNB()
model_nbayes.fit(X_train, y_train)

# Make predictions on the test data
y_pred = model_nbayes.predict(X_test)
accuracy = np.sum(y_pred == y_test)/len(y_pred)    # Calculate the accuracy (1 - error rate)

# obtain class probabilites
y_prob = model_nbayes.predict_proba(X_test)

print("\nNaive Bayes Classifier (independent variance)")
print(f"Test Accuracy: {accuracy}")
print(f"Per classs priors:   (0):{model_nbayes.class_prior_[0]}  |  (1):{model_nbayes.class_prior_[1]}")
print(f"Per classs mean:     (0):{model_nbayes.theta_[0]}  |  (1):{model_nbayes.theta_[1]}")
print(f"Per classs variance: (0):{model_nbayes.var_[0]}  |  (1):{model_nbayes.var_[1]}")

plotData(X_train, y_train, val_data=X_test, val_labels=y_test, model=model_nbayes, 
         fig=3, title="Naive Bayes Classifier (independent variance)")

#########################################################

# Correlated data generator
np.random.seed(42)
mean_class0 = [-1, 0]  # mean class 0
mean_class1 = [2, 2]  # mean class 1
cov_matrix = [[1, 0.8], [0.8, 1]]  # non diagonal covariance matrix

# Generate random samples for each class
m = 100  # number of samples
X0 = np.random.multivariate_normal(mean_class0, cov_matrix, m)
X1 = np.random.multivariate_normal(mean_class1, cov_matrix, m)

# Build dataset
X = np.vstack((X0, X1))
y = np.hstack((np.zeros(m), np.ones(m)))  # Labels {0  1}

# Split data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) 

###############################################################

# Generalized Bayesian Classifier for general covariance matrix
model_qda = QuadraticDiscriminantAnalysis(store_covariance=True)
model_qda.fit(X_train, y_train)

# Make predictions on the test data
y_pred = model_qda.predict(X_test)
accuracy = np.sum(y_pred == y_test)/len(y_pred)    # Calculate the accuracy (1 - error rate)

# obtain class probabilites
y_prob = model_qda.predict_proba(X_test)

print("\nGeneralized Bayesian Classifier (correlated data)")
print(f"Test Accuracy: {accuracy}")
print(f"Per classs priors:     (0):{model_qda.priors_[0]}  |  (1):{model_qda.priors_[1]}")
print(f"Per classs mean:     (0):{model_qda.means_[0]}  |  (1):{model_qda.means_[1]}")
for i, cov in enumerate(model_qda.covariance_):
    print(f"Covariance matrix class {i}:\n{cov}\n")

plotData(X_train, y_train, val_data=X_test, val_labels=y_test, model=model_qda, 
          fig=4, title="(QDA) Generalized Bayesian Classifier (correlated data)")

input('\n...Press return to finish program...')