md_analysis.py

"""
Function for the analysis of the simulations
"""

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


import pandas as pd
import numpy as np
import seaborn as sns
import re
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from typing import List, Union, Tuple, Dict


def rdf_input(file:str):
    """
    Read and parse a radial distribution function (RDF) data file given by GROMOS software

    The function extracts `radius` and `rdf` values from GROMOS output rdf function file, 
    Ignore comment lines that start with `#` of the warning.
    The data is assumed to have two columns:
        - the first column = radius, 
        - the second column = RDF values.

    Parameters
    ----------
    file : str
        Path to the RDF data file.

    Returns
    -------
    radius : list of float
        A list containing the radius values.
    rdf : list of float
        A list containing the corresponding RDF values.

    Notes
    -----
    - The file is expected to have numerical values separated by whitespace.
    - Lines starting with `#` are ignored.
    - The function assumes that all non-comment lines contain exactly two numerical values.
    """

    radius = []
    rdf = []
    
    with open(file, "r") as source:
        for line in source:
            if not line.startswith("#"):  # Ignore comment lines
                line = line.strip().split()
                radius.append(float(line[0]))
                rdf.append(float(line[1]))

    return radius, rdf


def parse_multi_rdf(rdf_files: List[str]) -> Tuple[List[float], List[List[float]]]:
    """
    Parses multiple RDF files and aggregates the radius and RDF data.

    Parameters
    ----------
    rdf_files : List[str]
        List of file paths to RDF data files.

    Returns
    -------
    radius : List[float]
        The list of the radius values
    rdf_list : List[List[float]]]
        A list of list of RDF values.
    """
    radius, rdf_list = [], []

    for file in rdf_files:
        try:
            rad, rdf = rdf_input(file)
            if not radius:
                radius = rad
            else:
                if radius != rad:
                    print("WARNING: Radius values are different across files.")
            rdf_list.append(rdf)
        except Exception as e:
            print(f"Error reading file {file}: {e}")

    return radius, rdf_list


def two_shield(rdf, radius):
    """
    Identify the two largest local maxima in the given radial distribution function (RDF).

    A local maximum is determined where the first derivative changes sign from positive to negative. 
    The function computes the numerical derivative and detects the first two such maxima.

    Parameters
    ----------
    rdf : array-like
        Values of the radial distribution function.
    radius : array-like
        Corresponding radius values.

    Returns
    -------
    max_rdf : float
        The radius at which the first local maximum occurs.
    max2_rdf : float or None
        The radius at which the second local maximum occurs, or None if only one maximum is found.

    Notes
    -----
    The function uses numerical differentiation via finite differences. It assumes that `radius` is 
    monotonically increasing and `rdf` has at least one local maximum.
    """
    
    # Compute the first derivative using finite differences
    dy_dx = np.diff(rdf) / np.diff(radius)
    
    # Identify indices where the derivative changes from positive to negative
    maxima_indices = np.where(np.diff(np.sign(dy_dx)) == -2)[0] + 1
    
    # Extract the first and second maxima, if available
    first_max_index = maxima_indices[0]
    second_max_index = maxima_indices[1] if len(maxima_indices) > 1 else None
    
    max_rdf = radius[first_max_index]
    max2_rdf = radius[second_max_index] if second_max_index is not None else None
    
    print(f'The first shield: {max_rdf} nm')
    print(f'The second shield: {max_rdf} nm')

    return max_rdf, max2_rdf


