import datetime
import pathlib
from typing import List, Set

# Needed to import `tools` for local testing
from femr_test_tools import EventsWithLabels, run_test_for_labeler

from femr.labelers import TimeHorizon
from femr.labelers.omop import (
    AKICodeLabeler,
    AnemiaCodeLabeler,
    CodeLabeler,
    HyperkalemiaCodeLabeler,
    HypoglycemiaCodeLabeler,
    HyponatremiaCodeLabeler,
    LupusCodeLabeler,
    MortalityCodeLabeler,
    NeutropeniaCodeLabeler,
    OMOPConceptCodeLabeler,
    ThrombocytopeniaCodeLabeler,
)

#############################################
#############################################
#
# Generic CodeLabeler
#
#############################################
#############################################


def test_outcome_codes(tmp_path: pathlib.Path):
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=10))
    # One outcome
    labeler = CodeLabeler(["2"], time_horizon)
    events_with_labels: EventsWithLabels = [
        (((2015, 1, 3), 2, None), "duplicate"),
        (((2015, 1, 3), 4, None), "duplicate"),
        (((2015, 1, 3), 1, None), "duplicate"),
        (((2015, 1, 3), 3, None), "skip"),
        (((2015, 10, 5), 1, None), False),
        (((2018, 1, 3), 2, None), "skip"),
        (((2018, 3, 1), 4, None), False),
        (((2018, 3, 3), 1, None), False),
        (((2018, 5, 2), 5, None), True),
        (((2018, 5, 3), 2, None), "skip"),
        (((2018, 5, 3, 11), 1, None), False),
        (((2018, 5, 4), 1, None), "duplicate"),
        (((2018, 5, 4), 4, None), False),
        (((2018, 11, 1), 5, None), False),
        (((2018, 12, 4), 1, None), False),
        (((2018, 12, 30), 4, None), "out of range"),
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_outcome_codes_one")

    # Zero outcome
    labeler = CodeLabeler([], time_horizon)
    events_with_labels = [
        (((2015, 1, 3), 2, None), "duplicate"),
        (((2015, 1, 3), 4, 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, 1), 4, None), False),
        (((2018, 3, 3), 1, None), False),
        (((2018, 5, 2), 5, None), False),
        (((2018, 5, 3), 2, None), False),
        (((2018, 5, 3, 11), 1, None), False),
        (((2018, 5, 4), 1, None), "duplicate"),
        (((2018, 5, 4), 4, None), False),
        (((2018, 11, 1), 5, None), False),
        (((2018, 12, 4), 1, None), False),
        (((2018, 12, 30), 4, None), "out of range"),
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_outcome_codes_zero")

    # Multiple outcomes
    labeler = CodeLabeler(["1", "4"], time_horizon)
    events_with_labels = [
        (((2015, 1, 3), 2, None), "duplicate"),
        (((2015, 1, 3), 4, None), "duplicate"),
        (((2015, 1, 3), 1, None), "duplicate"),
        (((2015, 1, 3), 3, None), "skip"),
        (((2015, 10, 5), 1, None), "skip"),
        (((2018, 1, 3), 2, None), False),
        (((2018, 3, 1), 4, None), "skip"),
        (((2018, 3, 3), 1, None), "skip"),
        (((2018, 5, 2), 5, None), False),
        (((2018, 5, 3), 2, None), False),
        (((2018, 6, 2), 0, None), True),
        (((2018, 6, 3, 11), 1, None), "skip"),
        (((2018, 6, 3, 23), 3, None), False),
        (((2018, 9, 1), 3, None), True),
        (((2018, 9, 4), 4, None), "skip"),
        (((2018, 11, 1), 5, None), False),
        (((2018, 12, 3), 0, None), True),
        (((2018, 12, 4), 4, None), "skip"),
        (((2018, 12, 30), 0, None), "out of range"),
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="test_outcome_codes_multiple")


def test_prediction_codes(tmp_path: pathlib.Path):
    # One outcome + multiple predictions
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=10))
    labeler = CodeLabeler(["2"], time_horizon, prediction_codes=["4", "5"])
    events_with_labels: EventsWithLabels = [
        (((2015, 1, 3), 2, None), "skip"),
        (((2015, 1, 3), 4, None), "skip"),
        (((2015, 1, 3), 1, None), "skip"),
        (((2015, 1, 3), 3, None), "skip"),
        (((2015, 10, 5), 1, None), "skip"),
        (((2018, 1, 3), 2, None), "skip"),
        (((2018, 3, 1), 4, None), False),
        (((2018, 3, 3), 1, None), "skip"),
        (((2018, 5, 2), 5, None), True),
        (((2018, 5, 3), 2, None), "skip"),
        (((2018, 5, 3, 11), 1, None), "skip"),
        (((2018, 5, 4), 4, None), False),
        (((2018, 5, 4), 1, None), "skip"),
        (((2018, 11, 1), 5, None), False),
        (((2018, 12, 4), 1, None), "skip"),
        (((2018, 12, 30), 4, None), "out of range"),
    ]
    run_test_for_labeler(labeler, events_with_labels, help_text="prediction_codes_one_outcomes")

    # Multiple outcomes + multiple predictions
    labeler = CodeLabeler(["2", "6", "7"], time_horizon, prediction_codes=["4", "5"])
    events_with_labels = [
        (((2010, 1, 1), 2, None), "skip"),
        (((2010, 1, 3), 4, None), True),
        (((2010, 1, 8), 6, None), "skip"),
        (((2010, 2, 1), 5, None), True),
        (((2010, 2, 9), 7, None), "skip"),
        (((2010, 2, 11), 4, None), False),
        (((2015, 1, 3), 2, None), "skip"),
        (((2015, 1, 3), 4, None), "skip"),
        (((2015, 1, 3), 1, None), "skip"),
        (((2015, 1, 3), 3, None), "skip"),
        (((2015, 10, 5), 1, None), "skip"),
        (((2018, 1, 3), 2, None), "skip"),
        (((2018, 3, 1), 4, None), True),
        (((2018, 3, 2), 7, None), "skip"),
        (((2018, 3, 3), 1, None), "skip"),
        (((2018, 5, 2), 5, None), True),
        (((2018, 5, 3), 2, None), "skip"),
        (((2018, 5, 3, 11), 1, None), "skip"),
        (((2018, 5, 4), 4, None), False),
        (((2018, 5, 4), 1, None), "skip"),
        (((2018, 11, 1), 5, None), False),
        (((2018, 12, 4), 1, None), "skip"),
        (((2018, 12, 30), 4, None), "out of range"),
    ]
    run_test_for_labeler(
        labeler,
        events_with_labels,
        help_text="prediction_codes_multiple_outcomes",
    )

    # Multiple outcomes + no predictions
    labeler = CodeLabeler(["2", "6", "7"], time_horizon, prediction_codes=[])
    events_with_labels = [
        (((2010, 1, 1), 2, None), "skip"),
        (((2010, 1, 3), 4, None), "skip"),
        (((2010, 1, 8), 6, None), "skip"),
        (((2010, 2, 1), 5, None), "skip"),
        (((2010, 2, 9), 7, None), "skip"),
        (((2010, 2, 11), 4, None), "skip"),
        (((2015, 1, 3), 2, None), "skip"),
        (((2015, 1, 3), 4, None), "skip"),
        (((2015, 1, 3), 1, None), "skip"),
        (((2015, 1, 3), 3, None), "skip"),
        (((2015, 10, 5), 1, None), "skip"),
        (((2018, 1, 3), 2, None), "skip"),
        (((2018, 3, 1), 4, None), "skip"),
        (((2018, 3, 2), 7, None), "skip"),
        (((2018, 3, 3), 1, None), "skip"),
        (((2018, 5, 2), 5, None), "skip"),
        (((2018, 5, 3), 2, None), "skip"),
        (((2018, 5, 3, 11), 1, None), "skip"),
        (((2018, 5, 4), 4, None), "skip"),
        (((2018, 5, 4), 1, None), "skip"),
        (((2018, 11, 1), 5, None), "skip"),
        (((2018, 12, 4), 1, None), "skip"),
        (((2018, 12, 30), 4, None), "skip"),
    ]
    run_test_for_labeler(
        labeler,
        events_with_labels,
        help_text="prediction_codes_zero_predictions",
    )


