"""
Build NSQL generation prompt.
Two main parts:
1) PromptBuilder makes prompt for calling codex to generate NSQL(Binder-SQL).
2) OpenAIQAPromptBuilder makes prompt for calling codex to generate QA answers.
"""

import random
from typing import Dict, Tuple
import pandas as pd
import copy

from utils.errors import DuplicateColumnsError
from utils.mmqa.image_stuff import get_caption_map
from retrieval.retrieve_pool import QAItem

from utils.normalizer import prepare_df_for_neuraldb_from_table


def _create_table_prompt(df: pd.DataFrame, title: str):
    """
    Return the CREATE TABLE clause as prompt.
    """
    string = "CREATE TABLE {}(\n".format(title)
    for header in df.columns:
        column_type = 'text'
        try:
            if df[header].dtype == 'int64':
                column_type = 'int'
            elif df[header].dtype == 'float64':
                column_type = 'real'
            elif df[header].dtype == 'datetime64':
                column_type = 'datetime'
        except AttributeError as e:
            raise DuplicateColumnsError(e)

        string += '\t{} {},\n'.format(header, column_type)
    string = string.rstrip(',\n') + ')\n'
    return string


class PromptBuilder(object):
    def __init__(self, args):
        self.args = args
        self.prompt_style = args.prompt_style
        random.seed(args.seed)

    def _select_x_prompt(self, df: pd.DataFrame, num_rows: int,
                         few_shot_demonstration=True):
        """
        Return the first X rows table contents as prompt.
        """
        if self.prompt_style == 'create_table_select_full_table':
            string = '/*\nAll rows of the table:\nSELECT * FROM w;\n'
        elif self.prompt_style == 'create_table_select_3':
            string = '/*\n{} example rows:\nSELECT * FROM w LIMIT {};\n'.format(num_rows, num_rows)
        elif self.prompt_style == 'create_table_select_3_hidden':
            string = '/*\n{} example rows:\n'.format(num_rows)
        elif few_shot_demonstration is True and self.prompt_style in \
                ["create_table_select_3_full_table",
                 "create_table_select_3_full_table_w_gold_passage_image",
                 "create_table_select_3_full_table_w_all_passage_image"]:
            string = '/*\n{} example rows:\nSELECT * FROM w LIMIT {};\n'.format(num_rows, num_rows)
        elif few_shot_demonstration is False and self.prompt_style in \
                ["create_table_select_3_full_table",
                 "create_table_select_3_full_table_w_gold_passage_image",
                 "create_table_select_3_full_table_w_all_passage_image"]:
            string = '/*\nAll rows of the table:\nSELECT * FROM w;\n'
        else:
            raise ValueError(f"Select x prompt style {self.prompt_style} is not supported.")

        for column_id, header in enumerate(df.columns):
            string += str(header)
            if column_id != len(df.columns) - 1:
                string += '\t'
        string += '\n'

        for row_id, row in df.iloc[:num_rows].iterrows():
            for column_id, header in enumerate(df.columns):
                string += str(row[header])
                if column_id != len(df.columns) - 1:
                    string += '\t'
            string += '\n'
        string += '*/\n'

        return string

    def _passage_prompt(self, passages, only_title, db_style_prompt=True):
        """
        Return the passage prompt.
        """
        if not db_style_prompt:
            string = "Passages: "
            for passage in passages:
                if only_title:
                    string += passage['title'] + ';; '
                else:
                    string += passage['title'] + f" ({passage['text']})" + ';; '
            string = string.rstrip(';; ')
            string += '\n'
            return string
        else:
            if len(passages) == 0:
                return ""
            passage_table_prompt = ""
            _header = []
            _rows = [[]]
            for passage in passages:
                _header.append(passage['title'])
                _rows[0].append(passage['text'])
            passage_table = prepare_df_for_neuraldb_from_table({"header": _header, "rows": _rows})
            passage_table_prompt += _create_table_prompt(passage_table, "Passages")
            if not only_title:
                passage_table_prompt += self._select_x_prompt(
                    df=passage_table,
                    num_rows=passage_table.shape[0]
                )
            return passage_table_prompt

    def _image_prompt(self, images, only_title, db_style_prompt=True):
        """
        Return the image prompt.
        """
        if not db_style_prompt:
            string = "Images: "
            for image in images:
                if only_title:
                    string += image['title'] + ';;'
                else:
                    string += image['title'] + f" ({image['caption']})" + ';; '
            string = string.rstrip(';; ')
            string += '\n'
            return string
        else:
            if len(images) == 0:
                return ""
            image_table_prompt = ""
            _header = []
            _rows = [[]]
            for image in images:
                _header.append(image['title'])
                _rows[0].append(image['caption'])
            image_table = prepare_df_for_neuraldb_from_table({"header": _header, "rows": _rows})
            image_table_prompt += _create_table_prompt(image_table, "Images")
            if not only_title:
                image_table_prompt += self._select_x_prompt(
                    df=image_table,
                    num_rows=image_table.shape[0]
                )
            return image_table_prompt

    def _pick_target_columns(self, df, strategy):
        """
        Pick the controllable target columns for generation.
        """
        if strategy == 'random':
            return random.choice(list(df.columns) + ['*'])
        elif strategy == 'traverse':
            raise NotImplementedError
        else:
            return ValueError

    def _pick_operators(self, df, strategy):
        """
        Pick the controllable operators for generation.
        """
        candidate_operators = ['none', 'count', 'max', 'min', 'sum']
        if strategy == 'random':
            return random.choice(candidate_operators)
        elif strategy == 'traverse':
            raise NotImplementedError
        else:
            return ValueError

    def _pick_nested_levels(self, df, strategy):
        """
        Pick the controllable(maybe) nested levels for generation.
        """
        if strategy == 'fixed':
            return 2
        elif strategy == 'random':
            raise NotImplementedError
        elif strategy == 'traverse':
            raise NotImplementedError
        else:
            raise ValueError

    def build_one_shot_prompt(
            self,
            prompt_type: Tuple,
            table: pd.DataFrame,
            question: str,
            answer_text: str,
            nsql: str,
            passages: Dict = None,
            images: Dict = None,
            title: str = None,
            only_title: bool = False,
            **kwargs
    ):
        """
        Build one-shot prompt with table-question-nsql.
        """
        one_shot_prompt = ""
        if self.prompt_style == 'create_table_select_full_table':
            one_shot_prompt += _create_table_prompt(table, title)
            one_shot_prompt += self._select_x_prompt(
                df=table,
                num_rows=table.shape[0]
            )
        elif self.prompt_style in ['create_table_select_3_full_table', 'create_table_select_3']:
            one_shot_prompt += _create_table_prompt(table, title)
            one_shot_prompt += self._select_x_prompt(
                df=table,
                num_rows=3,
            )
        elif self.prompt_style == 'create_table':
            one_shot_prompt += _create_table_prompt(table, title)
        elif self.prompt_style == 'no_table':
            # No table input, to test Codex QA with only internal knowledge
            pass
        elif self.prompt_style in ['create_table_select_3_full_table_w_all_passage_image']:
            assert passages is not None and images is not None
            one_shot_prompt += _create_table_prompt(table, title)
            one_shot_prompt += self._select_x_prompt(
                df=table,
                num_rows=3,
            )
            all_passages, all_images = [], []
            caption_map = get_caption_map()

            for passage_idx in range(len(passages['id'])):
                all_passages.append({
                    'id': passages['id'][passage_idx],
                    'title': passages['title'][passage_idx],
                    'url': passages['url'][passage_idx],
                    'text': passages['text'][passage_idx]
                })

            for image_idx in range(len(images['id'])):
                all_images.append({
                    "id": images['id'][image_idx],
                    "title": images['title'][image_idx],
                    "url": images['url'][image_idx],
                    "path": images['path'][image_idx],
                    "pic": images['pic'][image_idx],
                    "caption": caption_map[images['id'][image_idx]]
                })

            one_shot_prompt += self._passage_prompt(
                passages=all_passages,
                only_title=only_title
            )
            one_shot_prompt += self._image_prompt(
                images=all_images,
                only_title=only_title
            )
        else:
            raise ValueError('{} is not supported.'.format(self.prompt_style))

        # question and nsql pairs
        if prompt_type == ('question', 'nsql'):
            one_shot_prompt += 'Q: {}\n'.format(question)
            one_shot_prompt += 'NeuralSQL: {}\n'.format(nsql)
        elif prompt_type == ('question', 'sql'):
            one_shot_prompt += 'Q: {}\n'.format(question)
            one_shot_prompt += 'SQL: {}\n'.format(nsql)
        elif prompt_type == ('question', 'answer'):
            one_shot_prompt += 'Q: {}\n'.format(question)
            one_shot_prompt += 'A: {}\n'.format(', '.join(answer_text))
        else:
            raise ValueError(f'Prompt type {prompt_type} is not supported.')

        return one_shot_prompt

    def build_generate_prompt(
            self,
            generate_type: Tuple,
            table: pd.DataFrame,
            question: str = None,
            passages: Dict = None,
            images: Dict = None,
            title: str = None,
            only_title: bool = False,
            supporting_context: Dict = None,
            **kwargs
    ):
        """
        Build the prompt of the generation sample.
        """
        generate_prompt = ""

        # task instruction
        if generate_type == ('answer',):
            generate_prompt += """\n-- Answer the question based on the given table below.\n\n"""
        elif generate_type == ('nsql',):
            generate_prompt += """\n-- Parse the question into NeuralSQL based on the given table below.\n\n"""
        elif generate_type == ('sql',):
            generate_prompt += """\n-- Parse the question into SQL based on the given table below.\n\n"""
        elif generate_type == ('npython',):
            generate_prompt += """\n-- Parse the question into NeuralPython based on the given table below.\n\n"""
        elif generate_type == ('python',):
            generate_prompt += """\n-- Parse the question into Python based on the given table below.\n\n"""
        else:
            generate_prompt += """\n-- Generate NeuralSQL and question pairs based on the given table below.\n\n"""

        # table prompt
        if self.prompt_style in ['create_table_select_full_table', 'create_table_select_3_full_table']:
            generate_prompt += _create_table_prompt(table, title)
            generate_prompt += self._select_x_prompt(
                df=table,
                num_rows=table.shape[0],
                few_shot_demonstration=False
            )
        elif self.prompt_style in ['create_table_select_3']:
            generate_prompt += _create_table_prompt(table, title)
            generate_prompt += self._select_x_prompt(
                df=table,
                num_rows=3,
                few_shot_demonstration=False
            )
        elif self.prompt_style == 'create_table':
            generate_prompt += _create_table_prompt(table, title)
        elif self.prompt_style == 'no_table':
            # No table input, to test Codex QA with only internal knowledge
            pass
        elif self.prompt_style in ['create_table_select_3_full_table_w_all_passage_image']:
            assert passages is not None and images is not None
            generate_prompt += _create_table_prompt(table, title)
            generate_prompt += self._select_x_prompt(
                df=table,
                num_rows=table.shape[0],
                few_shot_demonstration=False
            )
            all_passages, all_images = [], []
            caption_map = get_caption_map()

            for passage_idx in range(len(passages['id'])):
                all_passages.append({
                    'id': passages['id'][passage_idx],
                    'title': passages['title'][passage_idx],
                    'url': passages['url'][passage_idx],
                    'text': passages['text'][passage_idx]
                })

            for image_idx in range(len(images['id'])):
                all_images.append({
                    "id": images['id'][image_idx],
                    "title": images['title'][image_idx],
                    "url": images['url'][image_idx],
                    "path": images['path'][image_idx],
                    "pic": images['pic'][image_idx],
                    "caption": caption_map[images['id'][image_idx]]
                })

            generate_prompt += self._passage_prompt(
                passages=all_passages,
                only_title=only_title
            )
            generate_prompt += self._image_prompt(
                images=all_images,
                only_title=only_title
            )
        elif self.prompt_style in ['create_table_select_3_full_table_w_gold_passage_image']:
            assert passages is not None and images is not None
            generate_prompt += _create_table_prompt(table, title)
            generate_prompt += self._select_x_prompt(
                df=table,
                num_rows=table.shape[0],
                few_shot_demonstration=False
            )
            gold_passages, gold_images = [], []
            caption_map = get_caption_map()
            for doc_id, doc_part in zip(supporting_context['doc_id'], supporting_context['doc_part']):
                if doc_part == 'text':
                    passage_idx = passages['id'].index(doc_id)
                    gold_passages.append({
                        'id': passages['id'][passage_idx],
                        'title': passages['title'][passage_idx],
                        'url': passages['url'][passage_idx],
                        'text': passages['text'][passage_idx]
                    })
                elif doc_part == 'image':
                    image_idx = images['id'].index(doc_id)
                    gold_images.append({
                        "id": images['id'][image_idx],
                        "title": images['title'][image_idx],
                        "url": images['url'][image_idx],
                        "path": images['path'][image_idx],
                        "pic": images['pic'][image_idx],
                        "caption": caption_map[doc_id]
                    })
            generate_prompt += self._passage_prompt(
                passages=gold_passages,
                only_title=only_title
            )
            generate_prompt += self._image_prompt(
                images=gold_images,
                only_title=only_title
            )
        else:
            raise ValueError('{} is not supported.'.format(self.prompt_style))

        # determine the target to generate
        if generate_type == ('answer',):
            generate_prompt += 'Q: {}\n'.format(question)
            generate_prompt += 'A: '
        elif generate_type == ('nsql',):
            generate_prompt += 'Q: {}\n'.format(question)
            generate_prompt += 'NeuralSQL: '
        elif generate_type == ('sql',):
            generate_prompt += 'Q: {}\n'.format(question)
            generate_prompt += 'SQL: '
        elif generate_type == ('npython',):
            generate_prompt += 'Q: {}\n'.format(question)
            generate_prompt += 'NeuralPython: '
        elif generate_type == ('python',):
            generate_prompt += 'Q: {}\n'.format(question)
            generate_prompt += 'Python: '
        else:
            raise ValueError(f'Generate type {generate_type} is not supported.')

        return generate_prompt