def plot_rdf(radius: List[float], rdf: List[Union[float, List[float]]], mol_name: str,
             mol_atom: str = 'C', water_atom: str = 'O', max_rdf: float = None,
             max2_rdf: float = None) -> None:
    """
    Plots the Radial Distribution Function (RDF) for a given molecule in water.

    Parameters
    ----------
    radius : List[float]
        List of radius values for the molecule.
    rdf : List[Union[float, List[float]]]
        List of RDF values corresponding to the radius values. Can be a list of floats or a list of lists of floats.
    mol_name : str
        Name of the molecule.
    mol_atom : str, optional
        Atom type of the molecule used for the rdf, default is 'C'.
    water_atom : str, optional
        Atom type of water used for the rdf, default is 'O'.
    max_rdf : float, optional
        Radius value where the first peak of the RDF occurs.
    max2_rdf : float, optional
        Radius value where the second peak of the RDF occurs.

    Returns
    -------
    None
        Displays the plot.
    """
    fig, ax = plt.subplots()
    
    if isinstance(rdf[0], list):
        for i, atom_rdf in enumerate(rdf):
            ax.plot(radius, atom_rdf, label=f'{mol_atom}{i+1}')
    else:
        ax.plot(radius, rdf)

    if max_rdf is not None:
        ax.axvline(max_rdf, linestyle="--", color="red", label=f"First shell {max_rdf} nm")
    if max2_rdf is not None:
        ax.axvline(max2_rdf, linestyle="--", color="green", label=f"Second shell {max2_rdf} nm")

    ax.set_title(f"RDF of the {mol_name} ({mol_atom}) in water ({water_atom})")
    ax.set_ylabel("RDF")
    ax.set_xlabel("Radius (nm)")
    ax.legend(frameon=False)

    plt.show()
    

def rdf_replicate(rdf_files: List[str]) -> Tuple[List[float], List[List[float]]]:
    """
    Replicate and process Radial Distribution Function (RDF) data from multiple files.

    This function processes multiple RDF files, extracts radius and RDF data, and plots them.
    It aggregates the RDF data across all files and checks for consistency in radius values.

    Parameters
    ----------
    rdf_files : List[str]
        A list of file paths to RDF data files.

    Returns
    -------
    Tuple[List[float], List[List[float]]]
        A tuple containing:
        - A list of radius values.
        - A list of aggregated RDF data from all files
            - lines: replicate * number of atoms
            - Columns: radius

    Notes
    -----
    Warns if radius values differ across replicates but the x axis remain the same.
    """
    global_rdf_list, global_radius = [], []

    for i, sub_file in enumerate(rdf_files):
        print(f'Replicate # {i+1}')
        radius, rdf_list = parse_multi_rdf(sub_file)
        plot_rdf(radius, rdf_list, mol_name='benzene')

        if not global_radius:
            global_radius = radius
        else:
            if global_radius != radius:
                print("WARNING: Radius values are different across replicates.\n")

        global_rdf_list.extend(rdf_list)

    return global_radius, global_rdf_list


def rdf_avg_std(global_rdf_list: List[List[float]]) -> Tuple[np.ndarray, np.ndarray]:
    """
    Calculate the average and standard deviation of Radial Distribution Function (RDF) data.

    Parameters
    ----------
    global_rdf_list : List[List[float]]
        A list of lists containing RDF data. Each inner list represents RDF values for a single replicate.

    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        A tuple containing:
        - An array of average RDF values.
        - An array of standard deviation values for the RDF data.

    Notes
    -----
    The function assumes that all inner lists in `global_rdf_list` have the same length.
    """
    arr = np.array(global_rdf_list, dtype='float32')
    arr = arr.T
    avg = arr.mean(axis=1)
    std = arr.std(axis=1)
    return avg, std


