# ---------------------------------------------------------------------
# IMPORTS
# ---------------------------------------------------------------------
# from otree.api import *
# We use explicit import instead of * in order to avoid import warnings
from otree.api import models
from otree.api import widgets
from otree.api import Page
from otree.api import WaitPage
from otree.api import BaseConstants
from otree.api import BaseSubsession
from otree.api import BaseGroup
from otree.api import BasePlayer

import random
from scipy.stats import gmean

# ---------------------------------------------------------------------
# CLASSES
# ---------------------------------------------------------------------
class C(BaseConstants):
    NAME_IN_URL = 'experiment'
    # Define the players per group (since the value can be set dynamically within the settings.py we can keep it as none (should keep it as none))
    PLAYERS_PER_GROUP = None
    # Define the questions and the corresponding answers - TODO: make flexible by reading in from csv or txt?
    # Within the dict we have an ID, the corresponding Question (label) and the correct answer
    QUESTIONS_DICT = {
        9: ('Wie viele Touristen wurden 2016 in Frankreich registriert?', 82600000),
        19: ('Wie groß ist der mittlere Abstand zwischen dem Planeten Merkur und der Sonne (in km)?', 57800000),
        23: ('Wie groß ist die Entfernung zwischen Erde und Mond (in km)?', 385000),
        11: ('Wie hoch ist das durchschnittliche Jahresgehalt eines Spielers in der Bundesliga (in Euro)?', 1456565),
        15: ('Wie viele E-Books wurden 2016 in Deutschland verkauft?', 28100000),
        21: ('Wie viel hat der Bau des Burj Khalifa Tower in Dubai gekostet (in Tausend Dollar)?', 1500000),
        13: ('Wie viele Bücher sind in der American Library of Congress?', 16000000),
        17: ('Wie groß ist der Durchmesser der Sonne (in km)?', 1391400),

        1: ('Wie viele Einwohner hat New York City und sein Ballungsgebiet?', 21000000),
        10: ('Wie viele Smartphones wurden 2017 in Deutschland verkauft?', 24100000),
        3: ('Wie viele Einwohner hat Tokyo und sein Ballungsgebiet?', 38000000),
        12: ('Wie hoch ist das Durchschnittsgehalt der Spieler beim FC Bayern München (in Euro)?', 5460000),
        18: ('Wie schwer ist die Cheops-Pyramide (in Tonnen)?', 5000000),
        14: ('Wie viele Menschen in Mexiko identifizieren sich als indigen?', 6000000),
        20: ('Wie groß ist die Gesamtmasse der Ozeane auf der Erde (in Tausend Milliarden Tonnen)?', 1400000),
        16: ('Wie viele Menschen starben 2015 weltweit an Krebs?', 8800000),

        24: ('Wie viele Galaxien beherbergt das sichtbare Universum (in Millionen Galaxien)?', 100000),
        2: ('Wie viele Einwohner hat Madrid und sein Ballungsgebiet?', 6500000),
        22: ('Wie viele Zellen gibt es im menschlichen Körper (in Milliarden Zellen)?', 100000),
        4: ('Wie viele Einwohner hat Shanghai und sein Ballungsgebiet?', 25000000),
        6: ('Wie viele Einwohner hat Melbourne und sein Ballungsgebiet?', 4500000),
        8: ('Wie viele Fahrräder gibt es Ihrer Meinung nach in Deutschland?', 62000000),
        5: ('Wie viele Einwohner hat Amsterdam und sein Ballungsgebiet?', 1600000),
        7: ('Wie viele Autos wurden 2016 in Deutschland zugelassen?',45071000)}
    # Randomize the order in which the questions will be asked (by shuffling the order of the keys (ID) of the dictionary)
    ORDER_L = list(QUESTIONS_DICT.keys())
    # random.shuffle(ORDER_L)
    print(ORDER_L)
    # Define the treatments - We add the exact amount of treatments we require to run the experiment 4 High, 4 Low, 4 Pressure, 4 Delay, 8 Control
    TREATMENTS_L = ['Low','Control',  'High', 'Low', 'Delay', 'High', 'Delay', 'High',
                    'Control','Pressure', 'Control', 'Pressure', 'Control','Pressure', 'Control', 'Pressure',
                    'Delay', 'Low', 'Delay','Control','Low','Control','High','Control']
    # Randomize the order of the treatments by shuffling we randomize the treatments
   # random.shuffle(TREATMENTS_L)
    print(TREATMENTS_L) # TODO delete after debug
    # Define the list of consonants for the high and low cognitive load treatments
    CONSONANTS_L = ['B', 'C', 'D', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'X', 'Y', 'Z']    

    # Define the default timer and the pressure timer as constants
    TIMER = 180 #TODO change to 120
    TIMER_DEL = 120
    PRESSURE_TIMER = 14

    # Define the number of rounds - number of rounds in this case matches the number of questions 
    NUM_ROUNDS = len(QUESTIONS_DICT.keys())

