import datetime
from typing import Any, List, Optional, Tuple, Union

import datasets
import meds

from femr.labelers import Labeler

# 2nd elem of tuple -- 'skip' means no label, None means censored
EventsWithLabels = List[Tuple[Tuple[Tuple, int, Any], Union[bool, str]]]


DUMMY_EVENTS = [
    ((1995, 1, 3), meds.birth_code, None),
    ((2010, 1, 1), 1, "test_value"),
    ((2010, 1, 1), 1, "test_value"),
    ((2010, 1, 5), 2, 1),
    ((2010, 6, 5), 3, None),
    ((2010, 8, 5), 2, None),
    ((2011, 7, 5), 2, None),
    ((2012, 10, 5), 3, None),
    ((2015, 6, 5, 0), 2, None),
    ((2015, 6, 5, 10, 10), 2, None),
    ((2015, 6, 15, 11), 3, None),
    ((2016, 1, 1), 2, None),
    ((2016, 3, 1, 10, 10, 10), 4, None),
]

NUM_EVENTS = len(DUMMY_EVENTS)
NUM_PATIENTS = 10


def create_patients_dataset(num_patients: int, events: List[Tuple[Tuple, int, Any]] = DUMMY_EVENTS) -> datasets.Dataset:
    """Creates a list of patients, each with the same events contained in `events`"""

    converted_events: List[meds.Event] = []

    for event in events:
        if isinstance(event[1], int):
            code = str(event[1])
        else:
            code = event[1]
        if event[2] is None:
            val = {}
        elif isinstance(event[2], str):
            val = {"text_value": event[2]}
        else:
            val = {"numeric_value": event[2]}
        converted_events.append({"time": datetime.datetime(*event[0]), "measurements": [{"code": code, **val}]})

    return datasets.Dataset.from_dict(
        {"patient_id": list(range(num_patients)), "events": [converted_events for _ in range(num_patients)]}
    )


def assert_labels_are_accurate(
    labeled_patients: List[meds.Label],
    patient_id: int,
    true_labels: List[Tuple[datetime.datetime, Optional[bool]]],
    help_text: str = "",
):
    """Passes if the labels in `labeled_patients` for `patient_id` exactly match the labels in `true_labels`."""
    generated_labels: List[meds.Label] = [a for a in labeled_patients if a["patient_id"] == patient_id]
    # Check that length of lists of labels are the same

    assert len(generated_labels) == len(
        true_labels
    ), f"len(generated): {len(generated_labels)} != len(expected): {len(true_labels)} | {help_text}"
    # Check that value of labels are the same
    for idx, (label, true_label) in enumerate(zip(generated_labels, true_labels)):
        assert label["boolean_value"] == true_label[1] and label["prediction_time"] == true_label[0], (
            f"patient_id={patient_id}, label_idx={idx}, label={label}  |  "
            f"{label} (Assigned) != {true_label} (Expected)  |  "
            f"{help_text}"
        )


def run_test_for_labeler(
    labeler: Labeler,
    events_with_labels: EventsWithLabels,
    true_outcome_times: Optional[List[datetime.datetime]] = None,
    true_prediction_times: Optional[List[datetime.datetime]] = None,
    help_text: str = "",
) -> None:
    patients: datasets.Database = create_patients_dataset(10, [x[0] for x in events_with_labels])
    true_labels: List[Tuple[datetime.datetime, Optional[bool]]] = [
        (datetime.datetime(*x[0][0]), x[1]) for x in events_with_labels if isinstance(x[1], bool)
    ]
    if true_prediction_times is not None:
        # If manually specified prediction times, adjust labels from occurring at `event.start`
        # e.g. we may make predictions at `event.end` or `event.start + 1 day`
        true_labels = [(tp, tl[1]) for (tl, tp) in zip(true_labels, true_prediction_times)]
    labeled_patients: List[meds.Label] = labeler.apply(patients)

    # Check accuracy of Labels
    for patient in patients:
        assert_labels_are_accurate(
            labeled_patients,
            patient["patient_id"],
            true_labels,
            help_text=help_text,
        )