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


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_c_MainExp'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 40

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


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    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)
    ObsB_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_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()

    input_submitted = models.BooleanField(initial=False)


    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,
        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)
    RT_inferred_choice = models.FloatField(initial=None)
    RT_predicted_reward = models.FloatField(initial=None)
    RT_predicted_choice = models.FloatField(initial=None)
    ObsB_RT_predicted_choice = models.FloatField(initial=None)
    ObsB_RT_predicted_reward = 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_inferred_reward = models.IntegerField(field_maybe_none=True)
    ObsB_predicted_choice_col = models.IntegerField(field_maybe_none=True)
    ObsB_predicted_reward = models.IntegerField(field_maybe_none=True)
    ObsB_predicted_choice_row = 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_TwoParticipants_MainExp/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_TwoParticipants_MainExp/predict_choice.html'

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

    def is_displayed(player: Player):
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObsB_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.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_prediction_query']) == 1) or (int(ObsB_this_trials_data['choice_prediction_query']) == 1)


    def vars_for_template(player):
        #environment = player.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)

        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group

        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number']) - 1 + player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number']) - 1 + player2_id) + '.csv', Row)

        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(player.session.config['subject_number'])) + '.csv', Row)  # THIS ALSO IS A PLACEHOLDER
        t = (player.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)'
            ObsA_show_form_field = 1
        else:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 = ''
            ObsA_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_show_form_field = 1
            ObsB_show_guess_text = 'Observer B\'s prediction is shown in purple.'

        else:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 = ''
            ObsB_show_form_field = 0
            ObsB_show_guess_text = ''

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

        return {'grid_data': grid_data,
                '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_show_form_field': ObsA_show_form_field, 'ObsB_show_form_field': ObsB_show_form_field,
                'ObsB_show_guess_text': ObsB_show_guess_text,
                'player_id':player1_id}

    @staticmethod
    def live_method(player: Player, data):
        print('live send worked')
        print('player', player.id_in_group, data)
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.csv', Row)
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]

        all_submitted = False

        if data['predicted_choice']:
            player.input_submitted = True
        if int(ObsA_this_trials_data['choice_prediction_query']) == 1 and int(ObsB_this_trials_data['choice_prediction_query']) == 1:
            all_submitted = all(p.input_submitted for p in player.group.get_players())
            print('all_submitted', all_submitted)

            if player.id_in_group == 1:
                predicted_choice_player1 = data["predicted_choice"].split(',')
                print('predicted_choice_player1', predicted_choice_player1)
                return {0: {"all_submitted": all_submitted, "predicted_choice_row_player1":int(predicted_choice_player1[0]), "predicted_choice_column_player1":int(predicted_choice_player1[1])}}
            elif player.id_in_group == 2:
                predicted_choice_player2 = data["predicted_choice"].split(',')
                print('predicted_choice_player2', predicted_choice_player2)
                return {0: {"all_submitted": all_submitted, "predicted_choice_row_player2":int(predicted_choice_player2[0]), "predicted_choice_column_player2":int(predicted_choice_player2[1])}}
        elif int(ObsA_this_trials_data['choice_prediction_query']) == 1 and int(ObsB_this_trials_data['choice_prediction_query']) == 0:
            if player.id_in_group == 1:
                predicted_choice_player1 = data["predicted_choice"].split(',')
                return {2: {"all_submitted": all_submitted, "predicted_choice_row_player1":int(predicted_choice_player1[0]), "predicted_choice_column_player1":int(predicted_choice_player1[1])}}
            elif player.id_in_group == 2:
                predicted_choice_player2 = data["predicted_choice"].split(',')
                return {1: {"all_submitted": all_submitted, "predicted_choice_row_player2":int(predicted_choice_player2[0]), "predicted_choice_column_player2":int(predicted_choice_player2[1])}}


            #predicted_choice_playerX = data["predicted_choice"].split(',')
            #if predicted_choice_playerX[0] != 'NA':
            #    predicted_choice_playerX_0 = int(predicted_choice_playerX[0])
            #else:
            #    predicted_choice_playerX_0 = 'NA'

            #if predicted_choice_playerX[1] != 'NA':
            #    predicted_choice_playerX_1 = int(predicted_choice_playerX[1])
            #else:
            #    predicted_choice_playerX_1 = 'NA'
            #return {0: {"all_submitted": all_submitted, "predicted_choice_row_player1":predicted_choice_playerX_0, "predicted_choice_column_player1":predicted_choice_playerX_1, "predicted_choice_row_player2":predicted_choice_playerX_0, "predicted_choice_column_player2":predicted_choice_playerX_1, "show_ObsA_results":int(ObsA_this_trials_data['choice_prediction_query']), "show_ObsB_results":int(ObsB_this_trials_data['choice_prediction_query'])}}



        #if data["predicted_choice"] == 'NA,NA':
        #    return {0: {"predicted_choice": '', "predicted_choice_row": '', "predicted_choice_column": ''}}  #
        #else:
        #    strings = data["predicted_choice"].split(',')
        #    return {0: {"predicted_choice":data["predicted_choice"], "predicted_choice_row":int(strings[0]), "predicted_choice_column":int(strings[1])}}




    def before_next_page(player: Player, timeout_happened):
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(player.session.config['subject_number'])) + '.csv', Row)  # THIS ALSO IS A PLACEHOLDER
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        if ObsA_this_trials_data['choice_prediction_query'] == 1:
            this_trials_demonstrator_data = demonstrator[t]

            # next couple lines will be for calculating performance
            demonstrator_choice_probs_list = this_trials_demonstrator_data['demonstrator_choice_probs'].replace("[","").replace("]", "")
            res = [ele for ele in demonstrator_choice_probs_list.split()]
            df_demonstrator_choice_probs = pd.DataFrame(data=res, columns=['demonstrator_choice_probs'])
            # save the old index
            df_demonstrator_choice_probs['old_index'] = df_demonstrator_choice_probs.index
            # sort pandas dataframe and reset the index
            df_demonstrator_choice_probs_sorted = df_demonstrator_choice_probs.sort_values(by=['demonstrator_choice_probs']).reset_index()

            # identify the new index for demonstrator's choice
            new_demo_choice_index = df_demonstrator_choice_probs[df_demonstrator_choice_probs['old_index'] == int(
                this_trials_demonstrator_data['demonstrator_choice_index'])].index.tolist()  # need the to list as it otherwise saves a pandasInt64 and not an int. and that's also why I have to select [0] below, because it's a list

            # identify the new index for participant's choice
            # participant_predicted_choice =  # pretty sure this is 3,5 format
            row_index, column_index = player.in_round(player.round_number).predicted_choice.split(",")
            participant_predicted_choice = [int(row_index), int(column_index)]

            # again some 1- and 2-dimensional back and forth
            array_11x11 = np.zeros((11, 11))
            array_11x11[participant_predicted_choice[0], participant_predicted_choice[1]] = 1
            flattened_array = array_11x11.flatten()
            participant_predicted_choice_index = np.where(flattened_array == 1)[0][0]
            new_participant_choice_index = df_demonstrator_choice_probs_sorted[
                df_demonstrator_choice_probs_sorted['old_index'] == int(
                    participant_predicted_choice_index)].index.tolist()  # need the to list as it otherwise saves a pandasInt64 and not an int. and that's also why I have to select [0] below, because it's a list

            index_difference = abs(new_demo_choice_index[0] - new_participant_choice_index[0])
            money = (121 - int(index_difference)) * 0.001


            Row.create(SN=player.session.config['subject_number'],
                       # participant_label= player.in_round(player.round_number).participant.label,
                       participant_label=player.in_round(player.round_number).id_in_group,
                       participant_code=player.participant.code,
                       IV_number_choice_only_trials=ObsA_this_trials_data['IV_number_choice_only_trials'],
                       trial=(player.round_number - 1),
                       trial_type=ObsA_this_trials_data['trial_type'],
                       predicted_choice=player.in_round(player.round_number).predicted_choice,
                       RT_predicted_choice=player.in_round(player.round_number).RT_predicted_choice,
                       demonstrator_choice_prob=this_trials_demonstrator_data['demonstrator_choice_prob'],
                       demonstrator_choice_index=this_trials_demonstrator_data['demonstrator_choice_index'],
                       demonstrator_noisy_rescaled_reward=this_trials_demonstrator_data[
                           'demonstrator_noisy_rescaled_reward'],
                       demonstrator_choice_probs=this_trials_demonstrator_data['demonstrator_choice_probs'],
                       demonstrator_noisy_rescaled_all_rewards=this_trials_demonstrator_data[
                           'demonstrator_noisy_rescaled_all_rewards'],
                       money=money,)


