from otree.api import *

doc = """
Your app description
"""

class C(BaseConstants):
    NAME_IN_URL = 'qlm'
    PLAYERS_PER_GROUP = 3
    NUM_ROUNDS = 2
    ENDOWMENT = 25
    CONVERSION_RATE = 0.25
    OFFER_MAX = 17
    PROB_FED_UP = 0.05
    CURRENT_EMPLOYERS = [1,1,2,2,3,3]
    OTHER_EMPLOYERS1 = [2,2,1,1,1,1]
    OTHER_EMPLOYERS2 = [3,3,3,3,2,2]
    WORKERS_PER_PLAYER = [[1,2,3,4,5,6],[3,4,1,2,5,6],[5,6,1,2,3,4]]
    COLORS_WORKERS_SAME = [0,4,1,6,2,8]
    S=0.5
    C1=12
    C2=4
    D1=0
    D2=6
    AH=2
    AL=1
    NUM_COLOR_FIELDS = 19
    BONUS_COUNTING = 5
    NUM_COMP_QUEST = 10
    BONUS_COMP_QUEST = 2

def creating_session(subsession):
    from itertools import cycle
    import random
    from random import choice

    session = subsession.session

    # Choosing payoff relevant round
    session.vars['paying_round'] = random.randint(1,C.NUM_ROUNDS)

    # Assigning half of the groups to both baseline and samecolor treatment
    same = cycle([True, False])
    for group in subsession.get_groups():
        group.samecolor = next(same)
    
    # Definition of colors
    subsession.session.colors = ['Red', 'Blue', 'Green', 'Orange', 'Gold', 'Pink', 'SaddleBrown', 'Purple', 'Aqua', 'Black']
    subsession.session.colors_de = ['rot', 'blau', 'grün', 'orange', 'gelb', 'pink', 'beige', 'violett', 'türkis', 'schwarz']

    # Creating 6 workers for each group
    for group in subsession.get_groups():
        for i in range(6):
            Worker.create(
                group = group,
                number = i+1,
                color_id = i+3 if group.samecolor == False else C.COLORS_WORKERS_SAME[i],                
                current_employer = group.get_player_by_id(C.CURRENT_EMPLOYERS[i]), 
                other_employer1 = group.get_player_by_id(C.OTHER_EMPLOYERS1[i]),
                other_employer2 = group.get_player_by_id(C.OTHER_EMPLOYERS2[i])
            )

    # Assigning colors to workers
    for group in subsession.get_groups():
        workers = Worker.filter(group = group)
        for worker in workers:
            worker.color_name = subsession.session.colors[worker.color_id]
            worker.color_de = subsession.session.colors_de[worker.color_id]
            worker.color_html = '<span style="color: {}">{}</span>'.format(worker.color_name, worker.color_de)

    # Color assignment, switch variables for random order of employees on the promotion page, workerids for players
    for player in subsession.get_players():
        player.color_id = player.id_in_group-1
        player.color_name = subsession.session.colors[player.color_id]
        player.color_de = subsession.session.colors_de[player.color_id]
        player.color_html = '<span style="color: {}">{}</span>'.format(player.color_name, player.color_de)
        add_e = '' if player.color_de[-1] == 'e' else 'e'
        player.color_html_akk = '<span style="color: {}">{}{}{}</span>'.format(player.color_name, player.color_de, add_e, 'n')
        player.color_html_fem = '<span style="color: {}">{}{}</span>'.format(player.color_name, player.color_de, add_e)

        player.switch_emp = choice([True, False])
        player.switch_worker12 = choice([True, False])
        player.switch_worker34 = choice([True, False])
        player.switch_worker56 = choice([True, False])
        player.worker1_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][int(player.switch_worker12)]
        player.worker2_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][1 - int(player.switch_worker12)]
        
        if not(player.switch_emp):
            player.other_emp1 = C.OTHER_EMPLOYERS1[player.id_in_group * 2 - 1]
            player.other_emp2 = C.OTHER_EMPLOYERS2[player.id_in_group * 2 - 1]
            player.worker3_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][2 + int(player.switch_worker34)]
            player.worker4_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][3 - int(player.switch_worker34)]
            player.worker5_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][4 + int(player.switch_worker56)]
            player.worker6_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][5 - int(player.switch_worker56)]
        else:
            player.other_emp1 = C.OTHER_EMPLOYERS2[player.id_in_group * 2 - 1]
            player.other_emp2 = C.OTHER_EMPLOYERS1[player.id_in_group * 2 - 1]
            player.worker3_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][4 + int(player.switch_worker56)]
            player.worker4_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][5 - int(player.switch_worker56)]
            player.worker5_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][2 + int(player.switch_worker34)]
            player.worker6_id = C.WORKERS_PER_PLAYER[player.id_in_group - 1][3 - int(player.switch_worker34)]
        
        player.worker1_html = Worker.filter(group=player.group, number=player.worker1_id)[0].color_html
        player.worker2_html = Worker.filter(group=player.group, number=player.worker2_id)[0].color_html
        player.worker3_html = Worker.filter(group=player.group, number=player.worker3_id)[0].color_html
        player.worker4_html = Worker.filter(group=player.group, number=player.worker4_id)[0].color_html
        player.worker5_html = Worker.filter(group=player.group, number=player.worker5_id)[0].color_html
        player.worker6_html = Worker.filter(group=player.group, number=player.worker6_id)[0].color_html

    for player in subsession.get_players():
        player.other_emp1_html = player.group.get_player_by_id(player.other_emp1).color_html
        player.other_emp2_html = player.group.get_player_by_id(player.other_emp2).color_html

