analysis.py

"""
    This module file contain functions regarding the performence analysis of the BuRNN training
"""

__author__  = "Pierre-Alexandre HO"
__date__    = "2025-02-03"
__version__ = "1.0"


# standard python modules
import os
from pathlib import Path
from glob import glob
from typing import List

# NN packages
import ase
import torch
import yaml
import schnetpack.transform as trn
import schnetpack as spk

# Data & vizualisation
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import seaborn as sns
from plotly.subplots import make_subplots
from matplotlib import pyplot as plt


class YamlParser:
    def __init__(self, yaml_file):
        """
        Initializes the parser by loading the YAML file.
        
        :param yaml_file: Path to the YAML file.
        """
        self.data = self._load_yaml(yaml_file)
    
    def _load_yaml(self, yaml_file):
        """
        Loads and parses the YAML file.
        
        :param yaml_file: Path to the YAML file.
        :return: Parsed YAML data as a dictionary.
        """
        with open(yaml_file, 'r') as file:
            return yaml.safe_load(file)
    
    def get_cutoff(self):
        """
        Retrieves the cutoff value from the parsed YAML data.
        
        :return: The cutoff value from the YAML file.
        """
        try:
            return self.data['globals']['cutoff']
        except KeyError as e:
            raise KeyError(f"Missing key in YAML file: {e}")
        

def get_calculator(model_path, device) -> spk.interfaces.SpkCalculator:
    """
    Create a SchnetPack calculator from the given model path.

    Args:
        model_path (Path): Path to the SchNet model.

    Returns:
        spk.interfaces.SpkCalculator: A configured SchnetPack calculator.
    """
    # get cutoff from configuration file
    model_args_path = model_path.parent / 'config.yaml'
    cutoff = YamlParser(model_args_path).get_cutoff()
    model_path = os.path.join(model_path.parent,'best_model')
    
    calculator = spk.interfaces.SpkCalculator(
        model_file = model_path, # path to model
        dtype=torch.float32, # 
        neighbor_list=trn.ASENeighborList(cutoff=cutoff), # neighbor list
        energy_key='V_BuRNN', # name of energy property in model
        force_key='F_BuRNN', # name of force property in model
        energy_unit="eV", # units of energy property predicted by ase
        device=device, # device for computation
    )
    return calculator


def predict_energy_and_forces(system: ase.Atoms, calculator):
    """
    Predict energy and forces for the given atomic system.

    Args:
        system (ase.Atoms): Atomic system for which predictions are required.

    Returns:
        tuple: Predicted energy (float) and forces (np.ndarray).
    """
    # Set the predictive calculator for the atomic system.
    system.set_calculator(calculator)

    # Perform predictions.
    energy = system.get_potential_energy()
    forces = system.get_forces()
    return energy, forces


def initialize_calculator(name_model: str, base_path: str, device: torch.device):
    """
    Initializes the predictive calculator using the specified model.

    Parameters
    ----------
    name_model : str
        The name of the model to load.
    base_path : str
        The base directory where the model is stored.
    device : torch.device
        The device to use for computation (e.g., 'cpu' or 'cuda').

    Returns
    -------
    Any
        The initialized predictive calculator.
    """
    model_path = Path(base_path) / name_model
    print(model_path)
    pred_calculator = get_calculator(model_path=model_path, device=device)

    return pred_calculator


def parse_energy_force(data_base, pred_calculator) -> dict[str, np.ndarray]:
    """
    Extracts reference and predicted energies and force magnitudes from an ASE database.

    The force magnitude is computed as the Euclidean norm of the force vector 
    for each atom in each molecular structure.

    Parameters
    ----------
    data_base : Any
        An ASE database containing molecular structures, atomic numbers, 
        reference energies, and force components.

    pred_calculator : Any
        A predictive model used to compute energy and force predictions.

    Returns
    -------
    dict[str, np.ndarray]
        A dictionary with:
        - 'idx' : np.ndarray
            A NumPy array containing the indices of processed structures.
        - 'energies' : np.ndarray of shape (n_structures, 2)
            - Column 0: Reference energy from the database.
            - Column 1: Predicted energy from the model.
        - 'magnitude' : np.ndarray of shape (n_structures, max_atoms, 2)
            - Column 0: Reference force magnitude for each atom.
            - Column 1: Predicted force magnitude for each atom.
            - NaN values are used for padding in structures with fewer atoms.
    """
    n_structure: int = len(data_base)
    list_idx = []
    energies = np.empty((n_structure, 2))
    numbers_atoms = []
    # Récupérer le nombre maximal d'atomes pour dimensionner le tableau 4D
    for db_item in data_base:
        numbers_atoms.append(int(db_item['_n_atoms']) )
        if db_item['_idx'] == n_structure-1:
            break
    max_atoms = max(numbers_atoms)
    
    # Initialisation du tableau 4D avec des NaN
    magnitude = np.full((n_structure, max_atoms, 2, 3), np.nan)

    for idx in range(n_structure):
        list_idx.append(idx)
        db_item = data_base[idx]
        molecule = ase.Atoms(numbers=db_item['_atomic_numbers'], positions=db_item['_positions'])
        energy, force = predict_energy_and_forces(molecule, pred_calculator)

        # Calcul des magnitudes des forces par atome
        ref = db_item['F_BuRNN'].numpy()
        pred = force

        # Insérer dans le tableau en conservant la structure (NaN pour padding)
        n_atoms = ref.shape[0]
        magnitude[idx, :n_atoms, 0] = ref
        magnitude[idx, :n_atoms, 1] = pred

        # Stockage des énergies
        energies[idx] = [db_item['V_BuRNN'][0].numpy(), energy]

    return {'idx': np.array(list_idx), 'energies': energies, 'magnitude': magnitude}


