diff --git a/main.py b/main.py index f2ef8bf..fee647d 100644 --- a/main.py +++ b/main.py @@ -1,130 +1,225 @@ -from datetime import datetime +import argparse +from datetime import date, datetime +import random from time import sleep -from sqlalchemy import func -from sqlalchemy.exc import IntegrityError, PendingRollbackError +from tqdm import tqdm -from data.database import dal, MegaMillionsDraw, PowerballDraw -import lottery_predictor.collect as predictor -from util.drawing import get_mega_millions_object, get_powerball_object -from util.environment import load_environment_variables +from data.database import dal2 +import lottery_predictor.analyze as analyze +from lottery_predictor.config import PROJECT_VARIABLES, GAME_INFO from util.scrape import scrape_game_history -project_variables = load_environment_variables() +BAR_FORMAT = PROJECT_VARIABLES["LP_BAR_FORMAT"] + +def dry_run_fill_game_data(game: str, year: int | None = None) -> None: + """ + + :param game: Name of the game to process + :param year: A starting year or None + :return: Nothing + + A dummy job to go through the motions of updating without doing anything + """ + if year is None: + year = GAME_INFO[game]['start_year'] + elif not isinstance(year, int): + raise ValueError(f"Year must be a valid year or None, received {year}") + elif year not in range(1992, 2030): + raise ValueError(f"Year must be within the range 1992 and 2030, received {year}") + + for year in range(year, datetime.today().year + 1): + for _ in tqdm( + range(random.randint(1, 106)), + desc=f"\tProcessing year {year}", + bar_format=BAR_FORMAT + ): + sleep(0.25) + + if year != datetime.now().year: + for _ in tqdm( + range(5), + desc='\tPseudo API backoff', + bar_format=BAR_FORMAT + ): + sleep(1) + else: + print() + + +def fill_game_data(game: str, year: int | None = None) -> None: + """ + :param game: Name of the game to process + :param year: A starting year or None + :return: Nothing + + If the year is None, all game data will be scraped for each year the game has data available. + Otherwise, the scrape will be limited to a range of years starting with the year specified + Not all data scraped will be used for predictions, but could be useful later. + """ + if year is None: + year = GAME_INFO[game]['start_year'] + elif not isinstance(year, int): + raise ValueError(f"Year must be a valid year or None, received {year}") + elif year not in range(1992, 2030): + raise ValueError(f"Year must be within the range 1992 and 2030, received {year}") + + # scrape records by year and insert into the database + for year in range(year, datetime.today().year + 1): + scrape_data = scrape_game_history(game=GAME_INFO[game]['scrape_name'], year=year) + # convert and insert records + for _ in tqdm( + range(len(scrape_data)), + desc=f"Processing year {year}", + bar_format=BAR_FORMAT + ): + for draw_date, numbers in scrape_data.items(): + record = GAME_INFO[game]['object_function'](draw_date=draw_date, draw_result=numbers) + dal2.add(record) + # sleep for 30 seconds to avoid overtaxing the endpoint, if the processed year is not the current one + if year != datetime.now().year: + for _ in tqdm( + range(30), + desc='\tAPI backoff', + bar_format=BAR_FORMAT + ): + sleep(1) + + +def update_games(dry_run: bool = False) -> None: + """ + :return: None + + Loop over games in the game_info dictionary and update the stored records + """ + date_format = PROJECT_VARIABLES["LP_DATE_INSERT_FORMAT"] + + for game, info in GAME_INFO.items(): + print(f"Updating {game} data...\n") + if dal2.count(table_name=info['table']) == 0: + if not dry_run: + fill_game_data(game=game) + else: + dry_run_fill_game_data(game=game) + else: + most_recent = dal2.most_recent(table_name=info['table']) + start_year = datetime.strptime(str(most_recent[0]), date_format).year + if not dry_run: + fill_game_data(game=game, year=start_year) + else: + dry_run_fill_game_data(game=game, year=start_year) + + +def get_table_counts(from_date: date | str | None) -> dict[str, int]: + """ + Loops over the tables in the database and gets a record count for each one. + + :return: A dictionary of table names and record counts + """ + counts = {} + for game, info in GAME_INFO.items(): + if isinstance(from_date, date): + counts[game] = dal2.count(table_name=info['table'], from_date=from_date) + elif from_date == 'rule-change': + rule_date = datetime.strptime(info['rule_change'], PROJECT_VARIABLES["LP_DATE_INSERT_FORMAT"]) + counts[game] = dal2.count(table_name=info['table'], from_date=rule_date) + else: + counts[game] = dal2.count(table_name=info['table']) + return counts + + +def get_prediction(game: str, window_size: int, test_case: bool = False) -> str: + # Get the list of drawings to make the prediction from + from_date = GAME_INFO[game]['rule_change'] + + if test_case: + if window_size == -1: + drawings = analyze.load_dataframe_by_dates(game=game, start_date=from_date) + else: + drawings = analyze.load_dataframe_most_recent(game=game, limit=window_size) + window_size = len(drawings) if window_size == -1 else window_size + test_case = drawings.iloc[0] + game_ball_out = 'powerball' if game == 'Powerball' else 'mega_ball' + case_string = (f"Draw Date: {test_case['draw_date']}, " + f"Main Balls: [{test_case['main_ball1']}, {test_case['main_ball2']}, {test_case['main_ball3']}, " + f"{test_case['main_ball4']}, {test_case['main_ball5']}], Game Ball: {test_case[game_ball_out]}") + + predicted = analyze.make_prediction(data_frame=drawings.iloc[1:], window_size=(window_size - 2)) + main_balls = ", ".join(predicted[0].astype(str)) + game_ball = "".join(predicted[1].astype(str)) + return f"{case_string}\n Prediction: Main Balls: [{main_balls}], Game Ball: {game_ball}\n" + else: + if window_size == -1: + drawings = analyze.load_dataframe_by_dates(game=game, start_date=from_date) + else: + drawings = analyze.load_dataframe_most_recent(game=game, limit=window_size) + window_size = len(drawings) if window_size == -1 else window_size + predicted = analyze.make_prediction(data_frame=drawings, window_size=(window_size - 1)) + main_balls = ", ".join(predicted[0].astype(str)) + game_ball = "".join(predicted[1].astype(str)) + return f"Prediction: Main Balls: {main_balls}, Game Ball: {game_ball}\n" + def main(): - print("Lottery Analyzer initializing...") - # convert specific variables for use - mm_start_year = int(project_variables["LP_MEGA_MILLIONS_INCEPTION"]) - pb_start_year = int(project_variables["LP_POWERBALL_INCEPTION"]) - end_year = datetime.today().year - date_format = project_variables["LP_DATE_INSERT_FORMAT"] - max_mm_main = int(project_variables["LP_MEGA_MILLIONS_MAX_MAIN_BALL"]) - max_megaball = int(project_variables["LP_MEGA_MILLIONS_MAX_MEGA_BALL"]) - max_pb_main = int(project_variables["LP_POWERBALL_MAX_MAIN_BALL"]) - max_powerball = int(project_variables["LP_POWERBALL_MAX_POWERBALL"]) + parser = argparse.ArgumentParser(description="A lottery prediction tool.") + subparsers = parser.add_subparsers(dest="command", required=True) + update_parser = subparsers.add_parser("update", help="Update database records") + update_parser.add_argument("-y", "--year", type=int, help="Year to start updates with") + update_parser.add_argument( + "-d", "--dryrun", + action="store_true", + help="Run the update process without actually making any database or API calls" + ) - # connect to the database - print("Connecting to the database...") - dal.connect() - session = dal.Session() + record_parser = subparsers.add_parser("record-count", help="Get record counts by table in the database") + record_parser.add_argument( + "-f", "--from-date", + type=str, + help="Date from which to base the count on (ex. 2026-01-01). \ + The default value is rule-change for the date when game rules last changed" + ) - print("Collecting and processing missing Mega Millions draws...") - # if there are no records in the MegaMillionsDraw table... - if predictor.get_record_count('MegaMillionsDraw') == 0: - # loop over all years and scrape the data for each - for year in range(mm_start_year, end_year + 1): - print(f"\tCollecting {year}...") - scrape_data = scrape_game_history(game='mega-millions', year=year) - # convert and insert the records - for key, value in scrape_data.items(): - record = get_mega_millions_object(draw_date=key, draw_result=value) - try: - session.add(record) - session.commit() - except IntegrityError: - pass - except PendingRollbackError: - session.rollback() - finally: - session.close() - # sleep between records and years to avoid querying the site too often - sleep(60) + predict_parser = subparsers.add_parser("predict", help="Predict the next drawing") + predict_parser.add_argument( + "--game", + type=str, + required=True, + help="The game to predict the next result for [MegaMillions, Powerball]" + ) + predict_parser.add_argument( + "--window-size", + type=int, + default=-1, + help="The number of records to use in the prediction model, defaults to -1 for all records" + ) + predict_parser.add_argument( + "--test", + action="store_true", + help="Use the most recent drawing as a test subject to find the right window_size" + ) - # commit all the records to the database - else: - most_recent = session.query(func.max(MegaMillionsDraw.draw_date)).first() - start_year = datetime.strptime(str(most_recent[0]), date_format).year - # loop over the years needed and gather records - for year in range(start_year, end_year + 1): - print(f"\tCollecting {year}...") - scrape_data = scrape_game_history(game='powerball', year=year) - # convert and insert the records - for key, value in scrape_data.items(): - record = get_powerball_object(draw_date=key, draw_result=value) - try: - session.add(record) - session.commit() - except IntegrityError: - pass - except PendingRollbackError: - session.rollback() - finally: - session.close() - # sleep between records and years to avoid querying the site too often - sleep(60) - - print("Collecting and processing missing Powerball draws...") - # if there are no records in the PowerBallDraw table... - if len(session.query(PowerballDraw).all()) == 0: - # loop over all years and scrape the data for each - for year in range(pb_start_year, end_year + 1): - print(f"\tCollecting {year}...") - scrape_data = scrape_game_history(game='powerball', year=year) - # convert and insert the records - for key, value in scrape_data.items(): - record = get_powerball_object(draw_date=key, draw_result=value) - try: - session.add(record) - session.commit() - except IntegrityError: - pass - except PendingRollbackError: - session.rollback() - finally: - session.close() - # sleep between records and years to avoid querying the site too often - sleep(60) - - # commit all the records to the database - session.commit() - else: - most_recent = session.query(func.max(PowerballDraw.draw_date)).first() - start_year = datetime.strptime(str(most_recent[0]), str(date_format)).year - # loop over the years needed and gather records - for year in range(start_year, end_year + 1): - print(f"\tCollecting {year}...") - scrape_data = scrape_game_history(game='powerball', year=year) - # convert and insert the records - for key, value in scrape_data.items(): - record = get_powerball_object(draw_date=key, draw_result=value) - try: - session.add(record) - session.commit() - except IntegrityError: - pass - except PendingRollbackError: - session.rollback() - finally: - session.close() - # sleep between records and years to avoid querying the site too often - sleep(60) - - mm_count = session.query(MegaMillionsDraw).count() - pb_count = session.query(PowerballDraw).count() - - print(f"There are {mm_count} Mega Millions draw records and {pb_count} Powerball records") + args = parser.parse_args() + match args.command: + case "record-count": + # load the date format for the database + date_format = PROJECT_VARIABLES["LP_DATE_INSERT_FORMAT"] + if args.from_date: + # if args.from_date is 'rule-change' use that date from vars, otherwise use the date provided + if args.from_date != 'rule-change': + from_date = datetime.strptime(args.from_date, date_format) + # get the table row counts from the date specified + table_counts = get_table_counts(from_date=from_date) + else: + table_counts = get_table_counts(from_date=args.from_date) + else: + table_counts = get_table_counts(from_date=None) + for table, count in table_counts.items(): + print(f"{table}:\t{count}") + case "update": + update_games(dry_run=args.dryrun) + case "predict": + print(get_prediction(game=args.game, window_size=args.window_size, test_case=args.test)) if __name__ == "__main__":