test_timeline.py

# Copyright (C) 2024, UChicago Argonne, LLC
# Licensed under the 3-clause BSD license.  See accompanying LICENSE.txt file
# in the top-level directory.

"""
Unit tests for timeline.
"""
import cProfile
import datetime
import pstats
from _decimal import Underflow, Overflow
from pprint import pformat
import random
from pstats import SortKey
from typing import List

import numpy as np
import numpy.testing as npt
import pandas as pd
from dateutil.parser import parse as date_parse
from pyspark.sql import SparkSession, DataFrame
import pyspark.sql.functions as F
from pyspark.sql.types import StructType, StructField, TimestampType, StringType

from Ocean.schema_lookup import schema_bitmask
from Octeres.bitmask import BASE_DTYPE_MAX
from Octeres.bitmask import eb
from Octeres.bitmask.bitmask_lite import LiteBitmask, LiteBitmaskSlots
from Octeres.data_generation import EventGeneration
from Octeres.forthwith import FORMAT_DATE_DAY
from Octeres.timeline import Dependency_Funcs, TLCollisionOverlap, TimelinePit, TimelineParallel
from Octeres.timeline import Event_Handler
from Octeres.timeline import (
    Timeline,
    sum_array,
    TLEvent,
    TLPointInTime,
    Event_Direction,
)
from Octeres.timeline import TimelineDict, TimelineMask
from Octeres.timeline import reduce_holes, Point_in_Time
from Octeres.util import superprint, df_to_lstofdct
import pytest

try:
    import dask
except ImportError:
    dask = None

pd.set_option("display.max_rows", None)
pd.set_option("display.max_columns", None)
pd.set_option("display.width", 512)
pd.set_option("display.max_colwidth", None)
pd.set_option("display.float_format", "{:.6f}".format)