def count_deviations(omd_files: List[str], round_index: int = 6, rep_index: int = 7) -> pd.DataFrame:
    """
    Counts the occurrences of deviation in the omd files and organizes the counts into a DataFrame.

    Parameters
    ----------
    omd_files : List[str]
        A list of file paths to the log files to be processed.
    round_index : int, optional
        The index position of the round name in the file path when split by '/'. Default is 6.
    rep_index : int, optional
        The index position of the replicate name in the file path when split by '/'. Default is 7.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing the counts of deviations for each round and replicate.

    Notes
    -----
    This function assumes that each file contains log messages and counts the occurrences
    of the specific message 'NOTICE NN Worker : Deviation from validation model above threshold'.
    """
    count_dev_dict = {}

    for file in omd_files:
        name_parts = file.split('/')
        round_name = name_parts[round_index]
        rep_name = name_parts[rep_index]

        with open(file, "r") as source:
            count_dev = 0
            for line in source:
                if 'NOTICE NN Worker : Deviation from validation model above threshold' in line:
                    count_dev += 1

        if round_name not in count_dev_dict:
            count_dev_dict[round_name] = {}

        if rep_name not in count_dev_dict[round_name]:
            count_dev_dict[round_name][rep_name] = 0

        count_dev_dict[round_name][rep_name] += count_dev

    count_dev_omd = pd.DataFrame(count_dev_dict)
    return count_dev_omd


def plot_deviation_statistics(count_dev_omd: pd.DataFrame, y_max: int, pc: bool = False) -> None:
    """
    Plots the average and standard deviation of model deviation counts for each round.

    Parameters
    ----------
    count_dev_omd : pd.DataFrame
        A DataFrame containing the counts of deviations for each round and replicate.
        Each column represents a round, and each row represents a replicate.
    y_max : int
        limit superior for the graph
    pc : bool
        parameter to show the graph in percentage
        
    Returns
    -------
    None
        Displays a bar plot with error bars representing the standard deviation.

    Notes
    -----
    This function calculates the mean and standard deviation of deviation counts for each round
    and plots them as a bar chart with error bars.
    """
    round_list, mean_list, std_list = [], [], []

    for col in count_dev_omd.columns:
        round_list.append(col)
        mean_list.append(count_dev_omd[col].mean())
        std_list.append(count_dev_omd[col].std())

    if pc:
        mean_list = (np.array(mean_list))/y_max*100
        std_list = (np.array(std_list))/y_max*100
        y_max = 100
        ylabel = 'Percentage of deviation'
    else:
        ylabel = 'Total number of deviations'
    
    fig, ax = plt.subplots()
    bars = ax.bar(round_list, mean_list, yerr=std_list, capsize=5, color='skyblue', alpha=0.7)

    for bar, mean, std in zip(bars, mean_list, std_list):
        if pc:
            text=f'{mean:.2f}±{std:.2f} %'
        else:
            text=f'{mean:.0f}±{std:.0f}'
        plt.text(bar.get_x() + bar.get_width() / 2.0, mean + std, text, ha='center', va='bottom')

    ax.set_ylabel(ylabel)
    ax.set_ylim(0, y_max)
    ax.set_title('Average of the total number of model deviation catches\nin each round of adaptive sampling')

    plt.show()


def read_dat_file(prop_file: str) -> Tuple[List[float], List[float]]:
    """
    Reads a .dat file containing time and the property evaluate by gromos software.

    Parameters
    ----------
    prop_file : str
        Path to the file containing the data.

    Returns
    -------
    Tuple[List[float], List[float]]
        A tuple containing two lists:
        - The first list contains time values as floats.
        - The second list contains prop values as floats.

    Notes
    -----
    The file is expected to have lines with two float values separated by whitespace.
    Lines starting with '#' are considered comments and are ignored.
    """
    time = []
    prop = []

    with open(prop_file, "r") as source:
        for line in source:
            if not line.startswith("#"):  # Ignore comment lines
                line = line.strip().split()
                time.append(float(line[0]))
                prop.append(float(line[1]))

    return time, prop


