from otree.api import *
from otree import settings
import numpy as np
import pandas as pd


doc = """Experiment 2 for Social Inference project
"""

# TO DOs ---------------------------------------------------------------------------------------------------------------
# Questions whether participants thought it was a real human partner
# Screen at the end of training phase that the experiment starts now

# NOTES TO SELF --------------------------------------------------------------------------------------------------------
# A player is initialized per session, i.e., in every session I will have x players
# Otree initializes a participant code that is valid across different subsessions
# I can create 1 room and in that 1 room I have 3 session configs that the experimenter can select. Then,
# once a session config is selected the experimenter can create however many links they need and share it with the participants.

# BUT with that set-up I would not have continuous participant numbers... however, I assign environment and demonstrator based on those participant numbers
# and I want to make sure that the same environments and demonstrators are used in all conditions...


# OTREE CLASSES --------------------------------------------------------------------------------------------------------
class C(BaseConstants):
    NAME_IN_URL = 'SI_Exp2_a_TrainingSequence'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 10

class Subsession(BaseSubsession):
    environment = models.IntegerField()
    subject_number = models.IntegerField()


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    Two_comprehension_check_spatial = models.IntegerField(label="Are points more similar to each other when the cells are closer to each other?", widget=widgets.RadioSelect, choices=[[1, 'a) Yes'], [2, 'b) No'],])
    Two_result = models.IntegerField()
    Three_comprehension_check_inference = models.IntegerField(label="Is the following statement True or False: I will see all selections that the player made and all points that the player received.", widget=widgets.RadioSelect, choices=[[1, 'a) True'], [2, 'b) False'],])
    Three_result = models.IntegerField()

    page_name = models.StringField()

    counter = models.IntegerField(default=0)

    IV_num_choice_only = models.IntegerField()

    inferred_reward = models.IntegerField(
        min=0,
        max=500,
        label='What reward did the player get on the previous trial?',
        blank=True,
        null=True,
        doc="""This player's decision""")

    inferred_choice = models.StringField(
        max_length=500,
        blank=True,
        null=True)

    predicted_reward = models.IntegerField(
        min=0,
        max=500,
        label='What reward will the player get on the next trial?',
        blank=True,
        null=True,
        doc="""This player's decision""")

    predicted_choice = models.StringField(
        max_length=500,
        blank=True,
        null=True)
    predicted_choice_row = models.IntegerField(
        max_length=500,
        blank=True,
        null=True)
    predicted_choice_column = models.IntegerField(
        max_length=500,
        blank=True,
        null=True)

    RT_inferred_reward = models.FloatField(initial=None, null=True, blank=True)
    RT_inferred_choice = models.FloatField(initial=None, null=True, blank=True)
    RT_predicted_reward = models.FloatField(initial=None, null=True, blank=True)
    RT_predicted_choice = models.FloatField(initial=None, null=True, blank=True)

    round_data_field = models.IntegerField()

    performance = models.FloatField()
    payout = models.FloatField()

    participant_code = models.StringField()

    ObsB_predicted_reward = models.IntegerField(field_maybe_none=True)
    ObsB_inferred_reward = models.IntegerField(field_maybe_none=True)
    ObsB_highlight_index_row = models.IntegerField(field_maybe_none=True)
    ObsB_highlight_index_column = models.IntegerField(field_maybe_none=True)
    ObsB_inf_highlight_index_row = models.IntegerField(field_maybe_none=True)
    ObsB_inf_highlight_index_column = models.IntegerField(field_maybe_none=True)
    ObsB_time_to_decide = models.FloatField(field_maybe_none=True)

    cheat_stay_on_page = models.IntegerField()

    # if we use Player model for read_csv these fields are required
    trial_index = models.IntegerField()
    trial_type = models.StringField()
    reward_inference_query = models.IntegerField()
    choice_inference_query = models.IntegerField()
    reward_prediction_query = models.IntegerField()
    choice_prediction_query = models.IntegerField()
    def store_round_data(self):
        self.round_data_field = self.round_number  # Storing round number as an example

    def before_next_page(self):
        self.store_round_data()

    def increment_counter(self):
        self.counter += 1
        return self.counter