class Test_Timeline:
    @classmethod
    def setup_class(cls):
        machine_name = "test"
        range_start = date_parse("2015-01-01")
        range_end = date_parse("2015-02-15")
        base_mask = eb.zeros(8)
        cls.tl = Timeline(machine_name, range_start, range_end, base_mask)

    def test_possible_unit_seconds(self):
        tl = self.tl
        possible_unit_seconds = tl.possible_unit_seconds
        correct_unit_seconds = (tl.range_end - tl.range_start).total_seconds() * len(tl.base_mask)
        assert possible_unit_seconds == correct_unit_seconds

    #     def test_prepare_event_timeline_00(self):
    #         #TODO
    #         pass
    #
    #     def test_group_timeline_00(self):
    #         #TODO
    #         pass
    #
    #     def test_normalize_timeline_00(self):
    #         #TODO
    #         pass
    #
    #     def test_mask_timeline_to_mag_timeline_00(self):
    #         #TODO
    #         pass
    #

    def test_search_timeline_for_holes_00(self):
        mask_timeline = TimelineMask()
        mask_timeline.append((date_parse("2017-01-01"), eb.ones(8)))
        mask_timeline.append((date_parse("2017-01-02"), eb.ones(8)))
        mask_timeline.append((date_parse("2017-01-03"), eb.array([1, 0, 1, 1, 1, 1, 1, 1])))
        mask_timeline.append((date_parse("2017-01-04"), eb.ones(8)))

        holes = self.tl.search_timeline_for_holes(mask_timeline)

        assert len(holes) == 1
        correct_ts = date_parse("2017-01-03")
        correct_te = date_parse("2017-01-04")
        correct_delta = eb.array([0, 1, 0, 0, 0, 0, 0, 0])
        hole_ts = holes[0][0]
        hole_te = holes[0][1]
        hole_delta = holes[0][2]
        superprint(correct_delta)
        superprint(hole_delta)
        assert correct_ts == hole_ts
        assert correct_te == hole_te
        assert (correct_delta == hole_delta).all()

    def test_search_timeline_for_holes_01(self):
        mask_timeline = TimelineMask()
        mask_timeline.append((date_parse("2017-01-01"), eb.array([1, 0, 1, 1, 1, 1, 1, 1])))
        mask_timeline.append((date_parse("2017-01-02"), eb.array([1, 0, 1, 1, 1, 1, 1, 1])))
        mask_timeline.append((date_parse("2017-01-03"), eb.array([1, 0, 1, 1, 1, 1, 1, 1])))
        mask_timeline.append((date_parse("2017-01-04"), eb.array([1, 0, 1, 1, 1, 1, 1, 1])))
        mask_timeline.append((date_parse("2017-01-05"), eb.array([1, 0, 1, 1, 1, 1, 1, 1])))
        mask_timeline.append((date_parse("2017-01-06"), eb.array([1, 0, 1, 1, 1, 1, 1, 1])))

        holes = self.tl.search_timeline_for_holes(mask_timeline)

        assert len(holes) == 1
        correct_delta = eb.array([0, 1, 0, 0, 0, 0, 0, 0])
        ts, te, hole_delta = holes[0]
        assert ts == date_parse("2017-01-01")
        assert te == date_parse("2017-01-06")
        assert (correct_delta == hole_delta).all()

    #     def test_calculate_area_00(self):
    #         #TODO
    #         pass
    def test_collision_detection_00(self):
        base_mask = self.tl.base_mask
        tl = self.tl

        event_lst = list()
        event_mask = base_mask.copy()
        event_mask[0:3] = 1
        dct = dict(
            pk="an_event",
            event_type_name="job",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 00:00:00"),
            time_end=date_parse("2015-02-04 00:00:00"),
        )
        event_lst.append(dct)
        event_mask = base_mask.copy()
        event_mask[1] = 1
        dct = dict(
            pk="an_event2",
            event_type_name="job",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 01:00:00"),
            time_end=date_parse("2015-02-03 02:00:00"),
        )
        event_lst.append(dct)
        timeline_lst = tl.prepare_event_timeline(event_lst)

        timeline_dct = tl.group_timeline(timeline_lst)
        timeline_sorted = tl.sort_timeline(timeline_dct)
        collisions = tl.find_collisions_timeline(timeline_sorted)
        for collision in collisions:
            superprint(collision)
        assert len(collisions) == 1

    def test_collision_detection_01(self):
        # https://pandas.pydata.org/pandas-docs/version/0.21.1/generated/pandas.Timestamp.to_datetime.html
        base_mask = self.tl.base_mask
        tl = self.tl

        event_lst = list()
        event_mask = base_mask.copy()
        event_mask[0:3] = 1
        dct: TLEvent = dict(
            pk="an_event",
            event_type_name="job",
            bitmask=event_mask,
            time_start=pd.Timestamp(date_parse("2015-02-03 00:00:00")),
            time_end=pd.Timestamp(date_parse("2015-02-04 00:00:00")),
        )
        event_lst.append(dct)
        event_mask = base_mask.copy()
        event_mask[1] = 1
        dct: TLEvent = dict(
            pk="an_event2",
            event_type_name="job",
            bitmask=event_mask,
            time_start=pd.Timestamp(date_parse("2015-02-03 01:00:00")),
            time_end=pd.Timestamp(date_parse("2015-02-03 02:00:00")),
        )
        event_lst.append(dct)
        timeline_lst: List[TLPointInTime] = tl.prepare_event_timeline(event_lst)

        timeline_dct: TimelineDict = tl.group_timeline(timeline_lst)
        timeline_sorted: TimelinePit = tl.sort_timeline(timeline_dct)
        tl.print_timeline(timeline_sorted)
        collisions: List[TLCollisionOverlap] = tl.find_collisions_timeline(timeline_sorted)
        for collision in collisions:
            superprint(collision)
        assert len(collisions) == 1

    def test_full_stack(self):
        base_mask = self.tl.base_mask
        tl = self.tl

        possible_unit_seconds = tl.possible_unit_seconds
        event_handler = Event_Handler()

        event_lst = list()
        event_mask = base_mask.copy()
        event_mask[0:3] = 1
        event_dct = event_handler.get_event(
            event_type_name="job",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03"),
            time_end=date_parse("2015-02-04"),
        )
        event_lst.append(event_dct)
        correct_unit_seconds = (event_dct["time_end"] - event_dct["time_start"]).total_seconds() * sum_array(event_mask)
        correct_possible_unit_seconds = (tl.range_end - tl.range_start).total_seconds() * len(tl.base_mask)

        timeline_lst = tl.prepare_event_timeline(event_lst)

        timeline_dct = tl.group_timeline(timeline_lst)
        timeline_sorted = tl.sort_timeline(timeline_dct)

        # mask_timeline = tl.normalize_timeline(timeline_sorted, test_negative=True)
        mask_timeline = tl.normalize_timeline_pit(timeline_sorted)
        mask_timeline = tl.sum_timeline(mask_timeline, test_negative=False)
        tl.print_timeline(mask_timeline, binary=True)
        mask_timeline = tl.flatten_timeline(mask_timeline)

        mask_timeline = tl.isolate_timeline_range(mask_timeline)

        # tl.print_timeline(mask_timeline)
        mag_timeline = tl.mask_timeline_to_mag_timeline(mask_timeline)

        consumed_unit_seconds = tl.calculate_area(mag_timeline)

        # return consumed_unit_seconds, possible_unit_seconds, mag_timeline
        assert round(abs(consumed_unit_seconds - correct_unit_seconds), 4) == 0
        assert possible_unit_seconds == correct_possible_unit_seconds

    def test_find_event_dependencies_00(self):
        base_mask = self.tl.base_mask
        tl = self.tl
        event_handler = Event_Handler()
        event_lst = list()
        event_mask = base_mask.copy()
        event_mask[0] = 1
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 00:00:00"),
            time_end=date_parse("2015-02-04 00:00:00"),
        )
        event_a_pk = event_dct["pk"]
        event_lst.append(event_dct)
        # joint events in time VVV ^^^
        event_mask = base_mask.copy()
        event_mask[0] = 1
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 12:00:00"),
            time_end=date_parse("2015-02-05 00:00:00"),
        )
        event_b_pk = event_dct["pk"]
        event_lst.append(event_dct)
        dependency_functions = list()
        dependency_functions.append(Dependency_Funcs.dep_time)
        dependency_functions.append(Dependency_Funcs.dep_space_all)
        event_dependancies = tl.find_event_dependencies(event_lst, dependency_functions)
        assert len(event_dependancies) == 2
        grouped_dependancies = tl.merge_event_dependencies(event_dependancies)
        assert len(grouped_dependancies) == 2
        assert event_a_pk in grouped_dependancies
        assert event_b_pk in grouped_dependancies[event_a_pk]
        assert event_b_pk in grouped_dependancies
        assert event_a_pk in grouped_dependancies[event_b_pk]

    def test_find_event_dependencies_01(self):
        base_mask = self.tl.base_mask
        tl = self.tl
        event_handler = Event_Handler()
        event_lst = list()
        event_mask = base_mask.copy()
        event_mask[0] = 2
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 00:00:00"),
            time_end=date_parse("2015-02-04 00:00:00"),
        )
        event_a_pk = event_dct["pk"]
        event_lst.append(event_dct)
        # joint events in time VVV ^^^
        event_mask = base_mask.copy()
        event_mask[0] = 1
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 12:00:00"),
            time_end=date_parse("2015-02-05 00:00:00"),
        )
        event_b_pk = event_dct["pk"]
        event_lst.append(event_dct)
        dependency_functions = list()
        dependency_functions.append(Dependency_Funcs.dep_time)
        dependency_functions.append(Dependency_Funcs.dep_space_all)
        event_dependancies = tl.find_event_dependencies(event_lst, dependency_functions)
        assert len(event_dependancies) == 2
        grouped_dependancies = tl.merge_event_dependencies(event_dependancies)
        assert len(grouped_dependancies) == 2
        assert event_a_pk in grouped_dependancies
        assert event_b_pk in grouped_dependancies[event_a_pk]
        assert event_b_pk in grouped_dependancies
        assert event_a_pk in grouped_dependancies[event_b_pk]

    def test_find_event_dependencies_02(self):
        base_mask = self.tl.base_mask
        tl = self.tl
        event_handler = Event_Handler()
        event_lst = list()
        event_mask = base_mask.copy()
        event_mask[0] = -1
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 00:00:00"),
            time_end=date_parse("2015-02-04 00:00:00"),
        )
        event_a_pk = event_dct["pk"]
        event_lst.append(event_dct)
        # joint events in time VVV ^^^
        event_mask = base_mask.copy()
        event_mask[0] = 1
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 12:00:00"),
            time_end=date_parse("2015-02-05 00:00:00"),
        )
        event_b_pk = event_dct["pk"]
        event_lst.append(event_dct)
        dependency_functions = list()
        dependency_functions.append(Dependency_Funcs.dep_time)
        dependency_functions.append(Dependency_Funcs.dep_space_all)
        event_dependancies = tl.find_event_dependencies(event_lst, dependency_functions)
        assert len(event_dependancies) == 2
        grouped_dependancies = tl.merge_event_dependencies(event_dependancies)
        assert len(grouped_dependancies) == 2
        assert event_a_pk in grouped_dependancies
        assert event_b_pk in grouped_dependancies[event_a_pk]
        assert event_b_pk in grouped_dependancies
        assert event_a_pk in grouped_dependancies[event_b_pk]

    def test_find_event_dependencies_03(self):
        base_mask = self.tl.base_mask
        tl = self.tl

        event_handler = Event_Handler()

        event_lst = list()

        event_mask = base_mask.copy()
        event_mask[0] = 1
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 00:00:00"),
            time_end=date_parse("2015-02-04 00:00:00"),
        )
        event_a_pk = event_dct["pk"]
        event_lst.append(event_dct)

        # disjoin events in time VVV ^^^

        event_mask = base_mask.copy()
        event_mask[0] = 1
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-05 00:00:00"),
            time_end=date_parse("2015-02-06 00:00:00"),
        )
        event_b_pk = event_dct["pk"]
        event_lst.append(event_dct)

        # joining event, but larger mask VVV

        event_mask = base_mask.copy()
        event_mask[0:2] = 1
        event_dct = event_handler.get_event(
            event_type_name="event",
            bitmask=event_mask,
            time_start=date_parse("2015-02-03 12:00:00"),
            time_end=date_parse("2015-02-05 12:00:00"),
        )
        event_c_pk = event_dct["pk"]
        event_lst.append(event_dct)

        dependency_functions = list()
        dependency_functions.append(Dependency_Funcs.dep_time)
        dependency_functions.append(Dependency_Funcs.dep_space_all)
        # we have time_start, time_end, bitmask, and event_type_name

        event_dependancies = tl.find_event_dependencies(event_lst, dependency_functions)
        assert len(event_dependancies) == 0

        dependency_functions = list()
        dependency_functions.append(Dependency_Funcs.dep_time)
        event_dependancies = tl.find_event_dependencies(event_lst, dependency_functions)
        superprint(pformat(event_dependancies))
        assert len(event_dependancies) == 4

        tl.generate_dependency_graph(
            event_handler.nodes,
            event_dependancies,
            filename="test_find_event_dependencies_01",
        )

        grouped_dependancies = tl.merge_event_dependencies(event_dependancies)
        assert len(grouped_dependancies) == 3
        superprint(pformat(grouped_dependancies))

        assert event_a_pk in grouped_dependancies
        assert event_b_pk in grouped_dependancies
        assert event_c_pk in grouped_dependancies


