import random

# Import the colrisk1 constants to use in the templates
from colrisk1.models import Constants as colrisk1Constants


def set_endowments(group):

    """
    We need to save the task and endowment variables in participant.vars
    because the player model is reset at the start of each round
    I assume that the players are put in a group randomly
    """
    print("Setting endowments")

    # Get participants
    players = group.get_players()
    participants = [player.participant for player in players]
    # Randomise participants
    random.shuffle(participants)

    if group.session.vars["group_type"] == "Luck":
        print("Luck Game")

        participants[0].vars.update(task_chosen="Luck", wealth="rich")
        participants[1].vars.update(task_chosen="Luck", wealth="rich")
        participants[2].vars.update(task_chosen="Luck", wealth="poor")
        participants[3].vars.update(task_chosen="Luck", wealth="poor")

    elif group.session.vars["group_type"] == "Effort":

        print("Effort Game")

        #print([p.vars["effort_score"] for p in participants])

        # Sort participants by their effort score
        participants.sort(key=lambda p:p.vars["effort_score"], reverse=True)

        # Set the top two scorers with high endowments, bottom two with low
        participants[0].vars.update(task_chosen="Effort", wealth="rich")
        participants[1].vars.update(task_chosen="Effort", wealth="rich")
        participants[2].vars.update(task_chosen="Effort", wealth="poor")
        participants[3].vars.update(task_chosen="Effort", wealth="poor")

    elif group.session.vars["group_type"] == "Mixed":

        print("Mixed Game")

        # First two participants are randomly given endowments
        participants[0].vars.update(task_chosen="Luck", wealth="rich")
        participants[1].vars.update(task_chosen="Luck", wealth="poor")
        
        # Second two given endowments based on who scored higher
        if participants[2].vars["effort_score"] > participants[3].vars["effort_score"]:
            participants[2].vars.update(task_chosen="Effort", wealth="rich")
            participants[3].vars.update(task_chosen="Effort", wealth="poor")
        else:
            participants[2].vars.update(task_chosen="Effort", wealth="poor")
            participants[3].vars.update(task_chosen="Effort", wealth="rich")

    else:
        raise ValueError("No valid group type selected!")


    # Set the participant payoffs to the endowment level
    for participant in participants:
        if participant.vars["wealth"] == "rich":
            participant.vars.update(endowment=colrisk1Constants.endowment_rich)
            participant.payoff = colrisk1Constants.endowment_rich
        elif participant.vars["wealth"] == "poor":
            participant.vars.update(endowment=colrisk1Constants.endowment_poor)
            participant.payoff = colrisk1Constants.endowment_poor