class Subsession(BaseSubsession):
    pass

class Group(BaseGroup):
    samecolor = models.BooleanField()

# Function to create promotion field variables
def var_prom():
    return models.BooleanField(
        choices=[[False, ''], [True, ''],],
        label='Beförderung',
        )

# Function to create offer field variables
def var_offer(label):
    return models.FloatField(
        min=0,
        max=C.OFFER_MAX,
        label=label,
        )

# Function to create comprehensve question variables
def var_comp(label):
    return models.BooleanField(
        choices=[[True, 'Richtig'], [False, 'Falsch']],
        label = label,
        widget = widgets.RadioSelectHorizontal
    )

# Function to create closeness questions
def var_close(label):
    return models.IntegerField(
        choices=[i for i in range(0,11)],
        label = label,
        widget = widgets.RadioSelectHorizontal
    )

class Player(BasePlayer):
    color_id = models.IntegerField()
    color_name = models.StringField()
    color_de = models.StringField()
    color_html = models.StringField()
    color_html_akk = models.StringField()
    color_html_fem = models.StringField()
    picture_count = models.IntegerField()
    comp_quest1 = var_comp(label='Aussage 1: Die Zuteilung der TeilnehmerInnen zu einer Farbgruppe, einem Arbeitsmarkt und einer Rolle (AN/AG) bleibt über alle {} Runden erhalten.'.format(C.NUM_ROUNDS))
    comp_quest2 = var_comp(label='Aussage 2: Die Farbgruppe jedes AN und jedes AG auf einem Arbeitsmarkt ist für alle AG desselben Arbeitsmarktes sichtbar.')
    comp_quest3 = var_comp(label='Aussage 3: Ein AG kann die Lohnangebote für seine aktuellen AN von deren Leistungsfähigkeit abhängig machen.')
    comp_quest4 = var_comp(label='Aussage 4: Ein AG kann seine Lohnangebote für die aktuellen AN der anderen AG nur von ihrer Zugehörigkeit zu einer bestimmten Farbgruppe abhängig machen.')
    comp_quest5 = var_comp(label='Aussage 5: Wenn das höchste Angebot der beiden abwerbenden AG gleich oder niedriger ist als das Lohnangebot des aktuellen AG, wechselt der AN mit einer Wahrscheinlichkeit von {}%.'.format(C.PROB_FED_UP))
    comp_quest6 = var_comp(label='Aussage 6: Die Leistungsfähigkeit eines AN wird in der ersten Runde zufällig bestimmt und bleibt dann für alle weiteren Runden gleich.')
    comp_quest7 = var_comp(label='Aussage 7: Ein AN arbeitet nach den Lohnangeboten auf jeden Fall für seinen aktuellen AG.')
    comp_quest8 = var_comp(label='Aussage 8: Die aktuellen AN eines AG sind zu Beginn jeder Runde dieselben.')
    comp_quest9 = var_comp(label='Aussage 9: Der erwirtschaftete Ertrag eines AN ist nur dann von seiner Leistungsfähigkeit abhängig, wenn er befördert wurde.')
    comp_quest10 = var_comp(label='Aussage 10: Der erwirtschaftete Ertrag eines AN ist höher, wenn er bei seinem aktuellen AG bleibt.')
    num_comp_correct = models.IntegerField()
    first_try = models.BooleanField(initial=True)
    promotion1 = var_prom()
    promotion2 = var_prom()
    promotion3 = var_prom()
    promotion4 = var_prom()
    promotion1e = var_prom()
    promotion2e = var_prom()
    promotion3e = var_prom()
    promotion4e = var_prom()
    offer11 = var_offer('Lohnangebot in Taler')
    offer12 = var_offer('Lohnangebot in Taler')
    offer13 = var_offer('Lohnangebot in Taler')
    offer14 = var_offer('Lohnangebot in Taler')
    offer21 = var_offer('')
    offer22 = var_offer('')
    offer23 = var_offer('')
    offer24 = var_offer('')
    offer31 = var_offer('')
    offer32 = var_offer('')
    offer41 = var_offer('')
    offer42 = var_offer('')
    offer51 = var_offer('')
    offer52 = var_offer('')
    offer61 = var_offer('')
    offer62 = var_offer('')
    quest1_1 = var_close('')
    quest1_2 = var_close('')
    quest1_3 = var_close('')
    quest1_4 = var_close('')
    quest1_5 = var_close('')
    quest2_1 = models.IntegerField(label='Ein Schläger und ein Ball kosten gemeinsam 1,10 Euro. Der Schläger kostet um 1,00 Euro mehr als der Ball. Wieviel Euro kostet der Ball?')
    quest2_2 = models.IntegerField(label='5 Maschinen brauchen 5 Minuten, um 5 Einheiten zu produzieren. Wieveile Minuten brauchen 100 Maschinen um 100 Einheiten zu produzieren?')
    quest2_3 = models.IntegerField(label='''In einem Teich befindet sich ein Seerosenbeet. Jeden Tag verdoppelt sich die Größe des Beets. 
                                   Es dauert 48 Tage bis der ganze Teich mit Seerosen bedeckt ist. Nach wievielen Tagen ist die Hälfte des Teiches mit Seerosen bedeckt?''')
    quest2_4 = var_close('Wenn es um Politik geht, werden oft die Begriffe "links" und "rechts" gebraucht. Wie würden Sie sich selbst auf einer Skala von 0 (links) bis 10 (rechts) einordnen?')
    quest2_5 = var_close('Wie wichtig ist es Ihnen auf einer Skala von 0 (überhaupt nicht wichtig) bis 10 (sehr wichtig), dass Sie in enem Land leben, das demokratisch regiert wird?')
    quest2_6 = var_close('Wie schätzen Sie die Neigung anderer Menschen ein Sie ausnutzen, wenn sie die Gelegenheit dazu haben? (0 = andere sind immer fair zu mir; 10 = andere nutzen mich immer aus)')
    quest3_1 = models.IntegerField(label = 'Ihr Geschlecht', choices = [[1, 'weiblich'], [2, 'männlich'], [3, 'divers'], [4, 'anderes']], widget = widgets.RadioSelect)
    quest3_2 = models.IntegerField(label = 'Ihr Alter')
    quest3_3 = models.IntegerField(choices = [[1, 'Betriebswirtschaftslehre'], [2, 'Formalwissenschaften'], [3, 'Geisteswissenschaften'], [4, 'Ingenieurwissenschaften'],
                                              [5, 'Kulturwisssenschaften'], [6, 'Medizin- Gesundheits- und Pflegewissenschaften'], [7, 'Naturwissenschaften'], 
                                              [8, 'Rechtswissenschaften'], [9, 'Sozialwissenschaften'], [10, 'Volkswirtschaftslehre'], [11, 'Sonstige'],
                                              [12, 'kein(e) Student(in)']], widget=widgets.RadioSelect, 
                                   label = 'Welcher Kategorie ist Ihr Studiengang am ehesten zuzuordnen?')
    case = models.IntegerField()
    switch_emp = models.BooleanField()
    other_emp1_id = models.IntegerField()
    other_emp1_html = models.StringField()
    other_emp2_id = models.IntegerField()
    other_emp2_html = models.StringField()
    switch_worker12 = models.BooleanField()
    switch_worker34 = models.BooleanField()
    switch_worker56 = models.BooleanField()
    worker1_id = models.IntegerField()
    worker1_html = models.StringField()
    worker2_id = models.IntegerField()
    worker2_html = models.StringField()
    worker3_id = models.IntegerField()
    worker3_html = models.StringField()
    worker4_id = models.IntegerField()
    worker4_html = models.StringField()
    worker5_id = models.IntegerField()
    worker5_html = models.StringField()
    worker6_id = models.IntegerField()
    worker6_html = models.StringField()
    payoff_round = models.FloatField(initial=0)
    payoff_pay_round = models.FloatField(initial=0)
    payoff_final = models.FloatField(initial=0)
    payout = models.FloatField(initial=0)
    bonus_counting = models.BooleanField(initial=True)

