Refactor code and update prediction output

main.py

@@ -1,17 +1,20 @@
 import argparse
-from datetime import date, datetime
 import random
+from datetime import date, datetime
 from time import sleep
 
 from tqdm import tqdm
 
-from data.database import dal2
 import lottery_predictor.analyze as analyze
-from lottery_predictor.config import PROJECT_VARIABLES, GAME_INFO
+import lottery_predictor.draw as ticket
+from data.database import dal2
+from lottery_predictor.config import GAME_INFO, PROJECT_VARIABLES
 from util.scrape import scrape_game_history
 
+
 BAR_FORMAT = PROJECT_VARIABLES["LP_BAR_FORMAT"]
 
+
 def dry_run_fill_game_data(game: str, year: int | None = None) -> None:
     """
 
@@ -26,21 +29,23 @@ def dry_run_fill_game_data(game: str, year: int | None = None) -> None:
     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}")
+        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)),
+                range(random.randint(1, 106)),
-            desc=f"\tProcessing year {year}",
+                desc=f"\tProcessing year {year}",
-            bar_format=BAR_FORMAT
+                bar_format=BAR_FORMAT,
         ):
             sleep(0.25)
 
         if year != datetime.now().year:
             for _ in tqdm(
-                range(5),
+                    range(5),
-                desc='\tPseudo API backoff',
+                    desc='\tPseudo API backoff',
-                bar_format=BAR_FORMAT
+                    bar_format=BAR_FORMAT,
             ):
                 sleep(1)
         else:
@@ -53,35 +58,45 @@ def fill_game_data(game: str, year: int | None = None) -> None:
     :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.
+    If the year is None, all game data will be scraped for each year the game
-    Otherwise, the scrape will be limited to a range of years starting with the year specified
+    has data available.
-    Not all data scraped will be used for predictions, but could be useful later.
+    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}")
+        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)
+        scrape_data = scrape_game_history(
+            game=GAME_INFO[game]['scrape_name'], year=year,
+        )
         # convert and insert records
         for _ in tqdm(
-            range(len(scrape_data)),
+                range(len(scrape_data)),
-            desc=f"Processing year {year}",
+                desc=f"Processing year {year}",
-            bar_format=BAR_FORMAT
+                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)
+                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
+        # 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),
+                    range(30),
-                desc='\tAPI backoff',
+                    desc='\tAPI backoff',
-                bar_format=BAR_FORMAT
+                    bar_format=BAR_FORMAT,
             ):
                 sleep(1)
 
@@ -103,7 +118,9 @@ def update_games(dry_run: bool = False) -> None:
                 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
+            start_year = datetime.strptime(
+                str(most_recent[0]), date_format,
+            ).year
             if not dry_run:
                 fill_game_data(game=game, year=start_year)
             else:
@@ -119,93 +136,174 @@ def get_table_counts(from_date: date | str | None) -> dict[str, int]:
     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)
+            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"])
+            rule_date = datetime.strptime(
-            counts[game] = dal2.count(table_name=info['table'], from_date=rule_date)
+                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:
+def get_prediction(
+        game: str,
+        window_size: int,
+        test_case: bool = False,
+) -> tuple[dict, dict] | dict:
     # 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)
+            drawings = analyze.load_dataframe_by_dates(
+                game=game,
+                start_date=from_date,
+            )
         else:
-            drawings = analyze.load_dataframe_most_recent(game=game, limit=window_size)
+            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']}, "
+        case_date = (test_case['draw_date'].
-              f"Main Balls: [{test_case['main_ball1']}, {test_case['main_ball2']}, {test_case['main_ball3']}, "
+                     strftime(PROJECT_VARIABLES['LP_DATE_INSERT_FORMAT']))
-              f"{test_case['main_ball4']}, {test_case['main_ball5']}], Game Ball: {test_case[game_ball_out]}")
+        case_main = [
+            test_case['main_ball1'].tolist(),
+            test_case['main_ball2'].tolist(),
+            test_case['main_ball3'].tolist(),
+            test_case['main_ball4'].tolist(),
+            test_case['main_ball5'].tolist(),
+        ]
+        case_game = test_case[game_ball_out].tolist()
+        case_output = {
+            'draw_date': case_date,
+            'main_balls': case_main,
+            'game_ball': case_game,
+        }
 
-        predicted = analyze.make_prediction(data_frame=drawings.iloc[1:], window_size=(window_size - 2))
+        predicted = analyze.make_prediction(
-        main_balls = ", ".join(predicted[0].astype(str))
+            data_frame=drawings.iloc[1:],
-        game_ball = "".join(predicted[1].astype(str))
+            window_size=(window_size - 2),
-        return f"{case_string}\n           Prediction: Main Balls: [{main_balls}], Game Ball: {game_ball}\n"
+        )
+        main_balls = [i.tolist() for i in predicted[0]]
+        game_ball = int(predicted[1].astype(str))
+        prediction_output = {'main_balls': main_balls, 'game_ball': game_ball}
+        return case_output, prediction_output
     else:
         if window_size == -1:
-            drawings = analyze.load_dataframe_by_dates(game=game, start_date=from_date)
+            drawings = analyze.load_dataframe_by_dates(
+                game=game,
+                start_date=from_date,
+            )
         else:
-            drawings = analyze.load_dataframe_most_recent(game=game, limit=window_size)
+            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))
+        predicted = analyze.make_prediction(
-        main_balls = ", ".join(predicted[0].astype(str))
+            data_frame=drawings,
-        game_ball = "".join(predicted[1].astype(str))
+            window_size=(window_size - 1),
-        return f"Prediction: Main Balls: {main_balls}, Game Ball: {game_ball}\n"
+        )
+        main_balls = [int(i) for i in predicted[0].astype(str)]
+        game_ball = int(predicted[1].astype(str))
+        prediction_output = {'main_balls': main_balls, 'game_ball': game_ball}
+        return prediction_output
 
 
 def main():
     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 = subparsers.add_parser(
-    update_parser.add_argument("-y", "--year", type=int, help="Year to start updates with")
+        "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"
+        help="Run the update process without actually making any "
+             "database or API calls",
     )
 
-    record_parser = subparsers.add_parser("record-count", help="Get record counts by table in the database")
+    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). \
+        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"
+             "The default value is rule-change for the date when game "
+             "rules last changed",
     )
 
-    predict_parser = subparsers.add_parser("predict", help="Predict the next drawing")
+    ticket_parser = subparsers.add_parser(
+        "generate-ticket",
+        help="Generate a random ticket for the specified game",
+    )
+    ticket_parser.add_argument(
+        "--game",
+        type=str,
+        required=True,
+        help="The game to predict the next result for [MegaMillions, Powerball]",
+    )
+
+    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]"
+        help="The game to predict the next result for [MegaMillions, Powerball]",
     )
     predict_parser.add_argument(
         "--window-size",
         type=int,
-        default=-1,
+        default=10,
-        help="The number of records to use in the prediction model, defaults to -1 for all records"
+        help="The number of records to use in the prediction model, "
+             "defaults to 10, use -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"
+        help="Use the most recent drawing as a test subject to find "
+             "the right window_size",
     )
 
     args = parser.parse_args()
     match args.command:
+        case "generate-ticket":
+            main_balls, game_ball = ticket.generate_random_ticket(
+                max_main=GAME_INFO[args.game]['max_main_ball'],
+                max_game=GAME_INFO[args.game]['max_game_ball'],
+                num_main=GAME_INFO[args.game]['main_count'],
+            )
+            print(
+                f"Random {args.game} play: Main Balls: {main_balls}, "
+                f"{GAME_INFO[args.game]['game_ball_name']}: {game_ball}",
+            )
         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 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
@@ -219,7 +317,36 @@ def main():
         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 args.test:
+                recent, prediction = get_prediction(
+                    game=args.game,
+                    window_size=args.window_size,
+                    test_case=args.test,
+                )
+                print(
+                    f"{args.game} prediction: Main Balls: "
+                    f"{prediction['main_balls']}, "
+                    f"{GAME_INFO[args.game]['game_ball_name']}: "
+                    f"{prediction['game_ball']}",
+                )
+                print(
+                    f"{recent['draw_date']} drawing  : "
+                    f"Main Balls: {recent['main_balls']}, "
+                    f"{GAME_INFO[args.game]['game_ball_name']}: "
+                    f"{recent['game_ball']}",
+                )
+            else:
+                prediction = get_prediction(
+                    game=args.game,
+                    window_size=args.window_size,
+                    test_case=args.test,
+                )
+                print(
+                    f"{args.game} prediction: Main Balls: "
+                    f"{prediction['main_balls']}, "
+                    f"{GAME_INFO[args.game]['game_ball_name']}: "
+                    f"{prediction['game_ball']}",
+                )
 
 
 if __name__ == "__main__":