from collections import OrderedDict
import numpy as np
from pytest import raises, warns
from ctapipe.utils.cutflow import CutFlow, UndefinedCut, PureCountingCut


def smaller2(x):
    return x < 2


def smaller3(x):
    return x < 3


def smaller5(x):
    return x < 5


def test_CutFlow():
    with warns(FutureWarning):
        flow = CutFlow("TestFlow")
        # set_cut and add_cut a aliases
        flow.set_cut("smaller5", smaller5)
        flow.add_cut("smaller3", lambda x: x < 3)

        for i in range(2, 6):
            flow.count("noCuts")
            # .keep counts if the function returns True,
            # i.e. when we "keep" the event
            if flow.keep("smaller5", i):
                # .cut counts if the function returns False,
                # i.e. when we do NOT "cut" the event
                if flow.cut("smaller3", i):
                    pass
                else:
                    # do something else that could fail or be rejected
                    try:
                        assert i == 3
                        flow.count("something")
                    except:
                        pass

        t = flow(sort_column=1)
        assert np.all(t["selected Events"] == [4, 3, 2, 1])

        with raises(UndefinedCut):
            flow.cut("undefined", 5)

        with raises(PureCountingCut):
            flow.cut("noCuts")


def test_set_cuts_clear():
    with warns(FutureWarning):
        flow = CutFlow("TestFlow")
        flow.set_cut("smaller5", smaller5)

        flow.set_cuts(
            OrderedDict([("smaller3", smaller3), ("smaller2", smaller2)]), clear=True
        )

        assert flow.cuts == OrderedDict(
            [("smaller3", [smaller3, 0]), ("smaller2", [smaller2, 0])]
        )


def test_set_cuts_no_clear():
    with warns(FutureWarning):
        flow = CutFlow("TestFlow")
        flow.set_cut("smaller5", smaller5)

        flow.set_cuts(
            OrderedDict([("smaller3", smaller3), ("smaller2", smaller2)]), clear=False
        )

        assert flow.cuts == OrderedDict(
            [
                ("smaller5", [smaller5, 0]),
                ("smaller3", [smaller3, 0]),
                ("smaller2", [smaller2, 0]),
            ]
        )


if __name__ == "__main__":
    test_CutFlow()
    test_set_cuts_clear()
    test_set_cuts_no_clear()