class Row(ExtraModel):
    player = models.Link(Player)

    SN = models.IntegerField()
    participant_label = models.IntegerField()
    participant_code = models.StringField()
    IV_number_choice_only_trials = models.IntegerField()
    trial_index = models.IntegerField()

    trial = models.IntegerField()
    trial_type = models.StringField()

    inferred_reward = models.IntegerField(
        min=0,
        max=500,
        label='What reward did the player get on the previous trial?',
        blank=True,
        doc="""This player's decision""")

    inferred_choice = models.StringField(
        max_length=500,
        blank=True)

    predicted_reward = models.IntegerField(
        min=0,
        max=500,
        label='What reward will the player get on the next trial?',
        blank=True,
        null=True,
        doc="""This player's decision""")

    predicted_choice = models.StringField(
        max_length=500,
        blank=True)

    RT_inferred_reward = models.FloatField(initial=None)
    RT_inferred_choice = models.FloatField(initial=None)
    RT_predicted_reward = models.FloatField(initial=None)
    RT_predicted_choice = models.FloatField(initial=None)

    reward_inference_query = models.BooleanField()
    choice_inference_query = models.BooleanField()
    reward_prediction_query = models.BooleanField()
    choice_prediction_query = models.BooleanField()

    demonstrator_choice_index = models.IntegerField()
    demonstrator_choice_index_row = models.IntegerField()
    demonstrator_choice_index_column = models.IntegerField()
    demonstrator_noisy_rescaled_reward = models.IntegerField()
    demonstrator_choice_probs = models.StringField()
    demonstrator_noisy_rescaled_all_rewards = models.StringField()

    mean_reward = models.FloatField()
    x1_index = models.IntegerField()
    x2_index = models.IntegerField()

    money = models.FloatField(initial=0)  # per trial

    # fields I have to create, but hopefully not save
    Demonstrator_for_SN = models.IntegerField()
    environment = models.IntegerField()
    scale = models.IntegerField()
    demonstrator_choice_prob = models.FloatField()
    demonstrator_true_all_rewards = models.LongStringField()
    demonstrator_noisy_reward = models.FloatField()

    ObsB_predicted_reward = models.IntegerField(field_maybe_none=True)
    ObsB_inferred_reward = models.IntegerField(field_maybe_none=True)
    ObsB_predicted_choice_row = models.IntegerField(field_maybe_none=True)
    ObsB_predicted_choice_col = models.IntegerField(field_maybe_none=True)
    ObsB_inferred_choice_row = models.IntegerField(field_maybe_none=True)
    ObsB_inferred_choice_col = models.IntegerField(field_maybe_none=True)
    ObsB_time_to_decide = models.FloatField(field_maybe_none=True)


# PAGES ----------------------------------------------------------------------------------------------------------------

class empty_grid(Page):
    template_name = 'SI_Exp2_NoPartner_TrainingSequence/empty_grid.html'
    counter = 0
    def is_displayed(player: Player):
        # "Next round" should not be shown on first trial
        if player.round_number > 1:
            return True

    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}
    timeout_seconds = 4


class predict_choice(Page):
    template_name = 'SI_Exp2_NoPartner_TrainingSequence/predict_choice.html'

    form_model = 'player'
    form_fields = ['predicted_choice', 'RT_predicted_choice'] #cheat_stay_on_page

    def is_displayed(player: Player):
        # import the csv
        # placerholder
        ObsA_trials = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Player)
        #print('ObsA_trials', ObsA_trials)
        #ObsB_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])+1) + '.csv', Row)
        # Change back
        #ObsA_trials = 'SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv'
        #ObsB_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])+1) + '.csv', Row)
        # check whether page sequence csv says we have to ask choice prediction query for either Observer
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObsA_trials[t]
        #print('ObsA_this_trials_data', ObsA_this_trials_data)
        #ObsB_this_trials_data = ObsB_trials[t]
        if int(ObsA_this_trials_data['choice_prediction_query']) == 1:
            return True

    def vars_for_template(self):
        #environment = self.session.config['environment']

        # this just gives the grid
        grid_numbers = np.empty((11, 11))
        grid_data = []
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'value': value, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)

        ObserverA_PageSequence = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])+1) + '.csv', Row) # THIS IS A PLACEHOLDER FOR NOW
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])-1+self.id_in_group+3) + '.csv', Row)  # THIS ALSO IS A PLACEHOLDER
        t = (self.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]
        demonstrator_this_trials_data = demonstrator[t]
        # What is being shown for Obs A
        if ObsA_this_trials_data['choice_prediction_query'] == 1:
            ObsA_query_row1 = 'Which cell do you think the player will select next?'
            ObsA_query_row2 = '(Please click into the respective cell and then submit)'
            show_form_field = 1
        else:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 =  ''
            show_form_field = 0
        # What is being shown for Obs B
        if ObsB_this_trials_data['choice_prediction_query'] == 1:
            ObsB_query_row1 = 'Which cell do you think the player will select next?'
            ObsB_query_row2 = '(Please click into the respective cell and then submit)'
            # The following is a placerholder because I take the real demonstrators value and just add some small noise
            ObsB_highlight_index_row = np.clip(int(np.random.normal(demonstrator_this_trials_data['demonstrator_choice_index_row'], scale=1)), 0, 10)
            ObsB_highlight_index_column = np.clip(int(np.random.normal(demonstrator_this_trials_data['demonstrator_choice_index_column'], scale=1)), 0, 10)
            ObsB_show_form_field = 1
            ObsB_show_guess_text = 'Observer B\'s prediction is shown in purple.'
            ObsB_time_to_decide = np.clip(np.random.lognormal(5000, 7000), 4000, 10000)
        else:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 =  ''
            ObsB_highlight_index_row = ''
            ObsB_highlight_index_column = ''
            ObsB_show_form_field = 0
            ObsB_show_guess_text = ''
            ObsB_time_to_decide = 0

        #print('ObsB_highlight_index_row', ObsB_highlight_index_row, type(ObsB_highlight_index_row))

        return {'grid_data': grid_data,
                'ObsB_highlight_index_row': ObsB_highlight_index_row, 'ObsB_highlight_index_column': ObsB_highlight_index_column,
                'ObsA_query_row1':ObsA_query_row1, 'ObsA_query_row2':ObsA_query_row2,
                'ObsB_query_row1':ObsB_query_row1, 'ObsB_query_row2':ObsB_query_row2,
                'show_form_field':show_form_field, 'ObsB_show_form_field':ObsB_show_form_field,
                'ObsB_show_guess_text': ObsB_show_guess_text, 'ObsB_time_to_decide':ObsB_time_to_decide}

    @staticmethod
    def live_method(player: Player, data):
        if player.id_in_group == 1:
            if data["predicted_choice"] == 'NA,NA':
                return {1: {"predicted_choice": '', "predicted_choice_row": '', "predicted_choice_column": ''}}  #
            else:
                strings = data["predicted_choice"].split(',')
                return {1: {"predicted_choice":data["predicted_choice"], "predicted_choice_row":int(strings[0]), "predicted_choice_column":int(strings[1])}} #

        if player.id_in_group == 2:
            if data["predicted_choice"] == 'NA,NA':
                return {2: {"predicted_choice": '', "predicted_choice_row": '', "predicted_choice_column": ''}}  #
            else:
                strings = data["predicted_choice"].split(',')
                #print('livesend success', data["predicted_choice"], type(data["predicted_choice"]))
                # for only Obs B query, [0] and [1] will return an N and an A
                return {2: {"predicted_choice":data["predicted_choice"], "predicted_choice_row":int(strings[0]), "predicted_choice_column":int(strings[1])}} #

        if player.id_in_group == 3:
            if data["predicted_choice"] == 'NA,NA':
                return {3: {"predicted_choice": '', "predicted_choice_row": '', "predicted_choice_column": ''}}  #
            else:
                strings = data["predicted_choice"].split(',')
                #print('livesend success', data["predicted_choice"], type(data["predicted_choice"]))
                # for only Obs B query, [0] and [1] will return an N and an A
                return {3: {"predicted_choice":data["predicted_choice"], "predicted_choice_row":int(strings[0]), "predicted_choice_column":int(strings[1])}} #

        if player.id_in_group == 4:
            if data["predicted_choice"] == 'NA,NA':
                return {4: {"predicted_choice": '', "predicted_choice_row": '', "predicted_choice_column": ''}}  #
            else:
                strings = data["predicted_choice"].split(',')
                #print('livesend success', data["predicted_choice"], type(data["predicted_choice"]))
                # for only Obs B query, [0] and [1] will return an N and an A
                return {4: {"predicted_choice":data["predicted_choice"], "predicted_choice_row":int(strings[0]), "predicted_choice_column":int(strings[1])}} #