class Worker(ExtraModel):
    group = models.Link(Group)
    number = models.IntegerField()
    color_id = models.IntegerField()
    color_name = models.StringField()
    color_de = models.StringField()
    color_html = models.StringField()
    current_employer = models.Link(Player)
    other_employer1 = models.Link(Player)
    other_employer2 = models.Link(Player)
    offer1 = models.FloatField()
    offer2 = models.FloatField()
    offer3 = models.FloatField()
    employer = models.Link(Player)
    stay = models.BooleanField(initial=False)
    ability = models.BooleanField()
    promoted = models.BooleanField()
    productivity = models.IntegerField()
    wage = models.FloatField()

class Offer(ExtraModel):
    group = models.Link(Group)
    employer = models.Link(Player)
    worker = models.Link(Worker)
    promoted = models.BooleanField()
    amount = models.FloatField()

# Function to create offers
def make_offer(player, worker_number, promoted, amount):
    Offer.create(group = player.group,
                employer = player,
                worker = Worker.filter(group = player.group, number = worker_number)[0], 
                promoted = promoted,
                amount = amount)

# PAGES

class Welcome(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1
    
    @staticmethod
    def vars_for_template(player: Player):
        session = player.session
        return dict(inv_exchange_rate = int(1/session.real_world_currency_per_point))

class ColorGroupAssignment(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1

class Counting(Page):
    form_model = 'player'
    form_fields = ['picture_count']

    @staticmethod
    def is_displayed(player):
        return player.round_number == 1

class AnnounceInstructions(Page):

    @staticmethod
    def is_displayed(player):
        return player.round_number == 1

class Instructions(Page):

    @staticmethod
    def is_displayed(player):
        return player.round_number == 1

class CompQuest(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.round_number == 1

    form_fields = ['comp_quest1', 'comp_quest2', 'comp_quest3', 'comp_quest4', 'comp_quest5', 'comp_quest6', 'comp_quest7', 'comp_quest8', 'comp_quest9', 'comp_quest10']

    @staticmethod
    def error_message(player: Player, values):

        solutions = dict(comp_quest1 = True, comp_quest2 = True, comp_quest3 = True, comp_quest4 = False, comp_quest5 = True, 
                         comp_quest6 = False, comp_quest7 = False, comp_quest8 = True, comp_quest9 = True, comp_quest10 = True)
        wrong_answers = []
        for i in range(len(solutions)):
            if list(values.values())[i] != list(solutions.values())[i]:
                wrong_answers.append(i+1)

        if player.first_try == True:
            player.num_comp_correct = C.NUM_COMP_QUEST - len(wrong_answers)
        player.first_try = False

        if values != solutions:
            return_string = ', Aussage '.join(str(x) for x in wrong_answers)
            return f"Ihre Bewertung folgender Aussagen ist nicht korrekt: Aussage {return_string}"
        
class AnnounceRound(Page):
    pass

class Promotion(Page):
    form_model = 'player'

    form_fields = ['promotion1', 'promotion2', 'promotion3', 'promotion4', 'offer11', 'offer12', 'offer13', 'offer14', 'offer21',
    'offer22', 'offer23', 'offer24', 'offer31', 'offer32', 'offer41', 'offer42', 'offer51', 'offer52', 'offer61', 'offer62']
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(workers = Worker.filter(group = player.group))

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        from random import choice
        
        # Collecting offers and promotions for current workers in lists
        offers1 = [player.offer11, player.offer12, player.offer13, player.offer14]
        offers2 = [player.offer21, player.offer22, player.offer23, player.offer24]
        promotions = [player.promotion1, player.promotion2, player.promotion3, player.promotion4]
        
        # Random decision on case
        player.case = choice([0,1,2,3])

        # Assigning offers and promotions from current employers to workers (worker1 and worker2)
        worker1 = Worker.filter(group = player.group, number = player.worker1_id)[0]
        worker1.offer1 = offers1[player.case]
        worker1.promoted = promotions[player.case]
        worker1.ability = False if player.case in [1,3] else True
        
        worker2 = Worker.filter(group = player.group, number = player.worker2_id)[0]
        worker2.offer1 = offers2[player.case]
        worker2.promoted = not(promotions[player.case])
        worker2.ability = False if player.case in [0,3] else True

        # Creating offers for workers of other employers (worker3,4,5,6)
        make_offer(player, player.worker3_id, True, player.offer31)
        make_offer(player, player.worker3_id, False, player.offer32)
        make_offer(player, player.worker4_id, True, player.offer41)
        make_offer(player, player.worker4_id, False, player.offer42)
        make_offer(player, player.worker5_id, True, player.offer51)
        make_offer(player, player.worker5_id, False, player.offer52)
        make_offer(player, player.worker6_id, True, player.offer61)
        make_offer(player, player.worker6_id, False, player.offer62)        

class ResultsWaitPage(WaitPage):
    def after_all_players_arrive(group: Group):
        from random import random

        workers = Worker.filter(group = group) 

        for worker in workers:
            # Identifying relevant offers from other employers (other_employer1 and other_employer2)
            worker.offer2 = Offer.filter(group=group, employer=worker.other_employer1, worker=worker, promoted = worker.promoted)[0].amount
            worker.offer3 = Offer.filter(group=group, employer=worker.other_employer2, worker=worker, promoted = worker.promoted)[0].amount       

            # Choosing an employer
            if random() > C.PROB_FED_UP and worker.offer1 >= max(worker.offer2, worker.offer3):
                worker.employer = worker.current_employer
                worker.stay = True
                worker.wage = worker.offer1
            else:
                if (worker.offer2 > worker.offer3) or (worker.offer2 == worker.offer3 and random() > 0.5):
                    worker.employer = worker.other_employer1
                    worker.wage = worker.offer2
                else:
                    worker.employer = worker.other_employer2
                    worker.wage = worker.offer3        
                       
            # Calculating worker productivity: (1+Si)*(ci + di*ai)
            worker.productivity = (1 + (C.S if worker.stay else 0)) \
                * ((C.C2 if worker.promoted else C.C1) + (C.D2 if worker.promoted else C.D1) * (C.AH if worker.ability else C.AL))
            
            # Updating employer's payoff: + worker.productivity - wage
            worker.employer.payoff_round += worker.productivity - worker.wage

class Results(Page):
    form_model = 'player'
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(workers = Worker.filter(group = player.group))

class Quest1(Page):
    form_model = 'player'
    form_fields = ['quest1_1', 'quest1_2', 'quest1_3', 'quest1_4', 'quest1_5']

    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            quest1_1_label = 'Wie verbunden fühlen Sie sich Ihrem Arbeitnehmer aus der Farbgruppe {}?'.format(player.worker1_html),
            quest1_2_label = 'Wie verbunden fühlen Sie sich Ihrem Arbeitnehmer aus der Farbgruppe {}?'.format(player.worker2_html),
            quest1_3_label = 'Wie verbunden fühlen Sie sich Ihrem Arbeitgeber aus der Farbgruppe {}?'.format(player.other_emp1_html),
            quest1_4_label = 'Wie verbunden fühlen Sie sich Ihrem Arbeitnehmer aus der Farbgruppe {}?'.format(player.worker3_html),
            quest1_5_label = 'Wie verbunden fühlen Sie sich Ihrem Arbeitnehmer aus der Farbgruppe {}?'.format(player.worker4_html)
        ) 

class Quest2(Page):

    form_model = 'player'
    form_fields = ['quest2_1', 'quest2_2', 'quest2_3', 'quest2_4', 'quest2_5', 'quest2_6']

    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS

class Quest3(Page):

    form_model = 'player'
    form_fields = ['quest3_1', 'quest3_2', 'quest3_3']

    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS

class FinalResultsWaitPage(WaitPage):

    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS
    
    wait_for_all_groups = True

    @staticmethod
    def after_all_players_arrive(subsession):
        session = subsession.session
        players = subsession.get_players()
        bonus_counting_list = [True for i in range(len(session.colors))]
        for player in players:
            # pick payoff from payoff relevant round
            player_in_pay_round = player.in_round(session.vars['paying_round'])
            player.payoff_pay_round = player_in_pay_round.payoff_round

            # set group bonus for counting to False if count of player is false
            player_in_first_round = player.in_round(1)
            if player_in_first_round.picture_count != C.NUM_COLOR_FIELDS:
                bonus_counting_list[player.color_id] = False

        # setting bonus for counting task
        for player in players:
            
            player.bonus_counting = bonus_counting_list[player.color_id]
            
            # calculate final payoff and EUR payoff amount
            player.payoff_final = player.payoff_pay_round + int(player.bonus_counting) * C.BONUS_COUNTING + player.num_comp_correct * C.BONUS_COMP_QUEST
            player.participant.payoff = cu(player.payoff_final)            

class FinalResults(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS

class End(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS


page_sequence = [Quest1, Quest2, Quest3] #, CompQuestResults, FinalResultsWaitPage, FinalResults]
# page_sequence = [Welcome, ColorGroupAssignment, Counting, AnnounceInstructions, Instructions, CompQuest, 
# AnnounceRound, Promotion, AnnounceExpectations, Expectations, ResultsWaitPage, Results, 
# Quest1, Quest2, Quest3, FinalResultsWaitPage, FinalResults]
# page_sequence = [Counting]

#cd Dropbox\Programming\otree\qlm