"""Entity Matching Module
"""
import statistics
import pandas as pd
from time import time

import matplotlib.pyplot as plt
import numpy as np
from networkx import Graph
from .string_matchers import WhitespaceTokenizer, Cosine, Jaccard, GeneralizedJaccard, Dice, OverlapCoefficient
from stringcompare import Levenshtein, Jaro
from stringcompare.preprocessing import NGramTokenizer
from tqdm.autonotebook import tqdm

from .datamodel import Data, PYJEDAIFeature
from .evaluation import Evaluation
from .utils import WordQgramTokenizer, cosine, get_qgram_from_tokenizer_name, FrequencyEvaluator


metrics_mapping = {
    'edit_distance': Levenshtein(),
    'cosine' : Cosine(),
    'jaro' : Jaro(),
    'jaccard' : Jaccard(),
    'generalized_jaccard' : GeneralizedJaccard(),
    'dice': Dice(),
    'overlap_coefficient' : OverlapCoefficient(),
}

vector_metrics_mapping = {
    'cosine': cosine
}

string_metrics = [
    'jaro', 'edit_distance'
]

set_metrics = [
    'cosine', 'dice', 'generalized_jaccard', 'jaccard', 'overlap_coefficient'
]

vector_metrics = [ 
    'cosine', 'dice', 'jaccard', 'sqeuclidean'
]

whoosh_index_metrics = [
    'TF-IDF', 'Frequency', 'PL2', 'BM25F'
] 

faiss_metrics = [
    'cosine', 'euclidean'
]

magellan_metrics = string_metrics + set_metrics
available_metrics = magellan_metrics + vector_metrics + whoosh_index_metrics + faiss_metrics

#
# Tokenizers
#
# char_qgram_tokenizers = { 'char_'+ str(i) + 'gram':i for i in range(1, 7) }
# word_qgram_tokenizers = { 'word_'+ str(i) + 'gram':i for i in range(1, 7) }
char_qgram_tokenizers = ['char_tokenizer']
word_qgram_tokenizers = ['word_tokenizer']
magellan_tokenizers = ['white_space_tokenizer']
joins_tokenizers = ["qgrams", "standard", "standard_multiset", "qgrams_multiset"]

# tfidf_tokenizers = [ 'tfidf_' + cq for cq in char_qgram_tokenizers.keys() ] + \
#                     [ 'tfidf_' + wq for wq in word_qgram_tokenizers.keys() ]

# tf_tokenizers = [ 'tf_' + cq for cq in char_qgram_tokenizers.keys() ] + \
#                     [ 'tf_' + wq for wq in word_qgram_tokenizers.keys() ]
                        
# boolean_tokenizers = [ 'boolean_' + cq for cq in char_qgram_tokenizers.keys() ] + \
#                         [ 'boolean_' + wq for wq in word_qgram_tokenizers.keys() ]

# vector_tokenizers = tfidf_tokenizers + tf_tokenizers + boolean_tokenizers

# available_tokenizers = [key for key in char_qgram_tokenizers] + [key for key in word_qgram_tokenizers] + magellan_tokenizers + vector_tokenizers
available_tokenizers = char_qgram_tokenizers + word_qgram_tokenizers + magellan_tokenizers + joins_tokenizers
available_vectorizers = ['tfidf', 'tf', 'boolean']

