
@author: Aniket Mitra
PostDoc ISTeP-CYU and AMGC-VUB
Modified after Miloikovitch, T. 2024
 

# === IMPORTS ===
import numpy as np
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import seaborn as sns
import os
from scipy.stats import gaussian_kde
import matplotlib.patches as mpatches
from scipy.stats import t
sns.set_theme(style="whitegrid")

# === FUNCTIONS ===

def load_data(filepath):
    filepath=()
    """Load quantitative and qualitative data from a CSV file."""
    if not os.path.exists(filepath):
        raise FileNotFoundError(
            f"The file {filepath} was not found. Please check the path.")
    
    
    
    data_total = pd.read_csv(filepath, sep=',')
    data_quant = (
        # Extract numerical data except the last column (qualitative)
        pd.concat([data_total.iloc[:, 8:12], data_total.iloc[:, 14:16]], axis=1)
        # Convert decimal separators if necessary
        # Ensure numerical format
        .apply(pd.to_numeric, errors="coerce")
        # Replace missing values with zero
        .fillna(0)
    )
    data_qual = data_total['Cluster']  # Extract the categorical column
    return data_quant, data_qual

def standardize_data(data):
    """Standardize data using StandardScaler (zero mean, unit variance)."""
    scaler = StandardScaler()
    data_scaled = scaler.fit_transform(data)
    return pd.DataFrame(data=data_scaled, columns=data.columns)

def perform_pca(data, n_components):
    """Perform Principal Component Analysis (PCA)."""
    pca = PCA(n_components=n_components)
    transformed_data = pca.fit_transform(data)
    return pca, transformed_data

def plot_heatmap_with_contours_all_data(pca_df, PC_x, PC_y, title, pca):
    """Plots heatmap with contours using ALL data points (no clusters)."""

    fig, ax = plt.subplots(figsize=(10, 8))

    # 1. Density Estimation & Contours (using ALL data):
    x = pca_df[f"PC{PC_x+1}"]
    y = pca_df[f"PC{PC_y+1}"]

    try:
        kde = gaussian_kde([x, y])  # Use ALL data points directly
        x_min, x_max = x.min() - 1, x.max() + 1
        y_min, y_max = y.min() - 1, y.max() + 1
        xx, yy = np.mgrid[x_min:x_max:50j, y_min:y_max:50j]  # Adjust grid density as needed
        positions = np.vstack([xx.ravel(), yy.ravel()])
        z = kde(positions)
        z = np.reshape(z.T, xx.shape)

        levels = np.linspace(0, z.max(), 7)  # Adjust levels
        

    except ValueError:
        print("Warning: Not enough data points to compute overall density.")
        xx, yy = np.mgrid[x.min() - 1:x.max() + 1:50j, y.min() - 1:y.max() + 1:50j]  # Default grid
        z = np.zeros_like(xx)  # Empty density
        levels = np.linspace(0, 1, 7)
        contours = ax.contour(xx, yy, z, levels=levels, colors='k', linewidths=0.5)  # Plot dummy contours

    # 2. Heatmap (using ALL data):
    x_edges = np.linspace(xx.min(), xx.max(), 51)  # 50 bins + 1 edge
    y_edges = np.linspace(yy.min(), yy.max(), 51)
    heatmap, _, _ = np.histogram2d(x, y, bins=(x_edges, y_edges))  # Use ALL data

    heatmap = heatmap / heatmap.sum()  # Normalize (optional)
    heatmap = heatmap.T
    extent = [x_edges[0], x_edges[-1], y_edges[0], y_edges[-1]]

    im = ax.imshow(heatmap, extent=extent, origin='lower', cmap='tab20b', aspect='auto', alpha=0.7)  # Added alpha


        # 4. Axis Labels and Title:
    var_x = pca.explained_variance_ratio_[PC_x] * 100  # Use pca here
    var_y = pca.explained_variance_ratio_[PC_y] * 100
    ax.set_xlabel(f'PC{PC_x+1} ({var_x:.1f}%)')
    ax.set_ylabel(f'PC{PC_y+1} ({var_y:.1f}%)')
    ax.set_title(title)

    # 5. Colorbar:
    cbar = fig.colorbar(im, ax=ax, label="Density")  # Use im and ax

    # 6. Legend:
    handles = [plt.Rectangle((0, 0), 1, 1, color=cluster_data.iloc[0]['ClusterColor']) for cluster_data in [pca_df[pca_df['Cluster'] == label] for label in pca_df['Cluster'].unique()]]
    labels = [label for label in pca_df['Cluster'].unique()]
    ax.legend(handles, labels, title="Groups")

    handles = []
    labels = []
    unique_labels = pca_df['Cluster'].unique()  # Get unique labels *once*
    colors_map = {'Blue': 'blue', 'Green': 'green', 'Red': 'red', 'Violet': 'violet', 'Blue': 'blue'} # Your color map
    for label in unique_labels:
        cluster_in_projection = pca_df[pca_df['Cluster'] == label]  # Check presence in projection
        if not cluster_in_projection.empty: # If the cluster exists in the projection
            color = colors_map.get(label, 'gray')  # Get color, default to gray
            handles.append(plt.Rectangle((0, 0), 1, 1, color=color))
            labels.append(label)
            mask = pca_df['Cluster'] == label
            x = pca_df.loc[mask, f"PC{PC_x+1}"]
            y = pca_df.loc[mask, f"PC{PC_y+1}"]
            confidence_ellipse(x, y, ax, color=color) # Add ellipse for each cluster and pass ax
  
    
    plt.tight_layout()
    plt.show()
    plt.grid()
    plt.close()
    
