"""
Created on Tue Feb 11 15:43:26 2025

@author: Aniket Mitra
PostDoc ISTeP-CYU and AMGC-VUB
Modified after: https://doi.org/10.5281/zenodo.14776445
"""


# === 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

sns.set_theme(style="whitegrid")

# === FUNCTIONS ===
def load_data(filepath):
    filepath = r'filepath'
    """Load quantitative data for a single group 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=',')

    # --- Select data for ONE group (e.g., 'Violet') ---
    group_name = 'Green'  # Change this to the group you want to analyze
    data_group = data_total[data_total['Cluster'] == group_name]

    data_quant = (
        data_group.iloc[:, :-1]  # Numerical data 
        .apply(pd.to_numeric, errors="coerce")
        .fillna(0)
    )
    return data_total, data_group, data_quant


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_density_with_heatmap_and_contours(pca_df, PC_x, PC_y, title, pca):
    """Plots density with heatmap, contours, and sample points for ONE group."""

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

    # 1. Heatmap (Background):
    x = pca_df[f"PC{PC_x+1}"]
    y = pca_df[f"PC{PC_y+1}"]
    heatmap, x_edges, y_edges = np.histogram2d(x, y, bins=50)
    heatmap = heatmap / heatmap.sum()
    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)

    # 2. Contours (Overlay):
    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)
        contour = ax.contour(xx, yy, z, levels=levels, colors='k', linewidths=0.5)

    except ValueError:
        print("Warning: Not enough data points to compute density.")

    # 3. Sample Points (Overlay):
    ax.scatter(pca_df[f"PC{PC_x+1}"], pca_df[f"PC{PC_y+1}"], c='none',  # Single color
               label=None, alpha=0.8, edgecolors='k', linewidths=0.2, zorder=2)  # No label needed

    # 4. Axis Labels and Title:
    var_x = pca.explained_variance_ratio_[PC_x] * 100
    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")

    plt.tight_layout()
    plt.show()
    
def filter_top_density(pca_df, original_data, PC_x, PC_y, threshold_fraction=2/3):
    """Filter the top 2/3 of samples based on density estimation."""
    X = pca_df[[f"PC{PC_x+1}", f"PC{PC_y+1}"]].values.T
    kde = gaussian_kde(X)
    density = kde(X)
        
    threshold = np.percentile(density, (1 - threshold_fraction) * 100)
    filtered_indices = density >= threshold
        
    return original_data.iloc[filtered_indices]

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_Loc123', fontsize=15)
    plt.xlabel('Component Number', fontsize=15)
    plt.ylabel('Proportion of Variance', fontsize=15)
    plt.grid()
    plt.show()
    

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_Loc123', fontsize=15)
    plt.xlabel('Component Number', fontsize=15)
    plt.ylabel('Eigenvalues', fontsize=15)
    plt.grid()
    plt.show()
    

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="coolwarm", 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_Loc123")
    plt.show()
    

def plot_biplot(PC, PC_x, PC_y, loadings, features, title, pca):
    """Generate a biplot to visualize PCA results 
    with vectors representing features."""
    
    plt.figure(figsize=(12, 8))
    
    
    var_x = pca.explained_variance_ratio_[PC_x] * 100
    var_y = pca.explained_variance_ratio_[PC_y] * 100
    
    # Add the feature vectors with scaling
    scale = 10  # Adjust this value to control the scaling
    for i, feature in enumerate(features):
        plt.arrow(0, 0, loadings[PC_x, i] * scale, loadings[PC_y, i] * scale,
              head_width=0.05, head_length=0.1, color='red')  # Increased head size
        plt.text(loadings[PC_x, i] * scale * 1.1, loadings[PC_y, i] * scale * 1.1,
             feature, color='black', fontsize=8)
        
    # Plot the points (all the same color now)
    plt.scatter(PC[:, PC_x], PC[:, PC_y], 
            facecolors='green',  # Fill color
            edgecolors='black',   # Border color
            linewidths=0.8,       # Border thickness
            s=30,                 
            alpha=0.8)            # Transparency

    

    
    
    # Labels and title
    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.grid()

    # Add the legend with formatted labels
    plt.legend(title="Groups", loc="best")

    plt.grid()
    plt.show()
    

def main():
    """Main function to execute the PCA analysis
    and generate visualizations."""
    filepath = r"C:\Users\amiar\Documents\Academic Phase 3\Calcite thin sections\20241129 LAser ablation raw data\Analysis\PCA-analysis_Copy2.csv"
    data_total, data_group, data_quant = load_data(filepath)
    DF_scaled = standardize_data(data_quant)
    
    pca, PC = perform_pca(DF_scaled, n_components=6)
    
    pca_df = pd.DataFrame(PC, columns=[f"PC{i+1}" for i in range(PC.shape[1])])
    

    
    plot_scree_plot(pca)
    plot_kaiser_rule(pca)
    plot_variance_bar(*compute_variance(pca, 6), 6,
                      "PCA on 6 components / Percentage of Variance Explained")
    plot_heatmap(pca.components_, data_quant.columns)
    
    pca, PC = perform_pca(DF_scaled, n_components=3)
    plot_variance_bar(*compute_variance(pca, 3), 3,
                      "PCA on 3 components / Percentage of Variance Explained")
    
  
    plot_density_with_heatmap_and_contours(pca_df, 0, 1, "Density with Heatmap and Contours PC1/PC2", pca)  # Pass pca
    plot_density_with_heatmap_and_contours(pca_df, 0, 2, "Density with Heatmap and Contours PC1/PC3", pca)  # Pass pca
    plot_density_with_heatmap_and_contours(pca_df, 1, 2, "Density with Heatmap and Contours PC2/PC3", pca)  # Pass pca
    
    
    
    plot_biplot(PC, 0, 1, pca.components_, data_quant.columns, "Biplot on PC1/PC2", pca)
    plot_biplot(PC, 0, 2, pca.components_, data_quant.columns, "Biplot on PC1/PC3", pca)
    plot_biplot(PC, 1, 2, pca.components_, data_quant.columns, "Biplot on PC2/PC3", pca)
    
    filtered_data = filter_top_density(pca_df, data_group, 0, 1)
    output_filepath = r'C:\Users\amiar\Documents\Academic Phase 3\Calcite thin sections\20241129 LAser ablation raw data\Analysis\filtered_data_Green.csv'
    filtered_data.to_csv(output_filepath, index=False)
    print(f"Filtered data saved to: {output_filepath}")
    
    # Perform PCA again on refined data
    DF_scaled_filtered = standardize_data(filtered_data.select_dtypes(include=[np.number]))
    pca_filtered, PC_filtered = perform_pca(DF_scaled_filtered, n_components=3)
    pca_df_filtered = pd.DataFrame(PC_filtered, columns=[f"PC{i+1}" for i in range(PC_filtered.shape[1])])
   
    # Generate plots for refined data
    plot_scree_plot(pca_filtered)
    plot_density_with_heatmap_and_contours(pca_df_filtered, 0, 1, "Density with Heatmap and Contours PC1 vs PC2 (Filtered)", pca)
    plot_density_with_heatmap_and_contours(pca_df_filtered, 0, 2, "Density with Heatmap and Contours PC1 vs PC3 (Filtered)", pca)
    plot_density_with_heatmap_and_contours(pca_df_filtered, 1, 2, "Density with Heatmap and Contours PC2 vs PC3 (Filtered)", pca)
    
    
    pca_df_filtered = pca_df_filtered.to_numpy()  # Convert DataFrame to NumPy array
    plot_scree_plot(pca_filtered)
    plot_biplot(pca_df_filtered, 0, 1, pca.components_, data_quant.columns, "Biplot on PC1/PC2 (Filtered)", pca)
    plot_biplot(pca_df_filtered, 0, 2, pca.components_, data_quant.columns, "Biplot on PC1/PC3 (Filtered)", pca)
    plot_biplot(pca_df_filtered, 1, 2, pca.components_, data_quant.columns, "Biplot on PC2/PC3 (Filtered)", pca)    