def parse_propr_multifiles(list_dat_files: List[str], round_index: int = 6, rep_index: int = 7) -> pd.DataFrame:
    """
    Extracts the desired property contained in the dat files and organizes the data into a DataFrame.

    Parameters
    ----------
    list_dat_files : List[str]
        A list of file paths to the dat property files to be processed.
    round_index : int, optional
        The index position of the round number in the file path when split by '/'. Default is 6.
    rep_index : int, optional
        The index position of the replicate number in the file path when split by '/'. Default is 7.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing the processed property data, organized by round and replicate numbers.
        Each cell contains a tuple of two lists corresponding to the time and the values of the desired property.

    Notes
    -----
    The file paths are expected to have a specific structure where the round number and replicate number
    can be extracted based on the provided indices.
    """
    property_data: Dict[str, Dict[str, Tuple[List[float], List[float]]]] = {}

    for file in list_dat_files:
        time, nnval = read_dat_file(file)  # Ensure `read_dat_file` is defined elsewhere
        name_parts = file.split('/')
        round_name = name_parts[round_index]
        rep_name = name_parts[rep_index]

        # Initialize the round dictionary if it doesn't exist
        if round_name not in property_data:
            property_data[round_name] = {}

        # Add the replicate to the round
        property_data[round_name][rep_name] = (time, nnval)

    # Convert the nested dictionary to a DataFrame
    property_df = pd.DataFrame(property_data)

    return property_df


def parse_time_cnf_file(extract_cnf_files: List[str], round_index: int = 6, rep_index: int = 7) -> pd.DataFrame:
    """
    Parses time data from cnf files and organizes it into a DataFrame.

    Parameters
    ----------
    extract_cnf_files : List[str]
        A list of file paths to the cnf files to be processed.
    round_index : int, optional
        The index position of the round number in the file path when split by '/'. Default is 6.
    rep_index : int, optional
        The index position of the replicate number in the file path when split by '/'. Default is 7.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing the parsed time data, organized by round and replicate numbers.
        Each cell contains a list of time values extracted from the files.

    Notes
    -----
    The file paths are expected to have a specific structure where the round number and replicate number
    can be extracted based on the provided indices. Each file is expected to contain lines starting with
    'TIMESTEP', followed by lines with time data.
    """
    extract_time: Dict[str, Dict[str, List[float]]] = {}

    for file in extract_cnf_files:
        with open(file, 'r') as source:
            read = False
            time = []
            for line in source:
                if line.startswith('TIMESTEP'):
                    read = True
                elif read:
                    time.append(float(line.strip().split()[1]))
                    read = False

        name_parts = file.split('/')
        round_num = name_parts[round_index]
        rep_num = name_parts[rep_index]

        # Initialize the round dictionary if it doesn't exist
        if round_num not in extract_time:
            extract_time[round_num] = {}

        # Add the replicate to the round
        extract_time[round_num][rep_num] = time

    # Convert the nested dictionary to a DataFrame
    extract_time_df = pd.DataFrame(extract_time)

    return extract_time_df


def take_energy(list_time_cnf: List[float], energy: Tuple[List[float], List[float]], col: int, row: int
                )-> Tuple[List[float], List[float]]:
    """
    Extracts energy values corresponding to given time_cnf points.

    Parameters
    ----------
    list_time_cnf : List[float]
        List of time_cnf points to extract energy values for.
    energy : Tuple[List[float], List[float]]
        A tuple containing two lists: (list of time_qm, list of energies).
    col : int
        Column identifier for logging purposes corresponding of the round
    row : int
        Row identifier for logging purposes corresponding of the replicate

    Returns
    -------
    Tuple[List[float], List[float]]
        A tuple containing two lists:
        - A list of energy values corresponding to the given time_cnf points.
        - A list of time_cnf points that did not have a corresponding time_qm value.
    """
    time_qm, energies = energy  # Unpack the tuple
    ener: List[float] = []
    no_match: List[float] = []

    for time_cnf in list_time_cnf:
        if time_cnf in time_qm:
            # Find the index of the time_cnf in time_qm
            index = time_qm.index(time_cnf)
            # Append the corresponding energy value to ener
            ener.append(energies[index])
        else:
            no_match.append(time_cnf)
            print(f'No correspondence in {col} {row} for time_cnf: {time_cnf}')

    return ener, no_match