def test_reduce_holes_00():
    """ordered timeline all ts butted up against each other"""
    timeline_holes = list()
    timeline_holes.append((1, 2, np.array([1, 0, 1, 1])))
    timeline_holes.append((2, 3, np.array([1, 0, 1, 1])))
    timeline_holes.append((3, 4, np.array([1, 0, 0, 1])))
    timeline_holes.append((4, 5, np.array([1, 1, 1, 1])))
    timeline_holes.append((5, 6, np.array([1, 1, 1, 1])))
    timeline_holes.append((6, 7, np.array([1, 1, 1, 1])))

    correct_holes = list()
    correct_holes.append((1, 3, np.array([1, 0, 1, 1])))
    correct_holes.append((3, 4, np.array([1, 0, 0, 1])))
    correct_holes.append((4, 7, np.array([1, 1, 1, 1])))

    reductions, timeline_holes = reduce_holes(timeline_holes)
    superprint("\n", pformat(timeline_holes))
    for idx, hole in enumerate(timeline_holes):
        ts, te, mask = hole
        cts, cte, cmask = correct_holes[idx]
        assert ts == cts
        assert te == cte
        npt.assert_equal(mask, cmask)
    assert reductions == 3


def test_reduce_holes_01():
    """ordered timeline all ts butted up against each other"""
    timeline_holes = list()

    timeline_holes.append((0, 1, np.array([1, 1, 1, 1])))
    timeline_holes.append((1, 2, np.array([1, 0, 1, 1])))
    timeline_holes.append((2, 3, np.array([1, 0, 1, 1])))
    timeline_holes.append((3, 4, np.array([1, 0, 0, 1])))
    timeline_holes.append((4, 5, np.array([1, 1, 1, 1])))
    timeline_holes.append((5, 6, np.array([1, 1, 1, 1])))
    timeline_holes.append((6, 7, np.array([1, 1, 1, 1])))

    correct_holes = list()
    correct_holes.append((0, 1, np.array([1, 1, 1, 1])))
    correct_holes.append((1, 3, np.array([1, 0, 1, 1])))
    correct_holes.append((3, 4, np.array([1, 0, 0, 1])))
    correct_holes.append((4, 7, np.array([1, 1, 1, 1])))

    reductions, timeline_holes = reduce_holes(timeline_holes)
    superprint("\n", pformat(timeline_holes))
    for idx, hole in enumerate(timeline_holes):
        ts, te, mask = hole
        cts, cte, cmask = correct_holes[idx]
        assert ts == cts
        assert te == cte
        npt.assert_equal(mask, cmask)
    assert reductions == 3