'''
    def get_timeout_seconds(player):
        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])+1) + '.csv', Row) # THIS IS A PLACEHOLDER FOR NOW
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]
        print(int(ObsB_this_trials_data['choice_prediction_query']), type(int(ObsB_this_trials_data['choice_prediction_query'])))
        if int(ObsB_this_trials_data['choice_prediction_query']) == 1 and int(ObsA_this_trials_data['choice_prediction_query']) == 0:
            # hope it works this way.
            # PLACERHOLDER. ideally I read this in
            return 2.5
        else:
            return 3600 #this should be 60 minutes
'''



class choice_phase(Page):
    template_name = 'SI_Exp2_TwoParticipants_MainExp/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)

        player1_id = self.id_in_group
        other_player = self.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])) + '.csv', Row)
        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1+player2_id) + '.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_TwoParticipants_MainExp/infer_reward.html'

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

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

    def is_displayed(player: Player):
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObsB_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.csv', Row)
        t = (player.round_number - 2)

        ObsA_previous_trials_data = ObsA_trials[t]
        ObsB_previous_trials_data = ObsB_trials[t]
        # this should be displayed on the current trial, if the reward inference query on the previous trial was...
        if int(ObsA_previous_trials_data['reward_inference_query']) == 1 or int(ObsB_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)


        player1_id = self.id_in_group
        other_player = self.get_others_in_group()[0]
        player2_id = other_player.id_in_group

        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])) + '.csv', Row)
        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1+player2_id) + '.csv', Row)

        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'])
            ObsA_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'])
            ObsA_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
            ObsA_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'])
            ObsA_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
            ObsA_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_show_guess = 'Observer B inferred a reward of: <span id=result_inferred_reward ></span>'
            ObsB_show_form_field = 1
        else:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 =  ''
            ObsB_show_guess = ''
            ObsB_show_form_field = 0

            # When to show background boxes
        # don't show box if 3
        print("Trial type 3?", int(ObsA_current_trials_data['trial_type']))
        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,
                'ObsA_show_form_field':ObsA_show_form_field, 'ObsB_show_form_field':ObsB_show_form_field,
                'ObsA_show_guess':ObsA_show_guess, 'ObsB_show_guess':ObsB_show_guess,
                'ObsA_sees':ObsA_sees, 'ObsB_sees':ObsB_sees, 'player_id':player1_id}

    @staticmethod
    def live_method(player: Player, data):
        print('live send worked')
        print('player', player.id_in_group, data)
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.csv', Row)
        t = (player.round_number - 2)
        ObsA_prev_trials_data = ObserverA_PageSequence[t]
        ObsB_prev_trials_data = ObserverB_PageSequence[t]

        all_submitted = False

        if data['inferred_reward']:
            print(player.id_in_group, 'submitted')
            player.input_submitted = True
        if int(ObsA_prev_trials_data['reward_inference_query']) == 1 and int(ObsB_prev_trials_data['reward_inference_query']) == 1:
            all_submitted = all(p.input_submitted for p in player.group.get_players())
            print('all_submitted', all_submitted)

            if player.id_in_group == 1:
                inferred_reward_player1 = data['inferred_reward']
                return {0: {"all_submitted": all_submitted, "inferred_reward_player1": inferred_reward_player1}}
            elif player.id_in_group == 2:
                inferred_reward_player2 = data['inferred_reward']
                return {0: {"all_submitted": all_submitted, "inferred_reward_player2": inferred_reward_player2, "show_ObsA_results":int(ObsA_prev_trials_data['reward_inference_query']), "show_ObsB_results":int(ObsB_prev_trials_data['reward_inference_query'])}}
        else:
            return {0: {"all_submitted": all_submitted, "inferred_reward_player1": data['inferred_reward'], "inferred_reward_player2": data['inferred_reward'], "show_ObsA_results":int(ObsA_prev_trials_data['reward_inference_query']), "show_ObsB_results":int(ObsB_prev_trials_data['reward_inference_query'])}}


    def before_next_page(player: Player, timeout_happened):
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(player.session.config['subject_number'])) + '.csv', Row)  # THIS ALSO IS A PLACEHOLDER
        t = (player.round_number - 2)
        ObsA_this_trials_data = ObsA_trials[t]
        if ObsA_this_trials_data['reward_inference_query'] == 1:
            this_trials_demonstrator_data = demonstrator[t]
            # next couple lines will be for calculating performance
            demonstrator_rewards_list = this_trials_demonstrator_data['demonstrator_noisy_rescaled_all_rewards'].replace("[","").replace("]", "")
            res = [int(ele) for ele in demonstrator_rewards_list.split()]
            df_demonstrator_rewards = pd.DataFrame(data=res, columns=['demonstrator_noisy_rescaled_all_rewards'])

            # sort pandas dataframe but keep original index?
            df_demonstrator_rewards.sort_values(by=['demonstrator_noisy_rescaled_all_rewards'], inplace=True)
            demonstrator_rewards_real_list =  df_demonstrator_rewards['demonstrator_noisy_rescaled_all_rewards'].tolist()

            # now take participants reward prediction
            participant_inferred_reward = player.in_round(player.round_number).inferred_reward
            # and insert it at the right position into the list
            # but first to the participants advantage I'm going to delete all rewards but one with the same value
            unique_values = list(set(demonstrator_rewards_real_list))

            for element in range(len(unique_values)):
                # if smaller than first element
                if participant_inferred_reward < unique_values[0]:
                    unique_values.insert(0, participant_inferred_reward)
                # if greater than last element
                elif participant_inferred_reward > unique_values[-1]:
                    unique_values.insert(len(unique_values), participant_inferred_reward)
                elif unique_values[element] < participant_inferred_reward and unique_values[element + 1] > participant_inferred_reward:
                    unique_values.insert(element + 1, participant_inferred_reward)

            # find demonstrator reward index in the unique_values list:
            demonstrator_reward_index = unique_values.index(int(this_trials_demonstrator_data['demonstrator_noisy_rescaled_reward']))
            participant_inferred_reward_index = unique_values.index(participant_inferred_reward)

            # compare "how many indicies" those two are away from each other
            index_difference = abs(demonstrator_reward_index - participant_inferred_reward_index)
            money= (121-index_difference)*0.001

            Row.create(SN=player.session.config['subject_number'],
                       participant_label=player.in_round(player.round_number).id_in_group,
                       participant_code=player.participant.code,
                       IV_number_choice_only_trials=ObsA_this_trials_data['IV_number_choice_only_trials'],
                       trial=(player.in_round(player.round_number).round_number-1),
                       trial_type=ObsA_this_trials_data['trial_type'],
                       inferred_reward=player.in_round(player.round_number).inferred_reward,
                       RT_inferred_reward=player.in_round(player.round_number).RT_inferred_reward,
                       demonstrator_choice_prob = this_trials_demonstrator_data['demonstrator_choice_prob'],
                       demonstrator_choice_index = this_trials_demonstrator_data['demonstrator_choice_index'],
                       demonstrator_noisy_rescaled_reward = this_trials_demonstrator_data['demonstrator_noisy_rescaled_reward'],
                       demonstrator_choice_probs = this_trials_demonstrator_data['demonstrator_choice_probs'],
                       demonstrator_noisy_rescaled_all_rewards = this_trials_demonstrator_data['demonstrator_noisy_rescaled_all_rewards'],
                       ObsB_inferred_reward =player.in_round(player.round_number).field_maybe_none('ObsB_inferred_reward'),
                       money=money)

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

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

    def error_message(self, values):
        player1_id = self.id_in_group
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1+player1_id) + '.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):
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObsB_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.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']))
        #print('ObsB_this_trials_data', int(ObsB_this_trials_data['reward_prediction_query']))
        return (int(ObsA_this_trials_data['reward_prediction_query']) == 1) or (int(ObsB_this_trials_data['reward_prediction_query']) == 1)

    def vars_for_template(player):
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group

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

        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
        # 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_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.csv', Row)

        #ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])) + '.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 = (player.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[player.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 = ''
            ObsA_show_form_field = 1
            ObsA_show_guess = 'Observer A predicted a reward of: <span id=result_predicted_reward_player2 ></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 = ''
            ObsA_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
            ObsA_show_form_field = 1
            ObsA_show_guess = 'Observer A predicted a reward of: <span id=result_predicted_reward_player1 ></span>'
        else:
            ObsA_query_row1 = 'Wait'
            ObsA_query_row2 =  ''
            current_highlight_index_row = np.NaN
            current_highlight_index_column = np.NaN
            ObsA_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_form_field = 1
            ObsB_show_guess = 'Observer B 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(ObsB_this_trials_data['reward_prediction_query']) == 0 and int(ObsB_this_trials_data['trial_type']) != 3:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 = ''
            ObsB_show_form_field = 0
            ObsB_show_guess = ''
        # if trial type reward-only (i.e. 3), no choice on screen
        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 = ''
            current_highlight_index_row = np.NaN
            current_highlight_index_column = np.NaN
            ObsB_show_form_field = 1
            ObsB_show_guess = 'Observer B predicted a reward of: '
        else:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 =  ''
            current_highlight_index_row = np.NaN
            current_highlight_index_column = np.NaN
            ObsB_show_form_field = 0
            ObsB_show_guess = ''

        # 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,
                'ObsA_show_form_field':ObsA_show_form_field, 'ObsB_show_form_field':ObsB_show_form_field,
                '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,
                'player_id':player1_id
                }

    # 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):
        print('live send worked')
        print('player', player.id_in_group, data)
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.csv', Row)
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]

        all_submitted = False

        if data['predicted_reward']:
            player.input_submitted = True
        if int(ObsA_this_trials_data['reward_prediction_query']) == 1 and int(ObsB_this_trials_data['reward_prediction_query']) == 1:
            all_submitted = all(p.input_submitted for p in player.group.get_players())
            print('all_submitted', all_submitted)

            if player.id_in_group == 1:
                predicted_reward_player1 = data['predicted_reward']
                return {0: {"all_submitted": all_submitted, "predicted_reward_player1": predicted_reward_player1}}
            elif player.id_in_group == 2:
                predicted_reward_player2 = data['predicted_reward']
                return {0: {"all_submitted": all_submitted, "predicted_reward_player2": predicted_reward_player2, "show_ObsA_results":int(ObsA_this_trials_data['reward_prediction_query']), "show_ObsB_results":int(ObsB_this_trials_data['reward_prediction_query'])}}
        else:
            return {0: {"all_submitted": all_submitted, "predicted_reward_player1": data['predicted_reward'], "predicted_reward_player2": data['predicted_reward'], "show_ObsA_results":int(ObsA_this_trials_data['reward_prediction_query']), "show_ObsB_results":int(ObsB_this_trials_data['reward_prediction_query'])}}

            # solution below doesn't work because form field is only saved after the page has been submitted
            #for p in player.group.get_players():
            #    print('p in player', p)
            #    if p.id_in_group == 1:
            #        predicted_reward_player1 = p.predicted_reward
            #    elif p.id_in_group == 2:
            #        predicted_reward_player2 = p.predicted_reward
                #return {p.id_in_group: {'all_submitted': all_submitted} for p in group.get_players()}
        #    if all_submitted == True:
        #        return {0: {"all_submitted":all_submitted, "predicted_reward_player1":predicted_reward_player1, "predicted_reward_player2":predicted_reward_player2}}
        #elif int(ObsA_this_trials_data['reward_prediction_query']) == 1 and int(ObsB_this_trials_data['reward_prediction_query']) == 0:
        #    if player.id_in_group == 1:
        #        predicted_reward_player1 = data['predicted_reward']
        #    return {0: {"all_submitted":all_submitted, "predicted_reward_player1":predicted_reward_player1, "predicted_reward_player2":None}}
        #elif int(ObsA_this_trials_data['reward_prediction_query']) == 0 and int(ObsB_this_trials_data['reward_prediction_query']) == 1:
        #    if player.id_in_group == 2:
        #        predicted_reward_player2 = data['predicted_reward']
        #    return {0: {"all_submitted":all_submitted, "predicted_reward_player1":None, "predicted_reward_player2":predicted_reward_player2}}

    def before_next_page(player: Player, timeout_happened):
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(player.session.config['subject_number'])) + '.csv', Row)  # THIS ALSO IS A PLACEHOLDER
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObsA_trials[t]
        this_trials_demonstrator_data = demonstrator[t]
        if ObsA_this_trials_data['reward_prediction_query'] == 1:
                # next couple lines will be for calculating performance
                demonstrator_rewards_list = this_trials_demonstrator_data['demonstrator_noisy_rescaled_all_rewards'].replace("[","").replace("]", "")
                res = [int(ele) for ele in demonstrator_rewards_list.split()]
                df_demonstrator_rewards = pd.DataFrame(data=res, columns=['demonstrator_noisy_rescaled_all_rewards'])
                # sort pandas dataframe but keep original index?
                df_demonstrator_rewards.sort_values(by=['demonstrator_noisy_rescaled_all_rewards'], inplace=True)
                demonstrator_rewards_real_list =  df_demonstrator_rewards['demonstrator_noisy_rescaled_all_rewards'].tolist()

                # now take participants reward prediction
                participant_predicted_reward = player.in_round(player.round_number).predicted_reward
                # and insert it at the right position into the list
                # but first to the participants advantage I'm going to delete all rewards but one with the same value
                unique_values = list(set(demonstrator_rewards_real_list))

                for element in range(len(unique_values)):
                    # if smaller than first element
                    if participant_predicted_reward < unique_values[0]:
                        unique_values.insert(0, participant_predicted_reward)
                    # if greater than last element
                    elif participant_predicted_reward > unique_values[-1]:
                        unique_values.insert(len(unique_values), participant_predicted_reward)
                    elif unique_values[element] < participant_predicted_reward and unique_values[element + 1] > participant_predicted_reward:
                        unique_values.insert(element + 1, participant_predicted_reward)

                # find demonstrator reward index in the unique_values list:
                demonstrator_reward_index = unique_values.index(int(this_trials_demonstrator_data['demonstrator_noisy_rescaled_reward']))
                participant_predicted_reward_index = unique_values.index(participant_predicted_reward)

                # compare "how many indicies" those two are away from each other
                index_difference = abs(demonstrator_reward_index - participant_predicted_reward_index)
                money= (121-index_difference)*0.001


                Row.create(SN=player.session.config['subject_number'],
                           participant_label=player.in_round(player.round_number).id_in_group,
                           participant_code=player.participant.code,
                           IV_number_choice_only_trials=ObsA_this_trials_data['IV_number_choice_only_trials'],
                           trial=(player.round_number-1),
                           trial_type=ObsA_this_trials_data['trial_type'],
                           predicted_reward=player.in_round(player.round_number).predicted_reward,
                           RT_predicted_reward=player.in_round(player.round_number).RT_predicted_reward,
                           demonstrator_choice_prob = this_trials_demonstrator_data['demonstrator_choice_prob'],
                           demonstrator_choice_index = this_trials_demonstrator_data['demonstrator_choice_index'],
                           demonstrator_noisy_rescaled_reward = this_trials_demonstrator_data['demonstrator_noisy_rescaled_reward'],
                           demonstrator_choice_probs = this_trials_demonstrator_data['demonstrator_choice_probs'],
                           demonstrator_noisy_rescaled_all_rewards = this_trials_demonstrator_data['demonstrator_noisy_rescaled_all_rewards'],
                           money=money)

    #timeout_seconds = 2.5

