# Copyright 2020 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
#
#     https://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.
""" Tests for results functions. """

import pytest
import statistics
from bigbench.api import task
from bigbench.api import results


def test_add_aggregate_scores_no_subtasks():
    """Tests that a task with no subtasks adds normalized_aggregate_score."""

    task_name = "test_task"
    scores = [
        task.ScoreData(
            score_dict={"multiple_choice_grade": i / 10},
            preferred_score="multiple_choice_grade",
            number_of_shots=i,
            low_score=i / 20,
            high_score=1.0,
            subtask_description=task_name,
        )
        for i in range(3)
    ]

    new_scores = results.add_aggregate_scores(task_name=task_name, scores=scores)

    norm_score = lambda x: (x / 10 - x / 20) / (1 - x / 20) * 100

    expected_scores = [
        task.ScoreData(
            score_dict={
                "multiple_choice_grade": i / 10,
                "normalized_aggregate_score": norm_score(i),
            },
            preferred_score="multiple_choice_grade",
            number_of_shots=i,
            low_score=i / 20,
            high_score=1.0,
            subtask_description=task_name,
        )
        for i in range(3)
    ]

    assert new_scores == expected_scores


@pytest.mark.parametrize(
    "preferred_metric", ("multiple_choice_grade", "random_metric", "subtask")
)
def test_add_aggregate_scores(preferred_metric):
    """Tests aggregate score computation."""
    task_name = "test_task"
    num_shots = 3
    num_subtasks = 2

    def multiple_choice_grade(shots: int, subtask: int) -> float:
        return (shots + 1) / (subtask + 10)

    def random_metric(shots: int, subtask: int) -> float:
        return (shots + 1) / (subtask + 8)

    def subtask_metric(shots: int, subtask: int) -> float:
        return (shots + 1) / (subtask + 6)

    def pref_metric(shots: int, subtask: int) -> float:
        if preferred_metric == "multiple_choice_grade":
            return multiple_choice_grade(shots, subtask)
        elif preferred_metric == "random_metric":
            return random_metric(shots, subtask)
        elif preferred_metric == "subtask":
            return subtask_metric(shots, subtask)
        else:
            assert False, f"invalid metric {preferred_metric}"

    def low_score(shots: int, subtask: int) -> float:
        return pref_metric(shots, subtask) / 2

    def high_score(shots: int, subtask: int) -> float:
        if preferred_metric == "multiple_choice_grade":
            return 1.0
        else:
            return pref_metric(shots, subtask) * 2

    def normalized_pref_metric(shots: int, subtask: int) -> float:
        return (pref_metric(shots, subtask) - low_score(shots, subtask)) / (
            high_score(shots, subtask) - low_score(shots, subtask)
        )

    scores = []
    for subtask_number in range(num_subtasks):
        subtask_scores = [
            task.ScoreData(
                score_dict={
                    "multiple_choice_grade": multiple_choice_grade(
                        shot, subtask_number
                    ),
                    "random_metric": random_metric(shot, subtask_number),
                    f"metric_{subtask_number}": subtask_metric(shot, subtask_number),
                },
                preferred_score=f"metric_{subtask_number}"
                if preferred_metric == "subtask"
                else preferred_metric,
                number_of_shots=shot,
                low_score=low_score(shot, subtask_number),
                high_score=high_score(shot, subtask_number),
                subtask_description=f"subtask{subtask_number}",
            )
            for shot in range(num_shots)
        ]

        subtask_scores[0].score_dict["zero_shot_metric"] = subtask_number

        scores += subtask_scores

    new_scores = results.add_aggregate_scores(task_name=task_name, scores=scores)

    if preferred_metric == "multiple_choice_grade":
        preferred_score = "multiple_choice_grade"
        low_fn = lambda shot_num: statistics.mean(
            low_score(shot_num, s) for s in range(num_subtasks)
        )
        high = 1.0
    else:
        preferred_score = "normalized_aggregate_score"
        low_fn = lambda shot_num: 0
        high = 100

    expected_aggregate_scores = [
        task.ScoreData(
            score_dict=dict(
                [
                    (
                        "multiple_choice_grade",
                        statistics.mean(
                            multiple_choice_grade(shot, t) for t in range(num_subtasks)
                        ),
                    ),
                    (
                        "random_metric",
                        statistics.mean(
                            random_metric(shot, t) for t in range(num_subtasks)
                        ),
                    ),
                    (
                        "normalized_aggregate_score",
                        100
                        * statistics.mean(
                            normalized_pref_metric(shot, t) for t in range(num_subtasks)
                        ),
                    ),
                ]
                + [
                    (f"metric_{subtask_number}", subtask_metric(shot, subtask_number))
                    for subtask_number in range(num_subtasks)
                ]
            ),
            preferred_score=preferred_score,
            number_of_shots=shot,
            low_score=low_fn(shot),
            high_score=high,
            subtask_description=task_name,
        )
        for shot in range(num_shots)
    ]

    expected_aggregate_scores[0].score_dict["zero_shot_metric"] = statistics.mean(
        range(num_subtasks)
    )

    expected_scores = scores + expected_aggregate_scores

    assert new_scores == expected_scores

    # make sure that operation is idempotent
    assert expected_scores == results.add_aggregate_scores(
        task_name=task_name, scores=new_scores
    )