from utils_ak.fluid_flow import *
from utils_ak.block_tree import *

import warnings

warnings.filterwarnings("ignore")


def test_water_line_flow_1():
    drenator = Container("Drenator", value=1000, max_pressures=[None, None])

    melting1 = Processor(
        "Melting1",
        items=["a", "a"],
        max_pressures=[1000, 1000],
        processing_time=0,
        limits=[750, 750],
    )
    melting2 = Processor(
        "Melting2",
        items=["b", "b"],
        max_pressures=[1000, 1000],
        processing_time=0,
        limits=[250, 250],
    )
    melting_queue = Queue(
        "MeltingQueue", [melting1, melting2], break_funcs={"in": lambda old, new: 1}
    )

    cooling1 = Processor(
        "Cooling1",
        items=["a", "a"],
        max_pressures=[None, None],
        processing_time=0.5,
        limits=[750, 750],
    )
    cooling2 = Processor(
        "Cooling2",
        items=["b", "b"],
        max_pressures=[None, None],
        processing_time=0.5,
        limits=[250, 250],
    )
    cooling_queue = Queue("CoolingQueue", [cooling1, cooling2])

    packing_hub = Hub("Hub")

    packing1 = Processor(
        "Packing1",
        items=["a", "a1"],
        max_pressures=[200, None],
        processing_time=0,
        limits=[750, None],
    )
    packing2 = Processor(
        "Packing2",
        items=["b", "b1"],
        max_pressures=[400, None],
        processing_time=0,
        limits=[100, None],
    )
    packing_queue1 = Queue("PackingQueue1", [packing1, packing2])

    # using sequence as packing
    container = Container(
        "Container3", item="b", max_pressures=[200, None], limits=[150, None]
    )
    # slow packing...
    processor = Processor(
        "Processor3", items=["b", "b2"], max_pressures=[50, None], processing_time=0
    )
    packing3 = Sequence("Packing3", [container, processor])
    packing_queue2 = Queue("PackingQueue1", [packing3])

    pipe_connect(drenator, melting_queue, "drenator-melting")
    pipe_connect(melting_queue, cooling_queue, "melting-cooling")
    pipe_connect(cooling_queue, packing_hub, "cooling-hub")

    pipe_connect(packing_hub, packing_queue1, "hub-packing_queue1")
    pipe_connect(packing_hub, packing_queue2, "hub-packing_queue2")

    flow = FluidFlow(drenator, verbose=True)
    run_flow(flow)

    maker, make = init_block_maker("root", axis=1)

    for node in drenator.iterate("down"):
        if node.active_periods():
            for period in node.active_periods():
                label = "-".join([str(node.name), period[0]])
                beg, end = period[1:]
                beg, end = (
                    custom_round(beg * 60, 5) // 5,
                    custom_round(end * 60, 5) // 5,
                )
                make(label, x=[beg, 0], size=(end - beg, 1))
    print(maker.root.tabular())


def test_water_line_flow_2_multiple_boilings():
    n_boilings = 2
    packing1_1 = Processor(
        "Packing1",
        items=["a", "a1"],
        max_pressures=[200, None],
        processing_time=0,
        limits=[750, None],
    )
    packing2_1 = Processor(
        "Packing2",
        items=["b", "b1"],
        max_pressures=[400, None],
        processing_time=0,
        limits=[100, None],
    )
    packing1_2 = Processor(
        "Packing1",
        items=["a", "a1"],
        max_pressures=[200, None],
        processing_time=0,
        limits=[750, None],
    )
    packing2_2 = Processor(
        "Packing2",
        items=["b", "b1"],
        max_pressures=[400, None],
        processing_time=0,
        limits=[100, None],
    )
    packing_queue1 = Queue(
        "PackingQueue1", [packing1_1, packing2_1, packing1_2, packing2_2]
    )

    packing3_1 = Processor(
        "Packing3",
        items=["b", "b2"],
        max_pressures=[200, None],
        processing_time=0,
        limits=[150, None],
    )
    packing3_2 = Processor(
        "Packing3",
        items=["b", "b2"],
        max_pressures=[200, None],
        processing_time=0,
        limits=[150, None],
    )
    packing_queue2 = Queue("PackingQueue1", [packing3_1, packing3_2])

    for _ in range(n_boilings):
        drenator = Container("Drenator", value=1000, max_pressures=[None, None])

        melting1 = Processor(
            "Melting1",
            items=["a", "a"],
            max_pressures=[1000, 1000],
            processing_time=0,
            limits=[750, 750],
        )
        melting2 = Processor(
            "Melting2",
            items=["b", "b"],
            max_pressures=[1000, 1000],
            processing_time=0,
            limits=[250, 250],
        )
        melting_queue = Queue(
            "MeltingQueue", [melting1, melting2], break_funcs={"in": lambda old, new: 1}
        )

        cooling1 = Processor(
            "Cooling1",
            items=["a", "a"],
            max_pressures=[None, None],
            processing_time=0.5,
            limits=[750, 750],
        )
        cooling2 = Processor(
            "Cooling2",
            items=["b", "b"],
            max_pressures=[None, None],
            processing_time=0.5,
            limits=[250, 250],
        )
        cooling_queue = Queue("CoolingQueue", [cooling1, cooling2])

        packing_hub = Hub("Hub")

        pipe_connect(drenator, melting_queue, "drenator-melting")
        pipe_connect(melting_queue, cooling_queue, "melting-cooling")
        pipe_connect(cooling_queue, packing_hub, "cooling-hub")

        pipe_connect(packing_hub, packing_queue1, "hub-packing_queue1")
        pipe_connect(packing_hub, packing_queue2, "hub-packing_queue2")

        flow = FluidFlow(drenator, verbose=True)
        run_flow(flow)

        maker, make = init_block_maker("root", axis=1)

        for node in drenator.iterate("down"):
            if node.active_periods():
                for period in node.active_periods():
                    label = "-".join([str(node.name), period[0]])
                    beg, end = period[1:]
                    beg, end = (
                        custom_round(beg * 60, 5) // 5,
                        custom_round(end * 60, 5) // 5,
                    )
                    make(label, x=[beg, 0], size=(end - beg, 1))
        print(maker.root.tabular())

        # clean-up
        for node in drenator.iterate("down"):
            node.reset()

        # remove packers from current boiling
        pipe_disconnect(packing_hub, packing_queue1)
        pipe_disconnect(packing_hub, packing_queue2)


if __name__ == "__main__":
    test_water_line_flow_1()
    # test_water_line_flow_2_multiple_boilings()
