from datetime import datetime


def split_sessions_into_train_test_val(all_sessions):
    """
    Splits session data into train/test/val lists based on session start time.
    1 March - 15 May -> Train
    16 May - 23 May -> Test
    24 May - 31 May -> Val
    Parameters
    ----------
    all_sessions: list of (session, min_session_time, max_session_time) tuples where session is deduplicated

    Returns train/test/val session lists
    -------

    """
    date_format = "%Y-%m-%d %H:%M:%S"
    train_time_range = datetime.strptime("2006-03-01 00:00:00", date_format), datetime.strptime(
        "2006-05-15 23:59:59", date_format
    )
    val_time_range = datetime.strptime("2006-05-16 00:00:00", date_format), datetime.strptime(
        "2006-05-23 23:59:59", date_format
    )
    test_time_range = datetime.strptime("2006-05-24 00:00:00", date_format), datetime.strptime(
        "2006-05-31 23:59:59", date_format
    )
    train_sessions, test_sessions, val_sessions = [], [], []
    for curr_session, min_sess_time, max_sess_time in all_sessions:

        assert (
            min_sess_time <= max_sess_time
        ), "Min session = {} time should be less than max session time = {}.".format(
            min_sess_time, max_sess_time
        )

        if train_time_range[0] <= min_sess_time <= train_time_range[1]:
            train_sessions.append(curr_session)
        elif val_time_range[0] <= min_sess_time <= val_time_range[1]:
            val_sessions.append(curr_session)
        elif test_time_range[0] <= min_sess_time <= test_time_range[1]:
            test_sessions.append(curr_session)
        else:
            raise Exception(
                "This session range is not handled = ({}, {})".format(min_sess_time, max_sess_time)
            )

    return train_sessions, test_sessions, val_sessions