import pandas as pd
import typing
from datamart_isi.joiners.joiner_base import JoinerBase
from datamart_isi.joiners.join_feature.feature_pairs import FeaturePairs
from datamart_isi.joiners.join_result import JoinResult
import rltk

"""
TODO
Implement RLTK joiner
"""


class RLTKJoiner_new(JoinerBase):
    """
        This new joiner just simply join all found pairs in rltk's block
    """
    def __init__(self):
        # if set to be true, we can conduct exact match with rltk joiner
        self.exact_match = False 
        self.join_target_column_names = ""

    def set_join_target_column_names(self, column_name):
        self.join_target_column_names = column_name

    def find_pair(self,
                  left_df: pd.DataFrame,
                  right_df: pd.DataFrame) -> typing.Tuple[pd.DataFrame, typing.List[tuple]]:
        class left(rltk.AutoGeneratedRecord):
            pass

        class right(rltk.AutoGeneratedRecord):
            pass

        left_df['id']=left_df.index.astype(str)
        right_df['id']=right_df.index.astype(str)
        if 'Unnamed: 0' in right_df.columns:
            right_df = right_df.drop(columns=['Unnamed: 0'])
        ds1 = rltk.Dataset(rltk.DataFrameReader(left_df), record_class=left)
        ds2 = rltk.Dataset(rltk.DataFrameReader(right_df), record_class=right)

        bg = rltk.HashBlockGenerator()
        block = bg.generate(bg.block(ds1, property_=self.join_target_column_names[0]),
                            bg.block(ds2, property_=self.join_target_column_names[1]))

        pairs = rltk.get_record_pairs(ds1, ds2, block=block)

        pairs_list = []
        pairs_column = [[] for _ in range(right_df.shape[0])]
        for r1, r2 in pairs:
            pairs_column[int(r2.id)].append(int(r1.id))
            pairs_list.append((r1.id, r2.id))

        right_df["joining_pairs"] = pairs_column
        return right_df, pairs_list

    def join(self,
             left_df: pd.DataFrame,
             right_df: pd.DataFrame,
             left_columns: typing.List[typing.List[int]],
             right_columns: typing.List[typing.List[int]],
             left_metadata: dict,
             right_metadata: dict
             ) -> JoinResult:

        class left(rltk.AutoGeneratedRecord):
            pass

        class right(rltk.AutoGeneratedRecord):
            pass

        left_df['id']=left_df.index.astype(str)
        right_df['id']=right_df.index.astype(str)
        if 'Unnamed: 0' in right_df.columns:
            right_df = right_df.drop(columns=['Unnamed: 0'])
        ds1 = rltk.Dataset(rltk.DataFrameReader(left_df), record_class=left)
        ds2 = rltk.Dataset(rltk.DataFrameReader(right_df), record_class=right)

        bg = rltk.HashBlockGenerator()
        block = bg.generate(bg.block(ds1, property_=self.join_target_column_names[0]),
                            bg.block(ds2, property_=self.join_target_column_names[1]))
        left_df = left_df.set_index('id')
        right_df = right_df.set_index('id')

        pairs = rltk.get_record_pairs(ds1, ds2, block=block)

        df_joined = pd.DataFrame()

        column_names_to_join = None
        for r1, r2 in pairs:
            left_res = left_df.loc[r1.id]
            right_res = right_df.loc[r2.id]
            if column_names_to_join is None:
                column_names_to_join = right_res.index.difference(left_res.index)
                matched_rows = right_res.index.intersection(left_res.index)
                columns_new = left_res.index.tolist()
                columns_new.extend(column_names_to_join.tolist())
            new = pd.concat([left_res, right_res[column_names_to_join]])
            df_joined = df_joined.append(new, ignore_index = True)
        # ensure that the original dataframe columns are at the first left part
        df_joined = df_joined[columns_new]

        return JoinResult(df_joined, matched_rows)