def parse_energy_dft(energy_file: List[str], unmatch_time_index: Dict[str, List[int]], converstion_unit: float,
                     round_index: int = 7) -> Dict[str, List[float]]:
    """
    Process energy files and store the energy values in a dictionary.
    Energy files corrspondent to the dft files .engrad

    Parameters
    ----------
    energy_file : List[str]
        List of file paths to process.
    unmatch_time_index : Dict[str, List[int]]
        Dictionary mapping round names to lists of unmatched time indices.
    converstion_unit : float
        value to convert of the energy into the desired unit
    round_index : int, optional
        The index position of the round number in the file path when split by '/'. Default is 7.

    Returns
    -------
    Dict[str, List[float]]
        A dictionary with round names as keys and lists of energy values as values.
    """
    energy = {}

    for file in energy_file:
        print(file)
        round_name_split = file.split('/')[round_index].split('_')
        round_number = int(round_name_split[-1])
        round_name = '_'.join(round_name_split[:-1]) + f'_{round_number}'
        sub_conf = int(re.sub(r'IR_BR_|BR_', '', file.split('/')[9].split('.')[0]))
                
        if round_name not in energy:
            energy[round_name] = []
       
        if sub_conf not in unmatch_time_index.get(round_name, []):  # return an empty list if round_name does not exist
            with open(file, 'r') as source:
                lines = source.readlines()
            ener = float(lines[7].strip()) * converstion_unit
            energy[round_name].append(ener)

    return energy


def multi_density_plot(
    property_df: pd.DataFrame,
    title: str,
    property_name: str,
    units: str
) -> None:
    """
    Generate a density plot for each column in the DataFrame.

    Parameters
    ----------
    property_df : pd.DataFrame
        DataFrame containing desired property to be plotted.
    title : str
        Title of the density plot.
    property_name : str
        Name of the property being plotted.
    units : str
        Units of the property being plotted.

    Returns
    -------
    None
    """
    palette = sns.color_palette("husl", len(property_df.columns))
    fig, ax = plt.subplots()
    for col in property_df.columns:
        for idx, row in property_df.iterrows():
            if not pd.isna(row[col]):
                sns.kdeplot(
                    row[col][1],
                    label=f'{col}/{idx}',
                    color=palette[int(col.split("_")[1])],
                    linewidth=0.5,
                )
    ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), frameon=False)
    ax.set_title(f'Density plot of {title} during the simulation')
    ax.set_xlabel(f'{property_name} ({units})')
    plt.show()
    
    
def read_spectrum_data(file_path: str, freq_conversion: float = 33.3564095198) -> Tuple[np.ndarray, np.ndarray]:
    """
    Reads spectrum data from the tcf output file from GRONOS.

    Parameters
    ----------
    file_path : str
        The path to the file containing the signal data.
    freq_conversion : float, optional
        The frequency conversion factor to apply to the frequency data.
        Default is 33.3564095198.

    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        A tuple containing two numpy arrays:
        - The first array contains the converted frequency data.
        - The second array contains the intensity data.

    Notes
    -----
    In the tcf file the spectrum com qfter the autocorelation and the line start
    with '#  frequency   intensity' to start reading and 'END' to stop reading. 
    Each data line should contain two float values: the frequency and the signal intensity.
    """
    signal: List[float] = []
    frequency: List[float] = []
    read = False

    with open(file_path, 'r') as source:
        for line in source:
            if 'frequency' in line:
                read = True
            elif 'END' in line:
                read = False
            elif read:
                line = line.strip().split()
                frequency.append(float(line[0]))
                signal.append(float(line[1]))

    frequency = np.array(frequency) * freq_conversion
    signal = np.array(signal)

    return frequency, signal