class reward_phase(Page):
    template_name = 'SI_Exp2_TwoParticipants_MainExp/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'])) + '.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))

        player1_id = self.id_in_group
        other_player = self.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        # 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_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1 + player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(self.session.config['subject_number'])-1 + player2_id) + '.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_TwoParticipants_MainExp/infer_choice.html'

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

    def is_displayed(player: Player):
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObsB_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.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) or (int(ObsB_this_trials_data['choice_inference_query']) == 1)

    def vars_for_template(player):
        # 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

        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group

        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number']) - 1 + player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number']) - 1 + player2_id) + '.csv', Row)

        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(player.session.config['subject_number'])) + '.csv',Row)  # THIS ALSO IS A PLACEHOLDER
        t = (player.round_number - 1)
        this_trials_demonstrator_data = demonstrator[t]
        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_show_guess_text = 'Observer B\'s prediction is shown in purple.'
            ObsB_show_form_field = 1
        else:
            ObsB_query_row1 = 'Wait'
            ObsB_query_row2 = ''
            ObsB_show_form_field = 0
            ObsB_show_guess_text = ''

        # 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
        print('ObsA sees', ObsA_sees)
        print('ObsB sees', ObsB_sees)

        #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,
                '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_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,
                'player_id':player1_id}

    @staticmethod
    def live_method(player: Player, data):
        print('live send worked')
        print('player', player.id_in_group, data)
        player1_id = player.id_in_group
        other_player = player.get_others_in_group()[0]
        player2_id = other_player.id_in_group
        ObserverA_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        ObserverB_PageSequence = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player2_id) + '.csv', Row)
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObserverA_PageSequence[t]
        ObsB_this_trials_data = ObserverB_PageSequence[t]

        all_submitted = False

        if data['inferred_choice']:
            player.input_submitted = True
        if int(ObsA_this_trials_data['choice_inference_query']) == 1 and int(ObsB_this_trials_data['choice_inference_query']) == 1:
            all_submitted = all(p.input_submitted for p in player.group.get_players())
            print('all_submitted', all_submitted)

            if player.id_in_group == 1:
                inferred_choice_player1 = data["inferred_choice"].split(',')
                print('inferred_choice_player1', inferred_choice_player1)
                return {0: {"all_submitted": all_submitted, "inferred_choice_row_player1":int(inferred_choice_player1[0]), "inferred_choice_column_player1":int(inferred_choice_player1[1])}}
            elif player.id_in_group == 2:
                inferred_choice_player2 = data["inferred_choice"].split(',')
                print('inferred_choice_player2', inferred_choice_player2)
                return {0: {"all_submitted": all_submitted, "inferred_choice_row_player2":int(inferred_choice_player2[0]), "inferred_choice_column_player2":int(inferred_choice_player2[1])}}
        else:
            inferred_choice_playerX = data["inferred_choice"].split(',')
            # this seems to be necessary to avoid a ValueError
            if inferred_choice_playerX[0] != 'NA':
                inferred_choice_playerX_0 = int(inferred_choice_playerX[0])
            else:
                inferred_choice_playerX_0 = 'NA'

            if inferred_choice_playerX[1] != 'NA':
                inferred_choice_playerX_1 = int(inferred_choice_playerX[1])
            else:
                inferred_choice_playerX_1 = 'NA'
            return {0: {"all_submitted": all_submitted, "inferred_choice_row_player1":inferred_choice_playerX_0, "inferred_choice_column_player1":inferred_choice_playerX_1, "inferred_choice_row_player2":inferred_choice_playerX_0, "inferred_choice_column_player2":inferred_choice_playerX_1, "show_ObsA_results":int(ObsA_this_trials_data['choice_inference_query']), "show_ObsB_results":int(ObsB_this_trials_data['choice_inference_query'])}}

    def before_next_page(player: Player, timeout_happened):
        participant_code = player.participant.code
        player1_id = player.id_in_group
        # this allows to read in a "variable" file name so that each player gets their own page sequence
        ObsA_trials = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_PageSequence_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(player.session.config['subject_number'])-1+player1_id) + '.csv', Row)
        t = (player.round_number - 1)
        ObsA_this_trials_data = ObsA_trials[t]
        this_trials_demonstrator_data = demonstrator[t]
        if ObsA_this_trials_data['choice_inference_query'] == 1:

            # next couple lines will be for calculating performance
            demonstrator_choice_probs_list = this_trials_demonstrator_data['demonstrator_choice_probs'].replace("[","").replace("]", "")
            res = [ele for ele in demonstrator_choice_probs_list.split()]
            df_demonstrator_choice_probs = pd.DataFrame(data=res, columns=['demonstrator_choice_probs'])
            df_demonstrator_choice_probs['old_index'] = df_demonstrator_choice_probs.index
            df_demonstrator_choice_probs_sorted = df_demonstrator_choice_probs.sort_values(by=['demonstrator_choice_probs']).reset_index()
            new_demo_choice_index = df_demonstrator_choice_probs[df_demonstrator_choice_probs['old_index'] == int(this_trials_demonstrator_data['demonstrator_choice_index'])].index.tolist()
            row_index, column_index = player.in_round(player.round_number).inferred_choice.split(",")
            participant_inferred_choice = [int(row_index), int(column_index)]
            array_11x11 = np.zeros((11, 11))
            array_11x11[participant_inferred_choice[0], participant_inferred_choice[1]] = 1
            flattened_array = array_11x11.flatten()
            participant_predicted_choice_index = np.where(flattened_array == 1)[0][0]
            new_participant_choice_index = df_demonstrator_choice_probs_sorted[df_demonstrator_choice_probs_sorted['old_index'] == int(participant_predicted_choice_index)].index.tolist()
            index_difference = abs(new_demo_choice_index[0] - new_participant_choice_index[0])
            money = (121 - int(index_difference)) * 0.001

            if t == (player.round_number-1):
                Row.create(SN=player.session.config['subject_number'],
                           participant_label=player.in_round(player.round_number).id_in_group,
                           participant_code=player.participant.code,
                           IV_number_choice_only_trials=ObsA_this_trials_data['IV_number_choice_only_trials'],
                           trial=(player.round_number-1),
                           trial_type=ObsA_this_trials_data['trial_type'],
                           inferred_choice=player.in_round(player.round_number).inferred_choice,
                           RT_inferred_choice=player.in_round(player.round_number).RT_inferred_choice,
                           demonstrator_choice_prob = this_trials_demonstrator_data['demonstrator_choice_prob'],
                           demonstrator_choice_index = this_trials_demonstrator_data['demonstrator_choice_index'],
                           demonstrator_noisy_rescaled_reward = this_trials_demonstrator_data['demonstrator_noisy_rescaled_reward'],
                           demonstrator_choice_probs = this_trials_demonstrator_data['demonstrator_choice_probs'],
                           demonstrator_noisy_rescaled_all_rewards = this_trials_demonstrator_data['demonstrator_noisy_rescaled_all_rewards'],
                           ObsB_inferred_choice_row=player.in_round(player.round_number).field_maybe_none('ObsB_inf_highlight_index_row'),
                           ObsB_inferred_choice_col=player.in_round(player.round_number).field_maybe_none('ObsB_inf_highlight_index_column'),
                           money=money)


