from otree.api import *


cu = cu
doc = ''


class C(BaseConstants):
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 21
    MIN_ROUNDS = 20 
    NAME_IN_URL = 'my_trust2'
    GAME_SEQ = 2


# swap role at first round
class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    contribution = models.IntegerField(widget=widgets.RadioSelect)


# FUNCTIONS
def creating_session(subsession: Subsession):
    if subsession.round_number == 1:
        role_1 = []
        role_2 = []
        for p in subsession.get_players():
            if p.participant.role == 1:
                role_1.append(p)
            if p.participant.role == 2:
                role_2.append(p)
        group_matrix = []
        for p1, p2 in zip(role_1, role_2):
            group_matrix.append([p2, p1])       
        subsession.set_group_matrix(group_matrix)   
    else:
        subsession.group_like_round(1)


def set_payoffs(group: Group):
    p1 = group.get_player_by_id(1)
    p2 = group.get_player_by_id(2)
    public_points = (
        cu(p1.contribution) * group.session.config["unit_more"]
        + cu(p2.contribution) * group.session.config["unit_less"]
    )
    profit = public_points/2
    p1.payoff = group.session.config["init_more"] - p1.contribution + profit
    p2.payoff = group.session.config["init_less"] - p2.contribution + profit
    if group.round_number == C.NUM_ROUNDS:
        p1.participant.sum_points2 = sum_payoff(p1)
        p2.participant.sum_points2 = sum_payoff(p2)


def is_successful(group: Group):
    p1 = group.in_round(group.round_number - 1).get_player_by_id(1)
    p2 = group.in_round(group.round_number - 1).get_player_by_id(2)
    public_points = (
        cu(p1.contribution) * group.session.config["unit_more"]
        + cu(p2.contribution) * group.session.config["unit_less"]
    )
    return "Total collective contribution in the common pool was {}, and each person received {} points.".format(public_points, public_points/2)


def contribution_choices(player: Player):
    init = (
        player.session.config["init_more"]
        if player.id_in_group == 1
        else player.session.config["init_less"]
    )
    return range(0, int(init + 1), 1)


def other_payoff(player: Player):
    group = player.group
    if player.id_in_group == 1:
        return group.in_round(player.round_number - 1).get_player_by_id(2).payoff
    else:
        return group.in_round(player.round_number - 1).get_player_by_id(1).payoff


def other_choice(player: Player):
    group = player.group
    if player.id_in_group == 1:
        return group.in_round(player.round_number - 1).get_player_by_id(2).contribution
    else:
        return group.in_round(player.round_number - 1).get_player_by_id(1).contribution


def sum_payoff(player: Player):
    return sum(map(lambda x: x.payoff, player.in_all_rounds()))


def my_choice(player: Player):
    return player.in_round(player.round_number - 1).contribution


def my_payoff(player: Player):
    return player.in_round(player.round_number - 1).payoff



# PAGES
class Introduction(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player: Player):
        player = player
        return player.round_number == 1


class Send(Page):
    form_model = 'player'
    form_fields = ['contribution']
    @staticmethod
    def vars_for_template(player :Player):
        if player.round_number == 1:
            total_payoff = 0
            my_contribution = 0
            other_contribution = 0
            my_last_payoff = 0
            other_last_payoff = 0
            get_or_not = 0
        else:
            total_payoff = sum_payoff(player)
            my_contribution = my_choice(player)
            other_contribution = other_choice(player)
            my_last_payoff = my_payoff(player)
            other_last_payoff = other_payoff(player)
            get_or_not = is_successful(player.group)
        return dict(
            sum_payoff = total_payoff,
            my_choice = my_contribution,
            other_choice = other_contribution,
            my_payoff = my_last_payoff,
            other_payoff = other_last_payoff,
            is_successful = get_or_not
        )


class ResultsWaitPage(WaitPage):
    after_all_players_arrive = 'set_payoffs'
    body_text = 'Please wait for the other player to make a decision.'
    title_text = ""


class End(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player: Player):
        player = player
        return player.round_number == C.NUM_ROUNDS

    @staticmethod
    def vars_for_template(player :Player):
        if player.round_number == 1:
            total_payoff = 0
        else:
            total_payoff = sum_payoff(player)
        return dict(
            sum_payoff = total_payoff,
        )


page_sequence = [Introduction, Send, ResultsWaitPage, End]