class choice_phase(Page):
    template_name = 'SI_Exp2_NoPartner_TrainingSequence/choice_phase.html'

    def vars_for_template(self):
        # this just gives the grid
        grid_numbers = np.empty((11, 11))
        #print('grid_numbers', grid_numbers)
        grid_data = []
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            #print('row_idx', row_idx)
            #print('row', row)
            for col_idx, value in enumerate(row):
                #print('col_idx', col_idx)
                #print('value', value)
                cell = {'value': value, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)
            #print('grid_data', grid_data)

        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])-1+self.id_in_group+3) + '.csv', Row)
        ObserverA_PageSequence = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])+1) + '.csv', Row) # THIS IS A PLACEHOLDER FOR NOW
        t = (self.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]
        if int(ObsA_this_trials_data['trial_type']) == 1 or int(ObsA_this_trials_data['trial_type']) == 2:
            ObsA_query_row1 = 'See player\'s selection'
            ObsA_query_row2 = ''
            this_trials_demonstrator_data = demonstrator[self.round_number - 1]
            highlight_index_row = this_trials_demonstrator_data['demonstrator_choice_index_row']
            highlight_index_column = this_trials_demonstrator_data['demonstrator_choice_index_column']
            ObsA_sees = True
        else:
            ObsA_query_row1 = '?'
            ObsA_query_row2 =  ''
            highlight_index_row = np.NaN
            highlight_index_column = np.NaN
            ObsA_sees = False
        if int(ObsB_this_trials_data['trial_type']) == 1 or int(ObsB_this_trials_data['trial_type']) == 2:
            ObsB_query_row1 = 'See player\'s selection'
            ObsB_query_row2 = ''
            ObsB_sees = True
        else:
            ObsB_query_row1 = '?'
            ObsB_query_row2 =  ''
            ObsB_sees = False

        return {'grid_data': grid_data, 'highlight_index_row': highlight_index_row,
                'highlight_index_column': highlight_index_column,
                'ObsA_query_row1':ObsA_query_row1, 'ObsA_query_row2':ObsA_query_row2,
                'ObsB_query_row1':ObsB_query_row1, 'ObsB_query_row2':ObsB_query_row2,
                'ObsA_sees':ObsA_sees, 'ObsB_sees':ObsB_sees}

    timeout_seconds = 6


