from otree.api import *
from random import randint, uniform, sample, shuffle


doc = """
Your app description
"""

class C(BaseConstants):
    NAME_IN_URL = 'helping_game'
    PLAYERS_PER_GROUP = 8
    NUM_BASE_ROUNDS = 50
    ENDOWMENT = 1
    PAYOFF_P_T = 2  # payoff to producer if trade happens
    PAYOFF_C_T = 20  # payoff to consumer if trade happens
    PAYOFF_P_N = 8  # payoff to producer if no trade happens
    PAYOFF_C_N = 6  # payoff to consumer if no trade happens
    PROB_EXTRA_ROUND = 0.75
    TOKEN_SUPPLY = 400

    # initialization
    NUM_ROUNDS = NUM_BASE_ROUNDS + 100
    TOKEN_BONUS = []
    with open("helping_game/helicopter_money.csv", "r") as f:
        f.readline()
        for line in f.readlines():
            line = line.strip().split(",")
            TOKEN_BONUS.append(int(line[1]))
    print(TOKEN_BONUS)


class Subsession(BaseSubsession):
    pass


def rematch(matrix):
    allowed_permutations = {
        0: [1, 0, 3, 2],
        1: [2, 0, 3, 1],
        2: [3, 0, 1, 2],
        3: [2, 3, 0, 1],
        4: [3, 2, 0, 1],
        5: [1, 3, 0, 2],
        6: [1, 2, 3, 0],
        7: [2, 3, 1, 0],
        8: [3, 2, 1, 0],
    }

    new_prods = [pair[1] for pair in matrix]
    new_cons = [pair[0] for pair in matrix]
    # print(new_prods, new_cons)

    permutation = allowed_permutations[randint(0, 8)]
    # print(permutation)

    new_matrix = []
    for i in range(4):
        new_matrix.append([new_prods[i], new_cons[permutation[i]]])
    # print(matrix)
    # print(new_matrix)
    return new_matrix


def flip_roles(group):
    for player in group.get_players():
        if player.participant.role == "producer":
            player.participant.role = "consumer"
        else:
            player.participant.role = "producer"


def get_matrix(group):
    matrix = []
    for player in group.get_players():
        if player.participant.role == "producer":
            matrix.append([player.id_in_group, player.participant.partner])
    return matrix


def set_matrix(group, matrix):
    # print(matrix)
    for pair in matrix:
        prod = group.get_player_by_id(pair[0])
        con = group.get_player_by_id(pair[1])
        prod.participant.partner = con.id_in_group
        con.participant.partner = prod.id_in_group


class ShuffleWaitPage(WaitPage):
    wait_for_all_groups = False
    def after_all_players_arrive(group):
        if group.round_number > 1:
            matrix = get_matrix(group)
            matrix = rematch(matrix)
            set_matrix(group, matrix)
            flip_roles(group)
        else:
            helicopter_money = bool(randint(0, 1))
            num_rounds = C.NUM_BASE_ROUNDS + num_extra_rounds()
            prods = sample(range(1, 9), k=4)
            cons = []
            for i in range(1, 9):
                if i not in prods:
                    cons.append(i)
            shuffle(prods)
            shuffle(cons)
            matrix = []
            for i in range(4):
                matrix.append([prods[i], cons[i]])
            for player in group.get_players():
                player.participant.helicopter_money = helicopter_money
                player.participant.num_rounds = num_rounds
                if player.id_in_group in prods:
                    player.participant.role = "producer"
                else:
                    player.participant.role = "consumer"
            set_matrix(group, matrix)


def num_extra_rounds():
    extra_rounds = 0
    while True:
        p = uniform(0, 1)
        if p <= C.PROB_EXTRA_ROUND:
            extra_rounds += 1
        else:
            return extra_rounds


class Group(BaseGroup):
    pass


def wtb_max(player):
    return player.participant.tokens


class Player(BasePlayer):
    tokens = models.IntegerField(initial=C.TOKEN_SUPPLY)
    points = models.IntegerField(initial=0)

    wts = models.IntegerField(
        min=0,
        label="Willingness to sell?"
    )
    wtb = models.IntegerField(
        min=0,
        label="Willingness to buy?"
    )


def supply_participants(group):
    if group.get_player_by_id(1).participant.helicopter_money:
        if group.round_number == 1:
            for player in group.get_players():
                player.participant.tokens = C.TOKEN_SUPPLY
                player.participant.points = 0
        else:
            for player in group.get_players():
                player.participant.tokens += C.TOKEN_BONUS[group.round_number-2]
    else:
        if group.round_number == 1:
            for player in group.get_players():
                player.participant.tokens = C.TOKEN_SUPPLY
                player.participant.points = 0


# PAGES
class InitWaitPage(WaitPage):
    wait_for_all_groups = False
    after_all_players_arrive = supply_participants


class WTS(Page):
    form_model = "player"
    form_fields = ["wts"]

    @staticmethod
    def is_displayed(player):
        return player.participant.role == "producer"


class WTB(Page):
    form_model = "player"
    form_fields = ["wtb"]

    @staticmethod
    def is_displayed(player):
        return player.participant.role == "consumer"


def set_payoffs(group):
    for prod in group.get_players():
        if prod.participant.role == "producer":
            con = group.get_player_by_id(prod.participant.partner)
            participant1 = prod.participant
            participant2 = con.participant
            if con.wtb >= prod.wts:
                participant1.points += C.PAYOFF_P_T
                participant2.points += C.PAYOFF_C_T
                participant1.tokens += prod.wts
                participant2.tokens -= prod.wts
            else:
                participant1.points += C.PAYOFF_P_N
                participant2.points += C.PAYOFF_C_N
            prod.tokens = participant1.tokens
            con.tokens = participant2.tokens
            prod.points = participant1.points
            con.points = participant2.points
    if group.round_number == group.get_player_by_id(1).participant.num_rounds:
        for player in group.get_players():
            player.participant.payoff = player.participant.points


class ResultsWaitPage(WaitPage):
    wait_for_all_groups = False
    after_all_players_arrive = set_payoffs


class Results(Page):
    @staticmethod
    def vars_for_template(player):
        msp = 0
        divisor = 0
        last_round = player.round_number == player.participant.num_rounds
        # print(player.round_number, player.participant.num_rounds, last_round)
        for prod in player.group.get_players():
            if prod.participant.role == "producer":
                if player.group.get_player_by_id(prod.participant.partner).wtb >= prod.wts:
                    msp += prod.wts
                    divisor += 1

        if divisor == 0:
            msp = None
        else:
            msp = int(msp / divisor)

        partner = player.group.get_player_by_id(player.participant.partner)
        if player.participant.role == "producer":
            trade_happened = partner.wtb >= player.wts
        else:
            trade_happened = player.wtb >= partner.wts

        if player.participant.helicopter_money:
            if player.round_number > 1:
                helicopter_bonus = C.TOKEN_BONUS[player.round_number-2]
            else:
                helicopter_bonus = 0
        else:
            helicopter_bonus = 0
        return dict(
            trade_happened = trade_happened,
            mean_selling_price = msp,
            last_round = last_round,
            helicopter_bonus = helicopter_bonus,
        )


def vars_for_admin_report(subsession):
    data = []
    for player in subsession.get_players():
        data.append(f"ID: {player.in_round(subsession.round_number).id},    Tokens: {player.in_round(subsession.round_number).tokens},    Points: {player.in_round(subsession.round_number).points}")
    return dict(
        data=data
    )


page_sequence = [ShuffleWaitPage, InitWaitPage, WTS, WTB, ResultsWaitPage, Results]