"""Unit tests for the apply_logic_response_criterion function."""

import copd
import numpy as np
import pandas as pd
import pytest


@pytest.fixture
def exacerbation_event():
    """Dataframe index (27) of the exacerbation event of interest."""
    return 27


@pytest.fixture
def first_pro_response():
    """Dataframe index (8) of the first weekly PRO response."""
    return 8


@pytest.fixture
def second_pro_response(first_pro_response):
    """Dataframe index of the second weekly PRO response. Seven days after first."""
    return first_pro_response + 7


@pytest.fixture
def input_df(exacerbation_event):
    """Sample input dataframe template - specific cases to be added in each test.

    This initial dataframe has no PRO responses between the initial exac and the event of
    interest with DaysSinceLastExac=25. Interim PRO responses should be added in tests.
    """
    df = pd.DataFrame({'PatientId': ['1'] * 31,
                       'DateOfEvent': pd.date_range('2022-01-01', '2022-01-31'),
                       'Q5Answered': [0] * 31,
                       'NegativeQ5': [np.nan] * 31,
                       'IsExac': [0] * 31,
                       'DaysSinceLastExac': [-1, -1, -1] + list(np.arange(1, 26)) +
                       list(np.arange(1, 4))})
    # Add initial event to simulate DaysSinceLastExac restart from 1
    df.loc[2, 'Q5Answered'] = 1
    df.loc[2, 'NegativeQ5'] = 0
    df.loc[2, 'IsExac'] = 1
    # Add event of interest (DaysSinceLastExac = 25)
    df.loc[exacerbation_event, 'Q5Answered'] = 1
    df.loc[exacerbation_event, 'NegativeQ5'] = 0
    df.loc[exacerbation_event, 'IsExac'] = 1

    # Add a negative response after the event of interest (should not be counted)
    df.loc[exacerbation_event + 2, 'Q5Answered'] = 1
    df.loc[exacerbation_event + 2, 'NegativeQ5'] = 1
    return df


def test_output_equals_expected_criterion_failed(input_df, exacerbation_event):
    """Test output is as expected for failed LOGIC response criterion."""
    # Output should be same as input with additional RemoveExac column
    expected_df = input_df.copy()
    # Insert RemoveExac column with exac flagged for removal
    expected_df['RemoveExac'] = np.nan
    expected_df.loc[exacerbation_event, 'RemoveExac'] = 1
    output_df = copd.apply_logic_response_criterion(input_df)
    pd.testing.assert_frame_equal(output_df, expected_df)


def test_output_equals_expected_criterion_passed(
        input_df, exacerbation_event, first_pro_response, second_pro_response):
    """Test output is as expected for passed LOGIC response criterion."""
    # Add PRO responses needed to pass LOGIC criterion
    input_df.loc[first_pro_response, 'Q5Answered'] = 1
    input_df.loc[first_pro_response, 'NegativeQ5'] = 1
    input_df.loc[second_pro_response, 'Q5Answered'] = 1
    input_df.loc[second_pro_response, 'NegativeQ5'] = 1

    # Output should be same as input with additional RemoveExac column
    expected_df = input_df.copy()
    # Insert RemoveExac column with exac flagged for removal
    expected_df['RemoveExac'] = np.nan
    expected_df.loc[exacerbation_event, 'RemoveExac'] = 0
    output_df = copd.apply_logic_response_criterion(input_df)
    pd.testing.assert_frame_equal(output_df, expected_df)