from otree.api import *
import random

doc = """
This experiment builds on Dal Bo et al. (2010) and uses the Prisoner's Dilemma and the Coordination Game
to study whether the Democracy Effect operates by changing subjects' beliefs or preferences.
"""


class C(BaseConstants):
    NAME_IN_URL = 'dem_pd'
    PLAYERS_PER_GROUP = 5
    ROUNDS_PER_PART = 10
    NUM_ROUNDS = ROUNDS_PER_PART * 3
    PAYOFF_A = cu(60)
    PAYOFF_A_mod = cu(48)
    PAYOFF_B = cu(50)
    PAYOFF_C = cu(30)
    PAYOFF_D = cu(20)
    PAYOFF_N = cu(0)
    TOTAL_BALLS = 10
    RED_BALLS = 2
    PROB_OVERRIDING = 0.6
    PROB_MOD_IF_OVER = 0.9
    GUESS_ROUNDS = [[1, 5], [2, 1], [3, 1]]  # Rounds with guessing task, each element is [part, round_in_this_part]


class Subsession(BaseSubsession):
    part = models.IntegerField()
    part_roman = models.StringField()
    round_in_part = models.IntegerField()
    guess_round = models.BooleanField(default=False)


def creating_session(subsession: Subsession):
    # Determine the part and the round number
    r = subsession.round_number
    n = C.ROUNDS_PER_PART
    p = C.PLAYERS_PER_GROUP
    subsession.part = (r - 1) // n + 1
    subsession.part_roman = 'I' * subsession.part
    subsession.round_in_part = (r - 1) % n + 1

    # Check if there is guessing task in this round
    part_round = [subsession.part, subsession.round_in_part]
    if part_round in C.GUESS_ROUNDS:
        subsession.guess_round = True

    # In the first round, create lists of guesses and partner's actions
    if subsession.round_number == 1:
        for participant in [player.participant for player in subsession.get_players()]:
            participant.guess_choices, participant.guess_partner_actions = [], []

    # If first round, group randomly, otherwise, stay in the same group
    if subsession.round_number == 1:
        subsession.group_randomly()
    else:
        subsession.group_like_round(r - 1)

    groups = subsession.get_groups()

    for g in groups:
        # Player who stays alone is set by player's id
        solo_id = (r-1) % p + 1
        solo = g.get_player_by_id(solo_id)
        solo.my_role = 'solo'

        # The rest of the players are paired randomly
        players = g.get_players()
        players.remove(solo)
        random.shuffle(players)
        pairs = [players[i:i + 2] for i in range(0, len(players), 2)]

        # Set partners for paired players
        for pair in pairs:
            pair[0].partner = pair[1].id_in_group
            pair[0].my_role = 'active'
            pair[1].partner = pair[0].id_in_group
            pair[1].my_role = 'passive'

        # Set solo player as own partner
        solo.partner = solo.id_in_group
        solo.my_role = 'solo'


class Group(BaseGroup):
    sum_vote = models.IntegerField(
        doc="Number of votes for modifying payoffs"
    )
    group_vote = models.BooleanField(
        doc="True: majority is for modifying; False: against"
    )
    overridden = models.BooleanField(
        default=False,
        doc="True: the group vote is overridden; False: not"
    )
    final_group_choice = models.BooleanField(
        doc="True: payoffs are modified in the group; False: not"
    )


class Player(BasePlayer):
    vote = models.BooleanField(
        doc="True: supporting payoff modification; False: not"
    )
    action = models.IntegerField(
        choices=[1, 2],
        doc="1: defect; 2: cooperate"
    )
    partner_action = models.IntegerField(
        choices=[1, 2]
    )
    partner = models.IntegerField()
    my_role = models.StringField()
    picked_ball = models.StringField(
        doc="Number of the ball that a subject picks"
    )
    picked_ball_color = models.StringField(
        doc="red: 1 (defect); black: 2: cooperate"
    )
    guess = models.IntegerField(
        label="Your guess:",
        min=0,
        max=100
    )
    comp_balls_1 = models.BooleanField()
    comp_balls_2 = models.BooleanField()
    comp_balls_3 = models.BooleanField()


# FUNCTIONS
def set_payoffs(group: Group):
    for p in group.get_players():
        set_payoff(p)