class OpenAIQAPromptBuilder(object):
    @staticmethod
    def table2codex_prompt(table, table_title=None, drop_row_id=True, ):
        _table = copy.deepcopy(table)
        header = _table['header']
        rows = _table['rows']
        if drop_row_id:
            if header[0] == "row_id":
                header = header[1:]
                rows = [_row[1:] for _row in rows]
        prompt_str = 'Table: {}\n'.format(table_title) if table_title else ''
        prompt_str += "/*\n"
        prompt_str += "\t".join(header) + "\n"
        prompt_str += '\n'.join(["\t".join([str(cell) for cell in row]) for row in rows]) + "\n"
        prompt_str += "*/"
        return prompt_str

    @staticmethod
    def build_one_shot_prompt(
            item: QAItem,
            answer_split_token: str = ';',
            verbose: bool = False,
            prompting_method='new_db',
            db_mapping_token="😅"
    ) -> str:
        """
        Build one-shot QA prompt.
        """
        assert prompting_method in ['basic', 'new_db']
        qa_type, qa_question = item.qa_question.split('@')
        prompt = ''
        db_prompt = OpenAIQAPromptBuilder.table2codex_prompt(item.table, item.title)
        prompt += "Give a database as shown below:\n{}\n\n".format(db_prompt)

        if prompting_method == 'basic':
            if qa_type == "map":
                prompt += "Q: Answer question \"{}\" row by row.".format(qa_question)
                assert answer_split_token is not None
                prompt += " The answer should be a list split by '{}' and have {} items in total.".format(
                    answer_split_token, len(item.table['rows']))
                prompt += "\nA: {}\n\n".format(f'{answer_split_token}'.join(item.qa_answer))
            elif qa_type == "ans":
                prompt += "Q: Answer question \"{}\" for the table.".format(qa_question)
                prompt += " "
                prompt += "\nA: {}\n\n".format(f'{answer_split_token}'.join(item.qa_answer))
            else:
                raise ValueError("The QA type is not supported!")

            return prompt

        elif prompting_method == "new_db":
            if qa_type == "map":
                prompt += "Q: Answer question \"{}\" row by row.".format(qa_question)
                assert answer_split_token is not None
                db_prompt_lines = db_prompt.split("\n")[2:-1]  # skip Title, /*, and */
                db_prompt_lines_with_answer = []
                db_prompt_lines_with_answer.append("/*")
                db_prompt_lines_with_answer.append(db_prompt_lines[0])
                assert len(db_prompt_lines[1:]) == len(
                    item.qa_answer), "answer items and table rows must be in the same number, check annotations"
                for db_prompt_line, qa_answer_item in zip(db_prompt_lines[1:], item.qa_answer):
                    db_prompt_lines_with_answer.append(
                        "{}{}{}".format(db_prompt_line, db_mapping_token, qa_answer_item))
                db_prompt_lines_with_answer.append("*/")
                prompt += "\n{}\n".format("\n".join(db_prompt_lines_with_answer))

            elif qa_type == "ans":
                prompt += "Q: Answer question \"{}\" for the table.".format(qa_question)
                prompt += " "
                prompt += "\nA: {}\n".format(f'{answer_split_token}'.join(item.qa_answer))
            else:
                raise ValueError("The QA type is not supported!")

            return prompt