class Subsession(BaseSubsession):
    pass

class Group(BaseGroup):
    # Define the variable for the social information
    averaged_estimates = models.FloatField()
    @property
    def rounded_averaged_estimates(self):
        return round(self.averaged_estimates)

class Player(BasePlayer):
    # Define question_ID - shows in the "Data" which question is being displayed
    question_ID = models.IntegerField()

    # Define the variables holding the answer and revised answer of the Player for a question
    answer =  models.FloatField(min=0, max=999999999999999)
    revised_answer = models.FloatField()

    @property
    def rounded_answer(self):
        return round(self.answer)


    # Define the variables for measuring time - time spent is measured within the html Pages: QuestionPage and ReiterationPage
    time_spent = models.FloatField()
    revised_time_spent = models.FloatField()
     
    # Define the current treatment variable and the flag if the treatment is fulfilled (used for reward calculation)
    curr_treatment = models.StringField()
    is_treatment_fulfilled = models.BooleanField()

    # Define the high and low cognitive load variables for displaying and submitting the randomly generated string
    low_cog_load_1 = models.StringField()
    low_cog_load_2 = models.StringField()
    low_cog_load_3 = models.StringField()
    high_cog_load_1 = models.StringField()
    high_cog_load_2 = models.StringField()
    high_cog_load_3 = models.StringField()
    low_cog_load_control_1 = models.StringField(label='Bitte geben Sie die Buchstabenkombination an.')
    low_cog_load_control_2 = models.StringField(label='Bitte geben Sie die Buchstabenkombination an.')
    low_cog_load_control_3 = models.StringField(label='Bitte geben Sie die Buchstabenkombination an.')
    high_cog_load_control_1 = models.StringField(label='Bitte geben Sie die Buchstabenkombination an.')
    high_cog_load_control_2 = models.StringField(label='Bitte geben Sie die Buchstabenkombination an.')
    high_cog_load_control_3 = models.StringField(label='Bitte geben Sie die Buchstabenkombination an.')
        
    # Define the confidence variable - using f'{"2":^{19}}' ensures that the labels are centered and that the spacing is even 
    initial_confidence = models.IntegerField(choices=[[1, '1 - sehr unsicher'], [2, f'{2:^19}'], [3, f'{3:^19}'],[4, f'{4:^19}'], [5, f'{5:^19}'], [6, f'{"6 - sehr sicher":^19}']],
                                     label='Bitte geben Sie an wie sicher Sie sich in Ihrer Schätzung sind.',
                                     widget=widgets.RadioSelectHorizontal)
    revised_confidence = models.IntegerField(choices=[[1, '1 - sehr unsicher'], [2, f'{2:^19}'], [3, f'{3:^19}'],[4, f'{4:^19}'], [5, f'{5:^19}'], [6, f'{"6 - sehr sicher":^19}']],
                                     label='Bitte geben Sie an wie sicher Sie sich in Ihrer Schätzung sind.',
                                     widget=widgets.RadioSelectHorizontal)  
    
    # Payoff does not need to be defined as it is a default variable it will be referred to as
    # the reward variable in further text(the value is only updated if the accuracy is within certain thresholds or if the treatment is not "fulfilled")
     

# ---------------------------------------------------------------------
# FUNCTIONS
# ---------------------------------------------------------------------
def averaging(group: Group):
    '''
        Calculates the averaged estimates of all players based on the geometrical means of their answers
        If the user does not enter a value, it is set to 0. These values are excluded from averaging (!= 0). 

        Args:
            group (Group):  the group holding all players
    '''
    answers_l = [p.answer for p in group.get_players() if p.answer != 0]
    
    # Only compute the geometrical mean if at least 1 person entered an answer (if answers_l is not empty), rounds the value to two places after comma
    if answers_l:
      #  group.averaged_estimates = round(int(gmean(answers_l)), 0)  # Convert to integer here
      group.averaged_estimates = gmean(answers_l)
    # This prevents the "maybe_none" error
    else:
        group.averaged_estimates = 0