def test_reduce_holes_02():
    """ordered timeline all ts separated a bit"""
    timeline_holes = list()

    timeline_holes.append((0, 1, np.array([1, 1, 1, 1])))
    timeline_holes.append((1, 2, np.array([1, 0, 1, 1])))
    timeline_holes.append((3, 4, np.array([1, 0, 0, 1])))
    timeline_holes.append((4, 5, np.array([1, 1, 1, 1])))
    timeline_holes.append((6, 7, np.array([1, 1, 1, 1])))

    correct_holes = list()
    correct_holes.append((0, 1, np.array([1, 1, 1, 1])))
    correct_holes.append((1, 2, np.array([1, 0, 1, 1])))
    correct_holes.append((3, 4, np.array([1, 0, 0, 1])))
    correct_holes.append((4, 5, np.array([1, 1, 1, 1])))
    correct_holes.append((6, 7, np.array([1, 1, 1, 1])))

    reductions, timeline_holes = reduce_holes(timeline_holes)
    superprint("\n", pformat(timeline_holes))
    for idx, hole in enumerate(timeline_holes):
        ts, te, mask = hole
        cts, cte, cmask = correct_holes[idx]
        assert ts == cts
        assert te == cte
        npt.assert_equal(mask, cmask)
    assert reductions == 0


