"""
Scripts to read, download or generate data
"""

from collections import Counter
from typing import List, Optional

import numpy as np
from collections import defaultdict
import graph_tool.all as gt


class Player:
    """
    Class representing a player in the Bradley-Terry model

    Attributes
    ----------
    id : int
        Player id
    rating : float (default=None)
        Player rating. If None, a score s is sampled from a logistic distribution. The rating is then computed as exp(s)
    """

    def __init__(self, id, rating=None):
        self.rating = rating
        if rating is None:
            # self.rating = np.random.rand() * max_rating
            # sample s from a logistic distribution
            loc, scale = 0, 1
            s = np.random.logistic(loc=loc, scale=scale)
            # Compute the rating from s
            self.rating = np.exp(s)
            # print(f"Player {id} rating: {self.rating} (s = {s})")
        self.id = id


def generate_matches(
    players,
    num_matches,
    replace=True,
    isolated=False,
    min_iso_match=False,
    self_edges=False,
    weighted=False,
):
    """
    Generate matches from the Bradley-Terry model

    Parameters
    ----------
    players : list
        List of players
    num_matches : int
        Number of matches to generate
    replace : bool
        Whether to sample player pairs with replacement
    isolated : bool
        Whether to allow for isolated players (players that do not play any match)
    min_iso_match : bool
        Whether to generate the minimum number of matches when avoiding isolated players.
        Matches are generated by randomly selecting N / 2 pairs of distinct players
        and placing an edge between them.
        The number of matches is then N / 2 (or (N + 1) / 2 if N is odd).
    self_edges : bool
        Whether to allow for self-edges (players playing against themselves)
    weighted : bool
        Whether to generate weighted matches (winner, looser, weight)

    Returns
    -------
    matches : list
        List of matches of the type [[winner, looser]] or [[winner, looser, weight]]
    """
    matches = []
    # Check if isolated players are allowed
    if not isolated:
        # Check if the number of matches is too low
        if num_matches < len(players) - 1:
            raise ValueError("Not enough matches to avoid isolated players")
        # Initialize set to track players that have played at least one match
        players_played = set()
        # Initialize tracker to count the total number of matches played
        total_matches = 0
        # While there are players that have not played yet...
        while len(players_played) < len(players):
            # Check if the number of matches is too low
            if total_matches >= num_matches:
                raise ValueError("Not enough matches to avoid isolated players")
            # Loop through all players
            for player1 in players:
                # If player1 has already played, skip
                if player1 in players_played:
                    continue
                # Randomly select another player that has not played yet and is different from player1
                # Check if the minimum number of matches is required
                if min_iso_match:
                    try:
                        player2 = np.random.choice(
                            [
                                player
                                for player in players
                                if player not in players_played and player != player1
                            ]
                        )
                    except (
                        ValueError
                    ):  # No more players to select. Occurs when there is an odd number of players.
                        # Randomly select a player from the set of players that have played
                        player2 = np.random.choice(list(players_played))
                else:
                    player2 = np.random.choice([player for player in players if player != player1])
                # Play a match
                match = []
                # Compute the probability of player1 winning
                prob = player1.rating / (player1.rating + player2.rating)
                # Randomly select the winner
                winner = player1 if np.random.rand() < prob else player2
                # Append the winner and looser to the match
                match.append(winner)
                match.append(player1 if winner == player2 else player2)
                # Append the match to the list of matches
                matches.append(match)
                # Update the set of players that have played
                players_played.update([player1, player2])
                # Update the total number of matches played
                total_matches += 1
        # print(f"Avoided isolated players with {total_matches} matches")
        # Update num_matches to the remaining number of matches to generate
        num_matches -= total_matches
    for i in range(num_matches):
        match = []
        # Randomly sample two players with repetition
        player1, player2 = np.random.choice(players, 2, replace=replace)
        # Check if self-edges are allowed
        if not self_edges:
            while player1 == player2:
                player2 = np.random.choice(players)
        # Compute the probability of player1 winning
        prob = player1.rating / (player1.rating + player2.rating)
        # Randomly select the winner
        winner = player1 if np.random.rand() < prob else player2
        # Append the winner and looser to the match
        match.append(winner)
        match.append(player1 if winner == player2 else player2)
        # Append the match to the list of matches
        matches.append(match)
    if weighted:
        # Count the number of times each match occurs
        match_count = Counter([tuple(match) for match in matches])
        # Generate a list of matches with weights
        matches = [[match[0], match[1], count] for match, count in match_count.items()]

    return matches


def generate_txt(matches, filename):
    """
    Generate a txt file from a list of matches

    Parameters
    ----------
    matches : list
        List of matches of the type [[winner, looser, (weight)]]
    filename : str
        Name of the file to write
    """
    # Determine whether match list is weighted
    weighted = False
    if len(matches[0]) == 3:
        weighted = True

    # Save match list to file
    with open(filename, "w") as f:
        for match in matches:
            try:
                if weighted:
                    f.write(f"{match[0].id} {match[1].id} {match[2]}\n")
                else:
                    f.write(f"{match[0].id} {match[1].id}\n")
            except AttributeError:  # If the players are not instances of the Player class
                if weighted:
                    f.write(f"{match[0]} {match[1]} {match[2]}\n")
                else:
                    f.write(f"{match[0]} {match[1]}\n")