class infer_reward(Page):
    template_name = 'SI_Exp2_NoPartner_TrainingSequence/infer_reward.html'

    form_model = 'player'
    form_fields = ['inferred_reward', 'RT_inferred_reward']

    def error_message(self, values):
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1+self.id_in_group) + '.csv', Row)
        t = (self.round_number - 2)
        ObsA_previous_trials_data = ObsA_trials[t]
        if int(ObsA_previous_trials_data['reward_inference_query']) == 1 :
            if values["inferred_reward"] is None or values["inferred_reward"] == '':
                return 'You must enter a value'

    def is_displayed(player: Player):
        #print("ROUND NUMBER:", player.round_number)
        # import the csv, #ALEX not if first trial
        ObsA_trials = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Player)
        t = (player.round_number - 2)
        #print('round number - 2', t)
        ObsA_previous_trials_data = ObsA_trials[t]

        #print('ObsA_previous_trials_data', ObsA_previous_trials_data)
        # this should be displayed on the current trial, if the reward inference query on the previous trial was...
        #print('show reward inf Obs A', int(ObsA_previous_trials_data['reward_inference_query']) == 1)
        #print('show reward inf Obs B', int(ObsB_previous_trials_data['reward_inference_query']) == 1)
        if int(ObsA_previous_trials_data['reward_inference_query']) == 1 : # if a reward inference should be made
            return True


    def vars_for_template(self):
        grid_numbers = np.empty((11, 11))
        # Restructure data to include indices
        grid_data = []
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'value': value, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)

        # need to identify the row and column that should be highlighted
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])-1+self.id_in_group+3) + '.csv', Row)
        ObserverA_PageSequence = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])+1) + '.csv', Row) # THIS IS A PLACEHOLDER FOR NOW
        t = (self.round_number - 2)
        ObsA_prev_trials_data = ObserverA_PageSequence[t]
        ObsB_prev_trials_data = ObserverB_PageSequence[t]
        c = (self.round_number - 1)
        ObsA_current_trials_data = ObserverA_PageSequence[c]
        ObsB_current_trials_data = ObserverB_PageSequence[c]

        # Get the "real" reward and grid index of demonstrator's choice and reward
        prev_trials_demonstrator_data = demonstrator[self.round_number - 2]
        prev_highlight_index_row = int(prev_trials_demonstrator_data['demonstrator_choice_index_row'])
        prev_highlight_index_column = int(prev_trials_demonstrator_data['demonstrator_choice_index_column'])
        demonstrator_reward = prev_trials_demonstrator_data['demonstrator_noisy_rescaled_reward']#grid_data[prev_highlight_index_row][prev_highlight_index_column]['value']

        # What is shown for Obs A
        if int(ObsA_prev_trials_data['reward_inference_query']) == 1:
            ObsA_query_row1 = 'How many points do you think the player received for their previous selection (highlighted in light red)?'
            ObsA_query_row2 = '(the current selection is highlighted in yellow)'
            ObsA_show_guess = 'Observer A inferred a reward of: <span id=result_inferred_reward ></span>'
            # This is important! The -2 indicates that we're getting the info from the trial before (-1 because counter starting at 1)
            previous_trials_demonstrator_data = demonstrator[self.round_number - 2]
            current_trials_demonstrator_data = demonstrator[self.round_number - 1]
            highlight_index_row = int(previous_trials_demonstrator_data['demonstrator_choice_index_row'])
            highlight_index_column = int(previous_trials_demonstrator_data['demonstrator_choice_index_column'])
            current_highlight_index_row = int(current_trials_demonstrator_data['demonstrator_choice_index_row'])
            current_highlight_index_column = int(current_trials_demonstrator_data['demonstrator_choice_index_column'])
            show_form_field = 1
        elif int(ObsA_prev_trials_data['reward_inference_query']) == 0 and int(ObsA_prev_trials_data['trial_type']) != 3 and int(ObsA_current_trials_data['trial_type']) != 3:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 =  ''
            ObsA_show_guess = ''
            previous_trials_demonstrator_data = demonstrator[self.round_number - 2]
            current_trials_demonstrator_data = demonstrator[self.round_number - 1]
            highlight_index_row = int(previous_trials_demonstrator_data['demonstrator_choice_index_row'])
            highlight_index_column = int(previous_trials_demonstrator_data['demonstrator_choice_index_column'])
            current_highlight_index_row = int(current_trials_demonstrator_data['demonstrator_choice_index_row'])
            current_highlight_index_column = int(current_trials_demonstrator_data['demonstrator_choice_index_column'])
            show_form_field = 0
        elif int(ObsA_prev_trials_data['reward_inference_query']) == 0 and int(ObsA_prev_trials_data['trial_type']) != 3 and int(ObsA_current_trials_data['trial_type']) == 3:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 =  ''
            ObsA_show_guess = ''
            previous_trials_demonstrator_data = demonstrator[self.round_number - 2]
            current_trials_demonstrator_data = demonstrator[self.round_number - 1]
            highlight_index_row = int(previous_trials_demonstrator_data['demonstrator_choice_index_row'])
            highlight_index_column = int(previous_trials_demonstrator_data['demonstrator_choice_index_column'])
            current_highlight_index_row = np.NaN
            current_highlight_index_column = np.NaN
            show_form_field = 0
        elif int(ObsA_prev_trials_data['reward_inference_query']) == 0 and int(ObsA_prev_trials_data['trial_type']) == 3 and int(ObsA_current_trials_data['trial_type']) != 3:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 =  ''
            ObsA_show_guess = ''
            previous_trials_demonstrator_data = demonstrator[self.round_number - 2]
            current_trials_demonstrator_data = demonstrator[self.round_number - 1]
            highlight_index_row = np.NaN
            highlight_index_column = np.NaN
            current_highlight_index_row = int(current_trials_demonstrator_data['demonstrator_choice_index_row'])
            current_highlight_index_column = int(current_trials_demonstrator_data['demonstrator_choice_index_column'])
            show_form_field = 0
        else:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 =  ''
            ObsA_show_guess = ''
            highlight_index_row = np.NaN
            highlight_index_column = np.NaN
            current_highlight_index_row = np.NaN
            current_highlight_index_column = np.NaN
            show_form_field = 0
        # What is shown for Obs B
        if int(ObsB_prev_trials_data['reward_inference_query']) == 1:
            ObsB_query_row1 = 'How many points do you think the player received for their previous selection (highlighted in light red)?'
            ObsB_query_row2 = '(the current selection is highlighted in yellow)'
            ObsB_inferred_reward = np.clip(int(np.random.normal(demonstrator_reward, scale=1)), 0, 100)
            ObsB_show_guess = 'Observer B inferred a reward of: '
            ObsB_time_to_decide = np.clip(np.random.lognormal(10000, 8000), 4000, 14000)
        else:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 =  ''
            ObsB_inferred_reward = ''
            ObsB_show_guess = ''
            ObsB_time_to_decide = 0

            # When to show background boxes
        # don't show box if 3
        if int(ObsA_current_trials_data['trial_type']) == 3:
            ObsA_sees = False
        else:
            ObsA_sees = True

        if int(ObsB_current_trials_data['trial_type']) == 3:
            ObsB_sees = False
        else:
            ObsB_sees = True


        return {'grid_data': grid_data, 'highlight_index_row': highlight_index_row, 'highlight_index_column': highlight_index_column,
                'current_highlight_index_row': current_highlight_index_row, 'current_highlight_index_column': current_highlight_index_column,
                'ObsA_query_row1':ObsA_query_row1, 'ObsA_query_row2':ObsA_query_row2,
                'ObsB_query_row1':ObsB_query_row1, 'ObsB_query_row2':ObsB_query_row2,
                'show_form_field':show_form_field, 'ObsB_inferred_reward':ObsB_inferred_reward,
                'ObsA_show_guess':ObsA_show_guess, 'ObsB_show_guess':ObsB_show_guess,
                'ObsA_sees':ObsA_sees, 'ObsB_sees':ObsB_sees, 'ObsB_time_to_decide':ObsB_time_to_decide}

    @staticmethod
    def live_method(player: Player, data):
        if player.id_in_group == 1:
            return {1: {"result":data["result"]}}

        if player.id_in_group == 2:
            return {2: {"result":data["result"]}}

        if player.id_in_group == 3:
            return {3: {"result":data["result"]}}

        if player.id_in_group == 4:
            return {4: {"result":data["result"]}}

