# ====================================================================================================================
# Imports
# ====================================================================================================================
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.neighbors import KernelDensity
from sklearn.metrics import pairwise
import warnings
from typing import Type
import matplotlib.pyplot as plt
from matplotlib import rcParams
from pathlib import Path
import sys
sys.path.append(str(Path(__file__).parents[2]))
from encoder.base_encoder import Encoder, EncodingDataFrameCreator
from defaults import DefaultPaths, DefaultNames


# ====================================================================================================================
# Similarity Measure
# ====================================================================================================================
class SimilarityMeasure:
    """
    Base class to measure the similarity between encodings of different classes.
    It calculates the intra and inter class similarity. The inter class similarity
    is calculated between the mean vectors of the classes.
    The intra class similarity is a list of similarity scores between the vectors of the same class.
    To get a single value for the intra class similarity, you can calculate e.g. the mean of the scores.

    Methods:
        - similarity: Calculate the intra and inter class similarity.

    Abstract:
        - property name: The name of the similarity measure.
        - calculate_similarity_matrix: Calculate the similarity matrix between the vectors.
    """
    @property
    def name(self) -> str:
        raise NotImplementedError('This property must be implemented by the subclass.')

    @property
    def merged_data(self) -> pd.DataFrame:
        return self._merged_data.copy()

    def __init__(
            self, encoder: Encoder,
            patho_data_path: Path = DefaultPaths().patho,
            clinical_data_path: Path = DefaultPaths().clinical,
            similarity_type: str = 'Matrix'
    ):
        self.encoder = encoder
        self._encoding_data_frame_creator = EncodingDataFrameCreator(
            encoder=encoder, patho_data_path=patho_data_path, clinical_data_path=clinical_data_path
        )
        self._merged_data = self._encoding_data_frame_creator.create_combined_data_frame()
        self._encoding_column_name = DefaultNames().encoding_column_name
        self._similarity_column_name = DefaultNames().similarity_column_name
        self.similarity_type = similarity_type

    # ================================================================================================================
    # Public Methods
    # ================================================================================================================
    def similarity(
            self, feature: str = 'primarily_metastasis', **kwargs
    ) -> tuple[pd.DataFrame, pd.DataFrame] | None:
        """
        Calculates the intra and inter class similarity for a given feature of
        the encodings generated by the encoder.

        Args:
            feature (str, optional): The feature for which the similarity should be calculated.

        Returns:
            tuple[pd.DataFrame, pd.DataFrame] | None: The intra and inter class similarity if the
            feature exists in the pathological or clinical data. Otherwise, None is returned.

            intra_class_similarity (pd.DataFrame): The similarity scores between the vectors of the same class.
            These scores are stored as list in the default name similarity_column_name column. The label
            column contains the class label. The similarities are all relevant scores between the encodings
            that are in that class (no vector to itself is compared).

            inter_class_similarity (pd.DataFrame): The similarity scores between the mean vectors of the classes.
        """
        merged_data = self.merged_data
        if not feature in merged_data.columns:
            warnings.warn('There is no feature with the name: {}'.format(feature))
            return None

        data_dict = self._vector_split_for_feature_value(feature=feature)
        intra_class_similarity = self._intra_class_similarity(data_dict)
        inter_class_similarity = self._inter_class_similarity(data_dict)

        return intra_class_similarity, inter_class_similarity

    # ================================================================================================================
    # Helper Methods
    # ================================================================================================================
    def _vector_split_for_feature_value(self, feature: str = 'primarily_metastasis') -> dict[any, pd.DataFrame]:
        merged_data = self._merged_data
        labels = self._merged_data[feature].unique()
        data_dict = {}
        for label in labels:
            if label and (type(label) != np.float64 or not np.isnan(label)):
                data = merged_data[merged_data[feature] == label][[self._encoding_column_name, 'patient_id']]
            else:
                data = merged_data[merged_data[feature].isnull()][[self._encoding_column_name, 'patient_id']]
                label = 'None'
            data_dict[label] = data
        return data_dict

    def _intra_class_similarity(self, data_dict: dict[any, pd.DataFrame]) -> pd.DataFrame:
        intra_class_dict = {}
        for key in data_dict.keys():
            vectors = self._data_dict_to_np_array(data_dict, key)
            sim_scores = self._intra_class_similarity_single_class(vectors)
            intra_class_dict[key] = sim_scores

        return pd.DataFrame(
            {
                'label': intra_class_dict.keys(),
                self._similarity_column_name: intra_class_dict.values()
            }
        )

    def _intra_class_similarity_single_class(self, vectors: np.ndarray) -> np.ndarray:
        sim_matrix = self.calculate_similarity_measure(vectors)
        sim_scores = self._return_flattened_upper_triangle_from_matrix(sim_matrix)
        return sim_scores

    def _inter_class_similarity(self, data_dict: dict[any, pd.DataFrame]) -> pd.DataFrame:
        mean_vector_dict = self._calculate_mean_vector(data_dict)
        sim_matrix = self.calculate_similarity_measure(np.array(list(mean_vector_dict.values())))
        return pd.DataFrame(sim_matrix,
                            index=list(mean_vector_dict.keys()),
                            columns=list(mean_vector_dict.keys()))

    def _calculate_mean_vector(
            self, data_dict: dict[any, pd.DataFrame]
    ) -> dict[any, np.ndarray]:
        mean_vector_dict = {}
        for key in data_dict.keys():
            vectors = self._data_dict_to_np_array(data_dict, key)
            mean_vector_dict[key] = np.mean(vectors, axis=0)
        return mean_vector_dict

    @staticmethod
    def _return_flattened_upper_triangle_from_matrix(sim_matrix: np.ndarray) -> np.ndarray:
        upper_triangular_indices = np.triu_indices_from(sim_matrix, k=1)
        upper_triangular_values = sim_matrix[upper_triangular_indices]
        return upper_triangular_values.reshape(-1)

    def _data_dict_to_np_array(self, data_dict: dict[any, pd.DataFrame], key: any) -> np.ndarray:
        data = data_dict[key]
        vectors = np.array(data[self._encoding_column_name].tolist())
        return vectors

    @staticmethod
    def _check_vector_shape(vectors: np.ndarray) -> None:
        if len(vectors.shape) != 2:
            raise ValueError('Can not calculate the similarity of a vector with shape: {}'.format(vectors.shape))

    def _create_all_vectors(self) -> np.ndarray:
        vectors = np.array(self.merged_data[self._encoding_column_name].tolist())
        return vectors

    # ================================================================================================================
    # Abstract Method
    # ================================================================================================================
    def calculate_similarity_measure(self, vectors: np.ndarray) -> np.ndarray:
        raise NotImplementedError('This method must be implemented by the subclass.')


class CosineSimilarity(SimilarityMeasure):
    @property
    def name(self) -> str:
        return 'Cosine Similarity'

    def calculate_similarity_measure(self, vectors: np.ndarray) -> np.ndarray:
        self._check_vector_shape(vectors)
        cosine_sim = pairwise.cosine_similarity(vectors)
        cosine_sim = np.abs(cosine_sim)
        return cosine_sim


class EntropySimilarity(SimilarityMeasure):
    @property
    def name(self) -> str:
        return 'Entropy Similarity'

    @property
    def pdf(self) -> KernelDensity:
        if not self._kde:
            self._kde = self._create_kde()
        return self._kde

    def __init__(self, encoder: Encoder,
                 patho_data_path: Path = DefaultPaths().patho,
                 clinical_data_path: Path = DefaultPaths().clinical
    ):
        super().__init__(encoder=encoder, patho_data_path=patho_data_path,
                         clinical_data_path=clinical_data_path,
                         similarity_type='Single')
        self._kde = None

    def _intra_class_similarity_single_class(self, vectors: np.ndarray) -> np.ndarray:
        sim_scores = self.calculate_similarity_measure(vectors)
        return sim_scores

    def _inter_class_similarity(self, data_dict: dict[any, pd.DataFrame]) -> pd.DataFrame:
        all_vectors = self._create_all_vectors()
        sim_measure = self.calculate_similarity_measure(all_vectors)
        return pd.DataFrame({'score': [sim_measure]})

    def _create_kde(self) -> KernelDensity:
        all_vectors = self._create_all_vectors()
        kde = KernelDensity(kernel='gaussian', bandwidth=1).fit(all_vectors)
        return kde

    def calculate_similarity_measure(self, vectors: np.ndarray):
        self._check_vector_shape(vectors)
        pdf = self.pdf
        log_probs = pdf.score_samples(vectors)
        probs = np.exp2(log_probs)
        entropy = -np.sum(probs * log_probs)
        return entropy


