import csv
import numpy as np
import random
from otree.api import *


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

# NOTES ----------------------------------------------------------------------------------------------------------------
# Training uses demonstrator with SN+2 to not have same demonstrator in training and experiment
# at the moment, the inference/prediction is saved ON the trial it was asked, not the trial it was asked FOR/about



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


class Subsession(BaseSubsession):
    pass

class Group(BaseGroup):
    pass

class Player(BasePlayer):
    One_comprehension_check_task = models.IntegerField(label="What does Observer B do?", widget=widgets.RadioSelect, choices=[[1, 'a) Observer B observers the player playing a game.'], [2, 'b) Observer B is the player.'],] )
    One_result = models.IntegerField()
    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,
        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,
        doc="""This player's decision""")

    predicted_reward_2 = models.IntegerField(
        min=0,
        max=500,
        label='What reward will the player get on the next trial?',
        blank=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)

    round_data_field = 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
        #self.save()



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

    #page_sequence = models.StringField()
    SN = models.IntegerField()
    participant_label = models.IntegerField()
    IV_number_choice_only_trials = models.IntegerField()
    trial_index = models.IntegerField()

    trial = models.IntegerField()
    trial_type = models.StringField()
    #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,
        initial=None,
        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,
        initial=None,
        doc="""This player's decision""")

    predicted_reward_2 = models.IntegerField(
        min=0,
        max=500,
        label='What reward will the player get on the next trial?',
        blank=True,
        initial=None,
        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(blank=True)
    choice_inference_query = models.BooleanField(blank=True)
    reward_prediction_query = models.BooleanField(blank=True)
    choice_prediction_query = models.BooleanField(blank=True)

    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_rewards = models.StringField()

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

    # 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_noisy_rescaled_all_rewards = models.LongStringField()
    demonstrator_true_all_rewards = models.LongStringField()
    demonstrator_noisy_reward = models.FloatField()



# PAGES ----------------------------------------------------------------------------------------------------------------
# Pages shows the html pages and corresponds to each html file within the folder (except the wait page, for that no specific html file is needed)
class AWelcomePage(Page):
    pass

class BConsentPage(Page):
    pass

class CInstructionsStartPage(Page):
    def is_displayed(player: Player):
        if player.round_number == 1:
            return True
        else:
            return player.participant.vars["ComprehensionCheckFailed"]

class TrainingAnnouncementPage(Page):
    pass

class DGrid(Page):
    template_name = 'SI_Exp2_Training/DGrid.html'
    #counter = 0
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}

class EGrid2(Page):
    template_name = 'SI_Exp2_Training/EGrid2.html'
    #counter = 0
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}

class HGrid3(Page):
    template_name = 'SI_Exp2_Training/HGrid3.html'
    #counter = 0
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}

class LGrid4(Page):
    template_name = 'SI_Exp2_Training/LGrid4.html'
    #counter = 0
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}

class OGrid5(Page):
    template_name = 'SI_Exp2_Training/OGrid5.html'
    pass

class RGrid6(Page):
    template_name = 'SI_Exp2_Training/RGrid6.html'
    pass

class TGrid7(Page):
    template_name = 'SI_Exp2_Training/TGrid7.html'
    pass

class WGrid8(Page):
    template_name = 'SI_Exp2_Training/WGrid8.html'
    pass

class ZAGrid9(Page):
    template_name = 'SI_Exp2_Training/ZAGrid9.html'
    pass

class ZDNextRound(Page):
    template_name = 'SI_Exp2_Training/ZDNextRound.html'
    timeout_seconds = 4

class FChoicePhase(Page):
    template_name = 'SI_Exp2_Training/FChoicePhase.html'

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

        ObsA_query_row1 = 'See player\'s selection'
        ObsA_query_row2 = ''
        highlight_index_row = 2
        highlight_index_column = 4
        ObsA_sees = True

        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,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6

class IChoicePhase2(Page):
    template_name = 'SI_Exp2_Training/IChoicePhase2.html'

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

        ObsA_query_row1 = 'See player\'s selection'
        ObsA_query_row2 = ''
        highlight_index_row = 2
        highlight_index_column = 4
        ObsA_sees = True

        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,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6

class MChoicePhase4(Page):
    template_name = 'SI_Exp2_Training/MChoicePhase4.html'

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

        ObsA_query_row1 = 'See player\'s selection'
        ObsA_query_row2 = ''
        highlight_index_row = 5
        highlight_index_column = 8
        ObsA_sees = True

        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,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6