def plot_multi_power_spectrum(datas: List[Union[str, Tuple[str, np.ndarray, np.ndarray]]],
                              average: bool = False,
                              rep: int = 6,
                              strips:bool = True)-> Union[Tuple[np.ndarray, np.ndarray], None]:
    """
    Plots the spectrum data from multiple files or tuples and optionally calculates the average signal.

    Parameters
    ----------
    datas : List[Union[str, Tuple[np.ndarray, np.ndarray]]]
        A list of file paths or tuples containing (frequency, signal) data to be processed and plotted.
    average : bool, optional
        If True, calculates and plots the average signal across all input files. Default is False.
    rep : int, optional
        Column identifier for logging purposes corresponding to the replicate name. Default is 6.
    strips : bool, optional
        aes of the experimental values, by default as strip else line.

    Returns
    -------
    Union[Tuple[np.ndarray, np.ndarray], None]
        If `average` is True, returns a tuple containing:
        - A numpy array representing the frequency data.
        - A numpy array representing the average signal across all input files.
        If `average` is False, returns None.

    Notes
    -----
    The function assumes that `read_spectrum_data` is defined elsewhere and returns
    frequency and signal data as numpy arrays or lists. The x-axis is limited to the
    range [500, 4000] for wavenumbers.
    """
    # Define a custom color palette
    custom_colors = ['#1f77b4', '#e41a1c', '#ff7f00', '#7a0403', '#999999', '#666666' ]
    color_map = ListedColormap(custom_colors)
    
    signal_list = []
    fig, ax = plt.subplots(figsize=(10,5))

    for idx, data in enumerate(datas):
        if isinstance(data, str):  # if data is a file path
            frequency, signal = read_spectrum_data(data)
            label = data.split('/')[rep]
        elif isinstance(data, tuple):
            label, frequency, signal = data

        ax.plot(frequency, signal, label=label, color=color_map(idx))
        signal_list.append(signal)

    if average:
        # Calculate the average signal
        signal_avg = np.mean(signal_list, axis=0)
        ax.plot(frequency, signal_avg, label='Signal Average')
        to_return = frequency, signal_avg
    else:
        to_return = None
    
    if strips:
        ax.axvspan(990, 1000, color='grey', alpha=0.3, label='Experiments')
        ax.axvspan(3056, 3066, color='grey', alpha=0.3)
    else:
        ax.axvline(3061, color='black', linestyle='--')
        ax.axvline(992, color='black', linestyle='--', label='Experiments')

    # Set the limits for the x-axis
    ax.set_xlim(500, 4000)

    # Add labels and title
    ax.set_xlabel('Wavenumber (cm$^{-1}$)')
    ax.set_ylabel('Intensity')
    ax.set_title('Benzene MD Power Spectra')

    # Add a legend
    ax.legend(frameon=False, loc='center left', bbox_to_anchor=(1, 0.5), ncol=1)

    
    # Display the plot
    plt.show()

    return to_return

    
def read_improper_data(file_path: str) -> pd.DataFrame:
    """
    Reads improper data from tser.out file and returns it as a DataFrame.

    Parameters
    ----------
    file_path : str
        The path to the file containing the improper data.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing the improper data with the first column set as the index.

    Notes
    -----
    The function assumes that the file has a header line starting with '#' and uses
    tab-separated values. The first column is set as the index of the DataFrame.
    """
    with open(file_path, 'r') as source:
        next(source)  # Skip the first line
        line1 = source.readline()

    column_names = line1.replace('#', '').strip().split()

    improper_df = pd.read_csv(file_path,
                            sep='\t+',
                            names=column_names,
                            comment='#',
                            dtype=np.float64,
                            engine='python')

    improper_df.set_index(improper_df.columns[0], inplace=True)

    return improper_df


