#!/usr/bin/env python
#
# file: 3D Multivariate Gaussian Surface Comparison
#
# revision history:
#
# 20260114 (AA): added file I/O for covariance matrices
# 20260114 (AA): reproduce slide 9 visualization with side-by-side comparison
#------------------------------------------------------------------------------

# import system modules
#
import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D 
from scipy.stats import multivariate_normal

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

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

# define default values for arguments
#
DEF_MEAN          = [0.0, 0.0]
DEF_GRID_SIZE     = int(100)
DEF_GRID_RANGE    = 15.0 
DEF_INPUT_FILE    = "cov_input.txt"
DEF_OUT_FILE_NAME = "gaussian_comparison.png"
DEF_PLOT_TITLE    = "Effect of Covariance on Distribution Shape"
DEF_SYMMETRY_TOL  = 1e-12

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

# read covariance matrices from a text file
#
def load_cov_matrices(filename):
    
    matrices = []
    current_rows = []

    # check if file exists
    #
    if not os.path.exists(filename):
        print("**> Error: input file not found (%s)" % filename)
        return None

    try:
        with open(filename, 'r') as fp:
            for line in fp:
                # strip whitespace and comments
                #
                line = line.strip()
                if (not line) or (line.startswith('#')):
                    continue
                
                # parse numbers
                #
                parts = line.split()
                if len(parts) != 2:
                    print("**> Error: expected 2 numbers per line, got: %s" 
                          % line)
                    return None
                
                # convert to float
                #
                row = [float(parts[0]), float(parts[1])]
                current_rows.append(row)

                # if we have collected 2 rows, we have a complete matrix
                #
                if len(current_rows) == 2:
                    matrices.append(current_rows)
                    current_rows = []
                    
                    # stop after finding 2 matrices
                    #
                    if len(matrices) == 2:
                        break
        
        # validate we found exactly two matrices
        #
        if len(matrices) != 2:
            print("**> Error: expected 2 matrices in file, found %d" 
                  % len(matrices))
            return None

        return matrices

    except ValueError:
        print("**> Error: file contains non-numeric data")
        return None
    except Exception as e:
        print("**> Error reading file: %s" % str(e))
        return None

# compute the PDF of a multivariate Gaussian on a 2D grid
#
def compute_gaussian_pdf(mu, cov, grid_range=DEF_GRID_RANGE, 
                         grid_size=DEF_GRID_SIZE):

    # check basic input dimensions
    #
    mu = np.asarray(mu, dtype=float).reshape(-1)
    cov = np.asarray(cov, dtype=float)

    if (mu.size != 2) or (cov.shape != (2, 2)):
        print("**> Error: mu must be length 2 and cov must be 2x2")
        return None, None, None

    # ensure the covariance matrix is symmetric
    #
    if not np.allclose(cov, cov.T, atol=DEF_SYMMETRY_TOL):
        print("**> Error: covariance matrix must be symmetric")
        return None, None, None

    # create the grid (x and y values)
    #
    x_val = np.linspace(-grid_range, grid_range, int(grid_size))
    y_val = np.linspace(-grid_range, grid_range, int(grid_size))
    X, Y = np.meshgrid(x_val, y_val)

    # stack x and y into a single (N, N, 2) array for the PDF function
    #
    pos = np.dstack((X, Y))

    # calculate the probability density function (PDF) values
    #
    try:
        # allow singular matrices for educational visualization purposes
        rv = multivariate_normal(mean=mu, cov=cov, allow_singular=True)
        Z = rv.pdf(pos)
    except Exception as e:
        print("**> Error computing PDF: %s" % str(e))
        return None, None, None

    # return the meshgrid and z values
    #
    return X, Y, Z

# helper function to format a subplot
#
def format_subplot(ax, X, Y, Z, title):
    
    # plot surface to match the slide style (Red surface, black mesh)
    #
    ax.plot_surface(X, Y, Z, 
                    color='red',
                    edgecolor='black',
                    linewidth=0.3,
                    alpha=0.9,
                    rstride=5, cstride=5,
                    antialiased=True)

    # label the plot and axes
    #
    ax.set_title(title, fontsize=10, fontfamily='monospace', pad=20)
    ax.set_xlabel("x-values", labelpad=10)
    ax.set_ylabel("y-values", labelpad=10)
    ax.set_zlabel("pdf values", labelpad=10)

    # set background color
    #
    ax.set_facecolor('white')
    
    # adjust initial view angle for better perspective
    #
    ax.view_init(elev=30, azim=-60)

# plot two 3D surface views side-by-side
#
def plot_dual_surfaces(data1, data2, title=DEF_PLOT_TITLE, 
                       outfile=DEF_OUT_FILE_NAME):
    
    # unpack data tuples
    #
    X1, Y1, Z1, cov1_str = data1
    X2, Y2, Z2, cov2_str = data2

    # check input data
    #
    if Z1 is None or Z2 is None:
        print("**> Error: no data to plot")
        return False

    # create the figure with a wider aspect ratio for two plots
    #
    fig = plt.figure(figsize=(16, 7))

    # determine common z-axis limit for fair comparison
    #
    z_max = max(np.max(Z1), np.max(Z2))
    # add a small buffer
    z_lim = (0, z_max * 1.1) 

    # --- Subplot 1 (Left) ---
    #
    ax1 = fig.add_subplot(121, projection='3d')
    title1 = "Covariance (Uncorrelated):\n%s" % cov1_str
    format_subplot(ax1, X1, Y1, Z1, title1)
    ax1.set_zlim(z_lim)

    # --- Subplot 2 (Right) ---
    #
    ax2 = fig.add_subplot(122, projection='3d')
    title2 = "Covariance (Correlated):\n%s" % cov2_str
    format_subplot(ax2, X2, Y2, Z2, title2)
    ax2.set_zlim(z_lim)

    # set main figure title
    #
    fig.suptitle(title, fontsize=14, y=0.98)

    # adjust layout to prevent overlapping labels
    #
    plt.tight_layout()

    # save to disk 
    #
    plt.savefig(outfile, bbox_inches='tight')
    print("Saved plot to: %s" % outfile)
    
    # show the plot to the user
    #
    plt.show()

    # exit gracefully
    #
    return True

# function: main
#
def main(argv):
    
    print("Generating comparison plot...")

    # parse command line arguments for input file
    #
    infile = DEF_INPUT_FILE
    if len(argv) > 1:
        infile = argv[1]
    
    print("Reading covariance matrices from: %s" % infile)

    # load matrices from file
    #
    matrices = load_cov_matrices(infile)
    
    if matrices is None:
        print("**> Error: Failed to load valid matrices.")
        return False

    cov1 = matrices[0]
    cov2 = matrices[1]

    # compute the data for both cases
    #
    X1, Y1, Z1 = compute_gaussian_pdf(DEF_MEAN, cov1)
    X2, Y2, Z2 = compute_gaussian_pdf(DEF_MEAN, cov2)

    # check for errors
    #
    if Z1 is None or Z2 is None:
        print("**> Error computing Gaussian data")
        return False

    # package data for the plotting function
    # passing the covariance as a string for the label
    #
    data1 = (X1, Y1, Z1, str(cov1))
    data2 = (X2, Y2, Z2, str(cov2))

    # plot the surfaces side-by-side
    #
    plot_dual_surfaces(data1, data2)

    # exit gracefully
    #
    return True

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

#
# end of file