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

author = 'Menglong Guan and Ryan Oprea'

doc = """
This is a repeated PD. In each math, two players decide which rule they want to use. 
The combination of chosen rules and a subject's implementation of his/her rule determine his/her payoff.
"""


class Constants(BaseConstants):
    name_in_url = 'Part2'
    players_per_group = 2
    conversion_rate = 1/100 # one experiment token for 0.01 dollar
    pay_method = 1 # pay_method=1 if pay the amount listed in strategy choice game; pay_method=0 if pay based on actual action paths

    time_limit = False # if ture, the experiment ends after time_limit_seconds. Not used in the current version
    time_limit_seconds = 10800

    #define payoffs of the stage game
    betray_payoff = 50
    betrayed_payoff_simple = 13 # use 13 for BOA=0.833; use 8 for BOA=0.779
    betrayed_payoff_complex = 23 # use 23.03279 for BOA=0.833; use 22.21311 for BOA=0.779, can use 22 with BOA=0.766
    both_cooperate_payoff=42
    both_defect_payoff = 25

    # deduction of payoff if not follow chosen rule
    payoff_deduct = -20

    delta=0.8 # the above stage game are chosen under delta=0.8
    # define payoffs of the strategy choice game (S1, C)=(AD, TF3T)
    AD_both = round(1/(1-delta)*both_defect_payoff)
    ADTF3T = round((1+delta+delta**2)*betray_payoff+delta**3/(1-delta)*both_defect_payoff)
    TF3TAD = round((1+delta+delta**2)*betrayed_payoff_complex+delta**3/(1-delta)*both_defect_payoff)
    TF3T_both = round(1/(1-delta)*both_cooperate_payoff)

    # define payoffs of the strategy choice game (S1, S2)=(AD, TFT)
    ADTFT=round(betray_payoff+delta/(1-delta)*both_defect_payoff)
    TFTAD=round(betrayed_payoff_simple+delta/(1-delta)*both_defect_payoff)
    TFT_both = round(1/(1-delta)*both_cooperate_payoff)

    # number of matches, randomly determined number of rounds in each match
    num_matches=10 #need to be an even integer

    ## if generate a sequence of game lengths in each session
    #match_duration = np.random.geometric(p=(1-delta), size=num_matches)# p is the probability the match ends after each round; size is the number of matchs

    ## use the pre-generated sequence
    match_duration =[4, 9, 5, 3, 7, 2, 12, 4, 1, 8]
    ## just for testing
    #match_duration = [4, 9]

    # the bolck design
    num_in_block=5
    last_rounds = []
    for k in range(0, len(match_duration)):
        if match_duration[k] % num_in_block == 0:
            last_rounds.append(int(match_duration[k]))
        else:
            last_rounds.append(int((match_duration[k] // num_in_block + 1) * num_in_block))

    last_rounds = np.cumsum(last_rounds)
    last_round=last_rounds[-1]
    num_rounds=last_round
    first_rounds=[1]
    for v in range(1, len(match_duration)):
        first_rounds.append(int(last_rounds[v-1]+1))

    lastrounds_plus1 = [x+1 for x in last_rounds]



class Subsession(BaseSubsession):

    def creating_session(self):

        if self.round_number==1:
            self.session.vars['start_time']=time.time()
            self.session.vars['alive']=True

        if self.round_number in Constants.first_rounds:
            self.group_randomly() # shuffle player matchings at the first round of a new match, see https://otreecb.netlify.app/rolesandmatching/matching.html
        else:
            self.group_like_round(self.round_number-1) # same match as in the previous round


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    strategy=models.StringField()

    decision=models.StringField(
        choices=['A','B'],
        doc="""This player's action""",
        widget=widgets.RadioSelect,
        label="Following your chosen rule, which action should you choose:"
    )
    match = models.IntegerField()
    round_in_match = models.IntegerField()
    chosenrule = models.StringField()
    correctaction = models.StringField() # the action that follows chosen rule in a round
    checkaction = models.IntegerField() # check if action follows strategy or not in a round
    matchpayoff = models.CurrencyField() # payoff in a match, value makes sense only in the last round of a match
    nummistake=models.IntegerField() # number of mistakes in the paid rounds of a match, value makes sense only in the last round of a match
    complex=models.IntegerField() # 1 if (S1,C); 0 if (S1,S2)
    #dieroll=models.IntegerField() # random number generated by computer at the end of each round
    counterpartdecision=models.StringField()
    QandC=models.StringField() #collect questions and comments at the end of Part 2
    useddecisionrule=models.StringField() #ask subjects what strategies they used to play the games, at the end of Part 2

    # define the function to identify math and round_in_match
    def matchround_num(self):
        l = 0
        while l < len(Constants.last_rounds):
            if self.round_number <= Constants.last_rounds[l]:
                self.match = l + 1
                l = len(Constants.last_rounds)
            else:
                l += 1  # short for k=k+1

        self.round_in_match = self.round_number - Constants.first_rounds[self.match - 1] + 1

    # half of matches play (S1,C), the other half play (S1,S2)
    def complex_treatment(self):
        if self.session.config['complex_first'] == 1:  # if (S1,C) is played first
            if self.match <= Constants.num_matches / 2:
                self.complex = 1
            else:
                self.complex = 0
        else:
            if self.match <= Constants.num_matches / 2:
                self.complex = 0
            else:
                self.complex = 1


    # define the function to pass chosen rule across rounds in a match
    def chosen_rule(self):
        if self.round_in_match == 1:
            self.chosenrule = self.strategy
        else:
            j = Constants.first_rounds[self.match - 1]
            self.chosenrule = self.in_round(j).strategy # only choose strategy in the 1st round of a match


    def other_player(self):
        return self.get_others_in_group()[0]  # get the counterpart of a player with 2-player groups

    def counterpart_decision(self):
        self.counterpartdecision = self.other_player().decision # get counterpart's decision

    # Define correct action
    def correct_action(self):
        m = self.round_number-1 # round_number of the previous round
        other_previous_3=[] # opponent's actions in previous 3 rounds: set initial values
        other_previous=[] # the list of opponent's previous 3 rounds: set initial values
        other_previous_1=[] # opponent's action in the previous round: set initial values
        if self.chosenrule == 'S1':
            self.correctaction = 'A'
        elif self.chosenrule == 'S2':
            if self.round_in_match==1:
                self.correctaction = 'B'
            else:
                other_previous_1=self.other_player().in_round(m).decision
                if other_previous_1=='A':
                    self.correctaction = 'A'
                else:
                    self.correctaction = 'B'
        else:
            if self.round_number - Constants.first_rounds[self.match - 1] < 3:
                self.correctaction = 'B'
            else:
                other_previous=self.other_player().in_rounds(m-2,m)#get the list of opponent's previous 3 rounds
                for t in other_previous:
                    other_previous_3 = other_previous_3+list((t.decision))
                #use t not p, otherwise there will be mistakes
                if other_previous_3==['A','A','A']:
                    self.correctaction = 'A'
                else:
                    self.correctaction ='B'

    # check if action follows the chosen strategy
    def check_action(self):
        if self.decision==self.correctaction:
            self.checkaction = 1
        else:
            self.checkaction = 0

    # define no implementation action
    def no_implementation(self):
        if self.chosenrule == 'S1':
            self.decision = 'A'
        elif self.chosenrule == 'S2':
            if self.round_in_match==1:
                self.decision='B'
            else:
                self.decision=self.correctaction
        else:
            if self.round_number - Constants.first_rounds[self.match - 1] < 3:
                self.decision = 'B'
            else:
                self.decision = self.correctaction

    # Define payoff calculation function, Dal Bo and Frechette (2011) paid every match
    def set_payoff(self):
        if self.complex==1:# if the (S1,C) treatment
            payoff_matrix_strategy = {
                'S1':
                    {
                        'S1': Constants.AD_both,
                        'C': Constants.ADTF3T
                    },
                'C':
                    {
                        'S1': Constants.TF3TAD,
                        'C': Constants.TF3T_both
                    }
            }
            payoff_matrix_round ={
                'B':
                    {
                        'B': Constants.both_cooperate_payoff,
                        'A': Constants.betrayed_payoff_complex
                    },
                'A':
                    {
                        'B': Constants.betray_payoff,
                        'A': Constants.both_defect_payoff
                    }
            }
        else:
            payoff_matrix_strategy = {
                'S1':
                    {
                        'S1': Constants.AD_both,
                        'S2': Constants.ADTFT
                    },
                'S2':
                    {
                        'S1': Constants.TFTAD,
                        'S2': Constants.TFT_both
                    }
            }
            payoff_matrix_round = {
                'B':
                    {
                        'B': Constants.both_cooperate_payoff,
                        'A': Constants.betrayed_payoff_simple
                    },
                'A':
                    {
                        'B': Constants.betray_payoff,
                        'A': Constants.both_defect_payoff
                    }
            }

        if Constants.pay_method == 1:
            if self.round_number in Constants.first_rounds:
                if self.checkaction == 1:
                    self.payoff = payoff_matrix_strategy[self.chosenrule][self.other_player().chosenrule]  # payoff in strategy choice game
                else:
                    self.payoff = payoff_matrix_strategy[self.chosenrule][self.other_player().chosenrule] + Constants.payoff_deduct
            elif self.round_in_match <= Constants.match_duration[self.match-1]:
                if self.checkaction == 1:
                    self.payoff = 0
                else:
                    self.payoff = Constants.payoff_deduct  # if not follow the chosen rule, lose a certain amount of payoff
            else:
                self.payoff = 0 # rounds not being taken into consideration for final payoff
        else:
            if self.round_in_match <= Constants.match_duration[self.match-1]:
                if self.checkaction == 1:
                    self.payoff = payoff_matrix_round[self.decision][self.other_player().decision]  # payoff in a stage game
                else:
                    self.payoff = payoff_matrix_round[self.decision][self.other_player().decision] + Constants.payoff_deduct
            else:
                self.payoff=0 # rounds not being taken into consideration for final payoff


    # Define the function to calculate match payoff and number of mistakes in a match
    def match_payoff(self):
        firstround=Constants.first_rounds[self.match - 1]
        lastround=Constants.last_rounds[self.match - 1]
        mpayoff=c(0) # set the initial value
        if self.round_number in Constants.last_rounds:
            previouspayoff=self.in_rounds(firstround,lastround)
            for s in previouspayoff:
                mpayoff=mpayoff+s.payoff
        else:
            mpayoff=c(0)

        self.matchpayoff=mpayoff

        numcorrect = 0  # number of correct action in a match, set initial value
        numpaidround=Constants.match_duration[self.match-1] # number of paid rounds in a match
        lastroundpaid=firstround+numpaidround-1
        if self.round_number in Constants.last_rounds:
            previouscorrect=self.in_rounds(firstround,lastroundpaid)
            for u in previouscorrect:
                numcorrect=numcorrect+u.checkaction
        else:
            numcorrect=0

        if self.round_number in Constants.last_rounds:
            self.nummistake = numpaidround - numcorrect # number of mistakes in paid rounds
        else:
            self.nummistake = 0