from wikifier.utils import wikifier_for_ethiopia
import typing
import rltk
from rltk.io.reader.dataframe_reader import DataFrameReader
import pandas as pd

class RLTKJoinerWikidata(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()
        str_key_l = self.join_target_column_names[0]
        str_key_r = self.join_target_column_names[1]
        block = bg.generate(
                bg.block(self.left_rltk_dataset, function_=lambda r: ''.join([str(getattr(r, pk)).lower()
                                                                              for pk in str_key_l])),
                bg.block(self.right_rltk_dataset, function_=lambda r: ''.join([str(getattr(r, pk)).lower()
                                                                               for pk in str_key_r]))
                )
        # 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))

        if 'id' in right_df.columns:
            right_df = right_df.drop(columns=['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)


def main():
    dd = pd.read_csv("~/Desktop/2017-2018_food_price_ethiopia.csv")
    phem_data = pd.read_csv("~/Desktop/phem_data_monthly_oromia.csv")
    res = wikifier_for_ethiopia(dd)
    joiner = RLTKJoinerWikidata()
    joiner.set_join_target_column_names([("dateTime", "WoredaName"), ("date", "mktname")])
    result, self.pairs = joiner.find_pair(left_df=phem_data, right_df=dd)