import datetime

from typing import Any, List, Mapping, cast

import femr_test_tools

import meds

import scipy.sparse

import femr

import femr.index

from femr.featurizers import FeaturizerList

from femr.featurizers.featurizers import AgeFeaturizer, CountFeaturizer

from femr.labelers import TimeHorizon

from femr.labelers.omop import CodeLabeler

def _assert_featurized_patients_structure(labels: List[meds.Label], features: Mapping[str, Any]):

    assert features["features"].dtype == "float32"

    assert features["patient_ids"].dtype == "int64"

    assert features["feature_times"].dtype == "datetime64[us]"

    assert features["feature_times"].shape[0] == len(labels)

    assert features["patient_ids"].shape[0] == len(labels)

    assert features["features"].shape[0] == len(labels)

    assert sorted(list(features["patient_ids"])) == sorted(list(label["patient_id"] for label in labels))

    assert sorted(list(features["feature_times"])) == sorted(list(label["prediction_time"] for label in labels))

def test_age_featurizer() -> None:

    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))

    dataset = femr_test_tools.create_patients_dataset(100)

    index = femr.index.PatientIndex(dataset)

    labeler = CodeLabeler(["2"], time_horizon, ["3"])

    patient: meds.Patient = dataset[0]

    labels = labeler.label(patient)

    featurizer = AgeFeaturizer(is_normalize=False)

    patient_features = featurizer.featurize(patient, labels)

    assert patient_features[0] == [(0, 15.43013698630137)]

    assert patient_features[1] == [(0, 17.767123287671232)]

    assert patient_features[-1] == [(0, 20.46027397260274)]

    all_labels = labeler.apply(dataset)

    featurizer = AgeFeaturizer(is_normalize=True)

    featurizer_list = FeaturizerList([featurizer])

    featurizer_list.preprocess_featurizers(dataset, index, all_labels)

    featurized_patients = featurizer_list.featurize(dataset, index, all_labels)

    _assert_featurized_patients_structure(all_labels, featurized_patients)

def test_count_featurizer() -> None:

    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))

    dataset = femr_test_tools.create_patients_dataset(100)

    index = femr.index.PatientIndex(dataset)

    labeler = CodeLabeler(["2"], time_horizon, ["3"])

    patient: meds.Patient = dataset[0]

    labels = labeler.label(patient)

    featurizer = CountFeaturizer()

    data = featurizer.generate_preprocess_data([patient], {patient["patient_id"]: labels})

    featurizer.encorperate_prepreprocessed_data([data])

    patient_features = featurizer.featurize(patient, labels)

    assert featurizer.get_num_columns() == 4, f"featurizer.get_num_columns() = {featurizer.get_num_columns()}"

    simple_patient_features = [{(featurizer.get_column_name(v.column), v.value) for v in a} for a in patient_features]

    assert simple_patient_features[0] == {

        ("SNOMED/184099003", 1),

        ("3", 1),

    }

    assert simple_patient_features[1] == {

        ("SNOMED/184099003", 1),

        ("3", 2),

        ("2", 2),

    }

    assert simple_patient_features[2] == {

        ("SNOMED/184099003", 1),

        ("3", 3),

        ("2", 4),

    }

    all_labels = labeler.apply(dataset)

    featurizer = CountFeaturizer()

    featurizer_list = FeaturizerList([featurizer])

    featurizer_list.preprocess_featurizers(dataset, index, all_labels)

    featurized_patients = featurizer_list.featurize(dataset, index, all_labels)

    _assert_featurized_patients_structure(all_labels, featurized_patients)