class MyWaitPage(WaitPage):
    def is_displayed(player: Player):
        return player.round_number==1

    title_text = "Please wait"
    #body_text = "Custom body text"


# These are calculations for the transfer phase, that are conducted on one trial that includes an empty_grid
def calcs_for_part3(player: Player):
    # Load the environment the participant's demonstrator experienced and determine which cells participants should be asked about in Part 3
    demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(player.session.config['subject_number'])-1+player.id_in_group) + '.csv', Row)
    environment_no = int(demonstrator[0]['environment'])
    environment = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/df_environment_' + str(environment_no) + '.csv', Row)

    rewards = []
    for cell in environment:
        reward_one_cell = cell['mean_reward']
        rewards.append(reward_one_cell)
    df = pd.DataFrame(rewards)

    # sort by first column, and I hope it keeps the index
    # if I don't save a copy of the df, inplace=True should keep the original index
    df.sort_values(by=0, inplace=True, ascending=False)

    start = 0
    end= 11

    list = []

    # there are 11 subgroups
    for subgroup in range(0,11):
        # from the sorted df with the original index, group into subframes of 11 and randomly select 3
        forced_choice = df.iloc[start:end].sample(n=3)

        list.append(forced_choice.index.tolist())

        start += 11
        end += 11

    order = random.sample(range(0, 11), 11)
    return list[order[0]], list[order[1]], list[order[2]], list[order[3]], list[order[4]], list[order[5]], list[order[6]], list[order[7]], list[order[8]], list[order[9]], list[order[10]], order