class predict_reward(Page):
    template_name = 'SI_Exp2_NoPartner_TrainingSequence/predict_reward.html'

    form_model = 'player'
    form_fields = ['predicted_reward', 'RT_predicted_reward']

    def error_message(self, values):
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1+self.id_in_group) + '.csv', Row)
        t = (self.round_number - 1)
        ObsA_previous_trials_data = ObsA_trials[t]
        if int(ObsA_previous_trials_data['reward_prediction_query']) == 1 :
            if values["predicted_reward"] is None or values["predicted_reward"] == '':
                return 'You must enter a value'

    def is_displayed(player: Player):
        ObsA_trials = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Player)
        #ObsB_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])+1) + '.csv', Row)
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObsA_trials[t]
        #ObsB_this_trials_data = ObsB_trials[t]
        #print('ObsA_this_trials_data', int(ObsA_this_trials_data['reward_prediction_query']) == 1)
        #print('ObsB_this_trials_data', int(ObsB_this_trials_data['reward_prediction_query']) == 1)
        return (int(ObsA_this_trials_data['reward_prediction_query']) == 1)

    def vars_for_template(self):
        grid_numbers = np.empty((11, 11))
        #highlight_index_row = 1  # Row index of the cell to highlight
        #highlight_index_column = 2  # Column index of the cell to highlight

        # This block gives which choice should be highlighted
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])-1+self.id_in_group+3) + '.csv', Row)
        trials_demonstrator_data = demonstrator[self.round_number - 1]
        highlight_index_row = trials_demonstrator_data['demonstrator_choice_index_row']
        highlight_index_column = trials_demonstrator_data['demonstrator_choice_index_column']

        # Restructure data to include indices
        grid_data = []
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'value': value, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)



        ObserverA_PageSequence = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])+1) + '.csv', Row) # THIS IS A PLACEHOLDER FOR NOW
        t = (self.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]

        # Get the "real" reward and grid index of demonstrator's choice and reward
        current_trials_demonstrator_data = demonstrator[self.round_number - 1]
        current_highlight_index_row = int(current_trials_demonstrator_data['demonstrator_choice_index_row'])
        current_highlight_index_column = int(current_trials_demonstrator_data['demonstrator_choice_index_column'])
        demonstrator_reward = trials_demonstrator_data['demonstrator_noisy_rescaled_reward']#int(grid_data[current_highlight_index_row][current_highlight_index_column]['value'])

        # What is shown for Obs A
        # if trial type either full info or choice-only and Obs A asked a query, show choice on screen
        if int(ObsA_this_trials_data['reward_prediction_query']) == 1 and int(ObsA_this_trials_data['trial_type']) != 3:
            ObsA_query_row1 = 'How many points do you think the player will receive for this selection?'
            ObsA_query_row2 = ''
            show_form_field = 1
            ObsA_show_guess = 'Observer A predicted a reward of: <span id=result_predicted_reward ></span>'
        # if trial type either full info or choice-only, but only Obs B makes prediction (but choice should still be shown on screen)
        elif int(ObsA_this_trials_data['reward_prediction_query']) == 0 and int(ObsA_this_trials_data['trial_type']) != 3:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 = ''
            show_form_field = 0
            ObsA_show_guess = ''
        # if trial type reward-only (i.e. 3), no choice on screen
        elif int(ObsA_this_trials_data['reward_prediction_query']) == 1 and int(ObsA_this_trials_data['trial_type']) == 3:
            ObsA_query_row1 = 'How many points do you think the player will receive for their selection?'
            ObsA_query_row2 = ''
            current_highlight_index_row = np.NaN
            current_highlight_index_column = np.NaN
            show_form_field = 1
            ObsA_show_guess = 'Observer A predicted a reward of: <span id=result_predicted_reward ></span>'
        else:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 =  ''
            current_highlight_index_row = np.NaN
            current_highlight_index_column = np.NaN
            show_form_field = 0
            ObsA_show_guess = ''
        # What is shown for Obs B
        if int(ObsB_this_trials_data['reward_prediction_query']) == 1 and int(ObsB_this_trials_data['trial_type']) != 3:
            ObsB_query_row1 = 'How many points do you think the player will receive for this selection?'
            ObsB_query_row2 = ''
            ObsB_show_guess = 'Observer B predicted a reward of: '
            show_form_field_ObsB = 1
            ObsB_time_to_decide = np.clip(np.random.lognormal(6000, 7000), 4000, 10000)
            # THIS IS A PLACEHOLDER. TAKING THE REAL REWARD VALUE AND PUTTING A NORMAL DISTRIBUTION ON IT
            ObsB_predicted_reward = np.clip(int(np.random.normal(demonstrator_reward, scale=1)), 0, 100)
        elif int(ObsB_this_trials_data['reward_prediction_query']) == 1 and int(ObsB_this_trials_data['trial_type']) == 3:
            ObsB_query_row1 = 'How many points do you think the player will receive for their selection?'
            ObsB_query_row2 = ''
            ObsB_show_guess = 'Observer B predicted a reward of: '
            show_form_field_ObsB = 1
            # THIS IS A PLACEHOLDER. TAKING THE REAL REWARD VALUE AND PUTTING A NORMAL DISTRIBUTION ON IT
            ObsB_predicted_reward = np.clip(int(np.random.normal(demonstrator_reward, scale=1)), 0, 100)
            ObsB_time_to_decide = np.clip(np.random.lognormal(6000, 7000), 4000, 10000)
        else:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 =  ''
            ObsB_show_guess = ' '
            show_form_field_ObsB = 0
            ObsB_predicted_reward = ''
            ObsB_time_to_decide = 0

        # Determine the background boxes. 1 full-inf trial, 2 choice only, 3 reward only
        # Should have background box if choice was shown
        if int(ObsA_this_trials_data['trial_type']) == 1 or int(ObsA_this_trials_data['trial_type']) == 2:
            ObsA_sees = True
        else:
            ObsA_sees = False

        if int(ObsB_this_trials_data['trial_type']) == 1 or int(ObsB_this_trials_data['trial_type']) == 2:
            ObsB_sees = True
        else:
            ObsB_sees = False

        ObsB_trial_type = int(ObsB_this_trials_data['trial_type'])

        return {'grid_data': grid_data, 'current_highlight_index_row': current_highlight_index_row,
                'current_highlight_index_column': current_highlight_index_column,
                'ObsA_query_row1':ObsA_query_row1, 'ObsA_query_row2':ObsA_query_row2,
                'ObsB_query_row1':ObsB_query_row1, 'ObsB_query_row2':ObsB_query_row2,
                'show_form_field':show_form_field, 'show_form_field_ObsB':show_form_field_ObsB,
                'ObsB_predicted_reward': ObsB_predicted_reward,
                'ObsA_show_guess':ObsA_show_guess, 'ObsB_show_guess':ObsB_show_guess,
                'ObsA_sees':ObsA_sees, 'ObsB_sees':ObsB_sees, 'ObsB_trial_type':ObsB_trial_type,
                'ObsB_time_to_decide':ObsB_time_to_decide}

    # live pages consists of 3 parts: py, once JS sends data, the other receives data
    # in the html file I'll have an input field for the form_fields.
    # <input type="hidden" id="money_after_trade_input" name="money_after_trade">
    # id is what will be used for JS, name has to be exactly the same as in python page file
    @staticmethod
    def live_method(player: Player, data):
        if player.id_in_group == 1:
            return {1: {"result":data["result"]}}

        if player.id_in_group == 2:
            return {2: {"result":data["result"]}}

        if player.id_in_group == 3:
            return {3: {"result":data["result"]}}

        if player.id_in_group == 4:
            return {4: {"result":data["result"]}}



