from otree.api import *
import random
import numpy as np
import itertools # for efficient cycling


doc = """
A rental game with one player proposing a price and the second
player accepting or rejecting
"""

# set all constant variables used throughout the game here
# eg the number of players per group, number of rounds
class Constants(BaseConstants):
    name_in_url       = 'RC'
    players_per_group = 2
    num_rounds        = 5 #adjust final page and call of calc_earnings func when changing this!
    seller_role       = 'Seller'
    buyer_role        = 'Buyer'
    value_high        = cu(100)
    value_low         = cu(50)
    price_min         = cu(0)
    price_max         = value_high # change as needed.

# subset players into subsessions
class Subsession(BaseSubsession):
    pass

# set buyer values here to ensure balanced treatment groups
# cycle through groups and assign value
# also set shuffling of groups
def creating_session(subsession):
    # set shuffling of groups
    subsession.group_randomly(fixed_id_in_group=True)
    #^ this should keep player roles fixed. NOTE reshuffles each round
    #set buyer values
    values = np.random.choice([Constants.value_low, Constants.value_high],
             100, p = [0.4, 0.6]) # 100 draws from given prob distribution
    values = itertools.cycle(values)
    for group in subsession.get_groups():
        group.value = next(values)
        # and start round count
        group.roundcount = 1