class AbstractEntityMatching(PYJEDAIFeature):
    """Calculates similarity from 0.0 to 1.0
    """

    _method_name: str = "Abstract Entity Matching"
    _method_info: str = "Calculates similarity from 0. to 1."

    def __init__(
            self,
            metric: str = 'dice',
            similarity_threshold: float = 0.5,
        ) -> None:
        self.pairs: Graph
        self.metric = metric
        self.similarity_threshold = similarity_threshold
        self.execution_time = 0
        self.qgram=None
        #
        # Selecting tokenizer
        #
        if metric not in available_metrics:
            raise AttributeError(
                'Metric ({}) does not exist. Please select one of the available. ({})'.format(
                    metric, available_metrics
                )
            )
        else:
            self._metric = metric

    def _predict_candidate_pairs(self, blocks: dict) -> None:
        """Similarity evaluation after comparison cleaning.

        Args:
            blocks (dict): Comparison cleaning blocks.
        """
        for entity_id, candidates in blocks.items():
            for candidate_id in candidates:
                similarity = self._similarity(entity_id, candidate_id)
                self._insert_to_graph(entity_id, candidate_id, similarity)
            self._progress_bar.update(1)

    def _insert_to_graph(self, entity_id1, entity_id2, similarity):
        if self.similarity_threshold is None or \
            (self.similarity_threshold is not None and similarity > self.similarity_threshold):
            self.pairs.add_edge(entity_id1, entity_id2, weight=similarity)

    def report(self) -> None:
        """Prints Block Building method configuration
        """
        print(
            "Method name: " + self._method_name +
            "\nMethod info: " + self._method_info +
            ("\nParameters: \n" + ''.join(['\t{0}: {1}\n'.format(k, v) for k, v in self._configuration().items()]) if self._configuration().items() else "\nParameters: Parameter-Free method\n") +
            "Attributes:\n\t" + ', '.join(c for c in (self.attributes if self.attributes is not None \
                else self.data.dataset_1.columns)) +
            "\nRuntime: {:2.4f} seconds".format(self.execution_time)
        )

    def get_weights_avg(self) -> float:
        return sum([w for _, _, w in self.pairs.edges(data='weight')])/len(self.pairs.edges(data='weight'))

    def get_weights_median(self) -> float:
        return [w for _, _, w in sorted(self.pairs.edges(data='weight'))][int(len(self.pairs.edges(data='weight'))/2)]    
    
    def get_weights_standard_deviation(self) -> float:
        return statistics.stdev([w for _, _, w in self.pairs.edges(data='weight')])
    
    def plot_distribution_of_all_weights(self, save_figure_path=None) -> None:
        title = "Distribution of scores with " + self.metric + " metric in graph from entity matching"
        plt.figure(figsize=(10, 6))
        all_weights = [w for _, _, w in self.pairs.edges(data='weight')]
        sorted_weights = sorted(all_weights, reverse=True)
        
        plt.hist(sorted_weights)
        plt.xlim(0, 1)
        # only one line may be specified; full height
        plt.axvline(x = self.get_weights_avg(), color = 'blue', label = 'Average weight')
        plt.axvline(x = self.get_weights_median(), color = 'black', label = 'Median weight')
        plt.axvline(x = self.get_weights_avg()+self.get_weights_standard_deviation(), color = 'green', label = 'Average + SD weight')
        plt.legend()
        if save_figure_path:
            plt.savefig(save_figure_path)
        plt.show()

    def plot_distribution_of_all_weights_2d(self, save_figure_path=None) -> None:
        title = "Distribution of scores with " + self.metric + " metric in graph from entity matching"
        plt.figure(figsize=(10, 6))
        all_weights = [w for _, _, w in self.pairs.edges(data='weight')]
        sorted_weights = sorted(all_weights, reverse=True)
        
        fig, ax = plt.subplots(tight_layout=True)
        hist = ax.hist2d(sorted_weights, sorted_weights)
        plt.axvline(x = self.get_weights_avg(), color = 'blue', label = 'Average weight')
        plt.axvline(x = self.get_weights_median(), color = 'black', label = 'Median weight')
        plt.axvline(x = self.get_weights_avg()+self.get_weights_standard_deviation(), color = 'green', label = 'Average + SD weight')
        plt.legend()
        if save_figure_path:
            plt.savefig(save_figure_path)
        plt.show()

    def plot_distribution_of_scores(self, save_figure_path=None) -> None:
        title = "Distribution of scores with " + self.metric + " metric in graph from entity matching"
        def weight_distribution(G):
            bins = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
            distribution = [0] * (len(bins)-1)
            for u, v, w in G.edges(data='weight'):
                for i in range(len(bins) - 1):
                    if bins[i] <= w < bins[i + 1]:
                        distribution[i] += 1
                        break
            return distribution, len(G.edges(data='weight'))

        labels = [f'{(i)/10:.1f} - {(i+1)/10:.1f}' for i in range(0, 10)]

        distribution, num_of_pairs = weight_distribution(self.pairs)
        width = 0.5
        x = np.arange(len(labels))  # the label locations
        distribution = list(map(lambda x: (x/num_of_pairs)*100, distribution))
        print("Distribution-% of predicted scores: ", distribution)

        fig, ax = plt.subplots(figsize=(10,6))
        r1 = ax.bar(x, distribution, width, align='center', color='red')
        ax.set_xticks(x)
        ax.set_xticklabels(labels)
    
        
        # Add some text for labels, title and custom x-axis tick labels, etc.
        ax.set_ylabel('Percentage of pairs in each range to all (%)')
        ax.set_title(title)
        ax.set_xlabel('Similarity score range')
        fig.tight_layout()
        
        # only one line may be specified; full height
        plt.axvline(x = self.get_weights_avg()*10, color = 'blue', label = 'Average weight')
        plt.axvline(x = self.get_weights_median()*10, color = 'black', label = 'Median weight')
        plt.axvline(x = self.get_weights_avg()*10+self.get_weights_standard_deviation()*10, color = 'green', label = 'Average + SD weight')
        plt.legend()
        if save_figure_path:
            plt.savefig(save_figure_path)
        plt.show()

    def plot_gt_distribution_of_scores(self, save_figure_path=None) -> None:
        title = "Distribution of scores with " + self.metric + " metric on ground truth pairs"
        def weight_distribution():
            bins = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
            distribution = [0] * (len(bins)-1)
            for _, (id1, id2) in self.data.ground_truth.iterrows():
                id1 = self.data._ids_mapping_1[id1]
                id2 = self.data._ids_mapping_1[id2] if self.data.is_dirty_er else self.data._ids_mapping_2[id2]
                w = self._calculate_vector_similarity(id1, id2)

                for i in range(len(bins) - 1):
                    if bins[i] <= w < bins[i + 1]:
                        distribution[i] += 1
                        break
            return distribution, len(self.data.ground_truth)

        labels = [f'{(i)/10:.1f} - {(i+1)/10:.1f}' for i in range(0, 10)]

        distribution, num_of_pairs = weight_distribution()
        width = 0.5
        x = np.arange(len(labels))  # the label locations
        distribution = list(map(lambda x: (x/num_of_pairs)*100, distribution))
        print("Distribution-% of predicted scores: ", distribution)

        fig, ax = plt.subplots(figsize=(10,6))
        r1 = ax.bar(x, distribution, width, align='center', color='blue')
        ax.set_xticks(x)
        ax.set_xticklabels(labels)

        # Add some text for labels, title and custom x-axis tick labels, etc.
        ax.set_ylabel('Percentage of pairs in each range to all (%)')
        ax.set_title(title)
        ax.set_xlabel('Similarity score range')
        fig.tight_layout()
        if save_figure_path:
            plt.savefig(save_figure_path)
        plt.show()

    def evaluate(self,
                 prediction,
                 export_to_df: bool = False,
                 export_to_dict: bool = False,
                 with_classification_report: bool = False,
                 verbose: bool = True) -> any:

        if self.data is None:
            raise AttributeError("Can not proceed to evaluation without data object.")

        if self.data.ground_truth is None:
            raise AttributeError("Can not proceed to evaluation without a ground-truth file. " +
                    "Data object has not been initialized with the ground-truth file")

        eval_obj = Evaluation(self.data)
        true_positives = 0
        total_matching_pairs = prediction.number_of_edges()
        for _, (id1, id2) in self.data.ground_truth.iterrows():
            id1 = self.data._ids_mapping_1[id1]
            id2 = self.data._ids_mapping_1[id2] if self.data.is_dirty_er \
                                                else self.data._ids_mapping_2[id2]
            if (id1 in prediction and id2 in prediction[id1]) or   \
                 (id2 in prediction and id1 in prediction[id2]):
                true_positives += 1

        eval_obj.calculate_scores(true_positives=true_positives,
                                  total_matching_pairs=total_matching_pairs)
        return eval_obj.report(self.method_configuration(),
                                export_to_df,
                                export_to_dict,
                                with_classification_report,
                                verbose)
        
    def stats(self) -> None:
        pass
        
    def export_to_df(self, prediction: Graph, tqdm_enable=False) -> pd.DataFrame:
        """Creates a dataframe with the predicted pairs.

        Args:
            prediction (Graph): Predicted graph
            tqdm_enable (bool): Whether to enable tqdm progress bar

        Returns:
            pd.DataFrame: Dataframe with the predicted pairs
        """
        pairs_list = []

        is_dirty_er = self.data.is_dirty_er
        dataset_limit = self.data.dataset_limit
        gt_to_ids_reversed_1 = self.data._gt_to_ids_reversed_1
        if not is_dirty_er:
            gt_to_ids_reversed_2 = self.data._gt_to_ids_reversed_2

        for edge in tqdm(prediction.edges, disable=not tqdm_enable, desc="Exporting to DataFrame"):
            node1, node2 = edge

            if not is_dirty_er:
                if node1 < dataset_limit:
                    id1 = gt_to_ids_reversed_1[node1]
                    id2 = gt_to_ids_reversed_2[node2]
                else:
                    id1 = gt_to_ids_reversed_1[node2]
                    id2 = gt_to_ids_reversed_2[node1]
            else:
                id1 = gt_to_ids_reversed_1[node1]
                id2 = gt_to_ids_reversed_1[node2]

            pairs_list.append((id1, id2))

        pairs_df = pd.DataFrame(pairs_list, columns=['id1', 'id2'])

        return pairs_df

