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       = 'RNC'
    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):
    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
    )
    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()
    market   = models.IntegerField()
    round    = 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 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 group.decision == 'Accept':
        p1.payoff = cu(max(0,group.price))
        p2.payoff = cu(max(0, group.value - group.price))
    else:
        p1.payoff = cu(0)
        p2.payoff = cu(0)
    Data.create(
        group    = group,
        player1  = p1.id_in_subsession,
        player2  = p2.id_in_subsession,
        value    = group.value,
        price    = group.price,
        response = group.decision,
        market   = group.round_number,
        round    = group.roundcount
        )
    # increase count. modulo operator returns remainder
    if group.roundcount % 10:
        group.roundcount = group.roundcount + 1
    else:
        group.roundcount = 1

# 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
    m = random.randint(1, group.round_number)
    if m == group.round_number:
        # assign to corresponding participant field
        if group.session.config['firstgame'] == 'rental_nc':
            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_nc':
            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_nc':
        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(10):
        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

# 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_nc'

# 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_nc' 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_nc' 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_nc'

# 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_nc' and player.round_number != 1) | (player.session.config['lastgame'] == 'rental_nc'))
    #^ display this page only if the players role is buyer and
    # it is not the first market

# the sellers page where he/she sets a price
class Seller(Page):
    form_model  = 'group'
    form_fields = ['price']
    @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 Buyer(Page):
    form_model  = 'group'
    form_fields = ['decision']
    @staticmethod
    def is_displayed(player):
        return player.role == Constants.buyer_role
    #^ display this page only if the players role is buyer

# waiting page to calculate and record round choices and payoffs
class RoundResultsWaitPage(WaitPage):
    after_all_players_arrive = set_payoffs_and_record_choices

# display results for the round
class Results(Page):
    pass

# waiting page to write data per group to the group level for aumatic export
# yes, its stupid. custom export function otree provides does not work
class MarketResultsWaitPage(WaitPage):
    after_all_players_arrive = write_data_to_group_level

# 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 = len(Data.filter(group = group, market = group.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 = group.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

# calculate earnings for all markets
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'] == 'rental_nc')

# sequence in which pages will be shown
# TODO: WHEN INCREASING ROUNDS, MAKE SURE TO ADJUST ROUND IN WHICH NEXTMARKET AND
#       TOTALRESULTS PAGES ARE SHOWN
page_sequence = [Info, WaitForNewInstructions, WaitAfterInstructions,
                 ComprehensionSeller, ComprehensionBuyer,
                 StartWaitPage, Role, BuyerValue,
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round1
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round2
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round3
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round4
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round5
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round6
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round7
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round8
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round9
                 Seller, WaitForP1, Buyer, RoundResultsWaitPage, Results, #round10
                 MarketResultsWaitPage, MarketResults, NextMarket,
                 WaitToCalculateResults, TotalResults]