# set variables at group level. since price, value, and accept/reject are
# relevant for all players in the group, define them at group level
# instead of player level
class Group(BaseGroup):
    price1 = models.CurrencyField(
        label = "Price in Period 1",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price2 = models.CurrencyField(
        label = "Price in Period 2",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price3 = models.CurrencyField(
        label = "Price in Period 3",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price4 = models.CurrencyField(
        label = "Price in Period 4",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price5 = models.CurrencyField(
        label = "Price in Period 5",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price6 = models.CurrencyField(
        label = "Price in Period 6",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price7 = models.CurrencyField(
        label = "Price in Period 7",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price8 = models.CurrencyField(
        label = "Price in Period 8",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price9 = models.CurrencyField(
        label = "Price in Period 9",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    price10 = models.CurrencyField(
        label = "Price in Period 10",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    #^ price can be called later as a form_field for input
    decision1 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision2 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision3 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision4 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision5 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision6 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision7 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision8 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision9 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    decision10 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    value = models.CurrencyField()
    #^ buyer value in this group
    roundcount = models.IntegerField()
    #^doesnt actually make sense to assign the count attribute at the group
    # level but couldnt get it to work at subsession level

# see group level above
class Player(BasePlayer):
    marketpayoff = models.CurrencyField()
    seller_matching = models.StringField(
                        choices=[['0%', '0%'], ['10%', '10%'], ['20%', '20%'],
                                 ['30%', '30%'], ['40%', '40%'], ['50%', '50%'],
                                 ['60%', '60%'], ['70%', '70%'], ['80%', '80%'],
                                 ['90%', '90%'], ['100%', '100%'], ['No idea', 'No idea']],
                            label='What is the probability of you being matched with a buyer that has a per-period value of 100 points for the good in a given market? ')

    seller_price = models.StringField(choices=[['First period', 'First period'],
                        ['All periods', 'All periods'], ['No idea', 'No idea']],
                    label='For which period(s) do you set a price when a market opens, if the market functions as described in these instructions?')

    seller_market = models.StringField(choices=[['After all 10 periods have been played', 'After all 10 periods have been played'],
                        ['After the buyer purchases the good', 'After the buyer purchases the good'],
                        ['After the buyer purchases the good or after all 10 periods have been played', 'After the buyer purchases the good or after all 10 periods have been played'],
                        ['No idea', 'No idea']],
                    label='When does a market close?')

    seller_payment = models.StringField(choices=[['CHF 4.00', 'CHF 4.00'],
                            ['CHF 6.00', 'CHF 6.00'], ['CHF 9.00', 'CHF 9.00'],
                            ['No idea', 'No idea']],
                        label='Suppose the computer selects a market that works as described in these instructions. If you sold your good while this market was open once at a price of 60 and three times at a price of 40, how much will you earn in addition to the guaranteed CHF 20 if the computer selects this market?')

    buyer_value = models.StringField(choices=[['Yes', 'Yes'], ['No', 'No'],
                                              ['No idea', 'No idea']],
                                     label='In each market you are randomly matched with a seller. Does the seller know about the value you assign to his/her good in a given market?')

    buyer_price = models.StringField(
                    choices=[['First period', 'First period'],
                             ['All periods', 'All periods'],
                             ['No idea', 'No idea']],
                    label='For which period(s) can the seller set a price when a market opens, if the market functions as described in these instructions?')

    buyer_market = models.StringField(
                    choices=[['After all 10 periods have been played',
                                'After all 10 periods have been played'],
                             ['After I have bought the good ',
                                'After I have bought the good '],
                             ['After I have bought the good or after all 10 periods have been played',
                              'After I have bought the good or after all 10 periods have been played'],
                             ['No idea', 'No idea']],
                             label='When does a market close?')

    buyer_payment = models.StringField(
                        choices=[['CHF 1.00', 'CHF 1.00'], ['CHF 2.50', 'CHF 2.50'],
                                 ['CHF 4.00', 'CHF 4.00'], ['CHF 7.00', 'CHF 7.00'],
                                 ['No idea', 'No idea']],
                        label='Assume that you value the good at 50 points in a market that functions as described in these instructions. If you bought the good while this market was open three times at a price of 40 and once at a price of 30, how much will you earn in addition to the guaranteed CHF 20 if the computer selects this market?')


# store groups prices, accepts/rejects, market and round IDs in an ExtraModel
class Data(ExtraModel):
    group      = models.Link(Group)
    player1    = models.IntegerField()
    player2    = models.IntegerField()
    value      = models.CurrencyField()
    price      = models.CurrencyField()
    response   = models.StringField()
    round      = models.IntegerField()
    market     = models.IntegerField()


#-------------------------- FUNCTIONS -----------------------------
# define function for assigning new values so that it can be called for next market
def assign_new_values(subsession: Subsession):
    #set buyer values anew
    values = np.random.choice([Constants.value_low, Constants.value_high],
             100, p = [0.4, 0.6]) # 100 draws from given prob distribution
    values = itertools.cycle(values)
    for group in subsession.get_groups():
        group.value = next(values)

# record all choices and calculate market payoffs
def record_choices_and_calculate_payoffs(group: Group):
    p1 = group.get_player_by_id(1)
    p2 = group.get_player_by_id(2)
    #record all choices in extramodel with one row per round
    prices = [group.price1, group.price2, group.price3, group.price4,
              group.price5, group.price6, group.price7, group.price8,
              group.price9, group.price10]
    decisions = [group.decision1, group.decision2, group.decision3, group.decision4,
                 group.decision5, group.decision6, group.decision7, group.decision8,
                 group.decision9, group.decision10]
    for i in range(1, 11):
        Data.create(
            group    = group,
            player1  = p1.id_in_subsession,
            player2  = p2.id_in_subsession,
            value    = group.value,
            price    = prices[i-1], #lists in python start at zero...
            response = decisions[i-1],
            round    = i,
            market   = group.round_number
        )
    #calculate and assign payoffs
    # seller first
    pay = cu(0)
    length = len(Data.filter(group = group, player1 = p1.id_in_subsession,
                            market = group.round_number, response = 'Accept'))
    for i in range(length):
        pay = pay + Data.filter(group = group, player1 = p1.id_in_subsession,
                                market = group.round_number, response = 'Accept')[i].price
    p1.payoff = pay
    # now buyer
    p = pay
    pay = group.value * length - p
    p2.payoff = cu(max(0, pay))

# calculate earnings after all markets have been played
def calculate_earnings(group: Group):
    p1 = group.get_player_by_id(1)
    p2 = group.get_player_by_id(2)
    # for the seller
    # pick a market at random
    m = random.randint(1, group.round_number)
    if m == group.round_number:
        # assign to corresponding participant field
        if group.session.config['firstgame'] == 'rental_c':
            p1.participant.exp1payoff = p1.marketpayoff
            p2.participant.exp1payoff = p2.marketpayoff
        else:
            p1.participant.exp2payoff = p1.marketpayoff
            p2.participant.exp2payoff = p2.marketpayoff
    else:
        p1_previous = p1.in_round(m)
        p2_previous = p2.in_round(m)
        # assign to corresponding participant field
        if group.session.config['firstgame'] == 'rental_c':
            p1.participant.exp1payoff = p1_previous.marketpayoff
            p2.participant.exp1payoff = p2_previous.marketpayoff
        else:
            p1.participant.exp2payoff = p1_previous.marketpayoff
            p2.participant.exp2payoff = p2_previous.marketpayoff
    # if this is the last market, roll the dice and assign payoff
    if group.session.config['lastgame'] == 'rental_c':
        e = random.randint(1,2)
        if e == 1:
            p1.participant.payoff = p1.participant.exp1payoff
        else:
            p1.participant.payoff = p1.participant.exp2payoff
        # roll separately for player 2
        e = random.randint(1,2)
        if e == 1:
            p2.participant.payoff = p2.participant.exp1payoff
        else:
            p2.participant.payoff = p2.participant.exp2payoff

# def shuffle_groups(subsession):
    #subsession.group_randomly(fixed_id_in_group=True)
    #^ this should keep player roles fixed. BUT reshuffles each round

# #function to reshuffle second entry of groups, and set as matrix
# if subsession.round_number == 2:
#     subsession.group_like_round(1)
# if subsession.round_number == 4:
#     subsession.group_like_round(3)
# if subsession.round_number == 3:
#     import random
#     matrix = self.get_group_matrix()
#     mcopy  = matrix[2:]
#     random.shuffle(mcopy)
#     matrix[2:] = mcopy
#     self.set_group_matrix(matrix)
#     print(self.get_group_matrix())#for debug


# custom export of data
# TODO: doesnt seem to work. maybe needs to be called?
def custom_export(players):
    # header row
    yield ['session', 'participant_code', 'group number', 'round_number',
           'id_in_group', 'payoff', 'price', 'response']
    for p in players:
        yield [p.session.code, p.participant.code, p.group_number,
        p.round_number, p.id_in_group, p.payoff, p.group.price,
        p.group.decision]


#----------------------------- PAGES ------------------------------
# define all pages that should be shown throughout the experiment

# start page if it is the first game
class Info(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1 and player.session.config['firstgame'] == 'rental_c'

# if this is the second game, ask participants to wait for new instructions
class WaitForNewInstructions(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1 and player.session.config['lastgame'] == 'rental_c' and player.session.config['firstgame'] != player.session.config['lastgame']

class WaitAfterInstructions(WaitPage):
    wait_for_all_groups = True
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1 and player.session.config['lastgame'] == 'rental_c' and player.session.config['firstgame'] != player.session.config['lastgame']

# comprehension questions
class ComprehensionSeller(Page):
    form_model = 'player'
    form_fields = ['seller_matching', 'seller_price', 'seller_market',
                    'seller_payment']
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1 and player.role == Constants.seller_role

class ComprehensionBuyer(Page):
    form_model = 'player'
    form_fields = ['buyer_value', 'buyer_price', 'buyer_market', 'buyer_payment']
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1 and player.role == Constants.buyer_role

class StartWaitPage(WaitPage):
    wait_for_all_groups = True

# page informs players of their role in the game
# constant over rounds so only show once
class Role(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1 and player.session.config['firstgame'] == 'rental_c'

# page informing the buyer of his/her new value
class BuyerValue(Page):
    @staticmethod
    def is_displayed(player):
        return (player.role == Constants.buyer_role) and ((player.session.config['firstgame'] == 'rental_c' and player.round_number != 1) | (player.session.config['lastgame'] == 'rental_c'))
    #^ display this page only if the players role is buyer and
    # it is not the first market

# the sellers page where he/she sets prices
class Seller(Page):
    form_model  = 'group'
    form_fields = ['price1', 'price2', 'price3', 'price4', 'price5', 'price6',
                    'price7', 'price8', 'price9', 'price10']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.seller_role
    #^ display this page only if the players role is seller

# page for buyer to wait while seller chooses prices
class WaitForP1(WaitPage):
    pass

# page for buyer to accept/reject
class Buyer1(Page):
    form_model  = 'group'
    form_fields = ['decision1']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer2(Page):
    form_model  = 'group'
    form_fields = ['decision2']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer3(Page):
    form_model  = 'group'
    form_fields = ['decision3']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer4(Page):
    form_model  = 'group'
    form_fields = ['decision4']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer5(Page):
    form_model  = 'group'
    form_fields = ['decision5']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer6(Page):
    form_model  = 'group'
    form_fields = ['decision6']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer7(Page):
    form_model  = 'group'
    form_fields = ['decision7']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer8(Page):
    form_model  = 'group'
    form_fields = ['decision8']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer9(Page):
    form_model  = 'group'
    form_fields = ['decision9']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# page for buyer to accept/reject
class Buyer10(Page):
    form_model  = 'group'
    form_fields = ['decision10']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role

# waiting page to calculate market payoffs
class MarketResultsWaitPage(WaitPage):
    after_all_players_arrive = record_choices_and_calculate_payoffs

# display results for the market at the end of 10 rounds
class MarketResults(Page):
    # calculate payoffs and pass as dictionary to the template
    @staticmethod
    def vars_for_template(player: Player):
        group = player.group
        length = 0 #for safety
        length = len(Data.filter(group = group, market = player.round_number,
                                 response = 'Accept'))
        # calculate sum of prices of all rounds where buyer accepted
        pricesum = cu(0)
        if length > 0:
            for i in range(length):
                pricesum = pricesum + Data.filter(group = group,
                                                  market = player.round_number,
                                                  response = 'Accept')[i].price
        # calculate player payoff based on role
        if player.role == Constants.seller_role:
            marketpayoff = pricesum
        else:
            marketpayoff = group.value * length - pricesum
        player.marketpayoff = cu(max(0, marketpayoff))
        # and return variables as a dictionary
        return dict(
            length = length,
            pricesum = pricesum,
            marketpayoff = marketpayoff
        )

# info page starting the next market
class NextMarket(Page):
    after_all_players_arrive = assign_new_values
    @staticmethod
    def is_displayed(player):
        return player.round_number != 5

class WaitToCalculateResults(WaitPage):
    after_all_players_arrive = calculate_earnings
    @staticmethod
    def is_displayed(player):
        return player.round_number == 5


# display total earnings over all rounds and markets
class TotalResults(Page):
    @staticmethod
    def is_displayed(player):
        return (player.round_number == 5) and (player.session.config['lastgame'] == 'rental_c')

# sequence in which pages will be shown
page_sequence = [Info, WaitForNewInstructions, WaitAfterInstructions,
                 ComprehensionSeller, ComprehensionBuyer,
                 StartWaitPage, Role, BuyerValue,
                 Seller, WaitForP1, Buyer1, Buyer2, Buyer3, Buyer4, Buyer5,
                 Buyer6, Buyer7, Buyer8, Buyer9, Buyer10, MarketResultsWaitPage,
                 MarketResults, NextMarket, WaitToCalculateResults, TotalResults]