def plot_scree_plot(pca):
    """Generate a scree plot to visualize the explained variance ratio."""
    PC_number = np.arange(pca.n_components_) + 1
    prop_var = pca.explained_variance_ratio_
    
    plt.figure(figsize=(10, 6))
    plt.plot(PC_number, prop_var, 'ro-')
    plt.title('Scree Plot (Elbow Method) PCA_Loc', fontsize=15)
    plt.xlabel('Component Number', fontsize=15)
    plt.ylabel('Proportion of Variance', fontsize=15)
    plt.grid()
    plt.show()
    plt.close()

def plot_kaiser_rule(pca):
    """Generate a scree plot with the Kaiser rule (eigenvalues > 1)."""
    PC_number = np.arange(pca.n_components_) + 1
    eigenvalues = pca.explained_variance_
    
    plt.figure(figsize=(10, 6))
    plt.plot(PC_number, eigenvalues, 'ro-')
    plt.axhline(y=1, color='r', linestyle='--')  # Kaiser threshold
    plt.title('Scree Plot (Kaiser Rule) PCA', fontsize=15)
    plt.xlabel('Component Number', fontsize=15)
    plt.ylabel('Eigenvalues', fontsize=15)
    plt.grid()
    plt.show()
    plt.close()

def compute_variance(pca, n_components):
    """Calculate explained variance for a given number of components."""
    prop_var = pca.explained_variance_ratio_[:n_components]
    variance_cumulative = np.cumsum(prop_var)
    return prop_var, variance_cumulative
    
def plot_variance_bar(prop_var, variance_cumulative, n_components, title):
    """Create a bar plot for explained variance."""
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.bar(range(1, n_components + 1), prop_var*100, alpha=0.5)
    ax.set_xlabel('Principal Component')
    ax.set_ylabel('Percentage of Explained Variance')
    ax.set_title(title)
    ax.axvline(x=3, color='red', linestyle='--')
    ax.text(3.2, variance_cumulative[2] * 0.002, 
            f'{variance_cumulative[2] * 100:.2f}% of explained variance', 
            color='red', ha='left', va='center', fontsize=12)
    plt.show()
    
def plot_heatmap(loadings, features):
    """Plot a heatmap showing the contributions of
    variables to principal components."""
    plt.figure(figsize=(12, 8))
    sns.heatmap(
        loadings, cmap="tab20b", annot=True, fmt=".2f", linewidths=0.5)
    plt.xticks(range(len(features)), features, rotation=45, ha='right')
    plt.yticks(
        range(len(loadings)), [f"PC{i+1}" for i in range(len(loadings))])
    plt.xlabel("Variables")
    plt.ylabel('Principal Components')
    plt.title("Variables' contributions to principal components PCA")
    plt.show()
    plt.close()
    
def confidence_ellipse(x, y, ax, color='black', n_std=1.96, **kwargs):
    """Draw a 95% confidence ellipse for a given cluster."""
    if len(x) < 2:
        return  # Not enough data points to draw an ellipse

    cov = np.cov(x, y)
    mean_x, mean_y = np.mean(x), np.mean(y)

    # Degrees of freedom
    n = len(x)
    dof = n - 1  # Degrees of freedom

    # t-statistic for 95% confidence
    t_crit = t.ppf(0.975, dof)  # Two-tailed test

    # Eigenvalues and eigenvectors
    eigenvalues, eigenvectors = np.linalg.eigh(cov)
    order = eigenvalues.argsort()[::-1]
    eigenvalues, eigenvectors = eigenvalues[order], eigenvectors[:, order]

    # Compute the angle and width-height of the ellipse
    theta = np.degrees(np.arctan2(*eigenvectors[:, 0][::-1]))
    width, height = 2 * t_crit * np.sqrt(eigenvalues)  # Use t-statistic

    ellipse = mpatches.Ellipse((mean_x, mean_y), width, height, angle=theta,
                                edgecolor=color, facecolor='none', lw=2, **kwargs)
    ax.add_patch(ellipse)