def parse_split(split: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Loads and extracts train, validation, and test indices from a saved NumPy file.

    Parameters
    ----------
    split : str
        Path to the `.npz` file containing the split indices.

    Returns
    -------
    tuple[np.ndarray, np.ndarray, np.ndarray]
        - train (np.ndarray): Array of training set indices.
        - val (np.ndarray): Array of validation set indices.
        - test (np.ndarray): Array of test set indices.
    """
    split = np.load(split)

    train = split['train_idx']
    val = split['val_idx']
    test = split['test_idx']

    return train, val, test


def split(idx, train, val, test):
    """
    Assigns a dataset split ('train', 'val', 'test') based on index membership.

    Parameters
    ----------
    idx : int
        The structure index.
    train : set
        Set of training indices.
    val : set
        Set of validation indices.
    test : set
        Set of test indices.

    Returns
    -------
    str
        The split category ('train', 'val', 'test') or 'error' if not found.
    """
    if idx in train:
        return 'train'
    elif idx in val:
        return 'val'
    elif idx in test:
        return 'test'
    else:
        return 'error'


def add_split(df: pd.DataFrame, train: set, val: set, test: set) -> pd.DataFrame:
    """
    Adds a 'Split' column to the DataFrame indicating the dataset partition.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing a 'structure_id' column.
    train : set
        Set of training indices.
    val : set
        Set of validation indices.
    test : set
        Set of test indices.

    Returns
    -------
    pd.DataFrame
        DataFrame with an added 'Split' column.
    """
    df['Split'] = df['structure_id'].apply(lambda idx: split(idx, train, val, test))
    return df


def compute_rmse(errors: np.ndarray) -> float:
    """
    Computes the Root Mean Square Error (RMSE).

    Parameters
    ----------
    errors : np.ndarray
        Array containing the prediction errors.

    Returns
    -------
    float
        The RMSE value.
    """
    return np.sqrt(np.mean(errors.explode()))


def compute_mae(y_ref: np.ndarray, y_pred: np.ndarray) -> float:
    """
    Computes the Mean Absolute Error (MAE).

    Parameters
    ----------
    y_ref : np.ndarray
        Array of references values.
    y_pred : np.ndarray
        Array of predicted values.

    Returns
    -------
    float
        The MAE value.
    """
    return np.mean(np.abs(y_ref - y_pred).explode())


def compute_r2(dft_error: np.ndarray, pred_error: np.ndarray) -> float:
    """
    Computes the R² score with division safety check.

    Parameters
    ----------
    dft_error : np.ndarray
        Array containing the dft errors.
    pred_error : np.ndarray
        Array containing the prediction errors.

    Returns
    -------
    float
        The R² score.
    """
    denominator = np.sum(dft_error.explode())
    return 1 - (np.sum(pred_error.explode()) / denominator) if denominator != 0 else float('nan')


def energy_analysis(energy_array: np.ndarray) -> pd.DataFrame:
    """
    Perform energy analysis on an array of reference and predicted energy values.

    Parameters
    ----------
    energy_array : np.ndarray
        A 2D numpy array where the first column contains reference energy values
        and the second column contains predicted energy values.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing the structure IDs, reference energy, predicted energy,
        dft error, and prediction error for each structure.

    Notes
    -----
    - NaN values in the reference and predicted energy columns are removed before analysis.
    - The dft error is calculated as the squared deviation from the mean of the reference values.
    - The prediction error is calculated as the squared difference between the reference and predicted values.
    """
    structure_list = list(range(len(energy_array)))

    y_ref = energy_array[:, 0]  # Reference values
    y_ref = y_ref[~np.isnan(y_ref)]
    y_pred = energy_array[:, 1]  # Predicted values
    y_pred = y_pred[~np.isnan(y_pred)]

    # Squared deviations and residuals
    dft_error = (y_ref - np.mean(y_ref)) ** 2  # Total variance
    pred_error = (y_ref - y_pred) ** 2  # Squared residuals

    # Creating DataFrame
    df = pd.DataFrame({
        'structure_id': structure_list,
        'dft_energy': y_ref,
        'pred_energy': y_pred,
        'dft_error_square': dft_error,
        'pred_error_square': pred_error
    })

    return df


def force_analysis(magnitude_array: list[np.ndarray]) -> pd.DataFrame:
    """
    Analyzes the force magnitudes from reference and predicted data.

    Parameters
    ----------
    magnitude_array : list[np.ndarray]
        A list of NumPy arrays where each entry corresponds to a structure.
        Each array has shape (n_atoms, 2), where:
        - Column 0: Reference magnitudes
        - Column 1: Predicted magnitudes

    Returns
    -------
    pd.DataFrame
        A DataFrame containing:
        - 'structure_id': The index of the structure.
        - 'atom_id': The index of each atom within the structure.
        - 'dft_force': The reference magnitude values.
        - 'pred_force': The predicted magnitude values.
        - 'dft_error_square': The dft error.
        - 'pred_error_square': The prediction error.
    Notes
    -----
    - NaN values in the reference and predicted energy columns are removed before analysis.
    - The dft error is calculated as the squared deviation from the mean of the reference values.
    - The prediction error is calculated as the squared difference between the reference and predicted values.
    """
    structure_list, atom_list = [], []
    y_ref, y_pred = [], []
    dft_error, pred_error = [], []

    for idx, structure in enumerate(magnitude_array):
        ref_values = structure[:, 0]
        ref_values = ref_values[~np.isnan(ref_values).any(axis=1)]
        pred_values = structure[:, 1]
        pred_values = pred_values[~np.isnan(pred_values).any(axis=1)]
        y_ref.extend(ref_values)  # Reference values
        y_pred.extend(pred_values)  # Predicted values
        
        n_atoms = len(ref_values)
        structure_list.extend([idx] * n_atoms)  # Repeat structure index for each atom
        atom_list.extend(range(n_atoms))  # Atom indices
        
        dft_error.extend((ref_values - np.mean(ref_values)) ** 2)  # Extract dft error
        pred_error.extend((ref_values - pred_values) ** 2)  # Extract prediction error
        
    # Creating DataFrame
    df = pd.DataFrame({
        'structure_id': structure_list,
        'atom_id': atom_list,
        'dft_force': y_ref,
        'pred_force': y_pred,
        'dft_error_square': dft_error,
        'pred_error_square': pred_error
    })

    return df


def compute_metrics(df: pd.DataFrame) -> dict[str, float | list]:
    """
    Computes prediction metrics (RMSE, MAE, R²) from a given DataFrame.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing columns:
        - 'y_ref' (np.ndarray): Reference force magnitudes or energy.
        - 'y_pred' (np.ndarray): Predicted force magnitudes or energy.
        - 'dft_error_square' (np.ndarray): Squared dft error values.
        - 'pred_error_square' (np.ndarray): Squared prediction error values.

    Returns
    -------
    dict[str, float if no split | list if split in df.colums]
        Dictionary containing:
        - 'RMSE' : Root Mean Squared Error.
        - 'MAE'  : Mean Absolute Error.
        - 'R²'   : Coefficient of determination.
    """
    
    # Give the property name (force or energy) of the ref/pred column
    ref_property = next((name for name in df.columns if name in ["dft_force", 'dft_energy']), None)
    pred_property = next((name for name in df.columns if name in ["pred_force", 'pred_energy']), None)
    
    if 'Split' in df.columns and not df['Split'].empty:
        split_list = df['Split'].unique().tolist()
        rmse = [compute_rmse(df['pred_error_square'])]
        mae = [compute_mae(df[ref_property], df[pred_property])]
        r2 = [compute_r2(df['dft_error_square'], df['pred_error_square'])]
        
        for split in split_list:
            sub_df = df[df['Split']==split]
            rmse.append(compute_rmse(sub_df['pred_error_square']))
            mae.append(compute_mae(sub_df[ref_property], sub_df[pred_property]))
            r2.append(compute_r2(sub_df['dft_error_square'], sub_df['pred_error_square']))
        
        split_list = ['Global'] + split_list
        
        print(f"Metrics:")
        print("+" + "-" * 12 + ("+" + "-" * 12) * len(split_list) + "+")
        print("| {:^10} ".format("Metric") + "".join(f"| {split:^10} " for split in split_list) + "|")
        print("+" + "-" * 12 + ("+" + "-" * 12) * len(split_list) + "+")
        print("| {:^10} ".format("RMSE") + "".join(f"| {val:^10.4f} " for val in rmse) + "|")
        print("| {:^10} ".format("MAE")  + "".join(f"| {val:^10.4f} " for val in mae)  + "|")
        print("| {:^10} ".format("R²")   + "".join(f"| {val:^10.4f} " for val in r2)   + "|")
        print("+" + "-" * 12 + ("+" + "-" * 12) * len(split_list) + "+")
                    
        return {"Split":split_list, "RMSE": rmse, "MAE": mae, "R²": r2}
    
    else:
        rmse = compute_rmse(df['pred_error_square'])
        mae = compute_mae(df[ref_property], df[pred_property])
        r2 = compute_r2(df['dft_error_square'], df['pred_error_square'])

        print(f"Global Metrics:\n"
            f"{'-'*20}\n"
            f"RMSE : {rmse:.4f}\n"
            f"MAE  : {mae:.4f}\n"
            f"R²   : {r2:.4f}")

        return {"RMSE": rmse, "MAE": mae, "R²": r2}
    

def box_ref_pred(model_df: pd.DataFrame) -> None:
    """
    Generates a boxplot comparing reference and predicted energies.

    Parameters
    ----------
    model_df : pd.DataFrame
        A DataFrame containing at least the following columns:
        - 'structure_id' : Index of the structures.
        - 'dft_energy' : Reference energy values.
        - 'pred_energy' : Predicted energy values.

    Returns
    -------
    None
        Displays an interactive boxplot comparing reference and predicted energy distributions.
    """
    # Convert to long format
    model_df_long = model_df.melt(id_vars=["structure_id"], 
                                  value_vars=["dft_energy", "pred_energy"], 
                                  var_name="Type", 
                                  value_name="Energy")

    # Box plot
    fig = px.box(model_df_long, x="Type", y="Energy", title="Boxplot of Reference vs Predicted Energy", hover_name='structure_id')
    fig.show()


def plot_ref_vs_pred_old(df: pd.DataFrame, ref_col: str, pred_col: str, idx_col: str, title: str, units: str, 
                 hover_info: dict = None, split_col: str = None, save:str = None, show: bool = False)-> None:
    """
    Generates a scatter plot comparing predicted and reference values, with optional coloring by dataset split.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing reference and predicted values.
    ref_col : str
        Column name for reference (dft) values.
    pred_col : str
        Column name for predicted values.
    idx_col : str
        Column name for index or structure identifier (used for hover text).
    title : str
        Title of the plot.
    units : str
        Units of the plot in MathJax format: <sup>-1</sup> 
    hover_info : dict, optional
        Dictionary containing additional columns to display in hover text.
        Example: {"label": "Atom", "col": "atom_id"}
    split_col : str, optional
        Column name indicating dataset split (e.g., "train", "val", "test").
        If provided, points will be colored based on this split.
    save : str, optional
        Path name of the file to save in html
    show : bool, optional
        Tell if the figure has to be show, by default False

    Returns
    -------
    None
    """
    fig = go.Figure()

    # Add identity line (x=y)
    fig.add_trace(go.Scatter(
        x=df[ref_col], y=df[ref_col],
        mode='lines',
        name='x=y',
        line=dict(color='gray', width=1, dash='dash')
    ))

    # Manage hover text and custom data
    hovertemplate = f"<b>Structure #: %{{text}}</b><br>Ref: %{{x}}<br>Pred: %{{y}}"
    customdata = None
    
    if hover_info:
        hover_label = hover_info.get("label", "Custom Data")
        hover_col = hover_info.get("col")
        if hover_col and hover_col in df.columns:
            hovertemplate = f"<b>Structure #: %{{text}}</b><br>{hover_label}: %{{customdata[0]}}<br>Ref: %{{x}}<br>Pred: %{{y}}"
            customdata = df[[hover_col]].to_numpy()

    # Define color mapping for splits
    split_colors = {'train': 'blue', 'val': 'orange', 'test': 'green'}
    
    if split_col and split_col in df.columns:
        unique_splits = df[split_col].unique()
        for split in unique_splits:
            subset = df[df[split_col] == split]
            fig.add_trace(go.Scatter(
                x=subset[ref_col].explode(), y=subset[pred_col].explode(),
                mode='markers',
                marker=dict(color=split_colors.get(split, 'gray'), size=6, opacity=0.7),
                hovertemplate=hovertemplate,
                text=subset[idx_col],
                customdata=subset[[hover_col]].to_numpy() if hover_info else None,
                name=f'Predictions ({split})'
            ))
    else:
        # Default scatter plot without split coloring
        fig.add_trace(go.Scatter(
            x=df[ref_col], y=df[pred_col],
            mode='markers',
            marker=dict(color='blue', size=6, opacity=0.7),
            hovertemplate=hovertemplate,
            text=df[idx_col],
            customdata=customdata,
            name='Predictions'
        ))

    # Styling axes
    fig.update_xaxes(title_text=f'<b>Referecence Values</b>', showline=True, linewidth=2, 
                     linecolor='black', mirror=True)
    fig.update_yaxes(title_text=f'<b>Predicted {title}</b>', showline=True, linewidth=2, 
                     linecolor='black', mirror=True)

    # Layout settings
    fig.update_layout(
        title=f'<b>{title} Prediction vs Reference in {units}</b>',
        plot_bgcolor='white',
        xaxis=dict(showgrid=True, gridwidth=1, gridcolor='lightgrey', zeroline=False),
        yaxis=dict(showgrid=True, gridwidth=1, gridcolor='lightgrey', zeroline=False)
    )
    
    if save:
        if not save.endswith('.html'):
            save = save + '.html'
        fig.write_html(save)
            
    if show:
        fig.show()


def plot_ref_vs_pred(df: pd.DataFrame, ref_col: str, pred_col: str, idx_col: str, title: str, units: str,
                     hover_info: dict = None, split_col: str = None, save: str = None, show: bool = False)-> None:
    """
    Generates a scatter plot comparing predicted and reference values, with optional coloring by dataset split.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing reference and predicted values.
    ref_col : str
        Column name for reference (dft) values.
    pred_col : str
        Column name for predicted values.
    idx_col : str
        Column name for index or structure identifier (used for hover text).
    title : str
        Title of the plot.
    units : str
        Units of the plot in MathJax format: <sup>-1</sup>
    hover_info : dict, optional
        Dictionary containing additional columns to display in hover text.
        Example: {"label": "Atom", "col": "atom_id"}
    split_col : str, optional
        Column name indicating dataset split (e.g., "train", "val", "test").
        If provided, points will be colored based on this split.
    save : str, optional
        Path name of the file to save in html

    Returns
    -------
    None
    """
    fig = go.Figure()

    # Manage hover text and custom data
    hovertemplate = f"<b>Structure #: %{{text}}</b><br>Ref: %{{x}}<br>Pred: %{{y}}"
    customdata = None

    if hover_info:
        hover_label = hover_info.get("label", "Custom Data")
        hover_col = hover_info.get("col")
        if hover_col and hover_col in df.columns:
            hovertemplate = f"<b>Structure #: %{{text}}</b><br>{hover_label}: %{{customdata[0]}}<br>Ref: %{{x}}<br>Pred: %{{y}}"
            customdata = df[[hover_col]].to_numpy()

    # Define color mapping for splits
    split_colors = {'train': 'blue', 'val': 'orange', 'test': 'green'}

    if split_col and split_col in df.columns:
        unique_splits = df[split_col].unique()
        for split in unique_splits:
            subset = df[df[split_col] == split]
            ref_subset = subset[ref_col].apply(pd.Series)
            pred_subset = subset[pred_col].apply(pd.Series)
            
            # loop for x y z point
            for i in range(3):
                if i == 0:
                    value = f"{'' if ref_subset.shape[1] == 1 else 'x'}"
                    # Add identity line (x=y)
                    if split == 'train':  # add only one time
                        fig.add_trace(go.Scatter(
                            x=ref_subset.iloc[:, i], y=ref_subset.iloc[:, i],
                            mode='lines',
                            name='x=y',
                            line=dict(color='gray', width=1, dash='dash')
                        )) 
                elif i == 1:
                    value = 'y'
                else:
                    value = 'z'
                fig.add_trace(go.Scatter(
                    x=ref_subset.iloc[:, i], y=pred_subset.iloc[:, i],
                    mode='markers',
                    marker=dict(color=split_colors.get(split, 'gray'), size=6, opacity=0.7),
                    hovertemplate=hovertemplate,
                    text=subset[idx_col],
                    customdata=subset[[hover_col]].to_numpy() if hover_info else None,
                    name=f'{split.capitalize()} {value}'
                ))
                if ref_subset.shape[1] == 1:
                    break
    else:
        # Default scatter plot without split coloring
        ref_values = df[ref_col].apply(pd.Series)
        pred_values = df[pred_col].apply(pd.Series)
        for i in range(ref_values.shape[1]):
            if i == 0:
                value = f"{'' if ref_values.shape[1] == 1 else 'x'}"
                # Add identity line (x=y)
                fig.add_trace(go.Scatter(
                    x=ref_values.iloc[:, i], y=ref_values.iloc[:, i],
                    mode='lines',
                    name='x=y',
                    line=dict(color='gray', width=1, dash='dash')
                )) 
            elif i == 1:
                value = 'y'
            else:
                value = 'z'
            fig.add_trace(go.Scatter(
                x=ref_values.iloc[:, i], y=pred_values.iloc[:, i],
                mode='markers',
                marker=dict(color='blue', size=6, opacity=0.7),
                hovertemplate=hovertemplate,
                text=df[idx_col],
                customdata=customdata,
                name=f'Predictions {value}'
            ))
            if ref_values.shape[1] == 1:
                break

    # Styling axes
    fig.update_xaxes(title_text=f'<b>Reference Values</b>', showline=True, linewidth=2,
                     linecolor='black', mirror=True)
    fig.update_yaxes(title_text=f'<b>Predicted {title}</b>', showline=True, linewidth=2,
                     linecolor='black', mirror=True)

    # Layout settings
    fig.update_layout(
        title=f'<b>{title} Prediction vs Reference in {units}</b>',
        plot_bgcolor='white',
        xaxis=dict(showgrid=True, gridwidth=1, gridcolor='lightgrey', zeroline=False),
        yaxis=dict(showgrid=True, gridwidth=1, gridcolor='lightgrey', zeroline=False)
    )

    if save:
        if not save.endswith('.html'):
            save = save + '.html'
        fig.write_html(save)

    if show:
        fig.show()


def plot_residuals(df: pd.DataFrame, ref_col: str, pred_col: str, idx_col: str, title: str, units: str,
                   split_col: str = None, hover_info: dict = None, standardized: bool = True,
                   save: str = None, show: bool = False)-> None:    
    """
    Plots residuals (standardized or not) versus predicted values.
    
    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing reference and predicted values.
    ref_col : str
        Column name for reference values.
    pred_col : str
        Column name for predicted values.
    idx_col : str
        Column name for structure identifiers.
    title : str
        Title of the plot.
    units : str
        Units of the plot in MathJax format: <sup>-1</sup> 
    split_col : str, optional
        Column name for data splits (train/val/test). Default is None.
    hover_info : dict, optional
        Dictionary with hover information. Should contain "label" and "col" keys. Default is None.
    standardized: bool, optional
        Default = True, enable to standardized the residuals
    save : str, optional
        Path name of the file to save in html
    show : bool, optional
        Tell if the figure has to be show, by default False
    
    Returns
    -------
    None
    """
    
    if not 'pred_error' in df.columns:
        df['pred_error'] = df[pred_col] - df[ref_col]
        
    flat_pred_error = df['pred_error'].explode()
    flat_pred_error = flat_pred_error.reset_index(drop=True)
    std_dev = np.std(flat_pred_error)

    fig = go.Figure()
    
    if standardized == True:
        df['std_residuals'] = df['pred_error'] / std_dev  
        residual_col = 'std_residuals'
        y_title = 'Standardized residuals'
        # Add threshold lines
        fig.add_hline(y=0, line_width=2, line_dash="dash", line_color="black", layer="below")
        fig.add_hline(y=-2, line_width=2, line_dash="dash", line_color="red", layer="below")
        fig.add_hline(y=2, line_width=2, line_dash="dash", line_color="red", layer="below")
        
    else:
        residual_col = 'pred_error'
        y_title = 'Residuals'
        avg = np.mean(flat_pred_error)
        # Add threshold lines
        fig.add_hline(y=np.mean(flat_pred_error), line_width=2, line_dash="dash", line_color="black", layer="below")
        fig.add_hline(y=avg-std_dev, line_width=2, line_dash="dash", line_color="red", layer="below")
        fig.add_hline(y=avg+std_dev, line_width=2, line_dash="dash", line_color="red", layer="below")
    
    # Manage hover text and custom data
    hovertemplate = f"<b>Structure #: %{{text}}</b><br>Ref: %{{customdata[0]:.3f}}<br>Pred: %{{x:.3f}}"
    customdata = df[[ref_col]].to_numpy()
    
    if hover_info:
        hover_label = hover_info.get("label", "Custom Data")
        hover_col = hover_info.get("col")
        if hover_col and hover_col in df.columns:
            hovertemplate = f"<b>Structure #: %{{text}}</b><br>{hover_label}: %{{customdata[0]}}<br>Ref: %{{customdata[1]:.3f}}<br>Pred: %{{x:.3f}}"
            customdata = df[[hover_col]].to_numpy()
    
    # Define color mapping for splits
    split_colors = {'train': 'blue', 'val': 'orange', 'test': 'green'}
    
    if split_col and split_col in df.columns:  #if not None AND split in df
        unique_splits = df[split_col].unique()
        for split in unique_splits:
            subset = df[df[split_col] == split]
            pred_subset = subset[pred_col].apply(pd.Series)
            residual_subset = subset[residual_col].apply(pd.Series)
            
            # loop for x y z point
            for i in range(3):
                if i == 0:
                    value = f"{'' if pred_subset.shape[1] == 1 else 'x'}"
                elif i == 1:
                    value = 'y'
                else:
                    value = 'z'            
                fig.add_trace(go.Scatter(
                    x=pred_subset.iloc[:, i], y=residual_subset.iloc[:, i],
                    mode='markers',
                    marker=dict(color=split_colors.get(split, 'gray'), size=6, opacity=0.7),
                    hovertemplate=hovertemplate,
                    text=subset[idx_col],
                    customdata = (
                        subset[['atom_id', 'dft_force']].to_numpy()
                        if 'atom_id' in subset.columns and 'dft_force' in subset.columns
                        else subset[['dft_energy']].to_numpy()
                    ),
                    name=f'{split.capitalize()} {value}'
                ))
                if pred_subset.shape[1] == 1:
                    break
    else:
        # Default scatter plot without split coloring
        residual_values = df[residual_col].apply(pd.Series)
        pred_values = df[pred_col].apply(pd.Series)
        for i in range(3):
            if i == 0:
                value = f"{'' if pred_values.shape[1] == 1 else 'x'}"
            elif i == 1:
                value = 'y'
            else:
                value = 'z'           
            fig.add_trace(go.Scatter(
                x=pred_values.iloc[:, i], y=residual_values.iloc[:, i],
                mode='markers',
                marker=dict(color='blue', size=6, opacity=0.7),
                hovertemplate=hovertemplate,
                text=df[idx_col],
                customdata=customdata,
                name=f'Residuals {value}'
            ))
            if pred_values.shape[1] == 1:
                break

    # Styling & legend part
    fig.update_xaxes(title_text=f'<b>{title} predicted ({units})</b>', showgrid=False, showline=True, linewidth=2, linecolor='black')
    fig.update_yaxes(title_text=f'<b>{y_title} ({units})</b>', showgrid=False, showline=True, linewidth=2, linecolor='black')
    fig.add_trace(go.Scatter(x=[None], y=[None], mode='lines', name='2 Standard Deviation',
                        line=dict(color='red', width=1, dash='dash')))
    
    fig.update_layout(
        title=dict(text=f'<b>{title} {y_title.lower()} versus predicted values</b>',
                   subtitle=dict(
                       text=f'Estimated Standard Deviation σ: {std_dev:.4f} {units}',
                       font=dict(color='gray', size =13)
                   )),        
        plot_bgcolor='white',
        legend=dict(
            itemsizing='constant',
            itemwidth=30,
            y=1.0,
            x=1.0,
            xanchor='right',
            yanchor='top'
        )
    )
    
    if save:
        if not save.endswith('.html'):
            save = save + '.html'
        fig.write_html(save)
        
    if show:
        fig.show()


def plot_residual_density(df: pd.DataFrame, ref_col: str, pred_col: str, units: str, split_col: str = 'Split',
                          save: str = None, show: bool = False) -> None:
    """
    Plots the density of residuals for a given DataFrame.

    This function calculates the residuals between predicted and reference values,
    then plots the density of these residuals using Kernel Density Estimation (KDE).
    It also adds vertical lines to indicate ±2 standard deviations from the mean.

    Parameters
    ----------
    df : pd.DataFrame
        The DataFrame containing the data. It must include the predicted and reference columns.
    ref_col : str
        The column name for reference values.
    pred_col : str
        The column name for predicted values.
    units : str
        Units of the plot in LaTeX format: $^{-1}$ 
    split_col : str, optional
        The column name for split categories (default is 'Split').
    save : str, optional
        Path name of the file to save in svg
    show : bool, optional
        Tell if the figure has to be show, by default False

    Returns
    -------
    None
        This function displays a plot and does not return any value.

    Notes
    -----
    - The function assumes that the DataFrame `df` contains the specified columns.
    - The 'pred_error' column is created if it does not already exist.
    """
    
    # Check if 'pred_error' column exists, if not, create it
    if 'pred_error' not in df.columns:
        df['pred_error'] = df[pred_col] - df[ref_col]
    residual_col = 'pred_error'
        
    # Check if split_col exists in the DataFrame
    if split_col not in df.columns:
        split_col = None
        unique_labels = ['density']
    else:
        unique_labels = list(reversed(sorted(df['Split'].unique().tolist())))

    # Flatten the 'pred_error' column
    flat_pred_error = df['pred_error'].explode()
    flat_pred_error = flat_pred_error.reset_index(drop=True)

    # Calculate standard deviation and mean
    std_dev = np.std(flat_pred_error)
    avg = np.mean(flat_pred_error)

    # Convert 'pred_error' to a DataFrame for plotting
    residuals = df[residual_col].apply(pd.Series)
    res_shape = residuals.shape[1]
    
    # Plot setup
    fig, ax = plt.subplots(1, res_shape, sharey=True, tight_layout=True, figsize=(5*res_shape, 5))
    fig.suptitle('Density Plot of Residuals', fontsize=16)
    
    if res_shape == 1:
        ax = [ax]

    for i in range(res_shape):
        if i == 0:
            value = f"{'' if res_shape == 1 else 'x'}"
            leg = True if res_shape == 1 else False
        elif i == 1:
            value = 'y'
            leg = False
        else:
            value = 'z'
            leg = False if res_shape == 1 else True

        # Plot KDE
        sns.kdeplot(data=df, x=residuals[i], hue=split_col, fill=False, alpha=0.6, linewidth=1.5, ax=ax[i], legend=leg)

        # Add vertical lines for avg ± 2*std_dev
        ax[i].axvline(avg - 2 * std_dev, color='grey', linestyle='--', linewidth=0.5, label='±2 std')
        ax[i].axvline(avg + 2 * std_dev, color='grey', linestyle='--', linewidth=0.5)
        ax[i].set_xlabel('')        
        ax[i].set_ylabel('')
        ax[i].set_xlim(-10,10)

        # Set individual titles
        if res_shape > 1 :
            ax[i].set_title(f'{value} force')

        # Plot the legend
        if leg:
            ax[i].legend(labels=unique_labels + ['±2 std'], frameon=False)

    # Set common x-label & y-label
    fig.supxlabel(f'Residuals {units}')
    fig.supylabel('Density')
    fig.supylabel('Density')
    fig.suptitle('Density Plot of Residuals', fontsize=16, fontweight='bold', x=0.05, ha='left')
    fig.text(0.05, 0.88, f'Estimated Standard Deviation σ: {std_dev:.4f} {units}',
            ha='left', fontsize=13, color='gray')
    plt.tight_layout(rect=[0, 0, 1, 0.95]) 
    
    if save:
        if not save.endswith('.svg'):
            save = save + '.svg'
        fig.savefig(save, format='svg')
    
    if show:
        plt.show()
    else:
        plt.close()


def plot_info_energy_prediction(model_energy_df: pd.DataFrame, num_model: int, standardized: bool = True, 
                                save:str = None, show_on: List[str]= None) -> None:
    """
    Generates and displays multiple plots related to energy predictions for a given model.
    
    This function performs the following steps:
    1. Displays a box plot comparing reference and predicted energy values.
    2. Plots the reference energy values versus predicted values with appropriate labels.
    3. Plots residuals or standardized residuals versus predicted values.
    
    Parameters
    ----------
    model_energy_df : pd.DataFrame
        DataFrame containing reference energy values, predicted energy values, 
        structure identifiers, and data split information.
    num_model : int
        Model number used for labeling the plots.
    standardized: bool, optional
        Default = True, enable to standerdized the residuals
    save : str, optional
        Path name of the file to save in html
    show_on: List[str], optinal
        Tell which plot has to be shown between
            - rp: Plot reference vs predicted energy values
            - rs: Plot residuals for model predictions
            - rd: Plot Density of the residues
    
    Returns
    -------
    None
        The function generates and displays the plots but does not return any value.
    """
    
    if save:
        save_pred = f'{save}_energy_prediction'
        save_res = f'{save}_energy_residuals'
        save_dens = f'{save}_energy_residuals_density'
    else:
        save_pred, save_res, save_dens = None, None, None
    
    # Ensure show_on is a list if it's None
    if show_on is None:
        show_on = []
    
    # Generate a box plot to compare reference and predicted energy values
    #box_ref_pred(model_energy_df)
    
    # Plot reference vs predicted energy values
    plot_ref_vs_pred(
        model_energy_df, 
        ref_col='dft_energy', 
        pred_col='pred_energy', 
        idx_col='structure_id', 
        title=f"Model {num_model} Energy", 
        units="kJ.mol<sup>-1</sup>", 
        split_col='Split',
        save=save_pred,
        show=True if 'rp' in show_on else False
    )
    
    # Plot residuals for model predictions
    plot_residuals(
        df=model_energy_df,
        ref_col='dft_energy',
        pred_col='pred_energy',
        idx_col='structure_id',
        title=f"Model {num_model} Energy",
        units="kJ.mol<sup>-1</sup>",
        split_col='Split',
        standardized = standardized,
        save=save_res,
        show=True if 'rs' in show_on else False
    )

    # Plot Density of the residues
    plot_residual_density(
        df = model_energy_df, 
        ref_col='dft_energy',
        pred_col='pred_energy',
        units="kJ.mol$^{-1}$.",
        split_col='Split',
        save=save_dens,
        show=True if 'rd' in show_on else False
    ) 
    
    
def plot_info_force_prediction(model_force_df: pd.DataFrame, num_model: int, standardized: bool = True,
                               save:str = None, show_on: List[str]= None) -> None:
    """
    Generates and displays multiple plots related to magnitude force predictions for a given model.

    This function performs the following steps:
    1. Plots the reference force values versus predicted force values with appropriate labels.
    2. Plots residuals or standardized residuals versus predicted values to assess model performance.

    Parameters
    ----------
    model_force_df : pd.DataFrame
        DataFrame containing reference magnitude force values, predicted force values, 
        structure identifiers, atom identifiers, and data split information.
    num_model : int
        Model number used for labeling the plots.
    standardized: bool, optional
        Default = True, enable to standerdized the residuals
    save : str, optional
        Path name of the file to save in html or do not save if None
    show_on: List[str], optinal
        Tell which plot has to be shown between
            - rp: Plot reference vs predicted energy values
            - rs: Plot residuals for model predictions
            - rd: Plot Density of the residues

    Returns
    -------
    None
        The function generates and displays the plots but does not return any value.
    """
    
    if save:
        save_pred = f'{save}_force_prediction'
        save_res = f'{save}_force_residuals'
        save_dens = f'{save}_force__density'
    else:
        save_pred, save_res, save_dens = None, None, None
    
    # Ensure show_on is a list if it's None
    if show_on is None:
        show_on = []
    
    # Plot reference vs predicted force values
    plot_ref_vs_pred(
        model_force_df, 
        ref_col='dft_force', 
        pred_col='pred_force', 
        idx_col='structure_id', 
        title=f"Model {num_model}: Forces", 
        units="kJ.mol<sup>-1</sup>.\u00C5<sup>-1</sup>", 
        split_col='Split',
        hover_info={"label": "Atom #", "col": "atom_id"},
        save=save_pred,
        show=True if 'rp' in show_on else False
    )

    # Plot standardized residuals for force predictions
    plot_residuals(
        df=model_force_df,
        ref_col='dft_force',
        pred_col='pred_force',
        idx_col='structure_id',
        title=f'Model {num_model}: Forces for each atom',
        units="kJ.mol<sup>-1</sup>.\u00C5<sup>-1</sup>",
        hover_info={'label': 'Atom #', 'col': 'atom_id'},
        split_col='Split',
        standardized = standardized,
        save=save_res,
        show=True if 'rs' in show_on else False
    )
    
    # Plot Density of the residues
    plot_residual_density(
        df = model_force_df, 
        ref_col='dft_force',
        pred_col='pred_force',
        units="kJ.mol$^{-1}$.\u00C5$^{-1}$",
        split_col='Split',
        save=save_dens,
        show=True if 'rd' in show_on else False
    ) 
        
def analyze_model(model_num: str, db: pd.DataFrame, data_train: str, base_path: str, torchdevice: str,
                  list_dic_metric_energy: list, list_dic_metric_force: list, list_model: list, split: bool = True,
                  standardized: bool = True, save:str = None, show_on: List[str]= None)->None:
    """
    Analyze the energy and force predictions of a model.

    Parameters
    ----------
    model_num : str
        The model number as a string.
    db : object
        The database object used for analysis.
    data_train : str
        The training data identifier.
    base_path : str
        The base path for model files.
    torchdevice : str
        The device to use for torch computations (e.g., 'cpu' or 'cuda').
    list_dic_metric_energy : list
        A list to store energy metrics DataFrames.
    list_dic_metric_force : list
        A list to store force metrics DataFrames.
    list_model : list
        A list to store model identifiers.
    split : bool
        Define if the split file exsite, by default True 
    standardized: bool, optional
        Default = True, enable to standerdized the residuals
    save : str, optional
        Path name of the file to save in html
        Do not put the extension file
    show_on: List[str], optinal
        Tell which plot has to be shown between
            - rp: Plot reference vs predicted energy values
            - rs: Plot residuals for model predictions
            - rd: Plot Density of the residues 

    Returns
    -------
    None 
        The dunction display figure and modify the list inplace
    """
    print()
    print('##########################')
    print(f'#      MODEL {model_num}           #')
    print('##########################')
    print()

    if split:
        train, val, test = parse_split(f'trainings/runs/{data_train}_{model_num}/split.npz')
    name_model = f'{data_train}_{model_num}/best_model'

    # Initialize the primary predictive calculator.
    model_calculator = initialize_calculator(name_model, base_path, torchdevice)
    model_energy_force = parse_energy_force(db, model_calculator)

    # Energy analysis
    print('\nEnergy\n')
    model_energy_df = energy_analysis(model_energy_force['energies'])
    if split:
        model_energy_df = add_split(model_energy_df, train, val, test)
    model_energy_metrics = compute_metrics(model_energy_df)
    plot_info_energy_prediction(model_energy_df, model_num, standardized=standardized, save=save, show_on=show_on)

    # Force analysis
    print('\nForces\n')
    model_force_df = force_analysis(model_energy_force['magnitude'])
    if split:
        model_force_df = add_split(model_force_df, train, val, test)
    model_force_metrics = compute_metrics(model_force_df)
    plot_info_force_prediction(model_force_df, model_num, standardized=standardized, save=save, show_on=show_on)

    list_dic_metric_energy.append(model_energy_metrics)
    list_dic_metric_force.append(model_force_metrics)
    list_model.append(f'{data_train.replace("_train", "")} {model_num}')
    

def add_bars(fig: go.Figure, model: List[str], rmse: List[float], mae: List[float], r2: List[float], col: int) -> None:
    """
    Add RMSE, MAE, and R² bar plots to a given subplot.

    Parameters
    ----------
    fig : go.Figure
        The figure to which the bars will be added.
    model : List[str]
        List of model names.
    rmse : List[float]
        List of RMSE values.
    mae : List[float]
        List of MAE values.
    r2 : List[float]
        List of R² values.
    col : int
        Column number for subplot placement.

    Returns
    -------
    None
        The function modifies `fig` in place.
    """
    legend: bool = (col == 1)  # Only show legend in the first column
    
    fig.add_trace(go.Bar(
        x=model,
        y=rmse,
        name='RMSE',
        marker_color='indianred',
        showlegend=legend,
    ), row=1, col=col)
    
    fig.add_trace(go.Bar(
        x=model,
        y=mae,
        name='MAE',
        marker_color='lightsalmon',
        showlegend=legend,
    ), row=1, col=col)
    
    fig.add_trace(go.Bar(
        x=model,
        y=r2,
        name='R²',
        marker_color='skyblue',
        showlegend=legend,
    ), row=2, col=col)


def performances_bar(df: pd.DataFrame, propertie: str, save: str = None, split: bool = True) -> None:
    """
    Generate a grouped bar chart showing RMSE, MAE, and R² values for different models
    across multiple datasets (Global, Train, Validation, Test).

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing model performance metrics with columns:
        - 'Model' (str): Names of models.
        - 'Split' (List[str]): List order of the split
        - 'RMSE' (List[float]): RMSE values.
        - 'MAE' (List[float]): MAE values.
        - 'R²' (List[float]): R² values.
    propertie : str
        Name of the property being predicted (used for the plot title).
    save : str, optional
        Path name of the file to save in html
    split : bool, optional
        Define if the split column exsite, by default True 
    Returns
    -------
    None
        The function displays the generated plot and save it.
    """
    if split:
        group_list = df['Split'][0]
    else:
        group_list = ['Global Prediction']
    models = df['Model'].tolist()
    fig = make_subplots(rows=2, cols=len(group_list), subplot_titles=group_list,
                        shared_yaxes=True, shared_xaxes = True,
                        vertical_spacing=0.1)

    for i, group in enumerate(group_list):
        rmse = []
        mae = []
        r2 = []
        
        for idx in range(len(df)):
            if split:
                split_index = df.iloc[idx]['Split'].index(group)
                rmse.append(float(df.iloc[idx]['RMSE'][split_index]) )
                mae.append(float(df.iloc[idx]['MAE'][split_index]) )        
                r2.append(float(df.iloc[idx]['R²'][split_index]) )        
            else:
                rmse.append(float(df.iloc[idx]['RMSE']) )
                mae.append(float(df.iloc[idx]['MAE']) )        
                r2.append(float(df.iloc[idx]['R²']) ) 
        
        add_bars(fig, models, rmse, mae, r2, i+1)

    fig.update_layout(
            title=f"Performance evaluation of the {propertie} prediction",
            showlegend=True,
            width=900,
            height=600,
            barmode='group',
            xaxis_tickangle=-45
        )
    
    if save:
        if not save.endswith('.html'):
            save = save + '.html'
        fig.write_html(save)
    
    fig.show()