class PChoicePhase5(Page):
    template_name = 'SI_Exp2_Training/PChoicePhase5.html'

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

        ObsA_query_row1 = '?'
        ObsA_query_row2 = ''
        highlight_index_row = np.NaN
        highlight_index_column = np.NaN
        ObsA_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,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6


class UChoicePhase7(Page):
    template_name = 'SI_Exp2_Training/UChoicePhase7.html'

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

        ObsA_query_row1 = 'See player\'s selection'
        ObsA_query_row2 = ''
        highlight_index_row = 9
        highlight_index_column = 10
        ObsA_sees = True

        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,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6

class GRewardPhase(Page):
    template_name = 'SI_Exp2_Training/GRewardPhase.html'

    def vars_for_template(self):
        # This block gives which choice should be highlighted
        highlight_index_row = 2
        highlight_index_column = 4

        grid_data = []

        grid_numbers = np.empty((11, 11))
        # here I need to get the value that corresponds to the choice row and column index
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'reward': 21, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)


        ObsA_query_row1 = 'See player\'s points'
        ObsA_query_row2 = ''
        ObsA_sees = True
        ObsA_reward = 21

        return {'grid_data': grid_data, 'highlight_index_row': highlight_index_row,
                'highlight_index_column': highlight_index_column,
                'ObsA_reward_below': ObsA_reward,
                'ObsA_query_row1': ObsA_query_row1, 'ObsA_query_row2': ObsA_query_row2,
                'ObsA_sees': ObsA_sees, 'nonsense':True}

    timeout_seconds = 6

class JRewardPhase2(Page):
    template_name = 'SI_Exp2_Training/JRewardPhase2.html'

    def vars_for_template(self):
        # This block gives which choice should be highlighted
        highlight_index_row = 2
        highlight_index_column = 4

        grid_data = []

        grid_numbers = np.empty((11, 11))
        # here I need to get the value that corresponds to the choice row and column index
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'reward': 19, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)


        ObsA_query_row1 = 'See player\'s points'
        ObsA_query_row2 = ''
        ObsA_sees = True
        ObsA_reward = 19

        return {'grid_data': grid_data, 'highlight_index_row': highlight_index_row,
                'highlight_index_column': highlight_index_column,
                'ObsA_reward_below': ObsA_reward,
                'ObsA_query_row1': ObsA_query_row1, 'ObsA_query_row2': ObsA_query_row2,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6

class NRewardPhase4(Page):
    template_name = 'SI_Exp2_Training/NRewardPhase4.html'

    def vars_for_template(self):
        # This block gives which choice should be highlighted
        highlight_index_row = 5
        highlight_index_column = 8

        grid_data = []

        grid_numbers = np.empty((11, 11))
        # here I need to get the value that corresponds to the choice row and column index
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'reward': '?', 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)


        ObsA_query_row1 = '?'
        ObsA_query_row2 = ''
        ObsA_sees = False
        ObsA_reward = '?'

        return {'grid_data': grid_data, 'highlight_index_row': highlight_index_row,
                'highlight_index_column': highlight_index_column,
                'ObsA_reward_below': ObsA_reward,
                'ObsA_query_row1': ObsA_query_row1, 'ObsA_query_row2': ObsA_query_row2,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6


class QRewardPhase5(Page):
    template_name = 'SI_Exp2_Training/QRewardPhase5.html'

    def vars_for_template(self):
        # This block gives which choice should be highlighted
        highlight_index_row = 5
        highlight_index_column = 8

        grid_data = []

        grid_numbers = np.empty((11, 11))
        # here I need to get the value that corresponds to the choice row and column index
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'reward': 40, 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)


        ObsA_query_row1 = 'See player\'s points'
        ObsA_query_row2 = ''
        ObsA_sees = True
        ObsA_reward = 40

        return {'grid_data': grid_data, 'highlight_index_row': np.NaN,
                'highlight_index_column': np.NaN,
                'ObsA_reward_below': ObsA_reward,
                'ObsA_query_row1': ObsA_query_row1, 'ObsA_query_row2': ObsA_query_row2,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6


class ZCRewardPhase9(Page):
    template_name = 'SI_Exp2_Training/ZCRewardPhase9.html'

    def vars_for_template(self):
        # This block gives which choice should be highlighted
        highlight_index_row = 2
        highlight_index_column = 4

        grid_data = []

        grid_numbers = np.empty((11, 11))
        # here I need to get the value that corresponds to the choice row and column index
        for row_idx, row in enumerate(grid_numbers):
            grid_row = []
            for col_idx, value in enumerate(row):
                cell = {'reward': '?', 'row_index': row_idx, 'column_index': col_idx}
                grid_row.append(cell)
            grid_data.append(grid_row)


        ObsA_query_row1 = '?'
        ObsA_query_row2 = ''
        ObsA_sees = False
        ObsA_reward = '?'

        return {'grid_data': grid_data, 'highlight_index_row': highlight_index_row,
                'highlight_index_column': highlight_index_column,
                'ObsA_reward_below': ObsA_reward,
                'ObsA_query_row1': ObsA_query_row1, 'ObsA_query_row2': ObsA_query_row2,
                'ObsA_sees': ObsA_sees}

    timeout_seconds = 6

class KExplainRewardCorrelation(Page):
    template_name = 'SI_Exp2_Training/KExplainRewardCorrelation.html'

    def vars_for_template(self):

        demonstrator = load_environment(self)
        this_trials_demonstrator_data = demonstrator[self.counter-1]
        grid_data = []
        demonstrator_rewards_list = this_trials_demonstrator_data['demonstrator_noisy_rescaled_all_rewards'].replace("[", "").replace("]","") # delete the brackets. need to ensure can re-create matrix list structure later
        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))

        for row_idx, row in enumerate(demonstrator_rewards_matrix):
            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)

        return {'grid_data': grid_data}