class Last(Page):
    template_name = 'SI_Exp2_AlgoPartner_MainExp/Last.html'

    form_model = 'player'
    form_fields = ['payout']

    # Is never displayed, just used for calculation on last page
    def before_next_page(player: Player, timeout_happened):

        if player.round_number == C.NUM_ROUNDS:
            payout = 0
            trials = Row.filter()

            for trial in trials:
                if int(trial.participant_label) == str(int(player.session.config['subject_number'])-1+player.id_in_group):
                    if trial.money == None:
                        money = 0
                    else:
                        money = trial.money
                    payout += money
            player.participant.vars['payout'] = round(payout, 2)
            player.performance = round(payout, 2)

        if player.counter == 0:
            forced_choice_0, forced_choice_1, forced_choice_2, forced_choice_3, forced_choice_4, forced_choice_5, forced_choice_6, forced_choice_7, forced_choice_8, forced_choice_9, forced_choice_10, order  = calcs_for_part3(player)
            player.participant.vars['forced_choice_0'] = forced_choice_0
            player.participant.vars['forced_choice_1'] = forced_choice_1
            player.participant.vars['forced_choice_2'] = forced_choice_2
            player.participant.vars['forced_choice_3'] = forced_choice_3
            player.participant.vars['forced_choice_4'] = forced_choice_4
            player.participant.vars['forced_choice_5'] = forced_choice_5
            player.participant.vars['forced_choice_6'] = forced_choice_6
            player.participant.vars['forced_choice_7'] = forced_choice_7
            player.participant.vars['forced_choice_8'] = forced_choice_8
            player.participant.vars['forced_choice_9'] = forced_choice_9
            player.participant.vars['forced_choice_10'] = forced_choice_10
            player.participant.vars['order'] = order

        # Access and increment the counter from the Player object
        player.counter = player.increment_counter()


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

    timeout_seconds = .0001



