From ca8bdb6290ef88505cc50d68c2165f6f688328e4 Mon Sep 17 00:00:00 2001
From: Chris Smith <ole088@gmail.com>
Date: Sun, 24 May 2026 21:49:24 -0400
Subject: [PATCH] Add prep and predict functions

Split data_frame functions into date and most most recent X
---
 lottery_predictor/analyze.py | 73 ++++++++++++++++++++++++++++++++++--
 1 file changed, 70 insertions(+), 3 deletions(-)

diff --git a/lottery_predictor/analyze.py b/lottery_predictor/analyze.py
index 76bf85f..773a4a8 100644
--- a/lottery_predictor/analyze.py
+++ b/lottery_predictor/analyze.py
@@ -3,16 +3,16 @@ from collections import Counter
 
 import numpy as np
 import pandas as pd
-from sklearn.ensemble import RandomForestClassifier
+from sklearn.ensemble import RandomForestRegressor
 from sqlalchemy import select
 
-from data.database import Base, dal, project_variables
+from data.database import Base, dal
 import util.drawing as drawutil
 
 RANDOM_SEED = 42
 
 
-def load_dataframe(
+def load_dataframe_by_dates(
         table_name: str, start_date: datetime | None = None, end_date: datetime | None = None
     ) -> pd.DataFrame:
     """
@@ -45,6 +45,73 @@ def load_dataframe(
         raise ValueError('An invalid table name was provided')
 
 
+def load_dataframe_most_recent(table_name: str, limit: int = 10) -> pd.DataFrame:
+    """
+        :param table_name: the name of the table to load data from
+        :param limit: The N most recent draws (default: 10)
+
+        :returns: a pandas DataFrame containing the data
+    """
+
+    dal.connect()
+    session = dal.Session()
+
+    game_dict = drawutil.check_table_name(table_name=table_name)
+    if game_dict:
+        target_table = Base.metadata.tables.get(game_dict['db_table_name'])
+        if limit is not None and limit > 0:
+            sql_statement = (select(target_table)
+                .order_by(target_table.columns.draw_date.desc())
+                .limit(limit))
+            return pd.read_sql(sql_statement, session.bind)
+        else:
+            raise ValueError('Limit must be a positive integer greater than zero.')
+    else:
+        raise ValueError('An invalid table name was provided.')
+
+
+def prepare_split_data(data: np.ndarray, window_size: int = 10) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+
+    # Calculate the indices for all windows at once
+    indices = np.arange(len(data) - window_size)
+
+    # Create X: flattened sliding windows
+    # Reshapes into (samples, window_size * columns)
+    X = np.array([data[i : i + window_size]. flatten() for i in indices])
+
+    # Create y_field: first 5 columns of the row which are the field balls
+    y_field = data[window_size:, 0:5]
+
+    # Create y_game: 6th column (index 5) of the row which is the game ball
+    y_game = data[window_size:, 5]
+
+    return X, y_field, y_game
+
+
+def make_prediction(data_frame: pd.DataFrame, window_size: int = 10) -> tuple[np.ndarray, np.ndarray]:
+    # Convert the data_frame to a numpy array for slicing
+    data = data_frame.values
+
+    # Prepare and split the data
+    X, y_field, y_game = prepare_split_data(data, window_size=window_size)
+
+    # Train the model for the field balls
+    field_model = RandomForestRegressor(n_estimators=200, random_state=RANDOM_SEED)
+    field_model.fit(X, y_field)
+
+    # Train the model for the game ball
+    game_model = RandomForestRegressor(n_estimators=200, random_state=RANDOM_SEED)
+    game_model.fit(X, y_game)
+
+    # Predict the next draw
+    current_window = data[-window_size:].flatten().reshape(1, -1)
+    # Get predictions and round to the nearest whole number
+    pred_field = np.sort(np.round(field_model.predict(current_window)).astype(int))
+    pred_game = np.round(game_model.predict(current_window)).astype(int)
+
+    return pred_field[0], pred_game[0]
+
+
 def get_most_common_number(
         data_frame: pd.DataFrame, columns: list | None = None, top: int = 1
     ) -> list[int]: