from datetime import datetime
import pytest

import lottery_predictor.analyze as analyze


def test_load_dataframe() -> None:
    pb_start = datetime(2015, 10, 4)
    mm_start = datetime(2025, 4, 5)
    date_end = datetime(2026, 3, 8)
    assert int(analyze.load_dataframe(
        table_name='PowerballDraw', start_date=pb_start, end_date=date_end).count()['draw_date']) == 1324
    assert int(analyze.load_dataframe(
        table_name='MegaMillionsDraw', start_date=mm_start, end_date=date_end).count()['draw_date']) == 96
    with pytest.raises(ValueError):
        analyze.load_dataframe(table_name='SomNonExistentTableName')


def test_least_and_most_common_number() -> None:
    pb_start = datetime(2015, 10, 4)
    mm_start = datetime(2025, 4, 5)
    date_end = datetime(2026, 3, 8)

    pb_df = analyze.load_dataframe(
        table_name='PowerballDraw', start_date=pb_start, end_date=date_end)
    mm_df = analyze.load_dataframe(
        table_name='MegaMillionsDraw', start_date=mm_start, end_date=date_end)

    assert analyze.get_most_common_number(pb_df, top=5) == [61, 21, 28, 23, 33]
    assert analyze.get_least_common_number(pb_df, bottom=5) == [13, 49, 46, 26, 34]
    assert analyze.get_most_common_number(mm_df, top=5) == [42, 40, 18, 10, 49]
    assert analyze.get_least_common_number(mm_df, bottom=5) == [35, 51, 3, 1, 20]

    assert analyze.get_most_common_number(pb_df, columns=['powerball'], top=1) == [4]
    assert analyze.get_least_common_number(pb_df, columns=['powerball'], bottom=1) == [16]
    assert analyze.get_most_common_number(mm_df, columns=['mega_ball'], top=1) == [1]
    assert analyze.get_least_common_number(mm_df, columns=['mega_ball'], bottom=1) == [3]