# Generated by CodiumAI
import unittest
from unittest.mock import patch

import numpy as np
import pandas as pd

from fiat_toolbox.metrics_writer.fiat_write_return_period_threshold import (
    ExceedanceProbabilityCalculator,
)


class TestExceedanceProbabilityCalculator(unittest.TestCase):
    # The class should be able to correctly calculate the exceedance probability for a given dataframe, threshold and time horizon.
    def test_calculate_exceedance_probability(self):
        # Arrange
        calculator = ExceedanceProbabilityCalculator("something")
        df = pd.DataFrame(
            {
                "something (2Y)": [0, 0.1, 0.2],
                "something (5Y)": [0, 0.2, 0.4],
                "something (10Y)": [0, 0.3, 0.6],
                "something (25Y)": [0.4, 0.6, 0.8],
                "something (50Y)": [0.9, 1.0, 1.1],
            }
        )
        threshold = 0.2
        T = 30

        # Act
        result = calculator.calculate(df, threshold, T)

        # Assert
        expected_result = pd.DataFrame({"Exceedance Probability": [82.0, 99.8, 100.0]})
        pd.testing.assert_frame_equal(result, expected_result)

    # The class should be able to correctly append the exceedance probability column to a given dataframe.
    def test_append_exceedance_probability(self):
        # Arrange
        calculator = ExceedanceProbabilityCalculator("something")
        df = pd.DataFrame(
            {
                "something (2Y)": [0, 0.1, 0.2],
                "something (5Y)": [0, 0.2, 0.4],
                "something (10Y)": [0, 0.3, 0.6],
                "something (25Y)": [0.4, 0.6, 0.8],
                "something (50Y)": [0.9, 1.0, 1.1],
            }
        )
        threshold = 0.2
        T = 30

        # Act
        result = calculator.append_probability(df, threshold, T)

        # Assert
        expected_result = pd.DataFrame(
            {
                "something (2Y)": [0, 0.1, 0.2],
                "something (5Y)": [0, 0.2, 0.4],
                "something (10Y)": [0, 0.3, 0.6],
                "something (25Y)": [0.4, 0.6, 0.8],
                "something (50Y)": [0.9, 1.0, 1.1],
                "Exceedance Probability": [82.0, 99.8, 100.0],
            }
        )
        pd.testing.assert_frame_equal(result, expected_result)

    # The class should be able to add the exceedance probability column to a given file.
    @patch("pandas.read_csv")
    @patch("pandas.DataFrame.to_csv")
    def test_append_exceedance_probability_to_file(self, mock_to_csv, mock_read_csv):
        # Arrange
        calculator = ExceedanceProbabilityCalculator("something")
        input_file = "tests/data/return_period_threshold_input.csv"
        output_file = "tests/data/return_period_threshold_output.csv"
        threshold = 0.2
        T = 30

        mock_read_csv.return_value = pd.DataFrame(
            {
                "something (2Y)": [0, 0.1, 0.2],
                "something (5Y)": [0, 0.2, 0.4],
                "something (10Y)": [0, 0.3, 0.6],
                "something (25Y)": [0.4, 0.6, 0.8],
                "something (50Y)": [0.9, 1.0, 1.1],
            }
        )

        # Act
        calculator.append_to_file(input_file, output_file, threshold, T)

        # Assert
        self.assertEqual(
            mock_to_csv.call_args[0][0], "tests/data/return_period_threshold_output.csv"
        )

    # The class should be able to handle a dataframe with all NaN values and return a dataframe with NaN values in the exceedance probability column.
    def test_handle_all_nan_values(self):
        # Arrange
        calculator = ExceedanceProbabilityCalculator("something")
        df = pd.DataFrame(
            {
                "something (2Y)": [np.nan, np.nan, np.nan],
                "something (5Y)": [np.nan, np.nan, np.nan],
                "something (10Y)": [np.nan, np.nan, np.nan],
                "something (25Y)": [np.nan, np.nan, np.nan],
                "something (50Y)": [np.nan, np.nan, np.nan],
            }
        )
        threshold = 0.2
        T = 30

        # Act
        result = calculator.calculate(df, threshold, T)

        # Assert
        expected_result = pd.DataFrame(
            {"Exceedance Probability": [np.nan, np.nan, np.nan]}
        )
        pd.testing.assert_frame_equal(result, expected_result)

    # The class should be able to handle a dataframe with multiple columns and some columns not starting with the specified column prefix.
    def test_handle_dataframe_with_multiple_columns(self):
        # Arrange
        calculator = ExceedanceProbabilityCalculator("something")

        # Create a dataframe with multiple columns and some columns not starting with the specified column prefix
        df = pd.DataFrame(
            {
                "something (2Y)": [0, 0.1, 0.2],
                "something (5Y)": [0, 0.2, 0.4],
                "something (10Y)": [0, 0.3, 0.6],
                "something (25Y)": [0.4, 0.6, 0.8],
                "something (50Y)": [0.9, 1.0, 1.1],
                "other_column": [
                    1,
                    2,
                    3,
                ],  # Column not starting with the specified column prefix
            }
        )

        # Expected result for the test data
        expected = pd.DataFrame(
            {
                "Exceedance Probability": [82.0, 99.8, 100.0],
            }
        )

        # Act
        result = calculator.calculate(df, threshold=0.2, T=30)

        # Assert
        pd.testing.assert_frame_equal(result, expected)

    # If the values of the row are all below the threshold, the exceedance probability should be NaN.
    # The other way around, if all values of the row are above the threshold, the exceedance probability
    # should use the smallest return period.
    def test_threshold_out_of_bounds(self):
        # Arrange
        df = pd.DataFrame(
            {
                "something (2Y)": [0, 0.1, 1.3],
                "something (5Y)": [0, 0.2, 1.5],
                "something (10Y)": [0, 0.5, 1.8],
                "something (25Y)": [0.4, 0.8, 2.0],
                "something (50Y)": [0.9, 1.5, 2.5],
            }
        )
        calculator = ExceedanceProbabilityCalculator("something")

        # Expected result for the test data
        expected = pd.DataFrame(
            {
                "Exceedance Probability": [np.nan, 14.4, 91.8],
            }
        )

        # Act
        result = calculator.calculate(df, threshold=1, T=5)

        # Assert
        pd.testing.assert_frame_equal(result, expected)

    # It shouldn't matter if the return period is out of bounds. The data is not representative as it will contain a lot of nan's, but that is the responsibility of the user.
    def test_return_period_out_of_bounds(self):
        # Arrange
        df = pd.DataFrame(
            {
                "something (2Y)": [0, 0, 1.3],
                "something (5Y)": [0, 0, 1.5],
                "something (10Y)": [0, 0, 1.8],
                "something (25Y)": [0.05, 0.1, 2.0],
                "something (50Y)": [0.1, 0.3, 2.5],
            }
        )
        calculator = ExceedanceProbabilityCalculator("something")

        # Expected result for the test data
        expected = pd.DataFrame(
            {
                "Exceedance Probability": [np.nan, 93.1, 100],
            }
        )

        # Act
        result = calculator.calculate(df, threshold=0.2, T=100)

        # Assert
        pd.testing.assert_frame_equal(result, expected)

    # The class should be able to handle a dataframe with non-numeric values and return NaN values in the exceedance probability column for the corresponding rows.
    def test_non_numeric_values(self):
        # Arrange
        df = pd.DataFrame(
            {
                "something (2Y)": [0, 0.1, 1.3],
                "something (5Y)": [0, 0.2, "hoi"],
                "something (10Y)": [0, "a", 1.8],
                "something (25Y)": [0.4, 0.8, 2.0],
                "something (50Y)": [0.9, 1.5, 2.5],
            }
        )
        threshold = 0.2
        T = 30
        calculator = ExceedanceProbabilityCalculator("something")

        # Act
        result = calculator.calculate(df, threshold, T)

        # Assert
        expected = pd.DataFrame({"Exceedance Probability": [82.0, np.nan, np.nan]})
        pd.testing.assert_frame_equal(result, expected)

    # The class should be able to handle a dataframe with negative values and return NaN values in the exceedance probability column for the corresponding rows.
    def test_negative_values(self):
        # Arrange
        calculator = ExceedanceProbabilityCalculator("something")
        df = pd.DataFrame(
            {
                "something (2Y)": [-0.5, -0.3, -0.1],
                "something (5Y)": [-0.4, -0.2, 0.1],
                "something (10Y)": [-0.3, 0, 0.2],
                "something (25Y)": [-0.1, 0.2, 0.4],
                "something (50Y)": [0, 0.4, 0.6],
            }
        )

        # Act
        result = calculator.calculate(df, threshold=0.2, T=30)

        # Assert
        expected = pd.DataFrame({"Exceedance Probability": [np.nan, 69.9, 95.0]})
        pd.testing.assert_frame_equal(result, expected)
