from otree.api import *
from otree.database import dbq
import numpy as np
import random

class C(BaseConstants):
    NAME_IN_URL = 'pre_survey'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1
    
    ENDOWMENT = cu(20)

    # Epper Allocation Setup (keeping your logic)
    Epper_alloc_03_me = np.ndarray.tolist(np.flip(np.round(np.linspace(550, 950, num=7) * (-1))).astype(int))
    Epper_alloc_03_other = np.ndarray.tolist(np.flip(np.round(np.linspace(850, 650, num=7))).astype(int))
    Epper_alloc_07_me = np.ndarray.tolist(np.flip(np.round(np.linspace(750, 750, num=7) * (-1))).astype(int))
    Epper_alloc_07_other = np.ndarray.tolist(np.flip(np.round(np.linspace(1050, 450, num=7))).astype(int))
    Epper_alloc_08_me = np.ndarray.tolist(np.flip(np.round(np.linspace(700, 800, num=7) * (-1))).astype(int))
    Epper_alloc_08_other = np.ndarray.tolist(np.flip(np.round(np.linspace(500, 1000, num=7))).astype(int))
    Epper_alloc_09_me = np.ndarray.tolist(np.flip(np.round(np.linspace(650, 850, num=7) * (-1))).astype(int))
    Epper_alloc_09_other = np.ndarray.tolist(np.flip(np.round(np.linspace(550, 950, num=7))).astype(int))
    Epper_alloc_13_me = np.ndarray.tolist(np.flip(np.round(np.linspace(700, 767, num=7) * (-1))).astype(int))
    Epper_alloc_13_other = np.ndarray.tolist(np.flip(np.round(np.linspace(800, 1133, num=7))).astype(int))

    LIKERT_CHOICES = [
        ('option1', 'Totally disagree'),
        ('option2', 'Rather disagree'),
        ('option3', 'Neither agree or disagree'),
        ('option4', 'Rather agree'),
        ('option5', 'Totally agree')
    ]

    # BRET Constants
    box_value = cu(30)
    num_rows = 10
    num_cols = 10
    box_height = '50px'
    box_width = '50px'
    random_payoff = True
    instructions = True
    feedback = True
    results = True
    dynamic = True
    time_interval = 0.50
    random = False
    devils_game = True
    undoable = True

    # --- ADD THESE MISSING LINES ---
    PD_PAYOFF_a = cu(750)
    PD_PAYOFF_b = cu(600)
    PD_PAYOFF_c = cu(500)
    PD_PAYOFF_d = cu(1000)
    # -------------------------------

class Subsession(BaseSubsession):
    pass

class Group(BaseGroup):
    pass

# Custom Model for Database Storage
class PreSurveyProfile(ExtraModel):
    user_id = models.StringField()
    
    # Allow null/blank for fields that might be skipped or fail to save
    cooperate_PD = models.BooleanField(blank=True, null=True)
    cooperate_Cond_PD_C = models.BooleanField(blank=True, null=True)
    cooperate_Cond_PD_D = models.BooleanField(blank=True, null=True)
    
    boxes_collected = models.IntegerField(blank=True, null=True)
    bomb = models.IntegerField(blank=True, null=True)
    
    epper_scenario_selected = models.StringField(blank=True, null=True)
    epper_choice_raw = models.IntegerField(blank=True, null=True)
    epper_payoff_me = models.IntegerField(blank=True, null=True)
    epper_payoff_other = models.IntegerField(blank=True, null=True)
    
    finished_part_1 = models.BooleanField(initial=True)