class reward_phase(Page):
    template_name = 'SI_Exp2_NoPartner_TrainingSequence/reward_phase.html'

    def vars_for_template(self):
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])-1+self.id_in_group+3) + '.csv', Row)
        this_trials_demonstrator_data = demonstrator[self.round_number - 1]

        # This block gives which choice should be highlighted
        highlight_index_row = this_trials_demonstrator_data['demonstrator_choice_index_row']
        highlight_index_column = this_trials_demonstrator_data['demonstrator_choice_index_column']

        grid_data = []
        # as input take the whole reward matrix from the current trial
        # PROBLEM: Otree imports the demonstrator_reward matrix into a string (but it seems like I can't import it as a list/array)
        demonstrator_rewards_list = this_trials_demonstrator_data['demonstrator_noisy_rescaled_all_rewards'].replace("[", "").replace("]","")
        # int() converts to required integers
        res = [int(ele) for ele in demonstrator_rewards_list.split()]
        demonstrator_rewards_matrix = np.array(res)
        demonstrator_rewards_matrix = np.reshape(demonstrator_rewards_matrix, (11, 11))

        # this is relevant to identify whether reward is shown in grid (full-info trial) or just below (in case of reward-info trial) or not at all
        ObserverA_PageSequence = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])+1) + '.csv', Row) # THIS IS A PLACEHOLDER FOR NOW
        t = (self.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]

        # here I need to get the value that corresponds to the choice row and column index
        for row_idx, row in enumerate(demonstrator_rewards_matrix):
            grid_row = []
            for col_idx, value in enumerate(row):
                if int(ObsA_this_trials_data['trial_type']) != 2:
                    reward = demonstrator_rewards_matrix[highlight_index_row][highlight_index_column]
                else:
                    reward = ''
                cell = {'reward': reward, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)

        # What is shown for Obs A
        if int(ObsA_this_trials_data['trial_type']) == 2: # if choice only trial, show 'wait'
            ObsA_query_row1 = '?'
            ObsA_query_row2 = ''
            ObsA_reward ='?'
            ObsA_sees = False
        else:
            ObsA_query_row1 = 'See player\'s points'
            ObsA_query_row2 =  ''
            ObsA_sees = True
            # I have this in there 2x because I don't know any better. I need it within cell to be referenced in grid, and just the one value to be referenced below grid
            ObsA_reward = demonstrator_rewards_matrix[highlight_index_row][highlight_index_column]
        # What is shown for Obs B
        if int(ObsB_this_trials_data['trial_type']) == 2: # if choice only trial, show 'wait'
            ObsB_query_row1 = '?'
            ObsB_query_row2 = ''
            ObsB_sees = False
        else:
            ObsB_query_row1 = 'See player\'s points'
            ObsB_query_row2 =  ''
            ObsB_sees = True

        ObsB_trial_type = int(ObsB_this_trials_data['trial_type'])
        #print('int(ObsA_this_trials_data[trial_type])', int(ObsA_this_trials_data['trial_type']))

        return {'grid_data': grid_data, 'highlight_index_row': highlight_index_row, 'highlight_index_column': highlight_index_column,
                'ObsA_reward_below':ObsA_reward, 'trial_type':int(ObsA_this_trials_data['trial_type']),
                'ObsA_query_row1':ObsA_query_row1, 'ObsA_query_row2':ObsA_query_row2,
                'ObsB_query_row1':ObsB_query_row1, 'ObsB_query_row2':ObsB_query_row2,
                'ObsA_sees':ObsA_sees, 'ObsB_sees':ObsB_sees, 'ObsB_trial_type':ObsB_trial_type}

    timeout_seconds = 6