def test_count_featurizer_with_ontology() -> None:

    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))

    dataset = femr_test_tools.create_patients_dataset(100)

    index = femr.index.PatientIndex(dataset)

    labeler = CodeLabeler(["2"], time_horizon, ["3"])

    patient: meds.Patient = dataset[0]

    labels = labeler.label(patient)

    class DummyOntology:

        def get_all_parents(self, code):

            if code in ("2", "SNOMED/184099003"):

                return {"parent", code}

            else:

                return {code}

    featurizer = CountFeaturizer(is_ontology_expansion=True, ontology=cast(femr.ontology.Ontology, DummyOntology()))

    data = featurizer.generate_preprocess_data([patient], {patient["patient_id"]: labels})

    featurizer.encorperate_prepreprocessed_data([data])

    patient_features = featurizer.featurize(patient, labels)

    assert featurizer.get_num_columns() == 5, f"featurizer.get_num_columns() = {featurizer.get_num_columns()}"

    simple_patient_features = [{(featurizer.get_column_name(v.column), v.value) for v in a} for a in patient_features]

    assert simple_patient_features[0] == {

        ("SNOMED/184099003", 1),

        ("3", 1),

        ("parent", 1),

    }

    assert simple_patient_features[1] == {

        ("SNOMED/184099003", 1),

        ("3", 2),

        ("2", 2),

        ("parent", 3),

    }

    assert simple_patient_features[2] == {

        ("SNOMED/184099003", 1),

        ("parent", 5),

        ("3", 3),

        ("2", 4),

    }

    all_labels = labeler.apply(dataset)

    featurizer = CountFeaturizer(is_ontology_expansion=True, ontology=cast(femr.ontology.Ontology, DummyOntology()))

    featurizer_list = FeaturizerList([featurizer])

    featurizer_list.preprocess_featurizers(dataset, index, all_labels)

    featurized_patients = featurizer_list.featurize(dataset, index, all_labels)

    _assert_featurized_patients_structure(all_labels, featurized_patients)

def test_count_featurizer_with_values() -> None:

    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))

    dataset = femr_test_tools.create_patients_dataset(100)

    index = femr.index.PatientIndex(dataset)

    labeler = CodeLabeler(["2"], time_horizon, ["3"])

    patient: meds.Patient = dataset[0]

    labels = labeler.label(patient)

    featurizer = CountFeaturizer(numeric_value_decile=True, string_value_combination=True)

    data = featurizer.generate_preprocess_data([patient], {patient["patient_id"]: labels})

    featurizer.encorperate_prepreprocessed_data([data])

    patient_features = featurizer.featurize(patient, labels)

    assert featurizer.get_num_columns() == 7

    simple_patient_features = [{(featurizer.get_column_name(v.column), v.value) for v in a} for a in patient_features]

    assert simple_patient_features[0] == {

        ("SNOMED/184099003", 1),

        ("3", 1),

        ("2 [1.0, inf)", 1),

        ("1 test_value", 2),

    }

    assert simple_patient_features[1] == {

        ("SNOMED/184099003", 1),

        ("3", 2),

        ("2", 2),

        ("2 [1.0, inf)", 1),

        ("1 test_value", 2),

    }

    assert simple_patient_features[2] == {

        ("SNOMED/184099003", 1),

        ("3", 3),

        ("2", 4),

        ("2 [1.0, inf)", 1),

        ("1 test_value", 2),

    }

    all_labels = labeler.apply(dataset)

    featurizer = CountFeaturizer(numeric_value_decile=True, string_value_combination=True)

    featurizer_list = FeaturizerList([featurizer])

    featurizer_list.preprocess_featurizers(dataset, index, all_labels)

    featurized_patients = featurizer_list.featurize(dataset, index, all_labels)

    _assert_featurized_patients_structure(all_labels, featurized_patients)

def test_count_featurizer_exclude_filter() -> None:

    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))

    dataset = femr_test_tools.create_patients_dataset(100)

    labeler = CodeLabeler(["2"], time_horizon, ["3"])

    patient: meds.Patient = dataset[0]

    labels = labeler.label(patient)

    # Test filtering all codes

    featurizer = CountFeaturizer(excluded_event_filter=lambda _: True)

    data = featurizer.generate_preprocess_data([patient], {patient["patient_id"]: labels})

    featurizer.encorperate_prepreprocessed_data([data])

    assert featurizer.get_num_columns() == 0

    # Test filtering no codes

    featurizer = CountFeaturizer(excluded_event_filter=lambda _: False)

    data = featurizer.generate_preprocess_data([patient], {patient["patient_id"]: labels})

    featurizer.encorperate_prepreprocessed_data([data])

    assert featurizer.get_num_columns() == 4

    # Test filtering single code

    featurizer = CountFeaturizer(excluded_event_filter=lambda e: e["code"] == "3")

    data = featurizer.generate_preprocess_data([patient], {patient["patient_id"]: labels})

    featurizer.encorperate_prepreprocessed_data([data])

    assert featurizer.get_num_columns() == 3