def plot_biplot(PC, PC_x, PC_y, loadings, features, data_qual, title, pca):
    """Generate a biplot to visualize PCA results with vectors representing features."""
    var_x = pca.explained_variance_ratio_[PC_x] * 100
    var_y = pca.explained_variance_ratio_[PC_y] * 100

    fig, ax = plt.subplots(figsize=(14, 9))  # Create figure and axes *ONLY ONCE*
    

    # Clean the labels in data_qual to remove spaces and standardize the case
    data_qual = data_qual.str.strip().str.capitalize()

    # Check the unique labels in data_qual
    unique_labels = data_qual.unique()
    print("Unique labels:", unique_labels)  

    # Adjust colors_map to match the labels
    colors_map = {'Green': 'green',
                  'Violet': 'violet',
                  'Blue': 'blue',
                  'Red': 'red'}
    
    # Prepare legend mapping
    legend_map = {'Green': 'Green: K/Fl/38',
                  'Violet': 'Violet: K/Pk/02',
                  'Blue': 'Blue: K/Fl/15',
                  'Red': 'Red: K/Fl/57'}

    # Check if all the labels are correctly mapped
    if not all(label in colors_map for label in unique_labels):
        print("Warning: some labels are not in colors_map.")
        print("Unmapped labels:",
              [label for label in unique_labels if label not in colors_map])
    
    # Assign colors to labels
    colors = data_qual.map(colors_map).fillna('gray')  

    # Check the color mapping
    print("Mapped colors:", colors.unique())

    # Plot the points with the legend, without black borders
    for label in unique_labels:
        mask = data_qual == label
        plt.scatter(PC[mask, PC_x], PC[mask, PC_y],
                    label=legend_map.get(label, label),
                    c=colors[mask])  # Remove edgecolors='k'

    # Add the feature vectors
    for i, feature in enumerate(features):
        plt.arrow(0, 0, loadings[PC_x, i]*28, loadings[PC_y, i]*8,
                  head_width=0.02*8, head_length=0.02*8, color='red')
        plt.text(loadings[PC_x, i] * 3*8, loadings[PC_y, i] * 1.1*8,
                 feature, color='black', fontsize=12)
    
    # Add confidence ellipses - CORRECTED CALL:
    unique_labels = data_qual.unique()
    colors_map = {'Green': 'green', 'Violet': 'violet', 'Blue': 'blue', 'Red': 'red'}  # Your color map
    for label in unique_labels:
        mask = data_qual == label
        x = PC[mask, PC_x]
        y = PC[mask, PC_y]
        color = colors_map.get(label, 'gray')  # Get color, default to gray
        confidence_ellipse(x, y, ax, color=color)  # Pass 'ax' here!

    plt.xlabel(f'PC{PC_x + 1} ({var_x:.2f}% of total variance)')
    plt.ylabel(f'PC{PC_y + 1} ({var_y:.2f}% of total variance)')
    plt.title(title)
    plt.legend(title="Groups", loc="best") # Add legend after ellipses are plotted
    plt.show()
    plt.grid()
    plt.close()



def main():
    """Main function to execute the PCA analysis
    and generate visualizations."""
    filepath = (filepath)
    data_quant, data_qual = load_data(filepath)
    DF_scaled = standardize_data(data_quant)
    
    pca, PC = perform_pca(DF_scaled, n_components=11)

    pca_df = pd.DataFrame(PC, columns=[f"PC{i+1}" for i in range(PC.shape[1])])
    pca_df['Cluster'] = data_qual

    # Color mapping (moved outside the plotting functions):
    colors_map = {'Green': 'green', 'Violet': 'violet', 'Blue': 'blue', 'Red': 'red'}
    pca_df['ClusterColor'] = pca_df['Cluster'].map(colors_map).fillna('gray')
    
    


    # Plot using ALL data points:
    plot_heatmap_with_contours_all_data(pca_df, 0, 1, "Density with Heatmap PC1/PC2 (All Data)", pca)
    plot_heatmap_with_contours_all_data(pca_df, 0, 2, "Density with Heatmap PC1/PC3 (All Data)", pca)
    plot_heatmap_with_contours_all_data(pca_df, 1, 2, "Density with Heatmap PC2/PC3 (All Data)", pca)

    
    plot_scree_plot(pca)
    plot_kaiser_rule(pca)
    plot_variance_bar(*compute_variance(pca, 11), 11,
                      "PCA on 16 components / Percentage of Variance Explained")
    plot_heatmap(pca.components_, data_quant.columns)
    
    pca, PC = perform_pca(DF_scaled, n_components=11)
    plot_variance_bar(*compute_variance(pca, 11), 11,
                      "PCA on 3 components / Percentage of Variance Explained")
    
  
    
    plot_biplot(PC, 0, 1, pca.components_,
                data_quant.columns, data_qual, "Biplot on PC1/PC2", pca)
    plot_biplot(PC, 0, 2, pca.components_,
                data_quant.columns, data_qual, "Biplot on PC1/PC3", pca)
    plot_biplot(PC, 1, 2, pca.components_,
                data_quant.columns, data_qual, "Biplot on PC2/PC3", pca)
    
    print("Unique cluster labels:", pca_df['Cluster'].unique())
    print("Missing ClusterColor:", pca_df[pca_df['ClusterColor'].isna()])
# === MAIN ===
if __name__ == "__main__":
    main()

