"""Unit tests for the logic_consecutive_negative_responses 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 third_pro_response(second_pro_response):
    """Dataframe index of the third weekly PRO response. Seven days after second."""
    return second_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 exacerbation at index
    2 and the event of interest with DaysSinceLastExac=25 at index exacerbation_event (set
    to 27). Interim PRO responses should be added in tests. Each row is a different day
    (in chronological order). Add/subtract N from exacerbation_event to refer to N days
    before or after the event by the dataframe index, e.g. exacerbation_event - 7 refers
    to the day a week prior.
    """
    df = pd.DataFrame({'PatientId': ['1'] * 31,
                       'DateOfEvent': pd.date_range('2022-01-01', '2022-01-31'),
                       'Q5Answered': [0] * 31,
                       'NegativeQ5': [np.nan] * 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
    # Add event of interest (DaysSinceLastExac = 25)
    df.loc[exacerbation_event, 'Q5Answered'] = 1
    df.loc[exacerbation_event, 'NegativeQ5'] = 0

    # Add a negative response 2 days 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_returns_one_when_no_responses(input_df, exacerbation_event):
    """Verify returns 1 (flag for removal) for no interim PRO responses."""
    assert copd.logic_consecutive_negative_responses(input_df, exacerbation_event) == 1


def test_returns_one_too_few_responses(input_df, exacerbation_event):
    """Verify returns 1 (flag for removal) for too few interim PRO responses."""
    # Add a single negative response 7 days before the exacerbation event. Should fail PRO
    # LOGIC because the negative response at index 29 is after the event of interest.
    input_df.loc[exacerbation_event - 7, 'Q5Answered'] = 1
    input_df.loc[exacerbation_event - 7, 'NegativeQ5'] = 1
    assert copd.logic_consecutive_negative_responses(input_df, exacerbation_event) == 1


def test_returns_one_too_few_negative_responses(
        input_df, exacerbation_event, second_pro_response, third_pro_response):
    """Verify returns 1 (flag for removal) for too few interim PRO responses."""
    # Add a positive response and a single negative response. Should return one because
    # the response at index 29 is after the period of interest.

    input_df.loc[second_pro_response, 'Q5Answered'] = 1
    input_df.loc[second_pro_response, 'NegativeQ5'] = 0

    input_df.loc[third_pro_response, 'Q5Answered'] = 1
    input_df.loc[third_pro_response, 'NegativeQ5'] = 1
    assert copd.logic_consecutive_negative_responses(input_df, exacerbation_event) == 1


def test_returns_one_too_few_consecutive_negative_responses_missing(
        input_df, exacerbation_event, first_pro_response, second_pro_response,
        third_pro_response):
    """Verify returns 1 (flag for removal) for too few consecutive -ve PRO responses.

    Input has a missing response between the two negative responses.
    """
    # Add negative responses at indices 8 and 22 (missing response at 15)
    input_df.loc[first_pro_response, 'Q5Answered'] = 1
    input_df.loc[first_pro_response, 'NegativeQ5'] = 1

    input_df.loc[third_pro_response, 'Q5Answered'] = 1
    input_df.loc[third_pro_response, 'NegativeQ5'] = 1

    assert copd.logic_consecutive_negative_responses(input_df, exacerbation_event) == 1


def test_returns_one_too_few_consecutive_negative_responses_positive(
        input_df, exacerbation_event, first_pro_response, second_pro_response,
        third_pro_response):
    """Verify returns 1 (flag for removal) for too few consecutive -ve PRO responses.

    Input has a positive response between the two negative responses.
    """
    # Add negative responses at indices 8 and 22, and a positive response at 15
    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'] = 0

    input_df.loc[third_pro_response, 'Q5Answered'] = 1
    input_df.loc[third_pro_response, 'NegativeQ5'] = 1
    assert copd.logic_consecutive_negative_responses(input_df, exacerbation_event) == 1


def test_returns_zero_enough_consecutive_negative_responses_default(
        input_df, exacerbation_event, first_pro_response, second_pro_response):
    """Verify returns 0 (pass LOGIC criterion) for required consecutive -ve PRO responses.

    Input has two consecutive negative responses. Should return 1 with default options.
    """
    # Add negative responses at indices 8 and 15
    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
    assert copd.logic_consecutive_negative_responses(input_df, exacerbation_event) == 0


def test_returns_one_too_few_consecutive_negative_responses_non_default(
        input_df, exacerbation_event, first_pro_response, second_pro_response):
    """Verify returns 1 (flag for removal) for too few consecutive -ve PRO responses.

    Input has two consecutive negative responses. Should return 0 with N=3.
    """
    # Add negative responses at indices 8 and 15
    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
    assert copd.logic_consecutive_negative_responses(
        input_df, exacerbation_event, N=3) == 1


def test_returns_zero_too_few_consecutive_negative_responses_non_default(
        input_df, exacerbation_event, first_pro_response, second_pro_response,
        third_pro_response):
    """Verify returns 0 (pass LOGIC criterion) for required consecutive -ve PRO responses.

    Input has three consecutive negative responses. Should return 0 with N=3
    """
    # Add negative responses at indices 8, 15, and 22
    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

    input_df.loc[third_pro_response, 'Q5Answered'] = 1
    input_df.loc[third_pro_response, 'NegativeQ5'] = 1
    assert copd.logic_consecutive_negative_responses(
        input_df, exacerbation_event, N=3) == 0