import pandas as pd

from src.tools.retriever import Retriever
from src.tools.llm import LlmAgent
from src.model.block import Block


class Controller:

    def __init__(self, retriever: Retriever, llm: LlmAgent, plan_language: str, content_language: str, specials: {}):
        self.plan_language = plan_language
        self.content_language = content_language
        self.retriever = retriever
        self.specials = specials
        self.llm = llm

    def get_response(self, query_fr: str, histo_fr: [(str, str)]) -> (str, [Block]):
        histo_conversation, histo_queries = self._get_histo(histo_fr)
        queries = self.llm.translate(text=histo_queries) if self.plan_language == 'en' else histo_queries
        block_sources = self.retriever.similarity_search(query=queries)
        block_sources = self._select_best_sources(block_sources)
        for block in block_sources:
            self._expand_block_with_specials(block,  histo_queries)
        sources_contents = [s.content for s in block_sources]
        context = '\n'.join(sources_contents)
        answer = self.llm.generate_paragraph(query=queries, histo=histo_conversation, context=context,
                                             language=self.content_language)
        sources_contents_fr = [s.content_fr for s in block_sources[:2]]
        context_fr = '\n'.join(sources_contents_fr)
        if self.content_language == 'en':
            answer = self.llm.generate_answer(answer_en=answer, query=query_fr,
                                              histo_fr=histo_conversation, context_fr=context_fr)
        answer = self._clean_answer(answer)
        return answer, block_sources

    @staticmethod
    def _get_histo(histo: [(str, str)]) -> (str, str):
        histo_conversation = ""
        histo_queries = ""

        for (query, answer) in histo[-5:]:
            histo_conversation += f'user: {query} \n bot: {answer}\n'
            histo_queries += query + '\n'
        return histo_conversation[:-1], histo_queries

    @staticmethod
    def _clean_answer(answer: str) -> str:
        answer = answer.strip('bot:')
        while answer and answer[-1] in {"'", '"', " ", "`"}:
            answer = answer[:-1]
        while answer and answer[0] in {"'", '"', " ", "`"}:
            answer = answer[1:]
        answer = answer.strip('bot:')
        if answer:
            if answer[-1] != ".":
                answer += "."
        return answer

    @staticmethod
    def _select_best_sources(sources: [Block], delta_1_2=0.15, delta_1_n=0.3, absolute=1.2, alpha=0.9) -> [Block]:
        """
        Select the best sources: not far from the very best, not far from the last selected, and not too bad per se
        """
        best_sources = []
        for idx, s in enumerate(sources):
            if idx == 0 \
                    or (s.distance - sources[idx - 1].distance < delta_1_2
                        and s.distance - sources[0].distance < delta_1_n) \
                    or s.distance < absolute:
                best_sources.append(s)
                delta_1_2 *= alpha
                delta_1_n *= alpha
                absolute *= alpha
            else:
                break
        return best_sources

    def _expand_block_with_specials(self, block: Block, query: str) -> Block:
        """
        Performs special treatments for blocks expanding the text in the block
        For example, it may add specific content extracted from a table based on elements of the query
        """

        def any_in(l1: [], l2: []) -> bool:
            """
            checks if any of el in l1 belongs to l2
            """
            return 0 < len([el for el in l1 if el in l2])

        def get_countries_names(df: pd.DataFrame) -> [str]:
            """
            extends the ortograph of countries: ex. Etats-Unis = USA = Etats Unis, etc.
            """
            countries_fr = list(df['pays'])
            countries_en = list(df['country'])
            countries_names = {c_fr: [c_fr, c_en] for c_fr, c_en in zip(countries_fr, countries_en)}
            countries_extensions = self.specials['countries_extensions']
            for c in set(countries_extensions.keys()).intersection(set(countries_names.keys())):
                countries_names[c] += countries_extensions[c]
            return countries_names

        def remote_rate_fn(ctrl: Controller, block: Block, query: str) -> Block:
            remote_rate_df = self.specials['remote_rate_df']
            remote_rate_known = self.specials['remote_rate_known']
            remote_rate_unknown = self.specials['remote_rate_unknown']
            countries_fr = list(remote_rate_df['pays'])
            countries_names = get_countries_names(remote_rate_df)
            countries_of_interest = [c for c in countries_fr if any_in(countries_names[c], query)]
            for c in countries_of_interest:
                rate = remote_rate_df[remote_rate_df['pays'] == c]['rate'].values[0]
                block.content += remote_rate_known + c + " is " + rate + '\n'
            if len(countries_of_interest) == 0:
                block.content += remote_rate_unknown
            return block

        def accommodation_meal_fn(ctrl: Controller, block: Block, query: str) -> Block:
            accommodation_meal_df = self.specials['accommodation_meal_df']
            accommodation_meal_known = self.specials['accommodation_meal_known']
            accommodation_meal_unknown = self.specials['accommodation_meal_unknown']
            countries_fr = list(accommodation_meal_df['pays'])
            countries_names = get_countries_names(df=accommodation_meal_df)
            countries_of_interest = [c for c in countries_fr if any_in(countries_names[c], query)]
            for c in countries_of_interest:
                rate = accommodation_meal_df[accommodation_meal_df['pays'] == c][['meal', 'accommodation']].values
                block.content += accommodation_meal_known + c + " is " + rate[0][0] + ' for meals and ' \
                                 + rate[0][1] + ' for accommodation\n'
            if len(countries_of_interest) == 0:
                block.content += accommodation_meal_unknown
            return block

        def expand_block(special: str, ctrl: Controller, block: Block, query: str) -> Block:
            routing_table = {'RemotenessRateTable': remote_rate_fn,
                             'AccommodationMealTable': accommodation_meal_fn, }
            if special in routing_table.keys():
                fn = routing_table[special]
                block = fn(ctrl, block, query)
            return block

        for special in block.specials:
            block = expand_block(special, self, block, query)
        return block