def set_payoff(player: Player):
    # Set payoff for unilateral defection
    payoff_a = C.PAYOFF_A
    k = C.ROUNDS_PER_PART * (player.subsession.part - 1) + 1
    if player.subsession.part > 1:
        if player.in_round(k).group.final_group_choice:
            payoff_a = C.PAYOFF_A_mod

    payoff_matrix = {
        (1, 1): C.PAYOFF_C,
        (1, 2): payoff_a,
        (2, 1): C.PAYOFF_D,
        (2, 2): C.PAYOFF_B,
    }

    # Set payoffs for paired players
    if player.partner == player.id_in_group:
        player.payoff = C.PAYOFF_N
    # Set payoffs for solo players
    else:
        other = player.group.get_player_by_id(player.partner)
        player.payoff = payoff_matrix[(player.action, other.action)]


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

    @staticmethod
    def is_displayed(player: Player):
        return (player.subsession.part == 3) & (player.subsession.round_in_part == 1)

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            nred=C.RED_BALLS,
            nblack=C.TOTAL_BALLS - C.RED_BALLS
        )


class CompBalls1(CompBallsBase):
    form_fields = ['comp_balls_1']


class CompBalls2(CompBallsBase):
    form_fields = ['comp_balls_2']


class CompBalls3(CompBallsBase):
    form_fields = ['comp_balls_3']


class Voting(Page):
    form_model = "player"
    form_fields = ["vote"]

    @staticmethod
    def is_displayed(player):
        return (player.subsession.part > 1) & (player.subsession.round_in_part == 1)


class ResultsWaitVoting(WaitPage):
    @staticmethod
    def is_displayed(player):
        return (player.subsession.part > 1) & (player.subsession.round_in_part == 1)

    @staticmethod
    def after_all_players_arrive(group: Group):

        # Determine majority vote
        group.sum_vote = sum([p.vote for p in group.get_players()])
        group.group_vote = (group.sum_vote * 2 > C.PLAYERS_PER_GROUP)

        # Random overriding
        p = C.PROB_OVERRIDING
        group.overridden = random.choices([True, False], weights=[p, 1-p], k=1)[0]

        # Determine final choice

        if group.overridden:
            # Randomly choose the outcome of the group
            q = C.PROB_MOD_IF_OVER
            group.final_group_choice = random.choices([True, False], weights=[q, 1-q], k=1)[0]
        else:
            group.final_group_choice = group.group_vote


class VoteResults1(Page):
    @staticmethod
    def is_displayed(player):
        return (player.subsession.part > 1) & (player.subsession.round_in_part == 1)

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            overridden=player.group.overridden
        )


class VoteResults2(Page):
    @staticmethod
    def is_displayed(player):
        return (player.subsession.part > 1) & (player.subsession.round_in_part == 1)

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            overridden=player.group.overridden,
            choice=player.group.final_group_choice
        )


class VoteResults3(Page):
    @staticmethod
    def is_displayed(player):
        return (player.subsession.part > 1) & (player.subsession.round_in_part == 1) & player.group.overridden

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            group_size=C.PLAYERS_PER_GROUP,
            majority_size=C.PLAYERS_PER_GROUP // 2 + 1,
            votes=player.group.sum_vote
        )


class DecisionBalls(Page):
    form_model = 'player'
    form_fields = ['picked_ball', 'picked_ball_color']

    @staticmethod
    def vars_for_template(player: Player):
        # Prepare a list of colors of the balls
        n_red = C.RED_BALLS
        n_black = C.TOTAL_BALLS - C.RED_BALLS
        colors = ['red'] * n_red + ['black'] * n_black
        random.shuffle(colors)

        # Information about the results of previous rounds
        p, r, a, f, s = [], [], [], [], []

        # Set nice role names
        round_num = player.round_number
        role_name = {'solo': 'Unmatched'}
        if player.subsession.part < 3:
            role_name['active'] = 'Matched'
            role_name['passive'] = 'Matched'
        else:
            role_name['active'] = 'Matched, A'
            role_name['passive'] = 'Matched, B'

        # If there were rounds in this part before, send information about them
        k = C.ROUNDS_PER_PART * (player.subsession.part - 1) + 1
        if player.subsession.round_in_part > 1:
            p = player.in_all_rounds()[0:round_num - 1][k-1:]
            r = [role_name[iplayer.my_role] for iplayer in p]
            a = ['' if iplayer.my_role == 'solo' else iplayer.action for iplayer in p]
            s = ['' if iplayer.my_role == 'solo' else iplayer.partner_action for iplayer in p]
            f = [iplayer.payoff.__int__() for iplayer in p]

        return dict(
            colors=colors,
            actions=a,
            partner_actions=s,
            roles=r,
            payoffs=f,
        )

    @staticmethod
    def is_displayed(player):
        return (player.my_role == 'passive') & (player.subsession.part == 3)


