from __future__ import annotations

import datetime

import meds

from femr.transforms import delta_encode, remove_nones
from femr.transforms.stanford import (
    move_billing_codes,
    move_pre_birth,
    move_to_day_end,
    move_visit_start_to_first_event_start,
)


def cleanup(patient):
    for event in patient["events"]:
        for measurement in event["measurements"]:
            if "metadata" not in measurement:
                measurement["metadata"] = {}

            for k in ("numeric_value", "text_value", "datetime_value"):
                if k not in measurement:
                    measurement[k] = None

            if "table" not in measurement["metadata"]:
                measurement["metadata"]["table"] = None


def test_pre_birth() -> None:
    patient = {
        "patient_id": 123,
        "events": [
            {"time": datetime.datetime(1999, 7, 2), "measurements": [{"code": 1234}]},
            {"time": datetime.datetime(1999, 7, 9), "measurements": [{"code": meds.birth_code}]},
            {"time": datetime.datetime(1999, 7, 11), "measurements": [{"code": 12345}]},
        ],
    }

    expected = {
        "patient_id": 123,
        "events": [
            {"time": datetime.datetime(1999, 7, 9), "measurements": [{"code": 1234}]},
            {"time": datetime.datetime(1999, 7, 9), "measurements": [{"code": meds.birth_code}]},
            {"time": datetime.datetime(1999, 7, 11), "measurements": [{"code": 12345}]},
        ],
    }

    cleanup(patient)
    cleanup(expected)

    assert move_pre_birth(patient) == expected


def test_move_visit_start_ignores_other_visits() -> None:
    patient = {
        "patient_id": 123,
        "events": [
            {  # A non-visit event with no explicit start time
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [{"code": 1234, "metadata": {"visit_id": 9999}}],
            },
            {  # A visit event with just date specified
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [
                    {
                        "code": 4567,
                        "metadata": {
                            "table": "visit_occurrence",
                            "visit_id": 9999,
                        },
                    }
                ],
            },
            {  # A non-visit event from a separate visit ID
                "time": datetime.datetime(1999, 7, 2, 11),
                "measurements": [{"code": 2345, "metadata": {"visit_id": 8888}}],
            },
            {  # First recorded non-visit event for visit ID 9999
                "time": datetime.datetime(1999, 7, 2, 12),
                "measurements": [{"code": 3456, "metadata": {"visit_id": 9999}}],
            },
        ],
    }

    # Note that events are implicitly sorted first by start time, then by code:
    # https://github.com/som-shahlab/femr/blob/main/src/femr/__init__.py#L69
    expected = {
        "patient_id": 123,
        "events": [
            {  # A non-visit event with no explicit start time
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [{"code": 1234, "metadata": {"visit_id": 9999}}],
            },
            {  # A non-visit event from a separate visit ID
                "time": datetime.datetime(1999, 7, 2, 11),
                "measurements": [{"code": 2345, "metadata": {"visit_id": 8888}}],
            },
            {  # Now visit event has date and time specified
                "time": datetime.datetime(1999, 7, 2, 12),
                "measurements": [
                    {
                        "code": 4567,
                        "metadata": {
                            "table": "visit_occurrence",
                            "visit_id": 9999,
                        },
                    }
                ],
            },
            {  # First recorded non-visit event for visit ID 9999
                "time": datetime.datetime(1999, 7, 2, 12),
                "measurements": [{"code": 3456, "metadata": {"visit_id": 9999}}],
            },
        ],
    }

    cleanup(patient)
    cleanup(expected)

    assert move_visit_start_to_first_event_start(patient) == expected


