# flake8: noqa: E402

import datetime
import os
import pathlib
import sys
import warnings
from typing import List

import meds
from femr_test_tools import EventsWithLabels, run_test_for_labeler

from femr.labelers import TimeHorizon, TimeHorizonEventLabeler


class DummyLabeler(TimeHorizonEventLabeler):
    """Dummy labeler that returns True if the event's `code` is in `self.outcome_codes`."""

    def __init__(self, outcome_codes: List[int], time_horizon: TimeHorizon, allow_same_time: bool = True):
        self.outcome_codes: List[str] = [str(a) for a in outcome_codes]
        self.time_horizon: TimeHorizon = time_horizon
        self.allow_same_time = allow_same_time

    def allow_same_time_labels(self) -> bool:
        return self.allow_same_time

    def get_prediction_times(self, patient: meds.Patient) -> List[datetime.datetime]:
        return sorted(list({e["time"] for e in patient["events"]}))

    def get_time_horizon(self) -> TimeHorizon:
        return self.time_horizon

    def get_outcome_times(self, patient: meds.Patient) -> List[datetime.datetime]:
        times: List[datetime.datetime] = []
        for e in patient["events"]:
            for m in e["measurements"]:
                if m["code"] in self.outcome_codes:
                    times.append(e["time"])
        return times


def test_no_outcomes(tmp_path: pathlib.Path):
    # No outcomes occur in this patient's timeline
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))
    labeler = DummyLabeler([100], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2015, 1, 3), 2, None), "duplicate"),
        (((2015, 1, 3), 1, None), "duplicate"),
        (((2015, 1, 3), 3, None), False),
        (((2015, 10, 5), 1, None), False),
        (((2018, 1, 3), 2, None), False),
        (((2018, 3, 3), 1, None), False),
        (((2018, 5, 3), 2, None), False),
        (((2018, 5, 3, 11), 1, None), False),
        (((2018, 5, 4), 1, None), False),
        (((2018, 12, 4), 1, None), "out of range"),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_no_outcomes")


def test_horizon_0_180_days(tmp_path: pathlib.Path):
    # (0, 180) days
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))
    labeler = DummyLabeler([2], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2015, 1, 3), 2, None), "duplicate"),
        (((2015, 1, 3), 1, None), "duplicate"),
        (((2015, 1, 3), 3, None), True),
        (((2015, 10, 5), 1, None), False),
        (((2018, 1, 3), 2, None), True),
        (((2018, 3, 3), 1, None), True),
        (((2018, 5, 3), 2, None), True),
        (((2018, 5, 3, 11), 1, None), False),
        (((2018, 5, 4), 1, None), False),
        (((2018, 12, 4), 1, None), "out of range"),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_0_180_days")


def test_horizon_0_180_days_no_same(tmp_path: pathlib.Path):
    # (0, 180) days
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))
    labeler = DummyLabeler([2], time_horizon, allow_same_time=False)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2015, 1, 3), 2, None), "duplicate"),
        (((2015, 1, 3), 1, None), "duplicate"),
        (((2015, 1, 3), 3, None), "same"),
        (((2015, 10, 5), 1, None), False),
        (((2018, 1, 3), 2, None), "same"),
        (((2018, 3, 3), 1, None), True),
        (((2018, 5, 3), 2, None), "same"),
        (((2018, 5, 3, 11), 1, None), False),
        (((2018, 5, 4), 1, None), False),
        (((2018, 12, 4), 1, None), "out of range"),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_0_180_days")


def test_horizon_1_180_days(tmp_path: pathlib.Path):
    # (1, 180) days
    time_horizon = TimeHorizon(datetime.timedelta(days=1), datetime.timedelta(days=180))
    labeler = DummyLabeler([2], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2015, 1, 3), 2, None), "duplicate"),
        (((2015, 1, 3), 1, None), "duplicate"),
        (((2015, 1, 3), 3, None), False),
        (((2015, 10, 5), 1, None), False),
        (((2018, 1, 3), 2, None), True),
        (((2018, 3, 3), 1, None), True),
        (((2018, 5, 3), 2, None), False),
        (((2018, 5, 3, 11), 1, None), False),
        (((2018, 5, 4), 1, None), False),
        (((2018, 12, 4), 1, None), "out of range"),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_1_180_days")