def test_reduce_holes_03():
    """ordered timeline all ts separated a bit"""
    timeline_holes = list()

    timeline_holes.append((0, 1, np.array([1, 1, 1, 1])))
    timeline_holes.append((1, 2, np.array([1, 0, 1, 1])))
    timeline_holes.append((3, 4, np.array([1, 0, 0, 1])))
    timeline_holes.append((4, 6, np.array([1, 1, 1, 1])))
    timeline_holes.append((6, 7, np.array([1, 1, 1, 1])))

    correct_holes = list()
    correct_holes.append((0, 1, np.array([1, 1, 1, 1])))
    correct_holes.append((1, 2, np.array([1, 0, 1, 1])))
    correct_holes.append((3, 4, np.array([1, 0, 0, 1])))
    correct_holes.append((4, 7, np.array([1, 1, 1, 1])))

    reductions, timeline_holes = reduce_holes(timeline_holes)
    superprint("\n", pformat(timeline_holes))
    for idx, hole in enumerate(timeline_holes):
        ts, te, mask = hole
        cts, cte, cmask = correct_holes[idx]
        assert ts == cts
        assert te == cte
        npt.assert_equal(mask, cmask)
    assert reductions == 1


def test_tlevent_00():
    event: TLEvent = dict(
        pk="1",
        event_type_name="fish",
        bitmask=eb.zeros(10),
        time_start=date_parse("2020-01-01"),
        time_end=date_parse("2020-01-02"),
    )
    superprint(type(event), event)
    assert type(event) == dict


def test_tlevent_01():
    event = TLEvent(
        pk="1",
        event_type_name="fish",
        bitmask=eb.zeros(10),
        time_start=date_parse("2020-01-01"),
        time_end=date_parse("2020-01-02"),
    )
    superprint(type(event), event)


def test_tlevent_02():
    event = TLEvent(
        pk="1",
        event_type_name="fish",
        bitmask=eb.zeros(10),
        time_start=pd.Timestamp(date_parse("2020-01-01")),
        time_end=pd.Timestamp(date_parse("2020-01-02")),
    )
    superprint(type(event), event)


