import random
import json

class EloRank:
    def __init__(self, initial_rating=1000, k_factor=32):
        """
        Initialize the EloRank class.
        :param initial_rating: Initial ELO rating for each model.
        :param k_factor: The K-factor that determines the sensitivity of rating changes.
        """
        self.ratings = {}
        self.initial_rating = initial_rating
        self.k_factor = k_factor
        self.wins = {}

    def add_model(self, model_id):
        """
        Add a new model with the initial rating.
        :param model_id: Unique identifier for the model.
        """
        self.ratings[model_id] = self.initial_rating
        self.wins[model_id] = 0

    def record_match(self, winner, loser):
        """
        Update the ratings based on a match result.
        :param winner: Model ID of the winner.
        :param loser: Model ID of the loser.
        """
        rating_winner = self.ratings[winner]
        rating_loser = self.ratings[loser]

        expected_winner = self.expected_score(rating_winner, rating_loser)
        expected_loser = self.expected_score(rating_loser, rating_winner)

        self.ratings[winner] += self.k_factor * (1 - expected_winner)
        self.ratings[loser] += self.k_factor * (0 - expected_loser)

        # Update win count
        self.wins[winner] += 1

    def expected_score(self, rating_a, rating_b):
        """
        Calculate the expected score for a model.
        :param rating_a: Rating of model A.
        :param rating_b: Rating of model B.
        :return: Expected score.
        """
        return 1 / (1 + 10 ** ((rating_b - rating_a) / 400))

    def get_rating(self, model_id):
        """
        Get the current rating of a model.
        :param model_id: Unique identifier for the model.
        :return: Current rating of the model.
        """
        return self.ratings.get(model_id, None)

    def get_wins(self, model_id):
        """
        Get the number of wins of a model.
        :param model_id: Unique identifier for the model.
        :return: Number of wins of the model.
        """
        return self.wins.get(model_id, 0)

    def get_top_models(self, n=2):
        """
        Get the top N models by rating.
        :param n: Number of top models to retrieve.
        :return: List of model IDs of the top models.
        """
        return sorted(self.ratings, key=self.ratings.get, reverse=True)[:n]

    def sample_next_match(self):
        """
        Sample the next match based on the probability proportional to the current rating.
        This approach helps accelerate the convergence of ranking.
        :return: Tuple of two model IDs for the next match.
        """
        model_ids = list(self.ratings.keys())
        probabilities = [self.ratings[model_id] for model_id in model_ids]
        total_rating = sum(probabilities)
        probabilities = [rating / total_rating for rating in probabilities]

        # Sample two different models for the next match
        next_match = random.choices(model_ids, probabilities, k=2)
        while next_match[0] == next_match[1]:
            next_match = random.choices(model_ids, probabilities, k=2)

        return tuple(next_match)

    def process_match_records(self, file_path):
        """
        Process match records from a JSON file and update ratings and win counts accordingly.
        :param file_path: Path to the JSON file containing match records.
        """
        with open(file_path, 'r') as file:
            match_records = json.load(file)

        for record in match_records:
            winner = record['winner']
            model_1 = record['model_1']
            model_2 = record['model_2']

            # Add models if they are not already added
            if model_1 not in self.ratings:
                self.add_model(model_1)
            if model_2 not in self.ratings:
                self.add_model(model_2)

            # Record the match result
            if winner == model_1:
                self.record_match(model_1, model_2)
            elif winner == model_2:
                self.record_match(model_2, model_1)

# # Example Usage
# e = EloRank()
# e.add_model('model_A')
# e.add_model('model_B')
# e.add_model('model_C')

# e.record_match('model_A', 'model_B')
# print(e.get_rating('model_A'))  # Should be greater than the initial rating
# print(e.get_rating('model_B'))  # Should be less than the initial rating

# print(e.get_top_models(2))  # Get the top 2 models
# print(e.sample_next_match())  # Sample the next match based on ratings

# # Process match records from a JSON file
# e.process_match_records('match_records.json')
# print(e.get_wins('model_A'))  # Get the number of wins for model_A