# -*- coding: utf-8 -*-
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Unittest for llmfn_outputs."""
from __future__ import annotations

from typing import Sequence

from absl.testing import absltest
from google.generativeai.notebook.lib import llmfn_output_row
from google.generativeai.notebook.lib import llmfn_outputs
from google.generativeai.notebook.lib import model as model_lib

LLMFnOutputEntry = llmfn_outputs.LLMFnOutputEntry
LLMFnOutputs = llmfn_outputs.LLMFnOutputs
LLMFnOutputRow = llmfn_output_row.LLMFnOutputRow


def _get_empty_model_results() -> model_lib.ModelResults:
    return model_lib.ModelResults(model_input="", text_results=[])


class LLMFnOutputsTest(absltest.TestCase):
    def _test_is_sequence(self, outputs: Sequence[LLMFnOutputEntry]):
        # Make sure `outputs` is iterable.
        count = 0
        for _ in outputs:
            count = count + 1

        # Make sure len(outputs) works.
        self.assertLen(outputs, count)

        # Make sure regular integer indices are accepted.
        self.assertIsInstance(outputs[0], LLMFnOutputEntry)

        # Make sure slices are accepted.
        self.assertLen(outputs[:-1], count - 1)
        self.assertIsInstance(outputs[:-1][0], LLMFnOutputEntry)

    def test_is_sequence(self):
        tmp_entry_list = []
        for prompt_num in range(0, 3):
            prompt = ["one", "two", "three"][prompt_num]
            output_entry = LLMFnOutputEntry(
                prompt_num=prompt_num,
                input_num=0,
                prompt=prompt,
                prompt_vars={},
                model_input=prompt,
                model_results=_get_empty_model_results(),
                output_rows=[],
            )
            tmp_entry_list.append(output_entry)

        outputs = LLMFnOutputs(tmp_entry_list)
        self._test_is_sequence(outputs)

    def test_as_dict_basic(self):
        outputs_list = []
        for prompt_num in range(0, 3):
            prompt = ["one", "two", "three"][prompt_num]
            text_results = ["red", "green", "blue"][prompt_num]
            output_entry = LLMFnOutputEntry(
                prompt_num=prompt_num,
                input_num=0,
                prompt=prompt,
                prompt_vars={},
                model_input=prompt,
                model_results=_get_empty_model_results(),
                output_rows=[
                    LLMFnOutputRow(
                        data={
                            "Result Num": 0,
                            "text_results": "{}_one".format(text_results),
                        },
                        result_type=str,
                    ),
                    LLMFnOutputRow(
                        data={
                            "Result Num": 1,
                            "text_results": "{}_two".format(text_results),
                        },
                        result_type=str,
                    ),
                ],
            )
            outputs_list.append(output_entry)

        expected_dict = {
            "Prompt Num": [0, 0, 1, 1, 2, 2],
            "Input Num": [0, 0, 0, 0, 0, 0],
            "Result Num": [0, 1, 0, 1, 0, 1],
            "Prompt": ["one", "one", "two", "two", "three", "three"],
            "text_results": [
                "red_one",
                "red_two",
                "green_one",
                "green_two",
                "blue_one",
                "blue_two",
            ],
        }

        outputs = LLMFnOutputs(outputs_list)
        self.assertEqual(expected_dict, outputs.as_dict())
        # Keys must be in the same order as well.
        self.assertEqual(list(expected_dict.keys()), list(outputs.as_dict().keys()))

    def test_as_dict_with_holes(self):
        outputs_list = []
        for prompt_num in range(0, 3):
            prompt = ["one", "two", "three"][prompt_num]
            text_results = ["red", "green", "blue"][prompt_num]
            output_entry = LLMFnOutputEntry(
                prompt_num=prompt_num,
                input_num=0,
                prompt=prompt,
                prompt_vars={},
                model_input=prompt,
                model_results=_get_empty_model_results(),
                output_rows=[
                    LLMFnOutputRow(
                        data={
                            "Result Num": 0,
                            text_results: True,
                            "text_results": "{}_one".format(text_results),
                        },
                        result_type=str,
                    ),
                    LLMFnOutputRow(
                        data={
                            "Result Num": 1,
                            text_results: True,
                            "text_results": "{}_two".format(text_results),
                        },
                        result_type=str,
                    ),
                ],
            )
            outputs_list.append(output_entry)

        expected_dict = {
            "Prompt Num": [0, 0, 1, 1, 2, 2],
            "Input Num": [0, 0, 0, 0, 0, 0],
            "Result Num": [0, 1, 0, 1, 0, 1],
            "Prompt": ["one", "one", "two", "two", "three", "three"],
            "red": [True, True, None, None, None, None],
            "green": [None, None, True, True, None, None],
            "blue": [None, None, None, None, True, True],
            "text_results": [
                "red_one",
                "red_two",
                "green_one",
                "green_two",
                "blue_one",
                "blue_two",
            ],
        }

        outputs = LLMFnOutputs(outputs_list)
        self.assertEqual(expected_dict, outputs.as_dict())
        # Keys must be in the same order as well.
        self.assertEqual(list(expected_dict.keys()), list(outputs.as_dict().keys()))


if __name__ == "__main__":
    absltest.main()