def test_tlpit_00():
    pit = TLPointInTime(
        pk="1",
        name="2",
        category="aaaa",
        bitmask=eb.zeros(10),
        ts=date_parse("2020-01-01"),
        direction=Event_Direction.POSITIVE,
    )
    superprint(pit)


def test_tlpit_01():
    pit = TLPointInTime(
        pk="1",
        name="2",
        category="aaaa",
        bitmask=eb.zeros(10),
        ts=date_parse("2020-01-01"),
        direction=Event_Direction.POSITIVE,
    )
    superprint(pit)


def test_tlpit_02():
    # noinspection PyTypeChecker
    pit: TLPointInTime = dict(
        pk="1",
        name="2",
        category="aaaa",
        bitmask=eb.zeros(10),
        ts=date_parse("2020-01-01"),
        direction=str(Event_Direction.POSITIVE),  # type: ignore
    )
    superprint(pit)


def test_get_mask_sum_00():
    bask_mask = eb.zeros(8)
    pit = Point_in_Time(bask_mask)

    bitmasks = [
        eb.array([0, 0, 0, 1, 1, 0, 0, 0]),
        eb.array([0, 0, 0, 1, 1, 0, 1, 0]),
        eb.array([0, 1, 1, 1, 1, 0, 0, 0]),
        eb.array([1, 0, 0, 1, 1, 0, 0, 0]),
    ]
    for idx, bitmask in enumerate(bitmasks):
        dct: TLPointInTime = TLPointInTime(
            pk=str(idx),
            name=str(idx),
            category="a",
            bitmask=bitmask,
            ts=date_parse("2020-01-01"),
            direction=Event_Direction.POSITIVE,
        )
        pit.positive_add(dct)
    value = pit.get_mask_sum()
    npt.assert_array_equal(value, eb.array([1, 1, 1, 4, 4, 0, 1, 0]))


def test_get_mask_sum_01():
    bask_mask = eb.zeros(8)
    pit = Point_in_Time(bask_mask)

    bitmasks = [
        eb.array([0, 0, 0, 1, 1, 0, 0, 0]),
        eb.array([0, 0, 0, 1, 1, 0, 1, 0]),
        eb.array([0, 1, 1, 1, 1, 0, 0, 0]),
        eb.array([1, 0, 0, 1, 1, 0, 0, 0]),
    ]
    for idx, bitmask in enumerate(bitmasks):
        dct: TLPointInTime = TLPointInTime(
            pk=str(idx),
            name=str(idx),
            category="a",
            bitmask=bitmask,
            ts=date_parse("2020-01-01"),
            direction=Event_Direction.POSITIVE,
        )
        pit.positive_add(dct)
    value = pit.get_mask_sum()
    npt.assert_array_equal(value, eb.array([1, 1, 1, 4, 4, 0, 1, 0]))

    bitmasks = [
        eb.array([0, 0, 1, 1, 0, 0, 0, 0]),
    ]
    for idx, bitmask in enumerate(bitmasks):
        dct: TLPointInTime = TLPointInTime(
            pk=str(idx),
            name=str(idx),
            category="a",
            bitmask=bitmask,
            ts=date_parse("2020-01-01"),
            direction=Event_Direction.NEGATIVE,
        )
        pit.negative_add(dct)
    value = pit.get_mask_sum()
    npt.assert_array_equal(value, eb.array([1, 1, 0, 3, 4, 0, 1, 0]))


def test_get_mask_sum_02():
    bask_mask = eb.zeros(8)
    pit = Point_in_Time(bask_mask)

    bitmasks = [
        eb.array([0, 0, 0, 1, 1, 0, 0, 0]),
        eb.array([0, 0, 0, 1, 1, 0, 1, 0]),
        eb.array([0, 1, 1, 1, 1, 0, 0, 0]),
        eb.array([1, 0, 0, 1, 1, 0, 0, 0]),
    ]
    for idx, bitmask in enumerate(bitmasks):
        dct: TLPointInTime = TLPointInTime(
            pk=str(idx),
            name=str(idx),
            category="a",
            bitmask=bitmask,
            ts=date_parse("2020-01-01"),
            direction=Event_Direction.POSITIVE,
        )
        pit.positive_add(dct)
    value = pit.get_mask_sum()
    npt.assert_array_equal(value, eb.array([1, 1, 1, 4, 4, 0, 1, 0]))

    bitmasks = [
        eb.array([0, 0, 1, 1, 0, 0, 0, 0]),
        eb.array([0, 0, 0, 0, 0, 0, 1, 0]),
    ]
    for idx, bitmask in enumerate(bitmasks):
        dct: TLPointInTime = TLPointInTime(
            pk=str(idx),
            name=str(idx),
            category="a",
            bitmask=bitmask,
            ts=date_parse("2020-01-01"),
            direction=Event_Direction.NEGATIVE,
        )
        pit.negative_add(dct)
    value = pit.get_mask_sum()
    npt.assert_array_equal(value, eb.array([1, 1, 0, 3, 4, 0, 0, 0]))