class Player(BasePlayer):
    id_created = models.StringField(label="Please enter your identification code here: ")
    task_counter = models.IntegerField(initial=0)
    
    # BRET
    # These are populated by JS, so we keep them as mandatory to ensure the task runs.
    bomb = models.IntegerField(initial=0)
    bomb_row = models.PositiveIntegerField()
    bomb_col = models.PositiveIntegerField()
    boxes_collected = models.IntegerField(initial=0)
    pay_this_round = models.BooleanField()
    round_result = models.CurrencyField()

    # PD - Mandatory
    cooperate_PD = models.BooleanField(
        choices=[[True, 'A'], [False, 'B']],
        widget=widgets.RadioSelect,
        label="1) Which action do you choose if you don't know the choice of the other player?"
    )
    cooperate_Cond_PD_C = models.BooleanField(
        choices=[[True, 'A'], [False, 'B']],
        widget=widgets.RadioSelect,
        label="2) Imagine the other player chose A. Which action do you want to choose in that case:"
    )
    cooperate_Cond_PD_D = models.BooleanField(
        choices=[[True, 'A'], [False, 'B']],
        widget=widgets.RadioSelect,
        label="3) Imagine the other player chose B. Which action do you want to choose in that case"
    )

    # Survey vars - Mandatory
    ambiguity_1 = models.StringField(choices=C.LIKERT_CHOICES, widget=widgets.RadioSelectHorizontal, label="1. People who insist on a yes or no answer don’t know how complicated things are.")
    ambiguity_2 = models.StringField(choices=C.LIKERT_CHOICES, widget=widgets.RadioSelectHorizontal, label="2. Many of our most important decisions are based on insufficient information.")
    ambiguity_3 = models.StringField(choices=C.LIKERT_CHOICES, widget=widgets.RadioSelectHorizontal, label="3. An expert who is unable to give a clear answer does not know much about the problem.")
    ambiguity_4 = models.StringField(choices=C.LIKERT_CHOICES, widget=widgets.RadioSelectHorizontal, label="4. Teachers who assign vague tasks give students a chance to show initiative and originality.")
    consequentialist_scale_1 = models.StringField(choices=C.LIKERT_CHOICES, widget=widgets.RadioSelectHorizontal, label="1. Rules and laws should only be followed when they maximize happiness.")
    consequentialist_scale_2 = models.StringField(choices=C.LIKERT_CHOICES, widget=widgets.RadioSelectHorizontal, label="2. When deciding what action to take, the only relevant factor to consider would be the outcome of the action.")
    consequentialist_scale_3 = models.StringField(choices=C.LIKERT_CHOICES, widget=widgets.RadioSelectHorizontal, label="3. Some rules should never be broken.")
    consequentialist_scale_4 = models.StringField(choices=C.LIKERT_CHOICES, widget=widgets.RadioSelectHorizontal, label="4. It is never morally justified to cause someone harm.")

    Epper_allocation_01 = models.IntegerField(choices=[1, 2, 3, 4, 5, 6, 7], label="")
    Epper_allocation_02 = models.IntegerField(choices=[1, 2, 3, 4, 5, 6, 7], label="")
    Epper_allocation_03 = models.IntegerField(choices=[1, 2, 3, 4, 5, 6, 7], label="")
    Epper_allocation_04 = models.IntegerField(choices=[1, 2, 3, 4, 5, 6, 7], label="")
    Epper_allocation_05 = models.IntegerField(choices=[1, 2, 3, 4, 5, 6, 7], label="")



def creating_session(subsession: Subsession):
    if subsession.round_number == 1:
        for p in subsession.get_players():
            # Define tasks: 1=PD, 2=Survey, 3=Epper, 4=BRET
            tasks = [1, 2, 3, 4]
            random.shuffle(tasks)
            p.participant.vars['task_order'] = tasks
            print(f"Player {p.id_in_group} task order: {tasks}")


# =============================================================================
# PAGES
# =============================================================================

class Introduction(Page):
    pass

class ID_Creation(Page):
    form_model = 'player'
    form_fields = ['id_created']
    
    @staticmethod
    def error_message(player, values):
        # 1. Normalize Input
        input_id = values['id_created'].strip()
        
        if not input_id:
            return "ID cannot be empty."

        # 2. Use dbq for global query
        # This will search the ENTIRE PreSurveyProfile table across all sessions
        existing_profile = dbq(PreSurveyProfile).filter(PreSurveyProfile.user_id == input_id).first()
        
        if existing_profile:
            return f"The ID '{input_id}' is already registered. Please use a different ID."