def test_move_visit_start_minute_after_midnight() -> None:
    patient = {
        "patient_id": 123,
        "events": [
            {
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [
                    {"code": 3456, "metadata": {"visit_id": 9999, "table": "visit_occurrence"}},
                    {"code": 1234, "metadata": {"visit_id": 9999}},
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 2, 0, 1),
                "measurements": [{"code": 2345, "metadata": {"visit_id": 9999}}],
            },
            {
                "time": datetime.datetime(1999, 7, 2, 12),
                "measurements": [{"code": 4567, "metadata": {"visit_id": 9999}}],
            },
        ],
    }

    expected = {
        "patient_id": 123,
        "events": [
            {
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [{"code": 1234, "metadata": {"visit_id": 9999}}],
            },
            {
                "time": datetime.datetime(1999, 7, 2, 0, 1),
                "measurements": [{"code": 3456, "metadata": {"visit_id": 9999, "table": "visit_occurrence"}}],
            },
            {
                "time": datetime.datetime(1999, 7, 2, 0, 1),
                "measurements": [{"code": 2345, "metadata": {"visit_id": 9999}}],
            },
            {
                "time": datetime.datetime(1999, 7, 2, 12),
                "measurements": [{"code": 4567, "metadata": {"visit_id": 9999}}],
            },
        ],
    }

    cleanup(patient)
    cleanup(expected)

    assert move_visit_start_to_first_event_start(patient) == expected


def test_move_visit_start_doesnt_move_without_event() -> None:
    patient = {
        "patient_id": 123,
        "events": [
            {
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [
                    {"code": 1234, "metadata": {"visit_id": 9999}},
                    {"code": 3456, "metadata": {"visit_id": 9999, "table": "visit_occurrence"}},
                    {"code": 2345, "metadata": {"visit_id": 9999}},
                ],
            },
        ],
    }

    # None of the non-visit events have start time > '00:00:00' so visit event
    # start time is unchanged, though order changes based on code under resort.
    expected = {
        "patient_id": 123,
        "events": [
            {
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [
                    {"code": 1234, "metadata": {"visit_id": 9999}},
                    {"code": 3456, "metadata": {"visit_id": 9999, "table": "visit_occurrence"}},
                    {"code": 2345, "metadata": {"visit_id": 9999}},
                ],
            }
        ],
    }

    cleanup(patient)
    cleanup(expected)

    assert move_visit_start_to_first_event_start(patient) == expected


def test_move_to_day_end() -> None:
    patient = {
        "patient_id": 123,
        "events": [
            {"time": datetime.datetime(1999, 7, 2), "measurements": [{"code": 1234}]},
            {"time": datetime.datetime(1999, 7, 2, 12), "measurements": [{"code": 4321}]},
            {"time": datetime.datetime(1999, 7, 9), "measurements": [{"code": meds.birth_code}]},
        ],
    }

    expected = {
        "patient_id": 123,
        "events": [
            {"time": datetime.datetime(1999, 7, 2, 12), "measurements": [{"code": 4321}]},
            {"time": datetime.datetime(1999, 7, 2, 23, 59), "measurements": [{"code": 1234}]},
            {"time": datetime.datetime(1999, 7, 9, 23, 59), "measurements": [{"code": meds.birth_code}]},
        ],
    }

    cleanup(patient)
    cleanup(expected)

    print(move_to_day_end(patient))
    print(expected)

    assert move_to_day_end(patient) == expected


def test_remove_nones() -> None:
    patient = {
        "patient_id": 123,
        "events": [
            {"time": datetime.datetime(1999, 7, 2), "measurements": [{"code": 1234}]},  # No value, to be removed
            {"time": datetime.datetime(1999, 7, 2, 12), "measurements": [{"code": 1234, "numeric_value": 3}]},
            {"time": datetime.datetime(1999, 7, 9), "measurements": [{"code": meds.birth_code}]},
        ],
    }

    expected = {
        "patient_id": 123,
        "events": [
            {"time": datetime.datetime(1999, 7, 2, 12), "measurements": [{"code": 1234, "numeric_value": 3}]},
            {"time": datetime.datetime(1999, 7, 9), "measurements": [{"code": meds.birth_code}]},
        ],
    }

    cleanup(patient)
    cleanup(expected)

    assert remove_nones(patient) == expected