def test_get_mask_sum_03():
    bask_mask = eb.zeros(8)
    pit = Point_in_Time(bask_mask)

    bitmasks = [
        eb.array([0, 0, 1, 1, 0, 0, 0, 0]),
        eb.array([0, 0, 0, 0, 0, 0, 1, 0]),
    ]
    for idx, bitmask in enumerate(bitmasks):
        dct: TLPointInTime = TLPointInTime(
            pk=str(idx),
            name=str(idx),
            category="a",
            bitmask=bitmask,
            ts=date_parse("2020-01-01"),
            direction=Event_Direction.NEGATIVE,
        )
        pit.negative_add(dct)
    value = pit.get_mask_sum()
    ecorr = eb.array([0, 0, -1, -1, 0, 0, -1, 0])
    npt.assert_array_equal(value, ecorr)


def test_get_mask_sum_03b():
    # underflow
    bask_mask = eb.zeros(8)
    pit = Point_in_Time(bask_mask)
    for idx in range(BASE_DTYPE_MAX * 2):  # negative needs two more, 127 to -128
        bitmask = eb.array([0, 0, -1, -1, 0, 0, 0, 0])
        dct: TLPointInTime = TLPointInTime(
            pk=str(idx),
            name=str(idx),
            category="a",
            bitmask=bitmask,
            ts=date_parse("2020-01-01"),
            direction=Event_Direction.NEGATIVE,
        )
        pit.negative_add(dct)
    with pytest.raises(Underflow):
        value = pit.get_mask_sum()


def test_get_mask_sum_03c():
    # underflow
    bask_mask = eb.zeros(8)
    pit = Point_in_Time(bask_mask)
    for idx in range(BASE_DTYPE_MAX * 2):  # negative needs two more, 127 to -128
        bitmask = eb.array([0, 0, 1, 1, 0, 0, 0, 0])
        dct: TLPointInTime = TLPointInTime(
            pk=str(idx),
            name=str(idx),
            category="a",
            bitmask=bitmask,
            ts=date_parse("2020-01-01"),
            direction=Event_Direction.NEGATIVE,
        )
        pit.negative_add(dct)
    with pytest.raises(Overflow):
        value = pit.get_mask_sum()


