from datetime import datetime
from collections import Counter

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sqlalchemy import select

from data.database import Base, dal
import util.drawing as drawutil

RANDOM_SEED = 42


def load_dataframe_by_dates(
        table_name: str, start_date: datetime | None = None, end_date: datetime | None = None
    ) -> pd.DataFrame:
    """
        :param table_name: the name of the table to load data from
        :param start_date: the start date for the data
        :param end_date: the end date for the data

        :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 start_date is not None and end_date is not None:
            sql_statement = (select(target_table)
                .where(target_table.columns.draw_date >= start_date)
                .where(target_table.columns.draw_date <= end_date))
        elif start_date is not None and end_date is None:
            sql_statement = (select(target_table).where(target_table.columns.draw_date >= start_date))
        elif start_date is None and end_date is not None:
            sql_statement = (select(target_table).where(target_table.columns.draw_date <= end_date))
        else:
            sql_statement = select(target_table)

        return pd.read_sql(sql_statement, session.bind)
    else:
        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]:
    """
    :param data_frame: a pandas DataFrame containing the data
    :param columns: a list of column names to use
    :param top: the number of top (most seen) numbers to return

    :returns: a list of the most common numbers
    """
    if columns is None:
        columns = ['main_ball1', 'main_ball2', 'main_ball3', 'main_ball4', 'main_ball5']
    flat_numbers = data_frame[columns].values.flatten()
    counts = Counter(flat_numbers)
    return [int(num) for num, _ in counts.most_common(top)]


def get_least_common_number(
        data_frame: pd.DataFrame, columns: list | None = None, bottom: int = 1
) -> list[int]:
    """
    :param data_frame: a pandas DataFrame containing the data
    :param columns: a list of column names to use
    :param bottom: the number of bottom (least seen) numbers to return

    :returns: a list of the least common numbers
    """
    if columns is None:
        columns = ['main_ball1', 'main_ball2', 'main_ball3', 'main_ball4', 'main_ball5']
    flat_numbers = data_frame[columns].values.flatten()
    counts = Counter(flat_numbers)
    return list(reversed([int(num) for num, _ in counts.most_common()[-bottom:]]))


def calculate_probabilities(
        data_frame: pd.DataFrame, max_number: int, columns: list | None = None
    ) -> dict[int, float]:
    """
    :param data_frame: A pandas DataFrame containing the data to calculate probabilities for
    :param max_number: The maximum number possible in the data_frame
    :param columns: The list of column names to use from the data_frame
    :returns dict: A dictionary containing the probabilities
    """
    if columns is None:
        columns = ['main_ball1', 'main_ball2', 'main_ball3', 'main_ball4', 'main_ball5']
    if data_frame.empty:
        return dict()

    # get all the numbers in the groups
    all_numbers = [num for group in data_frame[columns].values for num in group]
    # count all the occurrences of each number
    counts = Counter(all_numbers)
    # calculate the basic probability of each number occurring again
    probabilities = {num: counts.get(num, 0) / (max_number + 1) for num in range(1, max_number + 1)}

    # if any calculation is greater than one, use what is to the right of the decimal point as the value
    for key, value in probabilities.items():
        if value > 1:
            probabilities[key] = value - int(str(value).split('.')[0])

    # return the probability dict
    return probabilities


def get_hot_numbers(probabilities: dict[int, float], top: int = 5) -> list[tuple[int, float]]:
    """
    :param probabilities: A dictionary containing the probabilities
    :param top: The count of hottest items to return, defaults to 5
    :returns list of tuples: A list of the hot numbers and their raw score
    """
    return sorted(probabilities.items(), key=lambda item: item[1], reverse=True)[:top]


def get_cold_numbers(probabilities: dict[int, float], bottom: int = 5) -> list[tuple[int, float]]:
    """
    :param probabilities: A dictionary containing the probabilities
    :param bottom: The count of coldest items to return, defaults to 5
    :returns list of tuples: A list of the hot numbers and their raw score
    """
    return sorted(probabilities.items(), key=lambda item: item[1])[:bottom]