class EntityMatching(AbstractEntityMatching):
    """Calculates similarity from 0.0 to 1.0 for all blocks
    """

    _method_name: str = "Entity Matching"
    _method_info: str = "Calculates similarity from 0. to 1. for all blocks"

    def __init__(
            self,
            metric: str = 'dice',
            tokenizer: str = 'white_space_tokenizer',
            vectorizer : str = None,
            qgram : int = 1,
            similarity_threshold: float = 0.0,
            tokenizer_return_unique_values = False, # unique values or not,
            attributes: any = None,
        ) -> None:
        super().__init__()
        self.pairs: Graph
        self.metric = metric
        self.attributes: list = attributes
        self.similarity_threshold = similarity_threshold
        self.tokenizer = tokenizer
        self.execution_time = 0
        self.vectorizer = vectorizer
        self.qgram: int = -1
        #
        # Selecting tokenizer
        #
        if metric not in available_metrics:
            raise AttributeError(
                'Metric ({}) does not exist. Please select one of the available. ({})'.format(
                    metric, available_metrics
                )
            )
        else:
            self._metric = metric

        self.tokenizer_return_set = (metric in set_metrics) or tokenizer_return_unique_values    
        self.qgram : int = qgram
        
        if(vectorizer is not None):
            if self.vectorizer not in available_vectorizers:
                raise AttributeError(
                    'Weighting Scheme ({}) does not exist. Please select one of the available. ({})'.format(
                        vectorizer, available_vectorizers
                    )
                )
                
        elif(tokenizer is not None):
            if tokenizer == 'white_space_tokenizer':
                self._tokenizer = WhitespaceTokenizer()
            elif tokenizer == 'char_tokenizer':
                self._tokenizer = NGramTokenizer(n=self.qgram)
            elif tokenizer == 'word_tokenizer':
                self._tokenizer = WordQgramTokenizer(q=self.qgram)
            elif tokenizer not in available_tokenizers:
                raise AttributeError(
                    'Tokenizer ({}) does not exist. Please select one of the available. ({})'.format(
                        tokenizer, available_tokenizers
                    )
                )
        
    def predict(self,
                blocks: dict,
                data: Data,
                tqdm_disable: bool = False) -> Graph:
        """Main method of entity matching. Inputs a set of blocks and outputs a graph \
            that contains of the entity ids (nodes) and the similarity scores between them (edges).

            Args:
                blocks (dict): blocks of entities
                data (Data): dataset module
                tqdm_disable (bool, optional): Disables progress bar. Defaults to False.

            Returns:
                networkx.Graph: entity ids (nodes) and similarity scores between them (edges)
        """
        start_time = time()
        self.tqdm_disable = tqdm_disable
        
        if not blocks:
            raise ValueError("Empty blocks structure")
        self.data = data
        self.pairs = Graph()
        all_blocks = list(blocks.values())
        self._progress_bar = tqdm(total=len(blocks),
                                  desc=self._method_name+" ("+self.metric+ ", " + str(self.tokenizer) + ")",
                                  disable=self.tqdm_disable)

        if self.vectorizer is not None:
            self.initialize_vectorizer()

        if 'Block' in str(type(all_blocks[0])):
            self._predict_raw_blocks(blocks)
        elif isinstance(all_blocks[0], set):
            self._predict_candidate_pairs(blocks)
        else:
            raise AttributeError("Wrong type of Blocks")
        self.execution_time = time() - start_time
        self._progress_bar.close()

        return self.pairs
            
    def _predict_raw_blocks(self, blocks: dict) -> None:
        """Method for similarity evaluation blocks after Block building

        Args:
            blocks (dict): Block building blocks
        """
        if self.data.is_dirty_er:
            for _, block in blocks.items():
                entities_array = list(block.entities_D1)
                for index_1 in range(0, len(entities_array), 1):
                    for index_2 in range(index_1+1, len(entities_array), 1):
                        similarity = self._similarity(entities_array[index_1],
                                                      entities_array[index_2])
                        self._insert_to_graph(entities_array[index_1],
                                              entities_array[index_2],
                                              similarity)
                self._progress_bar.update(1)
        else:
            for _, block in blocks.items():
                for entity_id1 in block.entities_D1:
                    for entity_id2 in block.entities_D2:
                        similarity = self._similarity(entity_id1, entity_id2)
                        self._insert_to_graph(entity_id1, entity_id2, similarity)
                self._progress_bar.update(1)

    def initialize_vectorizer(self) -> None:
        self.frequency_evaluator : FrequencyEvaluator = FrequencyEvaluator(vectorizer=self.vectorizer,
                                                                            tokenizer=self.tokenizer,
                                                                            qgram=self.qgram)
        d1 = self.data.dataset_1[self.attributes] if self.attributes else self.data.dataset_1
        self._entities_d1 = d1 \
                    .apply(" ".join, axis=1) \
                    .apply(lambda x: x.lower()) \
                    .values.tolist()
        
        d2 = None
        if(not self.data.is_dirty_er):
            d2 = self.data.dataset_2
            if self.attributes:
                d2 = d2[self.attributes]

        self._entities_d2 = d2 \
                    .apply(" ".join, axis=1) \
                    .apply(lambda x: x.lower()) \
                    .values.tolist() if not self.data.is_dirty_er else self._entities_d1 
        
        
        _dataset_identifier : str = ('_'.join([self.data.dataset_name_1, self.data.dataset_name_2])) if(self.data.dataset_name_1 is not None and self.data.dataset_name_2 is not None) else ("dataset") 
        self.frequency_evaluator.fit(metric=self.metric,
                                    dataset_identifier=_dataset_identifier,
                                    indexing='inorder',
                                    d1_entities=self._entities_d1,
                                    d2_entities=self._entities_d2)

    def _similarity(self, entity_id1: int, entity_id2: int) -> float:

        similarity: float = 0.0
        if self.vectorizer is not None:
            return self.frequency_evaluator.predict(id1=entity_id1, id2=entity_id2)
        elif isinstance(self.attributes, dict):
            for attribute, weight in self.attributes.items():
                e1 = self.data.entities.iloc[entity_id1][attribute].lower()
                e2 = self.data.entities.iloc[entity_id2][attribute].lower()
                if self.tokenizer_return_set: 
                    similarity += weight*metrics_mapping[self._metric].compare(
                        set(self._tokenizer.tokenize(e1)) if self._metric in set_metrics else e1,
                        set(self._tokenizer.tokenize(e2)) if self._metric in set_metrics else e2
                    )
                else:     
                    similarity += weight*metrics_mapping[self._metric].compare(
                        self._tokenizer.tokenize(e1) if self._metric in set_metrics else e1,
                        self._tokenizer.tokenize(e2) if self._metric in set_metrics else e2
                    )
        elif isinstance(self.attributes, list):
            for attribute in self.attributes:
                e1 = self.data.entities.iloc[entity_id1][attribute].lower()
                e2 = self.data.entities.iloc[entity_id2][attribute].lower()
                if self.tokenizer_return_set: 
                    similarity += metrics_mapping[self._metric].compare(
                        set(self._tokenizer.tokenize(e1)) if self._metric in set_metrics else e1,
                        set(self._tokenizer.tokenize(e2)) if self._metric in set_metrics else e2
                    )
                else:
                    similarity += metrics_mapping[self._metric].compare(
                        self._tokenizer.tokenize(e1) if self._metric in set_metrics else e1,
                        self._tokenizer.tokenize(e2) if self._metric in set_metrics else e2
                    )
            similarity /= len(self.attributes)
        else:
            e1 = self.data.entities.iloc[entity_id1].str.cat(sep=' ').lower()
            e2 = self.data.entities.iloc[entity_id2].str.cat(sep=' ').lower()
            if self.tokenizer_return_set: 
                te1 = set(self._tokenizer.tokenize(e1)) if self._metric in set_metrics else e1
                te2 = set(self._tokenizer.tokenize(e2)) if self._metric in set_metrics else e2
            else:
                te1 = self._tokenizer.tokenize(e1) if self._metric in set_metrics else e1
                te2 = self._tokenizer.tokenize(e2) if self._metric in set_metrics else e2
            similarity = metrics_mapping[self._metric].compare(te1, te2)
        return similarity

    def _configuration(self) -> dict:
        return  {
            "Metric" : self.metric,
            "Attributes" : self.attributes,
            "Similarity threshold" : self.similarity_threshold,
            "Tokenizer" : self.tokenizer,
            "Vectorizer" : self.vectorizer if self.vectorizer is not None else "None",
            "Qgrams" : self.qgram
        }

