import pandas as pd
from datasets import Dataset

def calculate_elo(old_rating, opponent_rating, score, k_factor=32):
    """
    Calculate the new ELO rating for a player.
    
    :param old_rating: The current ELO rating of the player.
    :param opponent_rating: The ELO rating of the opponent.
    :param score: The score of the game (1 for win, 0.5 for draw, 0 for loss).
    :param k_factor: The K-factor used in ELO rating (default is 32).
    :return: The new ELO rating.
    """
    expected_score = 1 / (1 + 10 ** ((opponent_rating - old_rating) / 400))
    new_rating = old_rating + k_factor * (score - expected_score)
    return new_rating
    
def update_elo_ratings(ratings_dataset, winner, loser, k_factor=32):
    # Convert the Hugging Face dataset to a pandas DataFrame
    ratings_df = pd.DataFrame(ratings_dataset)

    # Check and add new players if they don't exist in the dataset
    for player in [winner, loser]:
        if player not in ratings_df['bot_name'].values:
            new_player = {'bot_name': player, 'elo_rating': 1200}
            ratings_df = pd.concat([ratings_df,pd.DataFrame([new_player])],ignore_index=True)

    # Extract old ratings
    winner_old_rating = ratings_df.loc[ratings_df['bot_name'] == winner, 'elo_rating'].iloc[0]
    loser_old_rating = ratings_df.loc[ratings_df['bot_name'] == loser, 'elo_rating'].iloc[0]

    # Calculate new ratings
    winner_new_rating = calculate_elo(winner_old_rating, loser_old_rating, 1, k_factor)
    loser_new_rating = calculate_elo(loser_old_rating, winner_old_rating, 0, k_factor)

    # Update the DataFrame
    ratings_df.loc[ratings_df['bot_name'] == winner, 'elo_rating'] = winner_new_rating
    ratings_df.loc[ratings_df['bot_name'] == loser, 'elo_rating'] = loser_new_rating

    # Convert the DataFrame back to a Hugging Face dataset
    updated_ratings_dataset = Dataset.from_pandas(ratings_df)

    return updated_ratings_dataset