# ====================================================================================================================
# Similarity Based Encoder Comparison
# ====================================================================================================================
class SimilarityVisualizer:
    def __init__(self):
        pass

    def visualize_similarity_single_encoder(
            self,  similarity: tuple[pd.DataFrame, pd.DataFrame],
            axes: list[plt.Axes],
            feature: str = 'primarily_metastasis',
            encoder_name: str = 'BioClinicalBert',
            similarity_type: str = 'Cosine Similarity',
    ) -> None:
        """
        Fills the given axes with the intra and inter class similarity visualization. The intra class similarity is
        visualized as a box plot and the inter class similarity as a heatmap.
        """
        rcParams.update({'font.size': 4})
        rcParams['svg.fonttype'] = 'none'
        if len(axes) != 2 or len(similarity) != 2:
            raise ValueError('The axes and similarity must have a length of 2.')
        intra_sim, inter_sim = similarity
        self._create_box_plot(axes[0], intra_sim, feature, encoder_name, similarity_type)
        self._create_heatmap(axes[1], inter_sim, feature, encoder_name, similarity_type)

    @staticmethod
    def _create_box_plot(
            ax: plt.Axes, intra_sim: pd.DataFrame,
            feature: str, encoder: str, similarity_type: str,
            vmin: float = 0.5, vmax: float = 1,
            color=sns.color_palette('flare')[4]
    ) -> None:
        df_transformed = intra_sim.explode('similarity', ignore_index=True)
        ax.set_xlim(vmin, vmax)
        sns.boxplot(
            data=df_transformed,
            x='similarity',
            y='label',
            color=color,
            ax=ax,
            orient='h',
            legend=None,
            flierprops={"marker": "x"},
            fliersize=2
        )
        sns.despine(ax=ax, top=True, left=True, right=True, bottom=False)
        ax.set_title(f'{encoder}: {feature}', fontsize=6)
        ax.set_ylabel('Label', fontsize=4)
        ax.set_xlabel(similarity_type, fontsize=4)
        for label in ax.get_yticklabels():
            label.set_fontsize(6)
        for label in ax.get_xticklabels():
            label.set_fontsize(4)

    @staticmethod
    def _create_heatmap(
            ax: plt.Axes, inter_sim: pd.DataFrame,
            feature: str, encoder: str, similarity_type: str,
            vmin: float = 0.9, vmax: float = 1,
            cmap = sns.color_palette('flare', as_cmap=True)
    ) -> None:
        sns.heatmap(
            data=inter_sim,
            cmap=cmap,
            annot=True,
            cbar=True,
            ax=ax,
            vmin=vmin,
            vmax=vmax
        )
        ax.set_title(f'{encoder}: {feature} with {similarity_type}', fontsize=6)
        for label in ax.get_yticklabels():
            label.set_rotation(0)
            label.set_fontsize(6)
        for label in ax.get_xticklabels():
            label.set_fontsize(6)