# --- Task 1: PD ---
class PD(Page):
    form_model = 'player'
    form_fields = ['cooperate_PD', 'cooperate_Cond_PD_C', 'cooperate_Cond_PD_D']
    
    @staticmethod
    def is_displayed(player: Player):
        # Show if current task in the sequence is 1
        # task_counter is 0-indexed (0, 1, 2, 3)
        if player.task_counter < 4:
            return player.participant.vars['task_order'][player.task_counter] == 1
        return False

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.task_counter += 1

# --- Task 2: Survey ---
class Survey(Page):
    form_model = 'player'
    form_fields = ['ambiguity_1', 'ambiguity_2', 'ambiguity_3', 'ambiguity_4',
                   'consequentialist_scale_1', 'consequentialist_scale_2', 'consequentialist_scale_3', 'consequentialist_scale_4']
    
    @staticmethod
    def is_displayed(player: Player):
        if player.task_counter < 4:
            return player.participant.vars['task_order'][player.task_counter] == 2
        return False

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.task_counter += 1

# --- Task 3: Epper ---
# Note: Epper has multiple pages. We only increment counter on the LAST page.
def is_epper(player):
    if player.task_counter < 4:
        return player.participant.vars['task_order'][player.task_counter] == 3
    return False

class epper_instructions(Page):
    @staticmethod
    def is_displayed(player: Player):
        return is_epper(player)

class epper_task_1(Page):
    form_model = 'player'
    form_fields = ['Epper_allocation_01']
    @staticmethod
    def is_displayed(player: Player):
        return is_epper(player)

class epper_task_2(Page):
    form_model = 'player'
    form_fields = ['Epper_allocation_02']
    @staticmethod
    def is_displayed(player: Player):
        return is_epper(player)

class epper_task_3(Page):
    form_model = 'player'
    form_fields = ['Epper_allocation_03']
    @staticmethod
    def is_displayed(player: Player):
        return is_epper(player)

class epper_task_4(Page):
    form_model = 'player'
    form_fields = ['Epper_allocation_04']
    @staticmethod
    def is_displayed(player: Player):
        return is_epper(player)