class RLTKJoiner(JoinerBase):

    def __init__(self):
        pass

    @staticmethod
    def find_pair(
             left_df: pd.DataFrame,
             right_df: pd.DataFrame,
             left_columns: typing.List[typing.List[int]],
             right_columns: typing.List[typing.List[int]],
             left_metadata: dict,
             right_metadata: dict
             ) -> typing.Tuple[pd.DataFrame, typing.List[tuple]]:

        fp = FeaturePairs(left_df, right_df, left_columns, right_columns, left_metadata, right_metadata)
        block = fp.get_rltk_block()
        record_pairs = rltk.get_record_pairs(fp.left_rltk_dataset, fp.right_rltk_dataset, block=block)
        pairs_list = []
        # for r1, r2 in record_pairs:
        #     pairs_list.append((r1.id, r2.id))

        pairs_column = [[] for _ in range(right_df.shape[0])]
        for r1, r2 in record_pairs:
            pairs_column[int(r2.id)].append(int(r1.id))
            pairs_list.append((r1.id, r2.id))

        right_df["joining_pairs"] = pairs_column

        return right_df, pairs_list

    def join(self,
             left_df: pd.DataFrame,
             right_df: pd.DataFrame,
             left_columns: typing.List[typing.List[int]],
             right_columns: typing.List[typing.List[int]],
             left_metadata: dict,
             right_metadata: dict
             ) -> JoinResult:
        # print(left_metadata)

        # step 1 : transform columns
        """
        1. merge columns if multi-multi relation, mark as "merged" so we use set-based functions only
        2. pull out the mapped columns and form to new datasets with same order to support
        """
        fp = FeaturePairs(left_df, right_df, left_columns, right_columns, left_metadata, right_metadata)
        block = fp.get_rltk_block()

        # test codes here for taking a look inside of block
        # for idx, b in enumerate(block.key_set_adapter):
        #     print(b)
        #     if idx == 10:
        #         break

        record_pairs = rltk.get_record_pairs(fp.left_rltk_dataset, fp.right_rltk_dataset, block=block)

        # similarity matrix, the shape is (left_df.shape[0] * right_df.shape[0])
        sim = [[0 for _ in range(len(right_df))] for _ in range(len(left_df))]
        # enumerate the record pairs and calculate the similarity
        for r1, r2 in record_pairs:
            similarities = []
            for f1, f2 in fp.pairs:
                v1 = f1.value_merge_func(r1)
                v2 = f2.value_merge_func(r2)
                if self.exact_match:
                    similarities.append(1 if v1 == v2 else 0)
                    continue
                # print(v1, v2, type(f1), type(f2))
                for similarity_func in f1.similarity_functions():
                    similarity = similarity_func(v1, v2)
                    similarities.append(similarity)
                    # print(f1.name, f2.name, v1, v2, similarity, similarity_func, type(f1))
                    # TODO: now only consider the first similarity function for now
                    break
                # print(v1, v2, similarities)
            sim[int(r1.id)][int(r2.id)] = sum(similarities)/len(similarities) if similarities else 0

        # choose the highest similarity row as the matched rows, the matched_rows have same length as left_df.shape[0]
        # TODO: why only choose the highest similarity? maybe considering add some other methods
        matched_rows = self.simple_best_match(sim)
        res = self.one_to_one_concat(matched_rows, left_df, right_df, right_columns)

        # step 2 : analyze target columns - get ranked similarity functions for each columns
        """
        see https://paper.dropbox.com/doc/ER-for-Datamart--ASlKtpR4ceGaj~6cN4Q7EWoSAQ-tRug6oRX6g5Ko5jzaeynT
        """

        # step 3 : check if 1-1, 1-n, n-1, or m-n relations,
        # based on the analyze in step 2 we can basically know if it is one-to-one relation

        return JoinResult(res, matched_rows)

    def one_to_one_concat(self, matched_rows, left_df, right_df, right_columns):
        right_remain = self.get_remain_list(right_df, right_columns)
        right_remain_headers = [right_df.columns.tolist()[j] for j in right_remain]
        series_arr = []
        for i in matched_rows:
            if isinstance(i, int):
                series_arr.append(right_df.iloc[i, right_remain])
            else:
                series_arr.append(pd.Series(index=right_remain_headers))
        to_join = pd.DataFrame(series_arr, index=range(len(matched_rows)))
        res = pd.concat([left_df, to_join], axis=1)
        return res

    def munkrus_match(self, sim: typing.List[typing.List[float]]):
        pass

    @staticmethod
    def simple_best_match(sim: typing.List[typing.List[float]], threshold=0.5):
        res = []
        for idx, v in enumerate(sim):
            max_val = threshold
            max_idx = None
            for idx_, v_ in enumerate(v):
                if v_ >= max_val:
                    max_idx = idx_
                    max_val = v_
            res.append(max_idx)
        return res

    @staticmethod
    def simple_best_matches(sim: typing.List[typing.List[float]], threshold=0.5):
        res = []
        for idx, v in enumerate(sim):
            cur = []
            for idx_, v_ in enumerate(v):
                if v_ >= threshold:
                    cur.append(idx_)
            res.append(cur)
        return res

    @staticmethod
    def get_remain_list(df: pd.DataFrame, columns_2d: typing.List[typing.List[int]]):
        all_columns = list(range(df.shape[1]))
        columns_1d = [item for sublist in columns_2d for item in sublist]
        remaining = [_ for _ in all_columns if _ not in columns_1d]
        return remaining