# CONTINUE WITH INFER CHOICE HTML
class infer_choice(Page):
    template_name = 'SI_Exp2_NoPartner_TrainingSequence/infer_choice.html'

    # Define form fields
    form_model = 'player'
    form_fields = ['inferred_choice', 'RT_inferred_choice']

    def is_displayed(player: Player):
        ObsA_trials = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Player)
        #ObsB_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])+1) + '.csv', Row)
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObsA_trials[t]
        #ObsB_this_trials_data = ObsB_trials[t]
        return (int(ObsA_this_trials_data['choice_inference_query']) == 1)

    def vars_for_template(self):
        # THIS WILL JUST GIVE THE GRID NUMBERS TO HTML
        grid_numbers = np.empty((11, 11))
        #grid_number = my_array.tolist()
        grid_data = []
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'value': value, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)

        # IF I WANT TO INCLUDE DEMONSTRATOR:
        # This gets the reward from the current trial and shows it on screen
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])-1+self.id_in_group+3) + '.csv', Row) #str(int(self.session.config['subject_number']))
        this_trials_demonstrator_data = demonstrator[self.round_number-1]

        ObserverA_PageSequence = read_csv('SI_Exp2_ObsB_TrainingSequence2/SI_Exp2_PageSequence_Training.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])+1) + '.csv', Row) # THIS IS A PLACEHOLDER FOR NOW
        t = (self.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]
        # What is shown for Obs A
        if int(ObsA_this_trials_data['choice_inference_query']) == 1:
            ObsA_query_row1 = 'Which cell do you think the player selected to receive ' + str(int(this_trials_demonstrator_data['demonstrator_noisy_rescaled_reward'])) + ' points?'
            ObsA_query_row2 = '(Please click into the respective cell and then submit)'
            show_form_field = 1
            demo_reward_for_html = {'value': this_trials_demonstrator_data['demonstrator_noisy_rescaled_reward']}
        else:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 =  ''
            show_form_field = 0
            demo_reward_for_html = {'value': '?'}

        # What is shown for Obs B
        if int(ObsB_this_trials_data['choice_inference_query']) == 1:
            if (int(ObsA_this_trials_data['choice_inference_query']) == 1) or (int(ObsA_this_trials_data['choice_inference_query']) == 3):
                ObsB_query_row1 = 'Which cell do you think the player selected to receive ' + str(int(this_trials_demonstrator_data['demonstrator_noisy_rescaled_reward'])) + ' points?'
            else:
                ObsB_query_row1 = 'Which cell do you think the player selected to receive ? points?'
            ObsB_query_row2 = '(Please click into the respective cell and then submit)'
            # The following is a placerholder because I take the real demonstrators value and just add some small noise
            ObsB_inf_highlight_index_row = np.clip(int(np.random.normal(this_trials_demonstrator_data['demonstrator_choice_index_row'], scale=1)), 0, 10)
            ObsB_inf_highlight_index_column = np.clip(int(np.random.normal(this_trials_demonstrator_data['demonstrator_choice_index_column'], scale=1)), 0, 10)
            ObsB_show_guess_text = 'Observer B\'s prediction is shown in purple.'
            ObsB_show_form_field = 1
            ObsB_time_to_decide = np.clip(np.random.lognormal(6000, 7000), 4000, 10000)
        else:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 =  ''
            ObsB_inf_highlight_index_row = ''
            ObsB_inf_highlight_index_column = ''
            ObsB_show_form_field = 0
            ObsB_show_guess_text = ''
            ObsB_time_to_decide = 0

            # Determine the background boxes. 1 full-inf trial, 2 choice only, 3 reward only
        # Should have background box if reward was presented as choice inference follows reward
        if int(ObsA_this_trials_data['trial_type']) == 2:
            ObsA_sees = False
        else:
            ObsA_sees = True

        if int(ObsB_this_trials_data['trial_type']) == 2:
            ObsB_sees = False
        else:
            ObsB_sees = True

        #return {'grid_number': grid_number, 'demo_reward_for_html':demo_reward_for_html, 'show_form_field':show_form_field,
        #        'ObsA_query_row1':ObsA_query_row1, 'ObsA_query_row2':ObsA_query_row2,
        #        'ObsB_query_row1':ObsB_query_row1, 'ObsB_query_row2':ObsB_query_row2}

        return {'grid_data': grid_data, 'demo_reward_for_html':demo_reward_for_html,
                'ObsB_inf_highlight_index_row': ObsB_inf_highlight_index_row, 'ObsB_inf_highlight_index_column': ObsB_inf_highlight_index_column,
                'ObsA_query_row1':ObsA_query_row1, 'ObsA_query_row2':ObsA_query_row2,
                'ObsB_query_row1':ObsB_query_row1, 'ObsB_query_row2':ObsB_query_row2,
                'show_form_field':show_form_field, 'ObsB_show_form_field':ObsB_show_form_field,
                'ObsB_show_guess_text': ObsB_show_guess_text,
                'ObsA_sees':ObsA_sees, 'ObsB_sees':ObsB_sees, 'ObsB_time_to_decide':ObsB_time_to_decide}

    @staticmethod
    def live_method(player: Player, data):
        if player.id_in_group == 1:
            if data["inferred_choice"] == 'NA,NA':
                return {1: {"inferred_choice": '', "inferred_choice_row": '', "inferred_choice_column": ''}}  #
            else:
                strings = data["inferred_choice"].split(',')
                #print('livesend success', data["inferred_choice"], type(data["inferred_choice"]))
                # for only Obs B query, [0] and [1] will return an N and an A
                return {1: {"inferred_choice":data["inferred_choice"], "inferred_choice_row":int(strings[0]), "inferred_choice_column":int(strings[1])}} #

        if player.id_in_group == 2:
            if data["inferred_choice"] == 'NA,NA':
                return {2: {"inferred_choice": '', "inferred_choice_row": '', "inferred_choice_column": ''}}  #
            else:
                strings = data["inferred_choice"].split(',')
                #print('livesend success', data["inferred_choice"], type(data["inferred_choice"]))
                # for only Obs B query, [0] and [1] will return an N and an A
                return {2: {"inferred_choice":data["inferred_choice"], "inferred_choice_row":int(strings[0]), "inferred_choice_column":int(strings[1])}} #

        if player.id_in_group == 3:
            if data["inferred_choice"] == 'NA,NA':
                return {3: {"inferred_choice": '', "inferred_choice_row": '', "inferred_choice_column": ''}}  #
            else:
                strings = data["inferred_choice"].split(',')
                return {3: {"inferred_choice":data["inferred_choice"], "inferred_choice_row":int(strings[0]), "inferred_choice_column":int(strings[1])}} #

        if player.id_in_group == 4:
            if data["inferred_choice"] == 'NA,NA':
                return {4: {"inferred_choice": '', "inferred_choice_row": '', "inferred_choice_column": ''}}  #
            else:
                strings = data["inferred_choice"].split(',')
                return {4: {"inferred_choice":data["inferred_choice"], "inferred_choice_row":int(strings[0]), "inferred_choice_column":int(strings[1])}} #


class finished(Page):
    template_name = 'SI_Exp2_NoPartner_TrainingSequence/finished.html'

    # Define form fields
    form_model = 'player'

    def is_displayed(player: Player):
        if player.round_number == 10:
            return True



page_sequence = [empty_grid, predict_choice, choice_phase, infer_reward, predict_reward, reward_phase, infer_choice, finished]