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 which a subject plays against a robot (computer) counterpart. The subject is required to follow a rule 
in each game, and the robot's choices are pre-determined.
"""


class Constants(BaseConstants):
    name_in_url = 'Part1'
    players_per_group = None
    conversion_rate = 1/100 # one experiment token for 0.01 dollar

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

    #initial payoff in trainer app
    payoff_initial=100

    # define payoffs of the stage game; used in the instructions
    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 in a round
    payoff_deduct = -20

    delta=0.8 # the above stage game are chosen under delta=0.8

    # number of matches, randomly determined number of rounds in each match
    num_matches=3
    ## 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 a pre-generated sequence
    match_duration =[4, 10, 5]

    # 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))



class Subsession(BaseSubsession):

    def creating_session(self):

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


class Group(BaseGroup):
    pass


class Player(BasePlayer):

    name=models.StringField(label="Your name:")

    assignedID=models.IntegerField(label='Your assigned ID:', min=1, max=30)

    perm=models.StringField(label="What is your Perm number?")

    venmo=models.StringField(label="What is your Venmo account?")

    quiz_1=models.IntegerField(
        choices=[
            [1, 'Another person in the experiment.'],
            [2, 'A robot that is making its decision based on your decisions.'],
            [3, 'A robot whose decisions are independent of your decisions.']
        ],
        widget=widgets.RadioSelect,
        label="Your counterpart for Part 1 is:"
    )

    quiz_2 = models.IntegerField(
        choices=[
            [1, 'A random number.'],
            [2, 'The choices made by the robot.'],
            [3, 'Whether your choices follow the rule you are given.']
        ],
        widget=widgets.RadioSelect,
        label="Your earnings in Part 1 are determined by:"
    )

    quiz_3 = models.IntegerField(
        choices=[
            [1, 'Is equal to 4.'],
            [2, 'Is equal to 20.'],
            [3, 'Is determined randomly, with the game continuing with 80% chance after each round.'],
            [4, 'Is determined randomly, with the game continuing with 50% chance after each round.']
        ],
        widget=widgets.RadioSelect,
        label="The number of rounds in each game:"
    )

    quiz_TFT = models.IntegerField(
        choices=[
            [1, 'Requires you to always choose A.'],
            [2, 'Requires you to always choose B.'],
            [3, 'Requires you to do whatever your robot counterpart did last round, after the first round.'],
            [4, 'Requires you to do whatever your robot counterpart did this round, after the first round.']
        ],
        widget=widgets.RadioSelect,
        label="In the next part of the experiment, the rule you must follow:"
    )

    quiz_TF3T = models.IntegerField(
        choices=[
            [1, 'Requires you to do whatever your robot counterpart did last round, after the first round.'],
            [2, 'Requires you to always choose A.'],
            [3, 'Requires you to, after the first three rounds, choose A if your robot counterpart chose A in the three previous rounds and otherwise choose B.'],
            [4, 'Requires you to always choose B.'],
            [5, 'Requires you to, after the first three rounds, choose B if your robot counterpart chose B in the three previous rounds and otherwise choose A.']
        ],
        widget=widgets.RadioSelectHorizontal,
        label="In the next part of the experiment, the rule you must follow:"
    )

    quizP2_1_1s = models.IntegerField(
        choices=[25,50,13,42],
        widget=widgets.RadioSelect,
        label="In the payment table above, if you choose A in some round and your counterpart chooses B, your earnings for the round is:"
    )

    quizP2_1_2s = models.IntegerField(
        choices=[25, 50, 13, 42],
        widget=widgets.RadioSelect,
        label="In the payment table above, if you choose A in some round and your counterpart chooses B, your counterpart's earnings for the round is:"
    )

    quizP2_1_3s = models.IntegerField(
        choices=[25, 50, 13, 42],
        widget=widgets.RadioSelect,
        label="In the payment table above, if you choose B in some round and your counterpart chooses A, your earnings for the round is:"
    )

    quizP2_1_4s = models.IntegerField(
        choices=[25, 50, 13, 42],
        widget=widgets.RadioSelect,
        label="In the payment table above, if you choose B in some round and your counterpart chooses A, your counterpart's earnings for the round is:"
    )

    quizP2_1_1c = models.IntegerField(
        choices=[25, 50, 23, 42],
        widget=widgets.RadioSelect,
        label="In the payment table above, if you choose A in some round and your counterpart chooses B, your earnings for the round is:"
    )

    quizP2_1_2c = models.IntegerField(
        choices=[25, 50, 23, 42],
        widget=widgets.RadioSelect,
        label="In the payment table above, if you choose A in some round and your counterpart chooses B, your counterpart's earnings for the round is:"
    )

    quizP2_1_3c = models.IntegerField(
        choices=[25, 50, 23, 42],
        widget=widgets.RadioSelect,
        label="In the payment table above, if you choose B in some round and your counterpart chooses A, your earnings for the round is:"
    )

    quizP2_1_4c = models.IntegerField(
        choices=[25, 50, 23, 42],
        widget=widgets.RadioSelect,
        label="In the payment table above, if you choose B in some round and your counterpart chooses A, your counterpart's earnings for the round is:"
    )

    quizP2_1_5 =models.IntegerField(
        choices=[
            [1, 'Is equal to 4.'],
            [2, 'Is equal to 20.'],
            [3, 'Is determined randomly, with the game continuing with 80% chance after each round.'],
            [4, 'Is determined randomly, with the game continuing with 50% chance after each round.']
        ],
        widget=widgets.RadioSelect,
        label="The number of rounds in each game:"
    )

    quizP2_2_1=models.IntegerField(
        choices=[
            [1, "Your and your counterpart's direct choices of A versus B in each of a series of rounds."],
            [2, "The computer's choices of A versus B each round, based on the rules you and your counterpart chose for it to follow, in a single game."],
            [3, "The computer's choices of A versus B each round, based on the rules you and your counterpart chose for it to follow, averaged over millions of games with the same length."],
            [4, "The computer's choices of A versus B each round, based on the rules you and your counterpart chose for it to follow, averaged over millions of games with randomly different lengths."]
        ],
        widget=widgets.RadioSelectHorizontal,
        label="Your payment for each game in Part 2 is determined by:"
    )

    quizP2_2_2s = models.IntegerField(
        choices=[
            [1, "125"],
            [2, "150"],
            [3, "113"],
            [4, "210"],
            [5, "It is impossible to know for sure."]
        ],
        widget=widgets.RadioSelect
        #label="Suppose you chose the rule 'Always choose A' and your counterpart chose the rule 'In round 1 choose B. After round 1: choose whatever counterpart chose last round'. Your base payment would be:"
    )

    quizP2_2_2c = models.IntegerField(
        choices=[
            [1, "125"],
            [2, "186"],
            [3, "120"],
            [4, "210"],
            [5, "It is impossible to know for sure."]
        ],
        widget=widgets.RadioSelect
        #label="Suppose you chose the rule 'Always choose A' and your counterpart chose the rule 'In rounds 1, 2 and 3 choose B. After round 3: choose A if counterpart chose A in all of the previous 3 rounds; otherwise choose B'. Your base payment would be:"
    )

    quizP2_2_3s = models.IntegerField(
        choices=[
            [1, "125"],
            [2, "150"],
            [3, "113"],
            [4, "210"],
            [5, "It is impossible to know for sure."]
        ],
        widget=widgets.RadioSelect
        #label="Suppose you chose the rule 'In round 1 choose B. After round 1: choose whatever counterpart chose last round' and your counterpart chose the rule 'Always choose A'. Your base payment would be:"
    )

    quizP2_2_3c = models.IntegerField(
        choices=[
            [1, "125"],
            [2, "186"],
            [3, "120"],
            [4, "210"],
            [5, "It is impossible to know for sure."]
        ],
        widget=widgets.RadioSelect
        #label="Suppose you chose the rule 'In rounds 1, 2 and 3 choose B. After round 3: choose A if counterpart chose A in all of the previous 3 rounds; otherwise choose B' and your counterpart chose the rule 'Always choose A'. Your base payment would be:"
    )

    quizP2_3= models.IntegerField(
        choices=[
            [1, "Based entirely on the choices (A or B) you and your counterpart choose each round when implementing the rules yourself given the numbers in the payment table."],
            [2, "Based entirely on the rules you each chose."],
            [3, "Based on the rules you each chose (the base payment), minus 20 for each mistake you made when attempting to implement your chosen rule."],
            [4, "Based on the rules you each chose (the base payment), minus 10 for each mistake you made when attempting to implement your chosen rule."]
        ],
        widget=widgets.RadioSelectHorizontal,
        label="Your payments for Part 2 will be:"
    )

    decision=models.StringField(
        choices=['A','B'],
        doc="""This player's action""",
        widget=widgets.RadioSelect,
        label="Following your 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 a match, value makes sense only in the last round of a match
    #dieroll=models.IntegerField() # random number generated by computer at the end of each round
    counterpartdecision=models.StringField()

    # define the function to identify match 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

    # define the chosen rule, play each rule twice, can modify
    def chosen_rule(self):
        if self.match<=1:
            self.chosenrule='S1'
        elif self.match>1 and self.match<=2:
            self.chosenrule='S2'
        else:
            self.chosenrule='C'

    # define robot counterpart's action in each round
    def counterpart_decision(self):
        roundrand = np.random.randint(1, 100)
        if self.chosenrule !='C':
            if roundrand<=50:
                self.counterpartdecision = 'A'
            else:
                self.counterpartdecision = 'B'
        else:
            if self.round_in_match <=3 or self.round_in_match ==5:
                self.counterpartdecision = 'A'
            elif self.round_in_match ==4:
                self.counterpartdecision = 'B'
            else:
                if roundrand <= 50:
                    self.counterpartdecision = 'A'
                else:
                    self.counterpartdecision = 'B'


    # 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.in_round(m).counterpartdecision
                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.in_rounds(m-2,m)
                for t in other_previous:
                    other_previous_3 = other_previous_3+list((t.counterpartdecision))
                #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 payoff calculation function, Dal Bo and Frechette (2011) paid every match
    def set_payoff(self):
        if self.round_number in Constants.first_rounds:
            if self.checkaction == 1:
                self.payoff = Constants.payoff_initial
            else:
                self.payoff = Constants.payoff_initial + 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


    # 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
        else:
            self.nummistake = 0