def custom_export(players):
    yield [
        "SN",
        "participant_label",
        "participant_code",
        "IV_number_choice_only_trials",
        "trial",
        "trial_type",
        "inferred_reward",
        "inferred_choice",
        "predicted_reward",
        "predicted_choice",
        "RT_inferred_reward",
        "RT_inferred_choice",
        "RT_predicted_reward",
        "RT_predicted_choice",
        "demonstrator_choice_prob",
        "demonstrator_choice_index",
        "demonstrator_noisy_rescaled_reward",
        "demonstrator_choice_probs",
        "demonstrator_noisy_rescaled_all_rewards",
        "ObsB_predicted_reward",
        "ObsB_inferred_reward",
        "ObsB_highlight_index_row",
        "ObsB_highlight_index_column",
        "ObsB_inf_highlight_index_row",
        "ObsB_inf_highlight_index_col",
        "money"
    ]

    trials = Row.filter()
    #print('trials', trials)
    for trial in trials:
        #print("TRIAL", trial.trial)
        #print('trial.money', trial.money)
        yield [trial.SN, trial.participant_label, trial.participant_code, trial.IV_number_choice_only_trials, trial.trial, trial.trial_type,
               trial.inferred_reward, trial.inferred_choice, trial.predicted_reward, trial.predicted_choice,
               trial.RT_inferred_reward, trial.RT_inferred_choice, trial.RT_predicted_reward, trial.RT_predicted_choice,
               trial.demonstrator_choice_prob, trial.demonstrator_choice_index, trial.demonstrator_noisy_rescaled_reward, trial.demonstrator_choice_probs,
               trial.demonstrator_noisy_rescaled_all_rewards,trial.ObsB_predicted_reward, trial.ObsB_inferred_reward, trial.ObsB_predicted_choice_row, trial.ObsB_predicted_choice_col,
               trial.ObsB_inferred_choice_row, trial.ObsB_inferred_choice_col, trial.money]

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