from otree.api import *
c = cu

doc = ''
class C(BaseConstants):
    NAME_IN_URL = 'repeated_PD_punishment'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 50
    MATRIX = (1, -2, 3, -5, 1, 3, -2, 0, -3, -2, 1, -5, -2, -3, -5)
class Subsession(BaseSubsession):
    pass
def creating_session(subsession: Subsession):
    session = subsession.session
    import random
    
    for player in subsession.get_players():
        if player.round_number == 1:
            maxround = 5 # play at least 5 round
            while (maxround < 50):
                maxround += 1
                if (random.uniform(0, 1) < 0.25): break
            for g in subsession.get_groups():
                p1 = g.get_player_by_id(1)
                p1.participant.vars['maxround'] = maxround
    
class Group(BaseGroup):
    maxround = models.IntegerField()
def set_payoffs(group: Group):
    for p in group.get_players():
        set_payoff(p)
class Player(BasePlayer):
    decision = models.StringField()
def other_player(player: Player):
    group = player.group
    return player.get_others_in_group()[0]
def set_payoff(player: Player):
    cc = C.MATRIX[0]
    cdrow = C.MATRIX[1]
    cprow = C.MATRIX[3]
    dcrow = C.MATRIX[5]
    dd = C.MATRIX[7]
    dprow = C.MATRIX[8]
    pcrow = C.MATRIX[10]
    pdrow = C.MATRIX[12]
    pp = C.MATRIX[14]
    
    
    points_matrix = {
        'C':
            {
                'C': cc,
                'D': cdrow,
                'P': cprow
            },
        'D':
            {
                'C': dcrow,
                'D': dd,
                'P': dprow
            },
        'P':
            {
                'C': pcrow,
                'D': pdrow,
                'P': pp
            }
    }
    
    other = other_player(player)
    player.payoff = (points_matrix[player.decision][other.decision])
class Instructions(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1
    @staticmethod
    def vars_for_template(player: Player):
        cc = C.MATRIX[0]
        cdrow = C.MATRIX[1]
        cdcol = C.MATRIX[2]
        cprow = C.MATRIX[3]
        cpcol = C.MATRIX[4]
        dcrow = C.MATRIX[5]
        dccol = C.MATRIX[6]
        dd = C.MATRIX[7]
        dprow = C.MATRIX[8]
        dpcol = C.MATRIX[9]
        pcrow = C.MATRIX[10]
        pccol = C.MATRIX[11]
        pdrow = C.MATRIX[12]
        pdcol = C.MATRIX[13]
        pp = C.MATRIX[14]
        
        return dict( cc = cc,
                     cdrow = cdrow,
                     cdcol = cdcol, 
                     cprow = cprow, 
                     cpcol = cpcol,
                     dcrow = dcrow,
                     dccol = dccol,
                     dd = dd,
                     dprow = dprow,
                     dpcol = dpcol,
                     pcrow = pcrow,
                     pccol = pccol,
                     pdrow = pdrow,
                     pdcol = pdcol,
                     pp = pp )
class PreDecision(Page):
    form_model = 'player'
    timeout_seconds = 1
class Decision(Page):
    form_model = 'player'
    form_fields = ['decision']
    @staticmethod
    def vars_for_template(player: Player):
        
        cc = C.MATRIX[0]
        cdrow = C.MATRIX[1]
        cdcol = C.MATRIX[2]
        cprow = C.MATRIX[3]
        cpcol = C.MATRIX[4]
        dcrow = C.MATRIX[5]
        dccol = C.MATRIX[6]
        dd = C.MATRIX[7]
        dprow = C.MATRIX[8]
        dpcol = C.MATRIX[9]
        pcrow = C.MATRIX[10]
        pccol = C.MATRIX[11]
        pdrow = C.MATRIX[12]
        pdcol = C.MATRIX[13]
        pp = C.MATRIX[14]
        
        return dict(
            roundnumber = player.round_number,
            cc = cc,
            cdrow = cdrow,
            cdcol = cdcol, 
            cprow = cprow, 
            cpcol = cpcol,
            dcrow = dcrow,
            dccol = dccol,
            dd = dd,
            dprow = dprow,
            dpcol = dpcol,
            pcrow = pcrow,
            pccol = pccol,
            pdrow = pdrow,
            pdcol = pdcol,
            pp = pp
        )
        
class ResultsWaitPage(WaitPage):
    after_all_players_arrive = set_payoffs
class Results(Page):
    form_model = 'player'
    @staticmethod
    def vars_for_template(player: Player):
        cc = C.MATRIX[0]
        cdrow = C.MATRIX[1]
        cdcol = C.MATRIX[2]
        cprow = C.MATRIX[3]
        cpcol = C.MATRIX[4]
        dcrow = C.MATRIX[5]
        dccol = C.MATRIX[6]
        dd = C.MATRIX[7]
        dprow = C.MATRIX[8]
        dpcol = C.MATRIX[9]
        pcrow = C.MATRIX[10]
        pccol = C.MATRIX[11]
        pdrow = C.MATRIX[12]
        pdcol = C.MATRIX[13]
        pp = C.MATRIX[14]
        
        other = other_player(player)
        
        return dict(
            roundnumber = player.round_number,
            other_player_decision = other.decision,
            same_choice = player.decision == other.decision,
            earned = player.payoff >= 0,
            othearned = other.payoff >= 0,
            lostpayoff = abs(player.payoff),
            lostothpayoff = abs(other.payoff),
            cc = cc,
            cdrow = cdrow,
            cdcol = cdcol, 
            cprow = cprow, 
            cpcol = cpcol,
            dcrow = dcrow,
            dccol = dccol,
            dd = dd,
            dprow = dprow,
            dpcol = dpcol,
            pcrow = pcrow,
            pccol = pccol,
            pdrow = pdrow,
            pdcol = pdcol,
            pp = pp
        )
class FinalPage(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        participant = player.participant
        p1 = player.group.get_player_by_id(1)
        maxround = p1.participant.vars['maxround']
        return player.round_number == maxround
    @staticmethod
    def vars_for_template(player: Player):
        other = other_player(player)
        table = '<table cellpadding=5><tr><th align=center>Round</th><th align=center>Own<br>choice</th><th align=center>Other<br>choice</th><th align=center>Own<br>payoff</th><th align=center>Other<br>payoff</th></tr>'
        for i in range(1,player.round_number+1):
            table += '<tr><td align=center>' + str(i) + '</td>'
            table += '<td align=center>' + player.in_round(i).decision + '</td>'
            table += '<td align=center>' + other.in_round(i).decision + '</td>'
            table += '<td align=center>' + str(player.in_round(i).payoff).split(' ')[0] + '</td>'
            table += '<td align=center>' + str(other.in_round(i).payoff).split(' ')[0] + '</td></tr>'
        
        table += '</table>'
        
        return dict( table = table )
page_sequence = [Instructions, PreDecision, Decision, ResultsWaitPage, Results, FinalPage]