from otree.api import (
    models, widgets, BaseConstants, BaseSubsession, BaseGroup, BasePlayer,
    Currency as c, currency_range
)

doc = ''
class Constants(BaseConstants):
    name_in_url = 'pd'
    players_per_group = 2
    num_rounds = 10
    cc_payoff = c(5)
    cd_payoff = c(-10)
    dc_payoff = c(10)
    a_payoff = c(0)
    dd_payoff = c(0)
    punish_cost = c(2)
    punish_dammage = c(8)
    share_groups_public = 0.5
    continuation_probability = 95
class Subsession(BaseSubsession):
    def creating_session(self):
        self.session.vars['alive']=True
        self.session.vars['show']=True
        self.group_randomly()
        number = (self.session.num_participants/Constants.players_per_group*Constants.share_groups_public)
        for idx, g in enumerate(self.get_groups()):
            if idx<number:
                g.condition_public=True
            else:
                g.condition_public=False
    def random_continuation(self):
        import random
        self.session.vars['continuation_draw']=random.choice(range(1,100))
        if self.session.vars['continuation_draw']>Constants.continuation_probability:
            self.session.vars['alive']=False
    def shutdown(self):
        self.session.vars['show']=False
class Group(BaseGroup):
    condition_public = models.BooleanField(initial=False)
    def set_payoffs1(self):
        for p in self.get_players():
            p.payoff1()
    def set_payoffs2(self):
        for p in self.get_players():
            p.payoff2()
    def set_condition(self):
        if self.condition_public==True:
            for p in self.get_players():
                p.condition_public=True
        else:
            for idx, p in enumerate(self.get_players()):
                if idx==0:
                    p.condition_public=True
                else:
                    p.condition_public=False
class Player(BasePlayer):
    action = models.StringField(choices=[['Cooperate', 'Cooperate'], ['Defect', 'Defect'], ['Avoid', 'Avoid']], label='Choose your first round action:')
    punish = models.StringField(choices=[['punish', 'Punish'], ['nopunish', "Don't punish"]], label='Do you want to punish?')
    action_opponent = models.StringField()
    condition_public = models.BooleanField(initial=False)
    def payoff1(self):
        payoff_matrix = dict(
            Cooperate=dict(
                Cooperate=Constants.cc_payoff,
                Defect=Constants.cd_payoff,
                Avoid=Constants.a_payoff
            ),
            Defect=dict(
                Cooperate=Constants.dc_payoff, 
                Defect=Constants.dd_payoff,
                Avoid=Constants.a_payoff
            ),
            Avoid=dict(
                Cooperate=Constants.a_payoff, 
                Defect=Constants.a_payoff,
                Avoid=Constants.a_payoff
            ),
        )
        
        self.payoff = payoff_matrix[self.action][self.other_player().action]
        self.action_opponent = self.other_player().action
        
    def other_player(self):
        return self.get_others_in_group()[0]
    def payoff2(self):
        self.payoff = self.payoff-Constants.punish_cost*(self.punish=='punish')-Constants.punish_dammage*(self.other_player().punish=='punish')