class VectorBasedMatching(AbstractEntityMatching):

    _method_name: str = "Vector Based Matching"
    _method_info: str = "Calculates similarity from 0. to 1. for vectors"

    def __init__(
            self,
            metric: str = 'cosine',
            similarity_threshold: float = 0.5,
        ) -> None:
        self.pairs: Graph
        self.metric = metric
        self.similarity_threshold = similarity_threshold
        self.vectors_d1 = None
        self.vectors_d2 = None
        self.execution_time = 0

        #
        # Selecting tokenizer
        #
        if metric not in vector_metrics:
            raise AttributeError(
                'Metric ({}) does not exist. Please select one of the available. ({})'.format(
                    metric, available_metrics
                )
            )
        else:
            self._metric = metric
            
    def predict(self,
                blocks: dict,
                data: Data,
                vectors_d1: np.array,
                vectors_d2: np.array = None,
                tqdm_disable: bool = False,
        ) -> Graph:
        """Main method of entity matching. Inputs a set of blocks and outputs a graph \
            that contains of the entity ids (nodes) and the similarity scores between them (edges).

            Args:
                blocks (dict): blocks of entities
                data (Data): dataset module
                tqdm_disable (bool, optional): Disables progress bar. Defaults to False.

            Returns:
                networkx.Graph: entity ids (nodes) and similarity scores between them (edges)
        """
        start_time = time()
        self.tqdm_disable = tqdm_disable
        self.vectors_d1 = vectors_d1
        # self.vectors_d2 = vectors_d2
        

        if(vectors_d1 is None):
            raise ValueError("Embeddings of the first dataset not given")
        else:
            self.vectors = vectors_d1
            if(not data.is_dirty_er):
                if(vectors_d2 is None):
                    raise ValueError("Embeddings of the second dataset not given")
                self.vectors = np.concatenate((vectors_d1,vectors_d2), axis=0)
        self.data = data
        self.pairs = Graph()
        self._progress_bar = tqdm(total=len(blocks),
                                  desc=self._method_name+" ("+self.metric+ ", " + str(self.tokenizer) + ")",
                                  disable=self.tqdm_disable)
        self._predict_candidate_pairs(blocks)
        self.execution_time = time() - start_time
        self._progress_bar.close()

        return self.pairs

    def _similarity(self, entity_id1: int, entity_id2: int) -> float:
        return vector_metrics_mapping[self._metric](self.vectors[entity_id1], self.vectors[entity_id2])

    def _configuration(self) -> dict:
        conf =  {
            "Metric" : self.metric,
            "Similarity threshold" : self.similarity_threshold
        }
        
        return conf