def rewarding(group: Group):
    '''
        Calculates the reward for every player based on the accuracy of their answer

        Args:
            group (Group):  the group holding all players
    '''
    # Iterate over every player in the group
    for player in group.get_players():
        # Calculate the accuracy by calculating the ratio of the player's answer and the correct value
        current_accuracy = int(player.revised_answer) / C.QUESTIONS_DICT[C.ORDER_L[player.round_number - 1]][1]
        # Determine the reward based on the accuracy (the ranges are symmetrical e.g. 0.9-1.1 --> 4 points)
        if abs(1 - current_accuracy) <= 0.1:
            player.payoff = 4
        elif abs(1 - current_accuracy) <= 0.2:
            player.payoff = 2
        elif abs(1 - current_accuracy) <= 0.4:
            player.payoff = 1

def control_treatment(group: Group):
    '''
        Checks, if the treatment is "fulfilled"        

        Args:
            group (Group):  the group holding all players
    '''
    # Iterate over every player in the group and set the treatment flag accordingly
    for player in group.get_players():
        # For Time Pressure, Delay, and Control the treatment is always "fulfilled"
        if player.curr_treatment not in ['High', 'Low']:
            player.is_treatment_fulfilled = True
        # However, for High and Low Cognitive Load the input of the user has to match the randomly generated string to "fulfill" the treatment
        elif player.curr_treatment == 'High':
            player.is_treatment_fulfilled = (player.high_cog_load_1 == player.high_cog_load_control_1.upper() and player.high_cog_load_2 == player.high_cog_load_control_2.upper() and player.high_cog_load_3 == player.high_cog_load_control_3.upper())
        elif player.curr_treatment == 'Low':
            player.is_treatment_fulfilled = (player.low_cog_load_1 == player.low_cog_load_control_1.upper() and player.low_cog_load_2 == player.low_cog_load_control_2.upper() and player.low_cog_load_3 == player.low_cog_load_control_3.upper())

        # Since the reward is calculated before the cognition check, the reward only has to be adapted if the treatment is not "fulfilled"
        if not player.is_treatment_fulfilled:
            player.payoff = 0

def get_timeout_seconds_custom(player):
    '''
        Override the get_timeout_seconds function to return a different timer depending on the current treatment

        Args:
            player (Player):  the player to change the timer for

        Returns:
            timer (int): the value of the timer in seconds
    '''
    if player.curr_treatment == 'Pressure':
        return C.PRESSURE_TIMER
    elif player.curr_treatment == 'Delay':
        return C.TIMER_DEL
    else:
        return C.TIMER


def js_vars_custom(player):
    '''
        Override the js_vars function to pass certain variables to javascript

        Args:
            player (Player):  the player to pass the variables to javascript for

        Returns:
            (dict): with the timer value handed over as timer and the current treatment handed over as treatment
    '''
    return dict(timer=C.TIMER,
                treatment=player.curr_treatment)

# ---------------------------------------------------------------------
# CLASSES - PAGES
# ---------------------------------------------------------------------
class LoCLPresentPage_1(Page):
    form_model = 'player'
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom
    
    @staticmethod
    def vars_for_template(player):
        # For Low Cognition a random consonant is selected and repeated five times
        player.low_cog_load_1 = ''.join(random.sample(C.CONSONANTS_L, 1) * 5)
    
    @staticmethod
    def is_displayed(player):
        # As this is the first page every round, the current treatment is assigned here
        player.curr_treatment = C.TREATMENTS_L[player.round_number - 1]
        # The page is only displayed if the treatment is Low
        return player.curr_treatment == 'Low'
    
class HiCLPresentPage_1(Page):
    form_model = 'player'
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def vars_for_template(player):
        # For High Cognition five random consonants are selected
        player.high_cog_load_1 = ''.join(random.sample(C.CONSONANTS_L, 5))
    
    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is High
        return player.curr_treatment == 'High'
    
class LoCLPresentPage_2(Page):
    form_model = 'player'
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom
    
    @staticmethod
    def vars_for_template(player):
        # For Low Cognition a random consonant is selected and repeated five times
        player.low_cog_load_2 = ''.join(random.sample(C.CONSONANTS_L, 1) * 5)
    
    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is Low
        return player.curr_treatment == 'Low'
    
class HiCLPresentPage_2(Page):
    form_model = 'player'
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def vars_for_template(player):
        # For High Cognition five random consonants are selected
        player.high_cog_load_2 = ''.join(random.sample(C.CONSONANTS_L, 5))
    
    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is High
        return player.curr_treatment == 'High'
    
class LoCLPresentPage_3(Page):
    form_model = 'player'
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom
    
    @staticmethod
    def vars_for_template(player):
        # For Low Cognition a random consonant is selected and repeated five times
        player.low_cog_load_3 = ''.join(random.sample(C.CONSONANTS_L, 1) * 5)
    
    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is Low
        return player.curr_treatment == 'Low'
    