@pytest.mark.skipif(dask is None, reason="could not import dask")
def test_the_gauntlet():
    # requires dask, dask[distributed]
    machine_name = "test_machine"
    dtype = "U2"
    empty_value = "  "
    test_value = ".."
    bit_total = 4000  # 100000
    bitmask_class = LiteBitmask
    time_seconds = 3600  # 86400
    range_start = datetime.datetime.strptime("2024-01-01", FORMAT_DATE_DAY)
    range_end = range_start + datetime.timedelta(seconds=time_seconds)
    eg = EventGeneration(
        time_seconds,
        bit_total,
        dtype=dtype,
        empty_value=empty_value,
        test_value=test_value,
        visualize=False,
        bitmask_class=bitmask_class,
        seed=42,
    )
    characters = []
    characters.extend(list(range(97, 122 + 1)))
    characters.extend(list(range(65, 90 + 1)))
    characters.extend(list(range(48, 57 + 1)))
    characters = list(map(chr, characters))
    event_names = eg.generate_names_n_level(characters, 3)
    events = eg.generate_non_overlapping_scattered_events_v0(event_names)
    event_lst = events.event_lst
    correct_unit_seconds = time_seconds * bit_total
    correct_possible_unit_seconds = (range_end - range_start).total_seconds() * bit_total

    import dask.dataframe as dd
    pdf = pd.DataFrame(event_lst)
    ddf = dd.from_pandas(pdf, npartitions=8)
    ddf = ddf.drop(['x', 'y'], axis=1)
    ddf['bitmask'] = ddf['bitmask'].apply(lambda b: b.to_dict(), meta=('bitmask', 'object'))
    ddf['time_start'] = ddf['ts'].apply(lambda tsi: pd.to_datetime(range_start + datetime.timedelta(seconds=tsi, microseconds=int(random.uniform(0, 1) * 1000000)), utc=False), meta=('ts', "object"))
    ddf['time_end'] = ddf['te'].apply(lambda tsi: pd.to_datetime(range_start + datetime.timedelta(seconds=tsi, microseconds=int(random.uniform(0, 1) * 1000000)), utc=False), meta=('te', 'object'))
    ddf['pk'] = ddf['name']
    ddf['event_type_name'] = 'job'
    pdf2 = ddf.compute()
    event_lst = df_to_lstofdct(pdf2)
    for event_dct in event_lst:
        event_dct['bitmask'] = LiteBitmask.from_dict(event_dct['bitmask'])

    # base_mask = bitmask_class.zeros(bit_total)
    # tl = Timeline(machine_name, range_start, range_end, base_mask)
    # possible_unit_seconds = tl.possible_unit_seconds
    # with cProfile.Profile() as pr:
    #     timeline_lst = tl.prepare_event_timeline(event_lst)
    #     timeline_dct = tl.group_timeline(timeline_lst)
    #     timeline_sorted = tl.sort_timeline(timeline_dct)
    #     mask_timeline = tl.normalize_timeline_pit(timeline_sorted)
    #     mask_timeline = tl.sum_timeline(mask_timeline, test_negative=False)
    #     # tl.print_timeline(mask_timeline, binary=True)
    #     mask_timeline = tl.flatten_timeline(mask_timeline)
    #     mask_timeline = tl.isolate_timeline_range(mask_timeline)
    #     # tl.print_timeline(mask_timeline)
    #     mag_timeline = tl.mask_timeline_to_mag_timeline(mask_timeline)
    #     consumed_unit_seconds = tl.calculate_area(mag_timeline)
    #     assert round(abs(consumed_unit_seconds - correct_unit_seconds), 4) == 0
    #     assert possible_unit_seconds == correct_possible_unit_seconds
    # profile_result = pstats.Stats(pr)
    # profile_result.sort_stats(SortKey.CUMULATIVE).print_stats(4)
    # print(f"{len(event_lst)=} {len(timeline_lst)=}")

    base_mask = bitmask_class.zeros(bit_total)
    tl = TimelineParallel(machine_name, range_start, range_end, base_mask, processes=16)
    possible_unit_seconds = tl.possible_unit_seconds
    with cProfile.Profile() as pr:
        tl.load(event_lst)
        tl.run()
        consumed_unit_seconds = tl.calculate_area()
        # assert round(abs(consumed_unit_seconds - correct_unit_seconds), 4) == 0
        # assert possible_unit_seconds == correct_possible_unit_seconds
    profile_result = pstats.Stats(pr)
    profile_result.sort_stats(SortKey.CUMULATIVE).print_stats(16)
    # print(f"{len(event_lst)=} {len(timeline_lst)=}")
    import ipdb; ipdb.set_trace()

    from dask.distributed import Client, Queue
    client = Client()
    spark: DataFrame = SparkSession.builder.appName("test_the_gauntlet").getOrCreate()
    spark.conf.set("spark.sql.session.timeZone", "UTC")
    schema = StructType(
        [
            StructField("name", StringType(), False),
            StructField("ts", TimestampType(), False),  # aurora_crayex_alerts, obtain from query
            StructField("te", TimestampType(), False),  # goes back to level1 -> table
            schema_bitmask,
        ])
    sdf = spark.createDataFrame(pdf2, schema=schema)  # if you don't provide a schema, the intervals in the bitmask won't cast right.
    sdf.show(32, truncate=False)
    pandas_partitions = Queue()

    def convert_to_pandas(iterator):
        for partition in iterator:
            print(partition)
            yield pd.DataFrame(list(partition))

    print(f"{sdf.rdd.getNumPartitions()=}")
    pdf3 = sdf.rdd.mapPartitions(convert_to_pandas).collect()

    def spark_to_multi_df(df: pd.DataFrame, address):
        with Client(address) as client:
            [future] = client.scatter([df])
            pandas_partitions.put(future)