def plot_improper_angle_density(data: Union[dict,pd.DataFrame], rep_name: str) -> None:
    """
    Plots the kernel density estimate (KDE) for each column in the DataFrame and stores the plot in a dictionary.

    Parameters
    ----------
    improper_df : pd.DataFrame
        A DataFrame containing improper angle data. Each column represents a different set of angles.
    rep_name : str
        The name of the replicate or simulation round to be included in the plot title and used as a key in the dictionary.

    Returns
    -------
    None
        Displays the KDE plot for each column in the DataFrame and stores the plot in the dictionary.

    Notes
    -----
    This function uses Seaborn to plot the KDE for each column in the DataFrame.
    The x-axis represents the improper angle in degrees, and the y-axis represents the density.
    """
    fig, ax = plt.subplots()

    for column in data:
        sns.kdeplot(data[column], label=column, ax=ax)

    ax.set_title(f'Improper angle density during the simulations {rep_name}')
    ax.set_xlabel('Improper angle $\phi$ (degree)')
    ax.set_ylabel('Density')
    ax.legend()


def process_improper_files(files: List[str], rep_index: int = 7) -> Dict[str, pd.Series]:
    """
    Processes a list of tser.out files containing improper angle data,
    plots the density, and stores the results.

    Parameters
    ----------
    files : List[str]
        A list of file paths to the improper angle data files to be processed.
    rep_index : int, optional
        The index position of the replicate name in the file path when split by '/'. Default is 7.

    Returns
    -------
    Dict[str, pd.Series]
        A dictionary where each key is a replicate name, and each value is a pandas Series
        containing the stacked improper angle data.

    """
    rep_improper = {}

    for file in files:
        improper_df = read_improper_data(file)
        rep_name = file.split('/')[rep_index]
        improper_df = improper_df.filter(like='PeriodicTorsion%')
        
        # Plot the improper angle density
        plot_improper_angle_density(improper_df, rep_name)

        # Store the stacked DataFrame in the dictionary
        rep_improper[rep_name] = improper_df.stack().reset_index(drop=True)

    return rep_improper


def aggregate_improper(files: List[str], improper_simulation: Dict[str, List[float]], global_name: str, rep_index: int = 7) -> None:
    """
    Aggregates improper angle data from multiple replicates into a global simulation dictionary.

    Parameters
    ----------
    files : List[str]
        A list of file paths to the improper angle data files to be processed.
    improper_simulation : Dict[str, List[float]]
        A dictionary to store aggregated improper angle data for each global simulation.
    global_name : str
        The name of the global simulation to aggregate data under.
    rep_index : int, optional
        The index position of the replicate name in the file path when split by '/'. Default is 7.
   
    Returns
    -------
    None
        Aggregates the improper angle data into the `improper_simulation` dictionary.
    """
    # Process the improper files to get a dictionary of replicate data
    rep_improper = process_improper_files(files,rep_index)

    # Plot the improper angle density for the global simulation
    plot_improper_angle_density(rep_improper, global_name)

    # Initialize the list for the global simulation name
    improper_simulation[global_name] = []

    # Aggregate the improper angle data from all replicates
    for angles in rep_improper.values():
        improper_simulation[global_name].extend(angles)
        
    