class HiCLPresentPage_3(Page):
    form_model = 'player'
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def vars_for_template(player):
        # For High Cognition five random consonants are selected
        player.high_cog_load_3 = ''.join(random.sample(C.CONSONANTS_L, 5))
    
    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is High
        return player.curr_treatment == 'High'

class QuestionPage(Page):
    form_model = 'player'
    form_fields = ['answer', 'time_spent']

    # Override the functions for the javascript variables and the timer with the custom functions
    js_vars = js_vars_custom
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def answer_error_message(player, value):
        if len(str(value).split('.')[0]) > 15: # Consider only the digits before the decimal
            return 'Your input should be a maximum of 15 digits.'

    @staticmethod
    def vars_for_template(player):
        # Assign the question ID according to the current predefined order and the current round
        player.question_ID = C.ORDER_L[player.round_number - 1]

        # Set the question label on the QuestionPage according to the currently selected question
        return dict(question_label = C.QUESTIONS_DICT[C.ORDER_L[player.round_number - 1]][0],
                    treatment = player.curr_treatment
                    )

class InitialConfidencePage(Page):
    form_model = 'player'
    form_fields = ['initial_confidence']
    
    # Override the functions for the javascript variables and the timer with the custom functions
    js_vars = js_vars_custom
    get_timeout_seconds = get_timeout_seconds_custom
    
class SocialInformationPage(Page):
    # Override the functions for the javascript variables and the timer with the custom functions
    js_vars = js_vars_custom
    get_timeout_seconds = get_timeout_seconds_custom


class ReiterationPage(Page):
    form_model = 'player'
    form_fields = ['revised_answer', 'revised_time_spent']

    # Override the functions for the javascript variables and the timer with the custom functions
    js_vars = js_vars_custom
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def vars_for_template(player):
        return dict(question_label=C.QUESTIONS_DICT[C.ORDER_L[player.round_number - 1]][0],
                    treatment=player.curr_treatment
                    )

class LoCLControlPage_1(Page):
    form_model = 'player'
    form_fields = ['low_cog_load_control_1']    
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is Low
        return player.curr_treatment == 'Low'

class HiCLControlPage_1(Page):
    form_model = 'player'
    form_fields = ['high_cog_load_control_1']    
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is High
        return player.curr_treatment == 'High'
    
class LoCLControlPage_2(Page):
    form_model = 'player'
    form_fields = ['low_cog_load_control_2']    
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is Low
        return player.curr_treatment == 'Low'

class HiCLControlPage_2(Page):
    form_model = 'player'
    form_fields = ['high_cog_load_control_2']    
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is High
        return player.curr_treatment == 'High'
    
class LoCLControlPage_3(Page):
    form_model = 'player'
    form_fields = ['low_cog_load_control_3']    
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is Low
        return player.curr_treatment == 'Low'

class HiCLControlPage_3(Page):
    form_model = 'player'
    form_fields = ['high_cog_load_control_3']    
    # Override the functions for the timer with the custom functions
    get_timeout_seconds = get_timeout_seconds_custom

    @staticmethod
    def is_displayed(player):
        # The page is only displayed if the treatment is High
        return player.curr_treatment == 'High'

class RevisedConfidencePage(Page):
    form_model = 'player'
    form_fields = ['revised_confidence']
    
    # Override the functions for the javascript variables and the timer with the custom functions
    js_vars = js_vars_custom
    get_timeout_seconds = get_timeout_seconds_custom
    
class QuestionWaitPage(WaitPage):
    pass

class ConfidenceWaitPage(WaitPage):
    after_all_players_arrive = averaging

class ReiterationWaitPage(WaitPage):
    after_all_players_arrive = rewarding

class TreatmentWaitPage(WaitPage):
    after_all_players_arrive = control_treatment

# ---------------------------------------------------------------------
# Define the order of the pages
# ---------------------------------------------------------------------
page_sequence = [LoCLPresentPage_1, HiCLPresentPage_1, QuestionPage, QuestionWaitPage, LoCLControlPage_1, HiCLControlPage_1, InitialConfidencePage, ConfidenceWaitPage, LoCLPresentPage_2, HiCLPresentPage_2, SocialInformationPage, LoCLControlPage_2, HiCLControlPage_2, LoCLPresentPage_3, HiCLPresentPage_3, ReiterationPage, LoCLControlPage_3, HiCLControlPage_3, ReiterationWaitPage, TreatmentWaitPage, RevisedConfidencePage]