import random
from otree.api import *

doc = """
Part 2: Robot task
"""


class C(BaseConstants):
    NAME_IN_URL = 'robot_task'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 6
    PERIODS_PER_BLOCK = 5
    PART2_PRIZE = cu(1000)

    SUCCESS_PROBS = {
        ('top', 'good'): 0.75,
        ('top', 'bad'): 0.50,
        ('bottom', 'good'): 0.50,
        ('bottom', 'bad'): 0.25,
    }


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    treatment = models.StringField()

    attention1 = models.IntegerField(
        label='What is 2 + 2?',
        blank=True
    )
    attention2 = models.IntegerField(
        label='What is 2 + 2?',
        blank=True
    )
    attention_fail = models.BooleanField(initial=False)

    # comprehension
    comp1 = models.BooleanField(
        choices=[[0, 'True'], [1, 'False']],
        widget=widgets.RadioSelect,
        blank=True,
        label='1. If you currently have a Good robot assistant, requesting a new robot assistant means you will get a Bad robot assistant for sure.'
    )
    comp2 = models.BooleanField(
        choices=[[0, 'True'], [1, 'False']],
        widget=widgets.RadioSelect,
        blank=True,
        label='2. Your quiz-score rank can change throughout the experiment.'
    )
    comp3 = models.BooleanField(
        choices=[[1, 'True'], [0, 'False']],
        widget=widgets.RadioSelect,
        blank=True,
        label='3. The probability of success is higher when you are ranked in the Top Half and your robot assistant is Good than when you are ranked in the Top Half and your robot assistant is Bad.'
    )
    comp4 = models.IntegerField(
        choices=[
            [0, 'The 5th outcome is more likely to be a failure than before.'],
            [0, 'The 5th outcome is more likely to be a success than before.'],
            [1, 'The 5th outcome is just as likely to be a success as before.'],
        ],
        widget=widgets.RadioSelect,
        blank=True,
        label='4. Suppose a participant has a 75% probability of success in each period. '
              'This probability is fixed across 5 periods because the participant’s quiz-score rank and robot\'s quality do not change. '
              'After observing 4 successes in a row, what can be said about the 5th outcome?'
    )

    comp1_wrong_once = models.BooleanField(initial=False)
    comp2_wrong_once = models.BooleanField(initial=False)
    comp3_wrong_once = models.BooleanField(initial=False)
    comp4_wrong_once = models.BooleanField(initial=False)

    comp1_wrong_count = models.IntegerField(initial=0)
    comp2_wrong_count = models.IntegerField(initial=0)
    comp3_wrong_count = models.IntegerField(initial=0)
    comp4_wrong_count = models.IntegerField(initial=0)

    comp_failed_questions_history = models.LongStringField(blank=True)

    # beliefs
    top_belief = models.IntegerField(min=0, max=100, blank=True)
    good_belief = models.IntegerField(min=0, max=100, blank=True)

    # switch decision
    switch_robot = models.BooleanField(
        choices=[[True, 'Yes'], [False, 'No']],
        widget=widgets.RadioSelect,
        blank=True,
    )

    # hidden state
    robot_type = models.StringField()

    # realized outcomes in this block
    block_outcomes = models.LongStringField()

    # questionnaire
    switch_reasons = models.LongStringField(blank=True)
    switch_reason_other_text = models.LongStringField(blank=True)

    most_important_switch_reason = models.StringField(
        blank=True,
        choices=[
            ['1', '1'],
            ['2', '2'],
            ['3', '3'],
            ['4', '4'],
            ['5', '5'],
            ['6', '6'],
        ],
        widget=widgets.RadioSelect,
    )

    risk_preference = models.StringField(
        choices=[
            ['A', 'Option A'],
            ['B', 'Option B'],
            ['C', 'Indifferent between the two options']
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )

    risk_CE = models.IntegerField(min=0, max=10, blank=True)

    ambiguity_preference = models.StringField(
        choices=[
            ['A', 'Option A'],
            ['B', 'Option B'],
            ['C', 'Indifferent between the two options']
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )

    ambiguity_prob = models.IntegerField(min=0, max=100, blank=True)


def creating_session(subsession):
    players = subsession.get_players()

    if subsession.round_number == 1:
        for player in players:
            # comes from iq_quiz participant fields
            player.treatment = player.participant.treatment
    else:
        for player in players:
            player.treatment = player.in_round(1).treatment


def assign_robot(player: Player):
    player.robot_type = random.choice(['good', 'bad'])
    player.participant.current_robot_type = player.robot_type


def carry_forward_robot(player: Player):
    player.robot_type = player.participant.current_robot_type


def success_probability(player: Player):
    rank_type = player.participant.rank_type
    robot_type = player.robot_type
    return C.SUCCESS_PROBS[(rank_type, robot_type)]


def generate_outcomes(player: Player):
    p = success_probability(player)
    outcomes = []
    for _ in range(C.PERIODS_PER_BLOCK):
        outcomes.append('success' if random.random() < p else 'failure')
    player.block_outcomes = ','.join(outcomes)


def prepare_block(player: Player):
    if player.round_number == 1:
        if not player.field_maybe_none('robot_type'):
            assign_robot(player)
        if not player.field_maybe_none('block_outcomes'):
            generate_outcomes(player)
    else:
        prev = player.in_round(player.round_number - 1)

        if not player.field_maybe_none('robot_type'):
            if prev.switch_robot:
                assign_robot(player)
            else:
                carry_forward_robot(player)

        if not player.field_maybe_none('block_outcomes'):
            generate_outcomes(player)


def cumulative_outcomes(player: Player):
    rows = []
    period = 1
    current_pattern = 1

    previous_reports = player.in_previous_rounds()

    for rd in player.in_rounds(1, player.round_number):
        block_outcomes = rd.field_maybe_none('block_outcomes')

        if block_outcomes not in [None, '']:
            outcomes = block_outcomes.split(',')

            for i, outcome in enumerate(outcomes):

                report_text = ''
                report_type = 'blank'

                if i > 0 and rd.round_number - 1 < len(previous_reports):

                    prev = previous_reports[rd.round_number - 1]

                    if i == 1:
                        report_text = f'In period {prev.round_number * 5}, you reported:'
                        report_type = 'header'

                    elif i == 2:
                        report_text = f'Likelihood of being in the Top half: {prev.top_belief}%'
                        report_type = 'top'

                    elif i == 3:
                        report_text = f'Likelihood of the quality of your robot being Good: {prev.good_belief}%'
                        report_type = 'good'

                    elif i == 4:
                        if prev.switch_robot:
                            report_text = 'You requested a new robot.'
                        else:
                            report_text = 'You did not request a new robot.'
                        report_type = 'switch'

                rows.append(dict(
                    period=period,
                    outcome=outcome,
                    pattern_class=f'pattern{current_pattern}',
                    report_text=report_text,
                    report_type=report_type,
                    is_block_end=(period % 5 == 0),
                ))

                period += 1

        if rd.round_number < player.round_number:
            if rd.field_maybe_none('switch_robot'):
                current_pattern += 1
                if current_pattern > 6:
                    current_pattern = 6

    return rows


def previous_reports(player: Player):
    rows = []
    for rd in player.in_previous_rounds():
        rows.append(dict(
            round_number=rd.round_number,
            top_belief=rd.top_belief,
            good_belief=rd.good_belief,
            switch_robot=rd.switch_robot,
        ))
    return rows


def bsr_payment(prob_percent, event_happened):
    q = prob_percent / 100
    if event_happened:
        win_prob = 1 - (1 - q) ** 2
    else:
        win_prob = 1 - q ** 2
    return 1000 if random.random() < win_prob else 0


class Part2Intro(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1


class Attention1(Page):
    form_model = 'player'
    form_fields = ['attention1']

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

    @staticmethod
    def error_message(player: Player, values):
        if values.get('attention1') is None:
            return 'Please answer the question.'

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        player.attention_fail = (player.attention1 != 10)


class AttentionPass(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            passed=not player.attention_fail
        )


class Part2Intro_2(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1 and player.attention_fail


class Attention2(Page):
    form_model = 'player'
    form_fields = ['attention2']

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1 and player.attention_fail

    @staticmethod
    def error_message(player: Player, values):
        if values.get('attention2') is None:
            return 'Please answer the question.'


class Comprehension(Page):
    form_model = 'player'
    form_fields = ['comp1', 'comp2', 'comp3', 'comp4']

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

    @staticmethod
    def error_message(player, values):
        solutions = dict(
            comp1=1, comp2=1, comp3=1, comp4=1
        )

        custom_error_messages = dict(
            comp1="Wrong. Each new assignment always gives a 50-50 chance of getting a Good or Bad robot assistant.",
            comp2="Wrong. Your quiz-score rank remains fixed throughout the experiment.",
            comp3="Wrong. If you are ranked in the Top Half and your robot assistant is Good, the probability of "
                  "getting a success is 75%, while it’s only 50% if you are ranked in the Top Half but your robot assistant is Bad.",
            comp4="Wrong. The probability of success remains the same as before because the participant's quiz-score rank and robot's quality do not change across the 5 periods."
        )

        error_messages = {}
        wrong_fields = []

        for field_name, correct_value in solutions.items():
            if values.get(field_name) != correct_value:
                error_messages[field_name] = custom_error_messages[field_name]
                wrong_fields.append(field_name)

                setattr(player, f'{field_name}_wrong_once', True)
                current_count = getattr(player, f'{field_name}_wrong_count')
                setattr(player, f'{field_name}_wrong_count', current_count + 1)

        if wrong_fields:
            old_history = player.field_maybe_none('comp_failed_questions_history') or ''
            this_attempt = ','.join(wrong_fields)

            if old_history:
                player.comp_failed_questions_history = old_history + ' | ' + this_attempt
            else:
                player.comp_failed_questions_history = this_attempt

        return error_messages

        # error_messages = dict()
        # for field_name in solutions:
        #     if values[field_name] != solutions[field_name]:
        #         error_messages[field_name] = custom_error_messages[field_name]
        # return error_messages

class RobotAssigned(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

    @staticmethod
    def vars_for_template(player: Player):
        prepare_block(player)
        return dict(
            initial_top_belief=player.participant.initial_top_belief,
        )


class BlockReport(Page):
    form_model = 'player'
    form_fields = ['top_belief', 'good_belief', 'switch_robot']

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number <= 5

    @staticmethod
    def vars_for_template(player: Player):
        prepare_block(player)
        return dict(
            outcome_rows=cumulative_outcomes(player),
            previous_reports=previous_reports(player),
            initial_top_belief=player.participant.initial_top_belief,
            periods_so_far=player.round_number * C.PERIODS_PER_BLOCK,
            total_periods=C.NUM_ROUNDS * C.PERIODS_PER_BLOCK,
        )

    @staticmethod
    def error_message(player: Player, values):
        if values.get('top_belief') is None:
            return 'Please report the probability of being ranked in the Top half.'
        if values.get('good_belief') is None:
            return 'Please report the probability of the robot quality being Good.'
        if values.get('switch_robot') is None:
            return 'Please choose whether you want to request a new robot.'

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # store potential payment objects, do not randomize final payment yet
        event_top = (player.participant.rank_type == 'top')
        top_points = bsr_payment(player.top_belief, event_top)
        setattr(player.participant, f'ability_belief_round{player.round_number}_points', top_points)

        event_good = (player.robot_type == 'good')
        good_points = bsr_payment(player.good_belief, event_good)
        setattr(player.participant, f'robot_belief_round{player.round_number}_points', good_points)

class FinalOutcomes(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 6

    @staticmethod
    def vars_for_template(player: Player):
        prepare_block(player)
        return dict(
            outcome_rows=cumulative_outcomes(player),
            periods_so_far=player.round_number * C.PERIODS_PER_BLOCK,
            total_periods=C.NUM_ROUNDS * C.PERIODS_PER_BLOCK,
            initial_top_belief=player.participant.initial_top_belief,
        )


class Questionnaire(Page):
    form_model = 'player'
    form_fields = [
        'switch_reasons',
        'switch_reason_other_text',
        'most_important_switch_reason',
        'risk_preference',
        'risk_CE',
        'ambiguity_preference',
        'ambiguity_prob',
    ]
    preserve_unsubmitted_inputs = True

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

    @staticmethod
    def vars_for_template(player: Player):
        switch_count = sum(
            1 for rd in player.in_rounds(1, 5)
            if rd.switch_robot
        )
        return dict(
            switch_count=switch_count,
        )

    @staticmethod
    def error_message(player: Player, values):
        selected_reasons_raw = values.get('switch_reasons') or ''
        selected_reasons = [
            x.strip() for x in selected_reasons_raw.split(',')
            if x.strip()
        ]

        if not selected_reasons:
            return 'Please select at least one reason for your switching decision.'

        if '6' in selected_reasons and not (values.get('switch_reason_other_text') or '').strip():
            return 'Please specify the other reason.'

        most_important = values.get('most_important_switch_reason')
        if not most_important:
            return 'Please indicate which selected reason was the most important.'

        if most_important not in selected_reasons:
            return 'Your most important reason must be one of the reasons you selected above.'

        if values.get('risk_preference') is None:
            return 'Please answer question 3.'
        if values.get('risk_CE') is None:
            return 'Please answer question 4.'
        if values.get('ambiguity_preference') is None:
            return 'Please answer question 5.'
        if values.get('ambiguity_prob') is None:
            return 'Please answer question 6.'


class FinalResults(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == C.NUM_ROUNDS

    @staticmethod
    def vars_for_template(player: Player):
        player.participant.finished = True

        # Draw the final payment only once.
        if player.participant.vars.get('final_selected_part') in [None, '']:

            # initialize stored values
            player.participant.final_selected_part = ''
            player.participant.final_payment_points = 0

            player.participant.selected_quiz_question = ''

            player.participant.selected_part2_subpart = ''
            player.participant.selected_ability_belief_period = ''
            player.participant.selected_robot_belief_period = ''
            player.participant.selected_task_period = ''
            player.participant.selected_task_outcome = ''
            player.participant.part2_payment_points = 0

            # -----------------------------
            # Step 1: select ONE subpart of Part 2
            # -----------------------------
            selected_part2_subpart = random.choice(['ability_belief', 'robot_belief', 'task_outcome'])
            player.participant.selected_part2_subpart = selected_part2_subpart

            if selected_part2_subpart == 'ability_belief':
                ability_options = [
                    ('initial', player.participant.initial_belief_payment_points),
                    (1, player.participant.ability_belief_round1_points),
                    (2, player.participant.ability_belief_round2_points),
                    (3, player.participant.ability_belief_round3_points),
                    (4, player.participant.ability_belief_round4_points),
                    (5, player.participant.ability_belief_round5_points),
                ]

                selected_period, payment_points = random.choice(ability_options)
                player.participant.selected_ability_belief_period = str(selected_period)
                player.participant.part2_payment_points = payment_points

            elif selected_part2_subpart == 'robot_belief':
                robot_options = [
                    (1, player.participant.robot_belief_round1_points),
                    (2, player.participant.robot_belief_round2_points),
                    (3, player.participant.robot_belief_round3_points),
                    (4, player.participant.robot_belief_round4_points),
                    (5, player.participant.robot_belief_round5_points),
                ]

                selected_period, payment_points = random.choice(robot_options)
                player.participant.selected_robot_belief_period = str(selected_period)
                player.participant.part2_payment_points = payment_points

            elif selected_part2_subpart == 'task_outcome':
                all_outcomes = []
                for rd in player.in_rounds(1, C.NUM_ROUNDS):
                    outcomes = rd.block_outcomes.split(',')
                    all_outcomes.extend(outcomes)

                selected_task_period = random.randint(1, 30)
                selected_task_outcome = all_outcomes[selected_task_period - 1]
                payment_points = 1000 if selected_task_outcome == 'success' else 0

                player.participant.selected_task_period = str(selected_task_period)
                player.participant.selected_task_outcome = selected_task_outcome
                player.participant.part2_payment_points = payment_points

            # -----------------------------
            # Step 2: select Part 1 or Part 2
            # -----------------------------
            final_selected_part = random.choice(['part1', 'part2'])
            player.participant.final_selected_part = final_selected_part

            if final_selected_part == 'part1':
                player.participant.selected_quiz_question = player.participant.quiz_payment_question
                player.participant.final_payment_points = player.participant.quiz_payment_points
            else:
                player.participant.final_payment_points = player.participant.part2_payment_points

        final_selected_part = player.participant.vars.get('final_selected_part', '')
        final_payment_points = player.participant.vars.get('final_payment_points', 0)
        selected_quiz_question = player.participant.vars.get('selected_quiz_question', '')

        selected_part2_subpart = player.participant.vars.get('selected_part2_subpart', '')
        selected_ability_belief_period = player.participant.vars.get('selected_ability_belief_period', '')
        selected_robot_belief_period = player.participant.vars.get('selected_robot_belief_period', '')
        selected_task_period = player.participant.vars.get('selected_task_period', '')
        selected_task_outcome = player.participant.vars.get('selected_task_outcome', '')

        selected_ability_belief_period_times_five = ''
        if selected_ability_belief_period not in ['', None, 'initial']:
            selected_ability_belief_period_times_five = int(selected_ability_belief_period) * 5

        selected_robot_belief_period_times_five = ''
        if selected_robot_belief_period not in ['', None]:
            selected_robot_belief_period_times_five = int(selected_robot_belief_period) * 5

        return dict(
            final_selected_part=final_selected_part,
            final_payment_points=final_payment_points,
            selected_quiz_question=selected_quiz_question,
            selected_part2_subpart=selected_part2_subpart,
            selected_ability_belief_period=selected_ability_belief_period,
            selected_ability_belief_period_times_five=selected_ability_belief_period_times_five,
            selected_robot_belief_period=selected_robot_belief_period,
            selected_robot_belief_period_times_five=selected_robot_belief_period_times_five,
            selected_task_period=selected_task_period,
            selected_task_outcome=selected_task_outcome,
        )






page_sequence = [
    Part2Intro,
    Attention1,
    AttentionPass,
    Part2Intro_2,
    Attention2,
    Comprehension,
    RobotAssigned,
    BlockReport,
    FinalOutcomes,
    Questionnaire,
    FinalResults
]