Rewrite main function and add support functions

Add argparse argument handling

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 tqdm import tqdm
-from sqlalchemy.exc import IntegrityError, PendingRollbackError
 
-from data.database import dal, MegaMillionsDraw, PowerballDraw
+from data.database import dal2
-import lottery_predictor.collect as predictor
+import lottery_predictor.analyze as analyze
-from util.drawing import get_mega_millions_object, get_powerball_object
+from lottery_predictor.config import PROJECT_VARIABLES, GAME_INFO
-from util.environment import load_environment_variables
 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...")
+    parser = argparse.ArgumentParser(description="A lottery prediction tool.")
-    # convert specific variables for use
+    subparsers = parser.add_subparsers(dest="command", required=True)
-    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"])
 
+    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
+    record_parser = subparsers.add_parser("record-count", help="Get record counts by table in the database")
-    print("Connecting to the database...")
+    record_parser.add_argument(
-    dal.connect()
+        "-f", "--from-date",
-    session = dal.Session()
+        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...")
+    predict_parser = subparsers.add_parser("predict", help="Predict the next drawing")
-    # if there are no records in the MegaMillionsDraw table...
+    predict_parser.add_argument(
-    if predictor.get_record_count('MegaMillionsDraw') == 0:
+        "--game",
-        # loop over all years and scrape the data for each
+        type=str,
-        for year in range(mm_start_year, end_year + 1):
+        required=True,
-            print(f"\tCollecting {year}...")
+        help="The game to predict the next result for [MegaMillions, Powerball]"
-            scrape_data = scrape_game_history(game='mega-millions', year=year)
+    )
-            # convert and insert the records
+    predict_parser.add_argument(
-            for key, value in scrape_data.items():
+        "--window-size",
-                record = get_mega_millions_object(draw_date=key, draw_result=value)
+        type=int,
-                try:
+        default=-1,
-                    session.add(record)
+        help="The number of records to use in the prediction model, defaults to -1 for all records"
-                    session.commit()
+    )
-                except IntegrityError:
+    predict_parser.add_argument(
-                    pass
+        "--test",
-                except PendingRollbackError:
+        action="store_true",
-                    session.rollback()
+        help="Use the most recent drawing as a test subject to find the right window_size"
-                finally:
+    )
-                    session.close()
-                # sleep between records and years to avoid querying the site too often
-            sleep(60)
 
-        # commit all the records to the database
+    args = parser.parse_args()
-    else:
+    match args.command:
-        most_recent = session.query(func.max(MegaMillionsDraw.draw_date)).first()
+        case "record-count":
-        start_year = datetime.strptime(str(most_recent[0]), date_format).year
+            # load the date format for the database
-        # loop over the years needed and gather records
+            date_format = PROJECT_VARIABLES["LP_DATE_INSERT_FORMAT"]
-        for year in range(start_year, end_year + 1):
+            if args.from_date:
-            print(f"\tCollecting {year}...")
+                # if args.from_date is 'rule-change' use that date from vars, otherwise use the date provided
-            scrape_data = scrape_game_history(game='powerball', year=year)
+                if args.from_date != 'rule-change':
-            # convert and insert the records
+                    from_date = datetime.strptime(args.from_date, date_format)
-            for key, value in scrape_data.items():
+                    # get the table row counts from the date specified
-                record = get_powerball_object(draw_date=key, draw_result=value)
+                    table_counts = get_table_counts(from_date=from_date)
-                try:
+                else:
-                    session.add(record)
+                    table_counts = get_table_counts(from_date=args.from_date)
-                    session.commit()
+            else:
-                except IntegrityError:
+                table_counts = get_table_counts(from_date=None)
-                    pass
+            for table, count in table_counts.items():
-                except PendingRollbackError:
+                print(f"{table}:\t{count}")
-                    session.rollback()
+        case "update":
-                finally:
+            update_games(dry_run=args.dryrun)
-                    session.close()
+        case "predict":
-                # sleep between records and years to avoid querying the site too often
+            print(get_prediction(game=args.game, window_size=args.window_size, test_case=args.test))
-            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")
 
 
 if __name__ == "__main__":