def test_horizon_180_365_days(tmp_path: pathlib.Path):
    # (180, 365) days
    time_horizon = TimeHorizon(datetime.timedelta(days=180), datetime.timedelta(days=365))
    labeler = DummyLabeler([2], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2000, 1, 3), 2, None), True),
        (((2000, 10, 5), 2, None), False),
        (((2002, 1, 5), 2, None), True),
        (((2002, 3, 1), 1, None), True),
        (((2002, 4, 5), 3, None), True),
        (((2002, 4, 12), 1, None), True),
        (((2002, 12, 5), 2, None), False),
        (((2002, 12, 10), 1, None), False),
        (((2004, 1, 10), 2, None), False),
        (((2008, 1, 10), 2, None), "out of range"),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_180_365_days")


def test_horizon_0_0_days(tmp_path: pathlib.Path):
    # (0, 0) days
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=0))
    labeler = DummyLabeler([2], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2015, 1, 3), 2, None), "duplicate"),
        (((2015, 1, 3), 1, None), True),
        (((2015, 1, 4), 1, None), False),
        (((2015, 1, 5), 2, None), True),
        (((2015, 1, 5, 10), 1, None), False),
        (((2015, 1, 6), 2, None), True),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_0_0_days")


def test_horizon_10_10_days(tmp_path: pathlib.Path):
    # (10, 10) days
    time_horizon = TimeHorizon(datetime.timedelta(days=10), datetime.timedelta(days=10))
    labeler = DummyLabeler([2], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2015, 1, 3), 2, None), False),
        (((2015, 1, 13), 1, None), True),
        (((2015, 1, 23), 2, None), True),
        (((2015, 2, 2), 2, None), False),
        (((2015, 3, 10), 1, None), True),
        (((2015, 3, 20), 2, None), False),
        (((2015, 3, 29), 2, None), "out of range"),
        (((2015, 3, 30), 1, None), "out of range"),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_10_10_days")


def test_horizon_0_1000000_days(tmp_path: pathlib.Path):
    # (0, 1000000) days
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=1000000))
    labeler = DummyLabeler([2], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2000, 1, 3), 2, None), True),
        (((2001, 10, 5), 1, None), True),
        (((2020, 10, 5), 2, None), True),
        (((2021, 10, 5), 1, None), True),
        (((2050, 1, 10), 2, None), True),
        (((2051, 1, 10), 1, None), False),
        (((5000, 1, 10), 1, None), "out of range"),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_0_1000000_days")


def test_horizon_5_10_hours(tmp_path: pathlib.Path):
    # (5 hours, 10.5 hours)
    time_horizon = TimeHorizon(datetime.timedelta(hours=5), datetime.timedelta(hours=10, minutes=30))
    labeler = DummyLabeler([2], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((2015, 1, 1, 0, 0), 1, None), True),
        (((2015, 1, 1, 10, 29), 2, None), False),
        (((2015, 1, 1, 10, 30), 1, None), False),
        (((2015, 1, 1, 10, 31), 1, None), False),
        #
        (((2016, 1, 1, 0, 0), 1, None), True),
        (((2016, 1, 1, 10, 29), 1, None), False),
        (((2016, 1, 1, 10, 30), 2, None), False),
        (((2016, 1, 1, 10, 31), 1, None), False),
        #
        (((2017, 1, 1, 0, 0), 1, None), False),
        (((2017, 1, 1, 10, 29), 1, None), False),
        (((2017, 1, 1, 10, 30), 1, None), False),
        (((2017, 1, 1, 10, 31), 2, None), False),
        #
        (((2018, 1, 1, 0, 0), 1, None), False),
        (((2018, 1, 1, 4, 59, 59), 2, None), False),
        (((2018, 1, 1, 5), 1, None), False),
        #
        (((2019, 1, 1, 0, 0), 1, None), True),
        (((2019, 1, 1, 4, 59, 59), 1, None), "out of range"),
        (((2019, 1, 1, 5), 2, None), "out of range"),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_5_10_hours")


def test_horizon_infinite(tmp_path: pathlib.Path):
    # Infinite horizon
    time_horizon = TimeHorizon(
        datetime.timedelta(days=10),
        None,
    )
    labeler = DummyLabeler([2], time_horizon)
    events_with_labels: EventsWithLabels = [
        # fmt: off
        (((1950, 1, 3), 1, None), True),
        (((2000, 1, 3), 1, None), True),
        (((2001, 10, 5), 1, None), True),
        (((2020, 10, 5), 1, None), True),
        (((2021, 10, 5), 1, None), True),
        (((2050, 1, 10), 2, None), True),
        (((2050, 1, 20), 2, None), False),
        (((2051, 1, 10), 1, None), False),
        (((5000, 1, 10), 1, None), False),
        # fmt: on
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_horizon_infinite")