def process_energy_data(energy_qm_df: pd.DataFrame, extract_time_df: pd.DataFrame
                        ) -> Tuple[Dict[str, List[float]], Dict[str, Dict[int, List[float]]]]:
    """
    Extract he QM energy from totqm file (already in a dataframe) for each snapshot extract from the simulation.

    Parameters
    ----------
    energy_qm_df : pd.DataFrame
        A DataFrame containing QM energy data from totam file. Each column represents a round of data.
    extract_time_df : pd.DataFrame
        A DataFrame containing time data corresponding to each snapshot. Each column represents a round of data.
    
    Returns
    -------
    Tuple[Dict[str, List[float]], Dict[str, Dict[int, List[float]]], List[str]]
        A tuple containing:
        - A dictionary with QM energy values for each snapshot for each round.
        - A dictionary with mismatched indices and their corresponding lists of mismatched values for each round.

    Notes
    -----
    This function iterates over each column in the energy DataFrame, extracts energy values,
    handles mismatches, and generates a summary of the parsing results.
    """
    ener, no_m = {}, {}
    resume_list = []
    for col in energy_qm_df.columns:
        ener_round = []
        no_match_round = {}
        count_no_match = 0

        for idx, row in energy_qm_df.iterrows():
            if not pd.isna(row[col]) and col in extract_time_df:
                energy_list, no_match_list = take_energy(extract_time_df[col][idx], energy_qm_df[col][idx], col, idx)
                ener_round += energy_list
                no_match_round[idx] = no_match_list
                count_no_match += len(no_match_list)

        ener[col] = ener_round
        no_m[col] = no_match_round

        s = 0
        for i in range(1, 6):
            if not isinstance(extract_time_df[col][f'rep_{i}'], float):
                s += len(extract_time_df[col][f'rep_{i}'])

        resume_list.append(f"Column: {col}: # Snapshots: {s} / # Energy values: {len(ener[col])} / # no match: {count_no_match}")

    print('\n ### Summary parsing  ###')
    for resume in resume_list:
        print(resume)

    return ener, no_m


def find_unmatched_time_indices(no_m: Dict[str, Dict[int, List[float]]], energy_qm_df: pd.DataFrame, 
                                extract_time_df: pd.DataFrame) -> Dict[str, List[int]]:
    """
    Finds the indices of unmatched times for each round in the data.

    Parameters
    ----------
    no_m : Dict[str, Dict[int, List[float]]]
        A dictionary containing mismatched time values for each round and index.
    energy_qm_df : pd.DataFrame
        A DataFrame containing QM energy data for each snapshot. Each column represents a round of data.
    extract_time_df : pd.DataFrame
        A DataFrame containing time data corresponding to each snapshot. Each column represents a round of data.

    Returns
    -------
    Dict[str, List[int]]
        A dictionary where each key is a round name, and each value is a list of indices
        corresponding to unmatched times in that round.

    Notes
    -----
    This function iterates over each column in the `no_m` dictionary, identifies the indices
    of unmatched times, and stores them in a new dictionary.
    """
    unmatch_time_index = {}

    for col in no_m:
        s = 0
        list_idexe = []
        for idx, row in energy_qm_df.iterrows():
            if not isinstance(extract_time_df[col][idx], float):
                for time in extract_time_df[col][idx]:
                    if time in no_m[col].get(idx, []):
                        list_idexe.append(s + extract_time_df[col][idx].index(time))
                s += len(extract_time_df[col][idx])
        unmatch_time_index[col] = list_idexe

    return unmatch_time_index


def calculate_energy_differences(energy_IR_BR: Dict[str, np.ndarray], energy_BR: Dict[str, np.ndarray], ref_energy: float) -> Dict[str, np.ndarray]:
    """
    Calculates the energy differences E(IR+BR)-E(BR).

    Parameters
    ----------
    energy_IR_BR : Dict[str, np.ndarray]
        A dictionary where each key is a round identifier and each value is a numpy array of energy values.
    energy_BR : Dict[str, np.ndarray]
        A dictionary where each key is a round identifier and each value is a numpy array of energy values.
    ref_energy : float
        A reference energy value to subtract from the calculated differences.

    Returns
    -------
    Dict[str, np.ndarray]
        A dictionary where each key is a round identifier from `energy_IR_BR`, and each value is a numpy array
        of the calculated energy differences.

    Notes
    -----
    This function assumes that the keys in `energy_IR_BR` and `energy_BR` correspond to each other in the same order.
    """
    final_energy_dft = {}

    for round_ir, round_br in zip(energy_IR_BR, energy_BR):
        diff = np.asarray(energy_IR_BR[round_ir]) - np.asarray(energy_BR[round_br]) - ref_energy
        final_energy_dft[round_ir] = diff

    return final_energy_dft