class EncoderSimilarityComparer:
    def __init__(
            self, encoder_list: list[Encoder],
            similarity_measure_cls: Type[SimilarityMeasure],
            patho_data_path: Path = DefaultPaths().patho,
            clinical_data_path: Path = DefaultPaths().clinical,
            image_base_name: str = "NameOfModalityToCompare",
            image_base_save_path: Path = DefaultPaths().encodings_comparison
    ):
        self.encoder_list = encoder_list
        self.similarity_measure_cls = similarity_measure_cls
        self.patho_data_path = patho_data_path
        self.clinical_data_path = clinical_data_path
        self._image_base_name = image_base_name
        self._image_base_save_path = image_base_save_path
        self._similarity_visualizer = SimilarityVisualizer()
        self._similarity_measurers = self._create_similarity_measures_list()
        # self.similarity_measures = self._create_similarity_measures()

    def visualize_similarity_inter_encoder_based(
            self, feature: str = 'primarily_metastasis', save_flag: bool = True, fig_row_height: float = 2.7
    ) -> None:
        """
        Creates a plot with n-encoders rows. Each row contains the intra and inter class similarity for one encoder
        based on the given feature. The intra class similarity is visualized as a box plot and the inter class similarity
        as a heatmap. The plot is saved in the image_base_save_path with a combination of the image_base_name and the
        given feature name.

        Args:
             feature (str, optional): The feature for which the intra and inter class similarity should be calculated.
             save_flag (bool, optional): If True, the plot is saved in the image_base_save_path. Default is False.
             fig_row_height (float, optional): The height of each row in the plot. Default is 2.7.
        """
        n_rows = len(self.encoder_list)
        fig, axes = plt.subplots(nrows=n_rows, ncols=2, figsize=(7, n_rows * fig_row_height))
        for i, sim_measurer in enumerate(self._similarity_measurers):
            self._similarity_visualizer.visualize_similarity_single_encoder(
                similarity=sim_measurer.similarity(feature=feature), axes=axes[i, :],
                encoder_name=self.encoder_list[i].encoder_model.name,
                similarity_type=self._similarity_measurers[i].name,
                feature=feature
            )
        plt.tight_layout()
        if save_flag:
            if not self._image_base_save_path.exists():
                self._image_base_save_path.mkdir(parents=True)
            plt.savefig(
                (self._image_base_save_path /
                 f'inter_{self._image_base_name}_{feature}_{self._similarity_measurers[0].name}.svg'),
                bbox_inches="tight",
                dpi=200
            )
        plt.show()


    def visualize_similarity_inter_encoder_based_all_relevant_features(
            self, save_flag: bool = True, fig_row_height: float = 2.7
    ):
        features = self._relevant_feature_list()
        for feature in features:
            print(f'Starting to visualize the feature similarity for {feature}')
            self.visualize_similarity_inter_encoder_based(
                feature=feature, save_flag=save_flag, fig_row_height=fig_row_height
            )

    # ================================================================================================================
    # Helper Methods
    # ================================================================================================================
    def _create_similarity_measures_list(self):
        sim_measure_list = []
        for encoder in self.encoder_list:
            sim_measure_list.append(
                self.similarity_measure_cls(
                    encoder=encoder, patho_data_path=self.patho_data_path, clinical_data_path=self.clinical_data_path
                )
            )

        return sim_measure_list

    def _create_similarity_measures(self) -> list[tuple[pd.DataFrame, pd.DataFrame]]:
        sim_measures = []
        for sim_measurer in self._similarity_measurers:
            sim_measures.append(sim_measurer.similarity())
        return sim_measures

    def _relevant_feature_list(self) -> list[str]:
        return [
            'sex', 'primarily_metastasis', 'smoking_status', 'age_at_initial_diagnosis',
            'pT_stage', 'grading', 'number_of_positive_lymph_nodes', 'infiltration_depth_in_mm',
            'perinodal_invasion', 'lymphovascular_invasion_L', 'vascular_invasion_V', 'perineural_invasion_Pn',
            'pN_stage', 'survival_status', 'hpv_association_p16', 'primary_tumor_site'
        ]


# ====================================================================================================================
# Testing
# ====================================================================================================================
if __name__ == '__main__':
    print('Similarity Measures started.')
    # my_encoder = SurgeryReportEncoderBioClinicalBert()
    # sim_measure = CosineSimilarity(encoder=my_encoder)
    # intra_sim, inter_sim = sim_measure.similarity()
