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

doc = """
A durable 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       = 'DC'
    players_per_group = 2
    num_rounds        = 5 #adjust final page when changing this!
    seller_role       = 'Seller'
    buyer_role        = 'Buyer'
    value_high        = cu(100)
    value_low         = cu(50)
    price_min         = cu(0)
    discountfactor    = float(0.6) # change as needed
    price_max         = cu(250) #NOTE CHANGE WHEN ADJUSTING PARAMETERS!!!
    #^unfortunately cant call a function or any parameters here

# 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
    # initialize empty value at player attribute
    for p in subsession.get_players():
        p.lastgroupdecision = 'None'

# 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):
    price = models.CurrencyField(
        doc   = """Sellers Price""",
        label = "What price between " + str(Constants.price_min) +
                " and " +
                 str(Constants.price_max*1.1) +
                 " do you want to set?",
        min = cu(0),
        max = cu(Constants.price_max)*1.1
    )
    #^ price can be called later as a form_field for input
    #^ TODO: unit shown as "points", ideal?
    decision = models.StringField(
        choices = ['Accept', 'Reject'],
        doc     = """Buyers decision""",
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect
    )
    #^ decision can be called later as a form_field for input
    # fields for all prices and decisions below for automatic export. ugly as hell
    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
    )
    decision1 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision2 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision3 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision4 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision5 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision6 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision7 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision8 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision9 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    decision10 = models.StringField(
        choices = ['Accept', 'Reject'],
        label   = "Do you accept or reject?",
        widget  = widgets.RadioSelect,
        initial = "None"
    )
    value = models.CurrencyField()
    #^ buyer value in this group
    roundcount = models.IntegerField(initial = 0)
    #^doesnt actually make sense to assign the count attribute at the group
    # level but couldnt get it to work at subsession level
    acceptanceperiod = models.IntegerField(initial = 0)

# see group level above
class Player(BasePlayer):
    lastgroupdecision = models.StringField()
    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 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 elapsed',
                                  'After all 10 periods have elapsed'],
                                 ['After the buyer purchases the good',
                                  'After the buyer purchases the good'],
                                 ['After the buyer purchases the good or after all 10 periods have elapsed',
                                  'After the buyer purchases the good or after all 10 periods have elapsed'],
                                 ['No idea', 'No idea']],
                        label='When does a market close?')
    seller_payment = models.StringField(
                        choices=[['CHF 1.50', 'CHF 1.50'], ['CHF 2.50', 'CHF 2.50'],
                                 ['CHF 5', 'CHF 5'], ['CHF 7.50', 'CHF 7.50'],
                                 ['CHF 15', 'CHF 15'], ['No idea', 'No idea']],
                        label='If you sold your good while the market was open at a price of 150 points, and the market functions as described in these instructions, 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 elapsed',
                              'After all 10 periods have elapsed'],
                             ['After I have bought the good ', 'After I have bought the good '],
                             ['After I have bought the good or after all 10 periods have elapsed',
                              'After I have bought the good or after all 10 periods have elapsed'],
                             ['No idea', 'No idea']],
                    label='When does a market close?')
    buyer_payment = models.StringField(
                        choices=[['CHF 0', 'CHF 0'], ['CHF 2.10', 'CHF 2.10'],
                                 ['CHF 4.20', 'CHF 4.20'], ['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 at a price of 80 in period 3, 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)

# function to calculate total value for a given buyer
def calculate_total_value(value, current_round):
    total_value = cu(0)
    if current_round == 10:
        total_value = value
    else:
        for t in range(1, 10-(current_round-2)): #include current round
            total_value = total_value + value*Constants.discountfactor**(t-1)
    return total_value

# function to set payoffs and record all choices made
# TODO: unclear if recording really needed or its all recorded anyways
def set_payoffs_and_record_choices(group: Group):
    p1 = group.get_player_by_id(1)
    p2 = group.get_player_by_id(2)
    if p2.lastgroupdecision != 'Accept':
        listd = [a for a in vars(group) if a.startswith('decision') and a[-1].isdigit()]
        listd.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))
        if getattr(group, listd[group.roundcount-1]) == 'Accept':
            listp = [a for a in vars(group) if a.startswith('price') and a[-1].isdigit()]
            listp.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))
            price = getattr(group, listp[group.roundcount-1])
            p1.payoff = cu(max(0, price))
            p2.payoff = cu(max(0, calculate_total_value(group.value, group.roundcount-1) - price))
            p1.lastgroupdecision = 'Accept'
            p2.lastgroupdecision = 'Accept'
            group.acceptanceperiod = group.roundcount
            #print(group.acceptanceperiod)
        else:
            p1.payoff = cu(0)
            p2.payoff = cu(0)
    group.roundcount = group.roundcount + 1
    if group.roundcount == 10:
        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
            )


# randomly pick a previous round and calculate payoff based on role
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
    p1_previous = p1.in_round(1)
    p2_previous = p2.in_round(1)
    m = random.randint(1, group.round_number)
    if m == group.round_number:
        # assign to corresponding participant field
        if group.session.config['firstgame'] == 'durable_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'] == 'durable_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'] == 'durable_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]

# write per-period data onto group model instead for automatic export
# ugly as hell, but no data is lost
# def write_data_to_group_level(group: Group):
#     listp = [a for a in vars(group) if a.startswith('price') and a[-1].isdigit()]
#     listp.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))
#     listd = [a for a in vars(group) if a.startswith('decision') and a[-1].isdigit()]
#     listd.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))
#     for i in range(group.roundcount-1):
#         setattr(group, listp[i], Data.filter(group = group, market = group.round_number)[i].price)
#         setattr(group, listd[i], Data.filter(group = group, market = group.round_number)[i].response)


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

#--------------1. Before the game begins
# 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'] == 'durable_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'] == 'durable_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'] == 'durable_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'] == 'durable_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'] == 'durable_c' and player.round_number != 1) | (player.session.config['lastgame'] == 'durable_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 a price
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) and (player.lastgroupdecision != 'Accept')

# waiting page until everyone is done and calculate payoffs
class RoundResultsWaitPage(WaitPage):
    after_all_players_arrive = set_payoffs_and_record_choices

# 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) and (player.lastgroupdecision != 'Accept')


# 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) and (player.lastgroupdecision != 'Accept')


# 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) and (player.lastgroupdecision != 'Accept')

# 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) and (player.lastgroupdecision != 'Accept')

# 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) and (player.lastgroupdecision != 'Accept')

# 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) and (player.lastgroupdecision != 'Accept')

# 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) and (player.lastgroupdecision != 'Accept')


# 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) and (player.lastgroupdecision != 'Accept')


# 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) and (player.lastgroupdecision != 'Accept')
    #^ display this page only if the players role is buyer

# # waiting page until everyone is done and calculate payoffs
# class RoundResultsWaitPage(WaitPage):
#     after_all_players_arrive = set_payoffs_and_record_choices

# # display results for the round
# class Results(Page):
#     @staticmethod
#     def is_displayed(player):
#         (player.lastgroupdecision != 'Accept') #and (player.round_number != 1)

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

# 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):
        length = 0
        price = 0
        val = 0
        group = player.group
        listd = [a for a in vars(group) if a.startswith('decision') and a[-1].isdigit()]
        listd.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))
        if player.lastgroupdecision == "Accept":
            listp = [a for a in vars(group) if a.startswith('price') and a[-1].isdigit()]
            listp.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))
            price = getattr(group, listp[group.acceptanceperiod-1])
            price = cu(max(0, price))
            val = cu(max(0, calculate_total_value(group.value, group.acceptanceperiod)))
            length = group.acceptanceperiod
        if player.role == Constants.seller_role:
            marketpayoff = price
        else:
            marketpayoff = val - price
        player.marketpayoff = max(cu(0), cu(marketpayoff))
        # print("player", player.role, "marketpayoff", marketpayoff)
        # print("assigned payoff", player.marketpayoff)
        # and return variables as a dictionary
        return dict(
            length = length,
            price = price,
            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):
    # calculate_earnings
    @staticmethod
    def is_displayed(player):
        return (player.round_number == 5) and (player.session.config['lastgame'] == 'durable_c')

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