from abc import ABC, abstractmethod


class BettingStrategy(ABC):
    def __init__(self, df: pl.DataFrame, target_columns: list[str]):
        self.df = df
        self.target_columns = target_columns

    @abstractmethod
    def add_bet(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Add a column to the dataframe which will be the bet number (1, 0 or 2)
        """
        pass

    def add_bet_won_column(self, df):
        """
        Add a column to the dataframe which will be True if the bet has won.
        """
        return df.with_columns(bet_won=pl.col("winner_num").eq(pl.col("bet")))

    def add_odds_to_use(self, df):
        """
        Add a column to the dataframe which will be the odds to use for the bet.
        """
        return df.with_columns(
            odds_to_use=pl.when(pl.col("winner_num").eq(1))
            .then(pl.col("1_avg_odds"))
            .when(pl.col("winner_num").eq(2))
            .then(pl.col("2_avg_odds"))
            .otherwise(pl.col("x_avg_odds"))
        )

    def calculate_return(self, df):
        """
        Add a column to the dataframe which will be the return on the bet.
        """
        return df.with_columns(
            return_on_bet=pl.when(pl.col("bet_won"))
            .then(pl.col("odds_to_use") - 1)
            .otherwise(pl.lit(-1))
        )

    def get_underdog(self, df):
        """
        Add a column to the dataframe which will be the underdog team (1=home, 2=away, or None)
        """
        return df.with_columns(
            underdog=pl.when(pl.col("1_avg_odds") < pl.col("2_avg_odds"))
            .then(pl.lit(2))
            .when(pl.col("2_avg_odds") < pl.col("1_avg_odds"))
            .then(pl.lit(1))
            .otherwise(pl.lit(None))
        )

    def get_favourite(self, df):
        """
        Add a column to the dataframe which will be the favourite team (1=home, 2=away, or None)
        """
        return df.with_columns(
            favourite=pl.when(pl.col("1_avg_odds") < pl.col("2_avg_odds"))
            .then(pl.lit(1))
            .when(pl.col("2_avg_odds") < pl.col("1_avg_odds"))
            .then(pl.lit(2))
            .otherwise(pl.lit(None))
        )

    def has_favourite_won(self, df):
        """
        Add a column to the dataframe which will be True if the favourite has won.
        """
        return df.with_columns(
            has_favourite_won=pl.col("favourite").eq(pl.col("winner_num"))
        )

    def has_underdog_won(self, df):
        """
        Add a column to the dataframe which will be True if the underdog has won.
        """
        return df.with_columns(
            has_underdog_won=pl.col("underdog").eq(pl.col("winner_num"))
        )

    def apply_strategy(self):
        """
        Apply the strategy to the dataframe.
        """
        prep_df = (
            self.df.pipe(self.get_underdog)
            .pipe(self.get_favourite)
        )
        bet_df = self.add_bet(prep_df)

        # check if the method has been implemented correctly
        if "bet" not in bet_df.columns:
            raise ValueError(
                "The add_bet method has not been implemented correctly."
                "Please add a column called 'bet'."
            )

        # check if the bet column only contains 1, 0 or 2
        required_bet_values = set([1, 0, 2])
        if not bet_df["bet"].is_in(required_bet_values).all():
            raise ValueError(
                "The add_bet method has not been implemented correctly."
                f"Please add a column called 'bet' with the values {required_bet_values}."
            )

        return (
            bet_df.pipe(self.add_bet_won_column)
            .pipe(self.add_odds_to_use)
            .pipe(self.has_favourite_won)
            .pipe(self.has_underdog_won)
            .pipe(self.calculate_return)

        )

class BetAlwaysOnFavourite(BettingStrategy):
    def add_bet(self, df: pl.DataFrame) -> pl.DataFrame:
        return df.with_columns(bet=pl.lit(1))

betting_strategy = BetAlwaysOnFavourite(df, target_columns)
result = betting_strategy.apply_strategy()

def run_strategy_n_times(
    strategy: BettingStrategy, df: pl.DataFrame, n_times: int, n_games: int, target_columns: list[str]
) -> list[float]:
    """
    Run a betting strategy multiple times on randomly sampled subsets of data.

    This function applies a given betting strategy to randomly sampled subsets of the input
    DataFrame multiple times and returns a list of the total returns for each run.

    Args:
        strategy (BettingStrategy): The betting strategy class to be applied.
        df (pl.DataFrame): The input DataFrame containing the full dataset of games and their information.
        n_times (int): The number of times to run the strategy.
        n_games (int): The number of games to sample for each run of the strategy.
        target_columns (list[str]): The target columns to be used in the strategy.

    Returns:
        list[float]: A list containing the total return on bets for each run of the strategy.

    Notes:
        - The function uses random sampling with replacement, so the same game may appear
          multiple times in a single run or across different runs.
        - The 'target_columns' variable is assumed to be defined in the outer scope.
    """
    results = []
    for i in range(n_times):
        result = strategy(
            df.sample(n_games, shuffle=True), target_columns
        ).apply_strategy()
        results.append(result["return_on_bet"].sum())
    return results