#############################################
#############################################
#
# Generic OMOPConceptCodeLabeler
#
#############################################
#############################################


class DummyOntology_Base:
    def get_children(self, code: str) -> Set[str]:
        return set()

    def get_all_children(self, code: str) -> Set[str]:
        val = {code}
        for child in self.get_children(code):
            val |= self.get_all_children(child)
        return val


class DummyOntology_OMOPConcept(DummyOntology_Base):
    def get_children(self, parent_code: str) -> Set[str]:
        if parent_code == "OMOP_CONCEPT_A":
            return {"OMOP_CONCEPT_A_CHILD", "OMOP_CONCEPT_A_CHILD2"}
        elif parent_code == "OMOP_CONCEPT_B":
            return {"OMOP_CONCEPT_B_CHILD"}
        elif parent_code == "OMOP_CONCEPT_A_CHILD":
            return {"OMOP_CONCEPT_A_CHILD_CHILD"}
        else:
            return set()


class DummyLabeler_OMOPConcept(OMOPConceptCodeLabeler):
    original_omop_concept_codes = [
        "OMOP_CONCEPT_A",
        "OMOP_CONCEPT_B",
    ]


def test_omop_concept_code_labeler(tmp_path: pathlib.Path):
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=10))
    ontology = DummyOntology_OMOPConcept()
    labeler = DummyLabeler_OMOPConcept(ontology, time_horizon, prediction_codes=["1", "2"])  # type: ignore
    assert set(labeler.outcome_codes) == {
        "OMOP_CONCEPT_A_CHILD_CHILD",
        "OMOP_CONCEPT_B",
        "OMOP_CONCEPT_B_CHILD",
        "OMOP_CONCEPT_A_CHILD2",
        "OMOP_CONCEPT_A",
        "OMOP_CONCEPT_A_CHILD",
    }
    assert labeler.prediction_codes == ["1", "2"]
    assert labeler.get_time_horizon() == time_horizon


#############################################
#############################################
#
# Specific instances of CodeLabeler
#
#############################################
#############################################


#############################################
# MortalityCodeLabeler
#############################################


class DummyOntology_Mortality(DummyOntology_Base):
    def get_children(self, parent_code: str) -> Set[str]:
        return set()


def test_MortalityCodeLabeler() -> None:
    """Create a MortalityCodeLabeler for codes 3 and 6"""
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))
    events_with_labels: EventsWithLabels = [
        (((1995, 1, 3), 0, 34.5), False),
        (((2000, 1, 1), 1, "test_value"), True),
        (((2000, 1, 5), 2, 1), True),
        (((2000, 6, 5), "SNOMED/419620001", True), "skip"),
        (((2005, 2, 5), 2, None), False),
        (((2005, 7, 5), 2, None), False),
        (((2010, 10, 5), 1, None), False),
        (((2015, 2, 5, 0), 2, None), False),
        (((2015, 7, 5, 0), 0, None), True),
        (((2015, 11, 5, 10, 10), 2, None), True),
        (((2015, 11, 15, 11), "SNOMED/419620001", None), "skip"),
        (((2020, 1, 1), 2, None), "out of range"),
        (((2020, 3, 1, 10, 10, 10), 2, None), "out of range"),
    ]

    ontology = DummyOntology_Mortality()

    # Run labeler
    labeler = MortalityCodeLabeler(ontology, time_horizon)  # type: ignore

    run_test_for_labeler(labeler, events_with_labels, help_text="MortalityLabeler")


#############################################
# LupusCodeLabeler
#############################################


class DummyOntology_Lupus(DummyOntology_Base):
    def get_children(self, parent_code: str) -> Set[str]:
        if parent_code == "SNOMED/55464009":
            return {"SNOMED_55464009", "Lupus_child_seven", "Lupus_child_nine", "Lupus_child_ten"}
        else:
            return set()


def test_LupusCodeLabeler() -> None:
    """Create a LupusCodeLabeler for codes 3 and 6"""
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))
    events_with_labels: EventsWithLabels = [
        (((1995, 1, 3), 0, 34.5), False),
        (((2000, 1, 1), 1, "test_value"), True),
        (((2000, 1, 5), 2, 1), True),
        (((2000, 5, 5), "SNOMED/201436003", None), "skip"),
        (((2005, 2, 5), 2, None), False),
        (((2005, 7, 5), 2, None), False),
        (((2010, 10, 5), 1, None), True),
        (((2010, 10, 8), "Lupus_child_seven", None), "skip"),
        (((2015, 2, 5, 0), 2, None), False),
        (((2015, 7, 5, 0), 0, None), True),
        (((2015, 11, 5, 10, 10), 2, None), True),
        (((2015, 11, 15, 11), "SNOMED/55464009", None), "skip"),
        (((2020, 1, 1), "Lupus_child_ten", None), "skip"),
        (((2020, 3, 1, 10, 10, 10), 2, None), "out of range"),
    ]

    ontology = DummyOntology_Lupus()
    labeler = LupusCodeLabeler(ontology, time_horizon)  # type: ignore
    # Check that we selected the right codes
    assert set(labeler.outcome_codes) == {
        "SNOMED_55464009",
        "SNOMED/201436003",
        "Lupus_child_nine",
        "SNOMED/55464009",
        "Lupus_child_ten",
        "Lupus_child_seven",
    }

    run_test_for_labeler(labeler, events_with_labels, help_text="LupusCodeLabeler")