class epper_task_5(Page):
    form_model = 'player'
    form_fields = ['Epper_allocation_05']
    @staticmethod
    def is_displayed(player: Player):
        return is_epper(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        # Increment counter only after the last Epper page
        player.task_counter += 1

# --- Task 4: BRET ---
def is_bret(player):
    if player.task_counter < 4:
        return player.participant.vars['task_order'][player.task_counter] == 4
    return False

class BRETInstructions(Page):
    @staticmethod
    def is_displayed(player: Player):
        return is_bret(player)
    
    def vars_for_template(player: Player):
        return {
            'num_rows': C.num_rows, 'num_cols': C.num_cols,
            'num_boxes': C.num_rows * C.num_cols,
            'num_nobomb': C.num_rows * C.num_cols - 1,
            'box_value': C.box_value, 'time_interval': C.time_interval,
        }

class BRETDecision(Page):
    form_model = 'player'
    form_fields = ['bomb', 'boxes_collected', 'bomb_row', 'bomb_col']
    
    @staticmethod
    def is_displayed(player: Player):
        return is_bret(player)

    def vars_for_template(self):
        reset = self.participant.vars.get('reset', False)
        if reset: del self.participant.vars['reset']
        input = not C.devils_game if not C.dynamic else False
        return {
            'otree_vars': {
                'reset': reset, 'input': input, 'random': C.random,
                'dynamic': C.dynamic, 'num_rows': C.num_rows, 'num_cols': C.num_cols,
                'feedback': C.feedback, 'undoable': C.undoable,
                'box_width': C.box_width, 'box_height': C.box_height,
                'time_interval': C.time_interval,
            }
        }

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.task_counter += 1

# --- Final Page ---
class Thank_you(Page):
    @staticmethod
    def is_displayed(player: Player):
        return True # Always display at the end

    @staticmethod
    def vars_for_template(player: Player):
        return {"p_id": player.id_created}

    @staticmethod
    def before_next_page(player, timeout_happened):
        p = player # No need for in_round(1) since num_rounds=1
        user_id = p.id_created

        # Epper Logic
        epper_scenarios = [
            {'name': '03', 'choice': p.Epper_allocation_01, 'vals_me': C.Epper_alloc_03_me, 'vals_other': C.Epper_alloc_03_other},
            {'name': '07', 'choice': p.Epper_allocation_02, 'vals_me': C.Epper_alloc_07_me, 'vals_other': C.Epper_alloc_07_other},
            {'name': '08', 'choice': p.Epper_allocation_03, 'vals_me': C.Epper_alloc_08_me, 'vals_other': C.Epper_alloc_08_other},
            {'name': '09', 'choice': p.Epper_allocation_04, 'vals_me': C.Epper_alloc_09_me, 'vals_other': C.Epper_alloc_09_other},
            {'name': '13', 'choice': p.Epper_allocation_05, 'vals_me': C.Epper_alloc_13_me, 'vals_other': C.Epper_alloc_13_other},
        ]
        selected_scenario = random.choice(epper_scenarios)
        
        # Safe Retrieval of Choice (handle None)
        choice_val = selected_scenario['choice']
        if choice_val is not None:
            idx = choice_val - 1
            epper_val_me = selected_scenario['vals_me'][idx]
            epper_val_other = selected_scenario['vals_other'][idx]
        else:
            epper_val_me = 0
            epper_val_other = 0

        # Safe DB Update
        existing = dbq(PreSurveyProfile).filter_by(user_id=user_id).first()
        
        # Use dictionary unpacking to handle fields safely
        defaults = dict(
            cooperate_PD=p.cooperate_PD,
            cooperate_Cond_PD_C=p.cooperate_Cond_PD_C,
            cooperate_Cond_PD_D=p.cooperate_Cond_PD_D,
            boxes_collected=p.boxes_collected,
            bomb=p.bomb,
            epper_scenario_selected=selected_scenario['name'],
            epper_choice_raw=choice_val,
            epper_payoff_me=abs(epper_val_me),
            epper_payoff_other=epper_val_other,
            finished_part_1=True
        )

        if existing:
            for key, value in defaults.items():
                setattr(existing, key, value)
        else:
            PreSurveyProfile.create(user_id=user_id, **defaults)
            
        print(f"SAVED PROFILE FOR {user_id}")

# --- Page Sequence ---
# Each block appears 4 times in the list because we have 4 tasks to display in random order.
# The `is_displayed` method checks `task_counter` to ensure only the CORRECT block shows up
# for the current step (0, 1, 2, 3).
page_sequence = [
    Introduction, 
    ID_Creation,
    
    # Block 1 (Task Counter = 0)
    PD, Survey, 
    epper_instructions, epper_task_1, epper_task_2, epper_task_3, epper_task_4, epper_task_5,
    BRETInstructions, BRETDecision,
    
    # Block 2 (Task Counter = 1)
    PD, Survey, 
    epper_instructions, epper_task_1, epper_task_2, epper_task_3, epper_task_4, epper_task_5,
    BRETInstructions, BRETDecision,
    
    # Block 3 (Task Counter = 2)
    PD, Survey, 
    epper_instructions, epper_task_1, epper_task_2, epper_task_3, epper_task_4, epper_task_5,
    BRETInstructions, BRETDecision,
    
    # Block 4 (Task Counter = 3)
    PD, Survey, 
    epper_instructions, epper_task_1, epper_task_2, epper_task_3, epper_task_4, epper_task_5,
    BRETInstructions, BRETDecision,
    
    Thank_you
]