from enum import Enum
from typing import List


class Treatment(Enum):
    FIXED_WAGE = "F"
    PIECE_RATE = "P"
    TOURNAMENT = "T"

    @staticmethod
    def set_fixed_wage_payoffs(players, fixed_wage: float):
        """Set fixed wage"""
        for p in players:
            p.fixed_wage_payoff = fixed_wage
            p.payoff += fixed_wage
            # Save to participant model
            pt = p.participant
            pt.fixed_wage_payoff = fixed_wage

    @staticmethod
    def set_piece_rate_payoffs(players, fixed_wage: float, piece_rate: float):
        """Set fixed wage plus piece rate"""
        for p in players:
            p.fixed_wage_payoff = fixed_wage
            piece_rate_payoff = piece_rate * p.score
            p.piece_rate_payoff = piece_rate_payoff
            p.payoff += fixed_wage + piece_rate_payoff
            # Save to participant model
            pt = p.participant
            pt.fixed_wage_payoff = fixed_wage
            pt.piece_rate_payoff = piece_rate_payoff

    @staticmethod
    def set_tournament_payoffs(players, fixed_wage: float, tournament_prizes: List[float]):
        """Set fixed wage plus tournament prize"""
        # Create list of tuples with (score, accuracy, player) for each player
        scores = [(p.score, p.score / p.worked if p.worked else 0, p) for p in players]
        scores.sort(key=lambda x: (x[0], x[1]), reverse=True)  # Sorts by score first, then by accuracy
        rankings = range(1, len(tournament_prizes) + 1)

        for ranking, score, prize in zip(rankings, scores, tournament_prizes):
            player = score[2]
            player.ranking = ranking
            player.fixed_wage_payoff = fixed_wage
            player.tournament_payoff = prize
            player.payoff += fixed_wage + prize
            # Save to participant model
            pt = player.participant
            pt.fixed_wage_payoff = fixed_wage
            pt.tournament_payoff = prize
            pt.ranking = ranking

    def set_payoffs(self, group, fixed_wage: float, piece_rate: float, tournament_prizes: List[float]):
        """Set fixed wage / piece rate / tournament prize respective to treatment"""
        players = group.get_players()

        if self == self.FIXED_WAGE:
            self.set_fixed_wage_payoffs(players, fixed_wage)
        elif self == self.PIECE_RATE:
            self.set_piece_rate_payoffs(players, fixed_wage, piece_rate)
        elif self == self.TOURNAMENT:
            self.set_tournament_payoffs(players, fixed_wage, tournament_prizes)