def test_count_bins_featurizer() -> None:

    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))

    dataset = femr_test_tools.create_patients_dataset(100)

    index = femr.index.PatientIndex(dataset)

    labeler = CodeLabeler(["2"], time_horizon, ["3"])

    patient: meds.Patient = dataset[0]

    labels = labeler.label(patient)

    time_bins = [

        datetime.timedelta(days=90),

        datetime.timedelta(days=180),

        datetime.timedelta(weeks=1e4),

    ]

    featurizer = CountFeaturizer(

        time_bins=time_bins,

    )

    data = featurizer.generate_preprocess_data([patient], {patient["patient_id"]: labels})

    featurizer.encorperate_prepreprocessed_data([data])

    patient_features = featurizer.featurize(patient, labels)

    assert featurizer.get_num_columns() == 12

    simple_patient_features = [{(featurizer.get_column_name(v.column), v.value) for v in a} for a in patient_features]

    assert simple_patient_features[0] == {

        ("SNOMED/184099003_70000 days, 0:00:00", 1),

        ("3_90 days, 0:00:00", 1),

    }

    assert simple_patient_features[1] == {

        ("3_90 days, 0:00:00", 1),

        ("SNOMED/184099003_70000 days, 0:00:00", 1),

        ("3_70000 days, 0:00:00", 1),

        ("2_70000 days, 0:00:00", 2),

    }

    assert simple_patient_features[2] == {

        ("2_70000 days, 0:00:00", 2),

        ("2_90 days, 0:00:00", 2),

        ("SNOMED/184099003_70000 days, 0:00:00", 1),

        ("3_90 days, 0:00:00", 1),

        ("3_70000 days, 0:00:00", 2),

    }

    all_labels = labeler.apply(dataset)

    time_bins = [

        datetime.timedelta(days=90),

        datetime.timedelta(days=180),

        datetime.timedelta(weeks=1e4),

    ]

    featurizer = CountFeaturizer(

        time_bins=time_bins,

    )

    featurizer_list = FeaturizerList([featurizer])

    featurizer_list.preprocess_featurizers(dataset, index, all_labels)

    featurized_patients = featurizer_list.featurize(dataset, index, all_labels)

    _assert_featurized_patients_structure(all_labels, featurized_patients)

def test_complete_featurization() -> None:

    time_horizon = TimeHorizon(datetime.timedelta(days=0), datetime.timedelta(days=180))

    dataset = femr_test_tools.create_patients_dataset(100)

    index = femr.index.PatientIndex(dataset)

    labeler = CodeLabeler(["2"], time_horizon, ["3"])

    all_labels = labeler.apply(dataset)

    age_featurizer = AgeFeaturizer(is_normalize=True)

    age_featurizer_list = FeaturizerList([age_featurizer])

    age_featurizer_list.preprocess_featurizers(dataset, index, all_labels)

    age_featurized_patients = age_featurizer_list.featurize(dataset, index, all_labels)

    time_bins = [

        datetime.timedelta(days=90),

        datetime.timedelta(days=180),

        datetime.timedelta(weeks=1e5),

    ]

    count_featurizer = CountFeaturizer(time_bins=time_bins)

    count_featurizer_list = FeaturizerList([count_featurizer])

    count_featurizer_list.preprocess_featurizers(dataset, index, all_labels)

    count_featurized_patients = count_featurizer_list.featurize(dataset, index, all_labels)

    age_featurizer = AgeFeaturizer(is_normalize=True)

    time_bins = [

        datetime.timedelta(days=90),

        datetime.timedelta(days=180),

        datetime.timedelta(weeks=1e5),

    ]

    count_featurizer = CountFeaturizer(time_bins=time_bins)

    featurizer_list = FeaturizerList([age_featurizer, count_featurizer])

    featurizer_list.preprocess_featurizers(dataset, index, all_labels)

    featurized_patients = featurizer_list.featurize(dataset, index, all_labels)

    assert featurized_patients["patient_ids"].shape == count_featurized_patients["patient_ids"].shape

    the_same = (

        featurized_patients["features"].toarray()

        == scipy.sparse.hstack((age_featurized_patients["features"], count_featurized_patients["features"])).toarray()

    )

    assert the_same.all()