class QueryTypes_FullInfo(Page):
    template_name = 'SI_Exp2_Training/QueryTypes_FullInfo.html'
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}

class QueryTypes_MissingChoice(Page):
    template_name = 'SI_Exp2_Training/QueryTypes_MissingChoice.html'
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}


class QueryTypes_MissingReward(Page):
    template_name = 'SI_Exp2_Training/QueryTypes_MissingReward.html'
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}

class QueryTypes_PredictChoices(Page):
    template_name = 'SI_Exp2_Training/QueryTypes_PredictChoices.html'
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}


class QueryTypes_PredictReward(Page):
    template_name = 'SI_Exp2_Training/QueryTypes_PredictReward.html'
    def vars_for_template(self):
        my_array = np.empty((11, 11))
        grid_number = my_array.tolist()
        return {'grid_number': grid_number}

class QueryTypes_InferChoices(Page):
    template_name = 'SI_Exp2_Training/QueryTypes_InferChoices.html'



class QueryTypes_InferRewards(Page):
    template_name = 'SI_Exp2_Training/QueryTypes_InferRewards.html'


class SChoicePredictionQuery(Page):
    template_name = 'SI_Exp2_Training/SChoicePredictionQuery.html'

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

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

        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

        return {'grid_data': grid_data,
                'ObsA_query_row1':ObsA_query_row1, 'ObsA_query_row2':ObsA_query_row2,
                'show_form_field':show_form_field}


class VRewardPredictionQuery(Page):
    template_name = 'SI_Exp2_Training/VRewardPredictionQuery.html'

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

    def error_message(self, values):
        if values["predicted_reward"] is None or values["predicted_reward"] == '':
            return 'You must enter a value'

    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)


        current_highlight_index_row = 9
        current_highlight_index_column = 10

        ObsA_query_row1 = 'How many points do you think the player will receive for this selection?'
        ObsA_query_row2 = ''
        show_form_field = 1
        ObsA_sees = True

        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,
                'show_form_field':show_form_field,
                'ObsA_sees':ObsA_sees}

    @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 XInferChoice(Page):
    template_name = 'SI_Exp2_Training/XInferChoice.html'

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

    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)


        ObsA_query_row1 = 'Which cell do you think the player selected to receive ' + str(40) + ' points?'
        ObsA_query_row2 = '(Please click into the respective cell and then submit)'
        show_form_field = 1
        demo_reward_for_html = {'value': str(40)}

        ObsA_sees = True

        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,
                'show_form_field':show_form_field, 'ObsA_sees':ObsA_sees}

    @staticmethod
    def live_method(player: Player, data):
        if player.id_in_group == 1:
            strings = data["inferred_choice"].split(',')
            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:
            strings = data["inferred_choice"].split(',')
            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:
            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:
            strings = data["inferred_choice"].split(',')
            return {4: {"inferred_choice": data["inferred_choice"], "inferred_choice_row": int(strings[0]), "inferred_choice_column": int(strings[1])}}

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

