"""
    @release_date  : $release_date
    @version       : $release_version
    @author        : Christos Matsingos, Ka Fu Man 
    
    This file is part of the TrIPP software
    (https://github.com/fornililab/TrIPP).
    Copyright (c) 2024 Christos Matsingos, Ka Fu Man and Arianna Fornili.

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, version 3.

    This program is distributed in the hope that it will be useful, but
    WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
    General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program. If not, see <http://www.gnu.org/licenses/>.
"""

import numpy as np
from tqdm import tqdm

def greedy_clustering(cutoff, rmsd_matrix, frames, trajectory_names):
    """
    Function to run greedy clustering as described by Micheletti et al.
    (DOI: 10.1002/1097-0134(20000901)40:4<662::aid-prot90>3.0.co;2-f).
    
    Parameters
    ----------
    cutoff : float
        The distance cutoff for clustering.
    rmsd_matrix : np.ndarray
        The distance matrix (Euclidean distance) generated by the calculate_rmsd_matrix function.
    frames : list
        A list of frames corresponding to the points to be clustered.
    trajectory_names : np.ndarray
        An array of trajectory names corresponding to the clustering points.
        Generated by the create_clustering_matrix function.
    Returns
    -------
    labels: np.ndarray
        An array of cluster labels for each point in the feature matrix.
    cluster_centers: list
        Indices of the cluster centers within their individual trajectories (local indices). 
        Indices start from 0 for each trajectory.
    cluster_center_indices: list
        Indices of the cluster centers within the full feature matrix (global indices).
    cluster_centers_trajectories: list
        Trajectory names of the cluster centers.
    """

    # A copy of the RMSD matrix is made to not modify the original array
    current_matrix = rmsd_matrix.copy()
    n_structures = len(frames)

    # A labels array with the same number of elements as the total numebr of
    # structures is generated. Initially all structures are labelled to
    # belong to cluster 0.
    labels = np.zeros(n_structures, dtype=int)
    cluster = 1
    cluster_centers = []
    cluster_center_indices = []
    cluster_centers_trajectories = []

    for i in tqdm(range(n_structures)):
        # A boolean array neighbor_mask is generated to determine what RMSDs
        # are under the cutoff.
        neighbor_mask = current_matrix <= cutoff

        # The current_matrix array is checked to determine if all elements are
        # set to infinity. If the condition is true the clustering is complete.
        if np.all(np.isinf(current_matrix)):
            break

        else:
            # The amount of neighbors (structures within the cutoff)
            # is determined for each structure.
            neighbor_counts = np.sum(neighbor_mask, axis=0)
            # The structure with the most neighbors is determined to
            # be the center of the cluster
            central_structure = np.argmax(neighbor_counts)
            cluster_centers.append(frames[central_structure])
            cluster_centers_trajectories.append(trajectory_names[central_structure])
            cluster_center_indices.append(central_structure)
            # All structures neighboring the central structure are determined
            # and are assigned to the same cluster
            neighbors = np.ravel(neighbor_mask[central_structure, :])
            labels[neighbors] = cluster
            # The current_matrix (copy of the RMSD matrix) is altered to
            # remove elements belonging to the cluster from future clustering
            # by setting the RMSD values to inf.
            current_matrix[central_structure, :] = np.inf
            current_matrix[:, central_structure] = np.inf
            current_matrix[neighbor_mask[central_structure], :] = np.inf
            current_matrix[:, neighbor_mask[central_structure]] = np.inf
            cluster += 1  # The following cluster label is generated

    cluster_centers = list(np.ravel(np.array(cluster_centers)))
    cluster_center_indices = list(np.ravel(np.array(cluster_center_indices)))
    cluster_centers_trajectories = list(np.ravel(np.array(cluster_centers_trajectories)))
    # All labels are changed to fit the numbering of the other clustering 
    # methods.
    labels = (labels - 1)

    return labels, cluster_centers, cluster_center_indices, cluster_centers_trajectories