class DecisionBallsResult(Page):
    form_model = 'player'
    form_fields = ['action']

    @staticmethod
    def vars_for_template(player: Player):

        # Information about the results of previous rounds
        p, r, a, f, s = [], [], [], [], []

        # Set nice role names
        round_num = player.round_number
        role_name = {'solo': 'Unmatched'}
        if player.subsession.part < 3:
            role_name['active'] = 'Matched'
            role_name['passive'] = 'Matched'
        else:
            role_name['active'] = 'Matched, A'
            role_name['passive'] = 'Matched, B'

        # If there were rounds in this part before, send information about them
        k = C.ROUNDS_PER_PART * (player.subsession.part - 1) + 1
        if player.subsession.round_in_part > 1:
            p = player.in_all_rounds()[0:round_num - 1][k-1:]
            r = [role_name[iplayer.my_role] for iplayer in p]
            a = ['' if iplayer.my_role == 'solo' else iplayer.action for iplayer in p]
            s = ['' if iplayer.my_role == 'solo' else iplayer.partner_action for iplayer in p]
            f = [iplayer.payoff.__int__() for iplayer in p]

        return dict(
            picked_ball=player.picked_ball,
            picked_ball_color=player.picked_ball_color,
            actions=a,
            partner_actions=s,
            roles=r,
            payoffs=f,
        )

    @staticmethod
    def is_displayed(player):
        return (player.my_role == 'passive') & (player.subsession.part == 3)


class Decision(Page):
    form_model = 'player'
    form_fields = ['action']

    @staticmethod
    def is_displayed(player):
        if player.subsession.part == 3:
            return player.my_role == 'active'
        else:
            return player.my_role != 'solo'

    @staticmethod
    def vars_for_template(player: Player):
        # Information about the results of previous rounds
        p, r, a, f, s = [], [], [], [], []

        # Set nice role names
        round_num = player.round_number
        role_name = {'solo': 'Unmatched'}
        if player.subsession.part < 3:
            role_name['active'] = 'Matched'
            role_name['passive'] = 'Matched'
        else:
            role_name['active'] = 'Matched, A'
            role_name['passive'] = 'Matched, B'

        # If there were rounds in this part before, send information about them
        k = C.ROUNDS_PER_PART * (player.subsession.part - 1) + 1
        if player.subsession.round_in_part > 1:
            p = player.in_all_rounds()[0:round_num - 1][k-1:]
            r = [role_name[iplayer.my_role] for iplayer in p]
            a = ['' if iplayer.my_role == 'solo' else iplayer.action for iplayer in p]
            s = ['' if iplayer.my_role == 'solo' else iplayer.partner_action for iplayer in p]
            f = [iplayer.payoff.__int__() for iplayer in p]

        # Default payoff for unilateral defection
        payoff_a = C.PAYOFF_A

        # If payoffs were modified, modify the payoff on the page
        if player.subsession.part > 1:
            if player.in_round(k).group.final_group_choice:
                payoff_a = C.PAYOFF_A_mod

        return dict(
            actions=a,
            partner_actions=s,
            roles=r,
            payoffs=f,
            payoff_a=payoff_a
        )


class NoDecision(WaitPage):
    title_text = 'You are unmatched in this round'
    body_text = 'Please wait for the other participants.'

    @staticmethod
    def is_displayed(player):
        return player.my_role == 'solo'


class ResultsWaitPage(WaitPage):
    after_all_players_arrive = set_payoffs


class Guess(Page):
    form_model = 'player'
    form_fields = ['guess']

    @staticmethod
    def is_displayed(player: Player):
        return player.subsession.guess_round & (player.my_role != 'solo')

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        p = player.participant

        # Save guess and last partner's action
        p.guess_choices.append(player.guess)
        p_action = player.group.get_player_by_id(player.partner).action
        p.guess_partner_actions.append(p_action)


class Results(Page):
    @staticmethod
    def is_displayed(player):
        return player.my_role != 'solo'

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        player.partner_action = player.group.get_player_by_id(player.partner).action

    @staticmethod
    def vars_for_template(player: Player):
        opponent = player.group.get_player_by_id(player.partner)
        return dict(
            opponent=opponent,
            same_choice=player.action == opponent.action,
            my_decision=player.field_display('action'),
            opponent_decision=opponent.field_display('action'),
        )


class WaitBetweenParts(Page):
    @staticmethod
    def is_displayed(player: Player):
        return (player.subsession.part in [1, 2]) & (player.subsession.round_in_part == C.ROUNDS_PER_PART)


page_sequence = [
    CompBalls1, CompBalls2, CompBalls3,
    Voting, ResultsWaitVoting,
    VoteResults1, VoteResults2, VoteResults3,
    DecisionBalls, DecisionBallsResult, Decision, NoDecision, ResultsWaitPage, Guess,
    Results, WaitBetweenParts
]