class ZBInferReward(Page):
    template_name = 'SI_Exp2_Training/ZBInferReward.html'

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

    def error_message(self, values):
        if values["inferred_reward"] is None or values["inferred_reward"] == '':
            return 'You must enter a value'


    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)

        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)

        highlight_index_row = 2
        highlight_index_column = 4
        current_highlight_index_row = 5
        current_highlight_index_column = 8
        show_form_field = 1
        ObsA_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,
                'show_form_field':show_form_field,
                'ObsA_show_guess':ObsA_show_guess,
                'ObsA_sees':ObsA_sees}

    @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"]}}




def load_environment(self):
    # ensure that demonstrator+environment during training is not the same as during the actual experiment --> always take next participant's environment and demonstrator
    if self.id_in_group < 50:
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(int(self.session.config['subject_number'])-1+self.id_in_group+2) + '.csv', Row)
    # need special solution for last participant as they couldn't refer to the next participant's demonstrator
    elif self.id_in_group == 50:
        demonstrator = read_csv('SI_Exp2_AlgoPartner_MainExp/01_ParticipantFiles/SI_Demonstrator_' + str(1) + '.csv', Row)
    else:
        print('error in finding right environment')
    return demonstrator



class ComprehensionCheck(Page):
    form_model = "player"
    form_fields = ["Two_comprehension_check_spatial", "Three_comprehension_check_inference"]

    def is_displayed(player: Player):
        if player.round_number == 1:
            return True
        else:
            return player.participant.vars["ComprehensionCheckFailed"]

    def before_next_page(player: Player, timeout_happened):
        if player.Two_comprehension_check_spatial == 1 and player.Three_comprehension_check_inference == 2:
            player.participant.vars["ComprehensionCheckFailed"] = False
        else:
            player.participant.vars["ComprehensionCheckFailed"] = True
        #else:
        #    if player.One_comprehension_check_task == 1 and player.Two_comprehension_check_spatial == 1 and player.Three_comprehension_check_inference == 2:
        #        player.participant.vars["ComprehensionCheckFailed"] = False
        #   else:
        #        player.participant.vars["ComprehensionCheckFailed"] = True

class ComprehensionCheckFailed(Page):
    def is_displayed(player: Player):
        return player.participant.vars["ComprehensionCheckFailed"]

    timeout_seconds = 5


class ComprehensionCheckPassed(Page):
    def is_displayed(player: Player):
        return player.participant.vars["ComprehensionCheckFailed"] == False

    def vars_for_template(player: Player):
        session_name = player.session.config['name']
        if session_name == 'SI_Exp2_a':
            end_of_app_text = 'Next, you will go through a short trial sequence so that you can get used to the pacing of the task.'
        else:
            end_of_app_text = 'You answered the questions correctly. Let\'s move on!'
        return {'end_of_app_text':end_of_app_text}

    # need to implement that you don't have to click the next button 30 or so times
    def app_after_this_page(player: Player, upcoming_apps):
        print(upcoming_apps)
        if player.participant.vars["ComprehensionCheckFailed"] == False:
            return upcoming_apps[0]


# EXECUTE -----------------------------------------------------------------------------------------------------

page_sequence = [
                AWelcomePage, BConsentPage, CInstructionsStartPage,
                DGrid, EGrid2, FChoicePhase, GRewardPhase, # Sequence
                HGrid3, ZDNextRound, IChoicePhase2, JRewardPhase2, #different reward
                KExplainRewardCorrelation, # full grid
                LGrid4, MChoicePhase4, NRewardPhase4, # missing reward
                OGrid5, PChoicePhase5, QRewardPhase5, # missing choice
                RGrid6, SChoicePredictionQuery, # choice prediction query
                TGrid7, UChoicePhase7, VRewardPredictionQuery, # reward prediction query
                WGrid8, PChoicePhase5, QRewardPhase5, XInferChoice, # choice inference query
                ZAGrid9, MChoicePhase4, NRewardPhase4, ZDNextRound, FChoicePhase,ZBInferReward,
                ComprehensionCheck, ComprehensionCheckFailed, ComprehensionCheckPassed,]

# Obs B in separate App! And training also in separate apps...