def test_delta_encode() -> None:
    patient = {
        "patient_id": 123,
        "events": [
            {"time": datetime.datetime(1999, 7, 2), "measurements": [{"code": 1234}]},
            {"time": datetime.datetime(1999, 7, 2), "measurements": [{"code": 1234}]},
            {"time": datetime.datetime(1999, 7, 2, 12), "measurements": [{"code": 1234, "numeric_value": 3}]},
            {"time": datetime.datetime(1999, 7, 2, 14), "measurements": [{"code": 1234, "numeric_value": 3}]},
            {"time": datetime.datetime(1999, 7, 2, 19), "measurements": [{"code": 1234, "numeric_value": 5}]},
            {"time": datetime.datetime(1999, 7, 2, 20), "measurements": [{"code": 1234, "numeric_value": 3}]},
        ],
    }

    expected = {
        "patient_id": 123,
        "events": [
            {"time": datetime.datetime(1999, 7, 2), "measurements": [{"code": 1234}]},
            {"time": datetime.datetime(1999, 7, 2, 12), "measurements": [{"code": 1234, "numeric_value": 3}]},
            {"time": datetime.datetime(1999, 7, 2, 19), "measurements": [{"code": 1234, "numeric_value": 5}]},
            {"time": datetime.datetime(1999, 7, 2, 20), "measurements": [{"code": 1234, "numeric_value": 3}]},
        ],
    }

    cleanup(patient)
    cleanup(expected)

    assert delta_encode(patient) == expected


def test_move_billing_codes() -> None:
    patient = {
        "patient_id": 123,
        "events": [
            {
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [
                    {
                        "code": 1234,
                        "metadata": {
                            "end": datetime.datetime(1999, 7, 20),
                            "visit_id": 10,
                            "clarity_table": "lpch_pat_enc",
                        },
                    }
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 9),
                "measurements": [
                    {
                        "code": meds.birth_code,
                        "metadata": {
                            "visit_id": 10,
                            "clarity_table": "lpch_pat_enc_dx",
                        },
                    }
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 10),
                "measurements": [
                    {
                        "code": 42165,
                        "metadata": {
                            "visit_id": 10,
                            "clarity_table": "shc_pat_enc_dx",
                        },
                    }
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 11),
                "measurements": [
                    {
                        "code": 12345,
                        "metadata": {
                            "visit_id": 10,
                        },
                    }
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 13),
                "measurements": [
                    {
                        "code": 123,
                        "metadata": {
                            "visit_id": 11,
                        },
                    }
                ],
            },
        ],
    }

    expected = {
        "patient_id": 123,
        "events": [
            {
                "time": datetime.datetime(1999, 7, 2),
                "measurements": [
                    {
                        "code": 1234,
                        "metadata": {
                            "end": datetime.datetime(1999, 7, 20),
                            "visit_id": 10,
                            "clarity_table": "lpch_pat_enc",
                        },
                    }
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 11),
                "measurements": [
                    {
                        "code": 12345,
                        "metadata": {
                            "visit_id": 10,
                        },
                    }
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 13),
                "measurements": [
                    {
                        "code": 123,
                        "metadata": {
                            "visit_id": 11,
                        },
                    }
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 20),
                "measurements": [
                    {
                        "code": meds.birth_code,
                        "metadata": {
                            "visit_id": 10,
                            "clarity_table": "lpch_pat_enc_dx",
                        },
                    }
                ],
            },
            {
                "time": datetime.datetime(1999, 7, 20),
                "measurements": [
                    {
                        "code": 42165,
                        "metadata": {
                            "visit_id": 10,
                            "clarity_table": "shc_pat_enc_dx",
                        },
                    }
                ],
            },
        ],
    }

    cleanup(patient)
    cleanup(expected)

    print(move_billing_codes(patient))
    print(expected)

    assert move_billing_codes(patient) == expected