def generate_gt(filename: str, players: Optional[List[Player]] = None) -> gt.Graph:
    """
    Generate a graph-tool graph from a txt file of matches

    Parameters
    ----------
    filename : str
        Path to the txt file. Matches are of the form "winner looser (weight)"
    players : list
        List of players. If None, the players are generated from the file.

    Returns
    -------
    g : graph-tool Graph
        Graph representing the matches
    """
    # If players is provided, generate an {id: rating} dictionary
    if players is not None:
        player_dict = {player.id: player.rating for player in players}
    else:
        player_dict = {}
    # Load the data
    data = np.loadtxt(filename, str, delimiter=" ")
    # Get the number of players by generating a set of all unique elements in the first two columns of data
    N = len(set(data[:, 0]).union(data[:, 1]))
    # Initialize the graph
    g = gt.Graph(directed=True)
    # Add vertex property for player ids
    v_id = g.new_vertex_property("int")
    # Add vertex property for player ratings
    v_rating = g.new_vertex_property("float")
    # Add edge property for match weights
    e_weight = g.new_edge_property("int")
    # Add vertex for each player
    vertices = list(g.add_vertex(N))
    # Build graph
    for match in data:
        i, j = int(match[0]), int(match[1])
        e = g.add_edge(vertices[i], vertices[j])
        if len(match) == 3:
            e_weight[e] = int(match[2])
        else:
            e_weight[e] = 1
        # Set properties
        v_id[vertices[i]] = i
        v_rating[vertices[i]] = player_dict[i] if players is not None else 0
        v_id[vertices[j]] = j
        v_rating[vertices[j]] = player_dict[j] if players is not None else 0
    # Set properties
    g.vertex_properties["id"] = v_id
    g.vertex_properties["rating"] = v_rating
    g.edge_properties["weight"] = e_weight

    return g


def read_matchlist(filename: str) -> np.ndarray:
    """
    Read a match list from a txt file and parse it into two dictionaries of dictionaries representing the out and in edges
    and the number of players and matches

    Parameters
    ----------
    filename : str
        Path to file containing match list of the form "winner, loser, (weight)". If no weight is provided, it is assumed to be 1.

    Returns
    -------
    match_list : ndarray
        Array of matches of the form [[i, j, w_ij], ...], where w_ij is the number of times i beats j. The types of the elements
        making up the tuple are str, str, int.
    """

    # Load
    data = np.loadtxt(filename, str, delimiter=" ")
    # Check for number of columns in data
    if data.shape[1] == 3:
        # Ensure correct data types
        data = np.array([[el[0], el[1], int(el[2])] for el in data], dtype=object)

    return data


def match_list_from_gt(g: gt.Graph, id: str, weight: Optional[str] = None) -> np.ndarray:
    """
    Generate a match list from a graph-tool graph

    Parameters
    ----------
    g : graph-tool Graph
        Graph representing the matches
    id : str
        Vertex property to use as id
    weight : Optional[str]
        Edge property to use as weight

    Returns
    -------
    match_list : ndarray
        Array of matches of the form [[i, j, w_ij], ...], where w_ij is the number of times i beats j. The types of the elements
        making up the tuple are str, str, int.
    """
    # Define dictionary to store player ids
    player_dict = {}
    # Get the number of edges
    E = g.num_edges()
    # Initialize match list
    match_list = np.zeros((E, 3), dtype=object)
    # Get the vertex property map for id
    id_property = g.vp[id]
    # Get the edge property map for weight if provided
    weight_property = g.ep[weight] if weight is not None else None

    # Initialize a dictionary to store the aggregated weights
    match_dict = defaultdict(int)

    # Loop through all edges
    for e in g.edges():
        # Get the source and target vertices
        source, target = e.source(), e.target()
        # Get the weight of the edge
        edge_weight = weight_property[e] if weight_property is not None else 1
        # Get the id of the source and target vertices
        source_id, target_id = id_property[source], id_property[target]
        # If there are any white spaces in the ids, replace them with underscores
        source_id = source_id.replace(" ", "_")
        target_id = target_id.replace(" ", "_")
        # Aggregate the weights
        match_dict[(source_id, target_id)] += edge_weight

    # Convert the dictionary to a list
    match_list = np.array(
        [[source, target, weight] for (source, target), weight in match_dict.items()], dtype=object
    )

    return match_list


def gt_from_match_list(matchlist: np.ndarray) -> gt.Graph:
    """
    Generate a graph-tool graph from a match list

    Parameters
    ----------
    matchlist : ndarray
        Array of matches of the form [[i, j, w_ij], ...], where w_ij is the number of times i beats j. The types of the elements
        making up the tuple are str, str, int.

    Returns
    -------
    g : graph-tool Graph
        Graph representing the matches
    """
    # Get the number of players
    N = len(set(matchlist[:, 0]).union(matchlist[:, 1]))
    # Initialize the graph
    g = gt.Graph(directed=True)
    # Add vertex property for player ids
    v_id = g.new_vertex_property("string")
    # Add edge property for match weights
    e_weight = g.new_edge_property("int")
    # Add vertex for each player
    vertices = list(g.add_vertex(N))
    # Build graph
    for match in matchlist:
        i, j = int(match[0]), int(match[1])
        e = g.add_edge(vertices[i], vertices[j])
        if len(match) == 3:
            e_weight[e] = int(match[2])
        else:
            e_weight[e] = 1
        # Set properties
        v_id[vertices[i]] = i
        v_id[vertices[j]] = j
    # Set properties
    g.vertex_properties["id"] = v_id
    g.edge_properties["weight"] = e_weight

    return g