#############################################
#############################################
#
# Specific instances of OMOPConceptCodeLabeler
#
#############################################
#############################################


class DummyOntology_OMOPConcept_Specific(DummyOntology_Base):
    def __init__(self, new_codes: List[str]):
        self.new_codes = new_codes + ["", ""]

    def get_children(self, parent_code: str) -> Set[str]:
        if parent_code == "child_1":
            return {"child_1_1"}
        elif parent_code == self.new_codes[0]:
            return {"child_1"}
        elif parent_code == self.new_codes[1]:
            return {"child_2"}
        return set()


def _assert_labvalue_constructor_correct(
    labeler: OMOPConceptCodeLabeler,
    time_horizon: TimeHorizon,
    outcome_codes: set,
):
    assert set(labeler.outcome_codes) == outcome_codes
    assert labeler.prediction_codes == ["1", "2"]
    assert labeler.get_time_horizon() == time_horizon


def _create_specific_labvalue_labeler(LabelerClass, outcome_codes: set):
    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=10))
    ontology = DummyOntology_OMOPConcept_Specific(LabelerClass.original_omop_concept_codes)
    labeler = LabelerClass(ontology, time_horizon, prediction_codes=["1", "2"])  # type: ignore
    _assert_labvalue_constructor_correct(labeler, time_horizon, outcome_codes)
    return labeler


def test_thrombocytopenia(tmp_path: pathlib.Path):
    outcome_codes = {"child_1_1", "child_2", "SNOMED/89627008", "child_1", "SNOMED/267447008"}
    _create_specific_labvalue_labeler(ThrombocytopeniaCodeLabeler, outcome_codes)


def test_hyperkalemia(tmp_path: pathlib.Path):
    outcome_codes = {"child_1", "SNOMED/14140009", "child_1_1"}
    _create_specific_labvalue_labeler(HyperkalemiaCodeLabeler, outcome_codes)


def test_hypoglycemia(tmp_path: pathlib.Path):
    outcome_codes = {
        "SNOMED/120731000119103",
        "child_2",
        "child_1",
        "SNOMED/52767006",
        "SNOMED/719216001",
        "SNOMED/302866003",
        "SNOMED/267384006",
        "SNOMED/421725003",
        "SNOMED/237637005",
        "SNOMED/237633009",
        "SNOMED/190448007",
        "child_1_1",
        "SNOMED/421437000",
        "SNOMED/230796005",
        "SNOMED/84371000119108",
    }
    _create_specific_labvalue_labeler(HypoglycemiaCodeLabeler, outcome_codes)


def test_hyponatremia(tmp_path: pathlib.Path):
    outcome_codes = {"child_2", "child_1", "SNOMED/89627008", "SNOMED/267447008", "child_1_1"}
    _create_specific_labvalue_labeler(HyponatremiaCodeLabeler, outcome_codes)


def test_anemia(tmp_path: pathlib.Path):
    outcome_codes = {
        "child_1",
        "SNOMED/713496008",
        "SNOMED/691411000119101",
        "SNOMED/691401000119104",
        "SNOMED/767657005",
        "child_2",
        "SNOMED/111570005",
        "SNOMED/271737000",
        "SNOMED/713349004",
        "child_1_1",
    }
    _create_specific_labvalue_labeler(AnemiaCodeLabeler, outcome_codes)


def test_neutropenia(tmp_path: pathlib.Path):
    outcome_codes = {"child_1", "SNOMED/165517008", "child_1_1"}
    _create_specific_labvalue_labeler(NeutropeniaCodeLabeler, outcome_codes)


def test_aki(tmp_path: pathlib.Path):
    outcome_codes = {"child_2", "child_1_1", "child_1", "SNOMED/298015003", "SNOMED/14669001", "SNOMED/35455006"}
    _create_specific_labvalue_labeler(AKICodeLabeler, outcome_codes)
