from otree.api import *
import random

doc = """
P3_Shock:
1. Belief elicitation about destruction probability (ambiguity treatment only)
2. Connection destruction (spinning wheel, same as Phase 4)
3. SWB II (mid-task mood)
4. Decision: Option A or Option B
5. Help decision (simultaneous, interval-based)
6. Belief about help received
"""


class C(BaseConstants):
    NAME_IN_URL = 'p3_shock'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 1

    # Destruction probabilities
    DESTRUCTION_PROB  = 0.35

    # Help intervals
    HELP_OPTIONS = [0, 15, 30, 45, 60, 75, 90]

    PHASE4_MIN = 13
    PRACTICE_MIN = 8
    # Belief bonus
    BELIEF_BONUS    = 15
    BELIEF_TOLERANCE = 5


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    connection_destroyed = models.BooleanField(initial=False)
    spin_degrees         = models.FloatField(initial=0)


class Player(BasePlayer):
    # ── Belief about destruction probability (ambiguity only) ─────────────────
    belief_destruction_lower = models.IntegerField(
        label='', min=0, max=100, initial=0
    )
    belief_destruction_upper = models.IntegerField(
        label='', min=0, max=100, initial=100
    )
    belief_tokens_0 = models.IntegerField(initial=0, min=0, max=100)
    belief_tokens_15 = models.IntegerField(initial=0, min=0, max=100)
    belief_tokens_30 = models.IntegerField(initial=0, min=0, max=100)
    belief_tokens_45 = models.IntegerField(initial=0, min=0, max=100)
    belief_tokens_60 = models.IntegerField(initial=0, min=0, max=100)
    belief_tokens_75 = models.IntegerField(initial=0, min=0, max=100)
    belief_tokens_90 = models.IntegerField(initial=0, min=0, max=100)

    belief_help_bonus = models.IntegerField(initial=0)

    # ── SWB II ────────────────────────────────────────────────────────────────
    swb_md_happy = models.IntegerField(
        choices=[0, 1, 2, 3, 4, 5, 6], label='',
        widget=widgets.RadioSelectHorizontal,
    )
    swb_md_friendly = models.IntegerField(
        choices=[0, 1, 2, 3, 4, 5, 6], label='',
        widget=widgets.RadioSelectHorizontal,
    )
    swb_md_relaxed = models.IntegerField(
        choices=[0, 1, 2, 3, 4, 5, 6], label='',
        widget=widgets.RadioSelectHorizontal,
    )
    swb_md_stressed = models.IntegerField(
        choices=[0, 1, 2, 3, 4, 5, 6], label='',
        widget=widgets.RadioSelectHorizontal,
    )
    swb_md_sad = models.IntegerField(
        choices=[0, 1, 2, 3, 4, 5, 6], label='',
        widget=widgets.RadioSelectHorizontal,
    )
    swb_md_angry = models.IntegerField(
        choices=[0, 1, 2, 3, 4, 5, 6], label='',
        widget=widgets.RadioSelectHorizontal,
    )
    swb_md_frustrated = models.IntegerField(
        choices=[0, 1, 2, 3, 4, 5, 6], label='',
        widget=widgets.RadioSelectHorizontal,
    )
    feeling_order_md = models.StringField()

    # ── Decision ──────────────────────────────────────────────────────────────
    task_decision = models.StringField(
        choices=[['A', 'Option A'], ['B', 'Option B']],
        label='', widget=widgets.RadioSelect,
    )

    # ── Help ──────────────────────────────────────────────────────────────────
    help_given    = models.IntegerField(initial=0, label='')
    help_received = models.IntegerField(initial=0)

    # ── Belief about help received ────────────────────────────────────────────
    belief_help_received = models.IntegerField(
        label='', min=0, max=20
    )
    belief_help_correct = models.BooleanField(initial=False)
    belief_help_bonus   = models.IntegerField(initial=0)


# ── FUNCTIONS ─────────────────────────────────────────────────────────────────

def determine_destruction(group: Group):
    p1 = group.get_player_by_id(1)
    initial_connection = p1.participant.vars.get('group_connection_initial', 'none')

    if initial_connection != 'none':
        destroyed      = random.random() < C.DESTRUCTION_PROB
        survive_prob   = 1 - C.DESTRUCTION_PROB
    else:
        destroyed    = False
        survive_prob = 0

    group.connection_destroyed = destroyed
    survive_angle = survive_prob * 360
    base_rotations = 5 * 360

    if initial_connection != 'none':
        if not destroyed:
            target_angle = random.uniform(20, max(survive_angle - 20, 21))
            final_angle  = (270 - target_angle) % 360
        else:
            target_angle = random.uniform(survive_angle + 20, 340)
            final_angle  = (270 - target_angle) % 360
        spin_degrees = base_rotations + final_angle
    else:
        spin_degrees = 0

    group.spin_degrees = spin_degrees

    for player in group.get_players():
        player.participant.vars['group_connection_initial'] = initial_connection
        player.participant.vars['group_connection_after_shock'] = (
            'none' if initial_connection == 'none' or destroyed else initial_connection
        )
        player.participant.vars['spin_degrees_shock'] = spin_degrees


def process_help(group: Group):
    p1 = group.get_player_by_id(1)
    p2 = group.get_player_by_id(2)

    connection = p1.participant.vars.get('group_connection_after_shock', 'none')

    if connection == 'none':
        p1.help_received = 0
        p2.help_received = 0
    else:
        p1.help_received = p2.help_given
        p2.help_received = p1.help_given

    p1.participant.vars['help_received_p3'] = p1.help_received
    p2.participant.vars['help_received_p3'] = p2.help_received
    p1.participant.vars['help_given_p3']    = p1.help_given
    p2.participant.vars['help_given_p3']    = p2.help_given


def score_belief_help(player: Player):
    actual       = player.help_received
    token_field  = f'belief_tokens_{actual}'
    tokens_correct = getattr(player, token_field, 0) or 0
    bonus = round((tokens_correct / 100) * 20)
    player.belief_help_bonus = bonus
    player.participant.vars['belief_help_bonus_p3'] = bonus


# ── PAGES ─────────────────────────────────────────────────────────────────────

class BeliefDestruction(Page):
    form_model  = 'player'
    form_fields = ['belief_destruction_lower', 'belief_destruction_upper']

    @staticmethod
    def is_displayed(player: Player):
        return player.participant.vars.get('task_type') == 'ambiguity'

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            initial_connection=player.participant.vars.get(
                'group_connection_initial', 'none')
        )

    @staticmethod
    def error_message(player: Player, values):
        if values['belief_destruction_lower'] > values['belief_destruction_upper']:
            if player.participant.language == 'de':
                return 'Die untere Grenze muss kleiner oder gleich der oberen Grenze sein.'
            return 'The lower bound must be less than or equal to the upper bound.'


class DestructionWaitPage(WaitPage):
    after_all_players_arrive = determine_destruction

    def title_text(self):
        return ('Please wait'
                if self.player.participant.language == 'en'
                else 'Bitte warten')

    def body_text(self):
        return ('Please wait until the other participant is ready.'
                if self.player.participant.language == 'en'
                else 'Bitte warten Sie bis der*die andere Teilnehmer*in bereit ist.')


class ConnectionShock(Page):
    @staticmethod
    def vars_for_template(player: Player):
        initial_connection = player.participant.vars.get('group_connection_initial', 'none')
        connection_after = player.participant.vars.get('group_connection_after_shock', 'none')
        return dict(
            initial_connection=initial_connection,
            connection_destroyed=player.group.connection_destroyed,
            spin_degrees=player.participant.vars.get('spin_degrees_shock', 0),
            language=player.participant.language,
            destruction_prob_pct=int(C.DESTRUCTION_PROB * 100),
            survival_prob_pct=int((1 - C.DESTRUCTION_PROB) * 100),
            is_ambiguity=player.participant.vars.get('task_type') == 'ambiguity',
            connection_bonus=36 if connection_after == 'strong' else 0,
        )


class SWBII(Page):
    form_model  = 'player'
    form_fields = [
        'swb_md_happy', 'swb_md_friendly', 'swb_md_relaxed', 'swb_md_stressed',
        'swb_md_sad', 'swb_md_angry', 'swb_md_frustrated',
    ]

    @staticmethod
    def vars_for_template(player: Player):
        import random as _random
        feelings = ['happy', 'friendly', 'relaxed', 'stressed', 'sad', 'angry', 'frustrated']
        _random.shuffle(feelings)
        player.feeling_order_md = ','.join(feelings)
        return dict(feeling_order=feelings)


class Decision(Page):
    form_model  = 'player'
    form_fields = ['task_decision']

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            connection=player.participant.vars.get(
                'group_connection_after_shock', 'none')
        )

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        player.participant.vars['task_decision_p3'] = player.task_decision


class WaitForBothDecisions(WaitPage):
    def title_text(self):
        return ('Please wait'
                if self.player.participant.language == 'en'
                else 'Bitte warten')

    def body_text(self):
        return ('Please wait until your partner has made their decision.'
                if self.player.participant.language == 'en'
                else 'Bitte warten Sie, bis Ihr*e Mitspieler*in eine Entscheidung getroffen hat.')

    @staticmethod
    def is_displayed(player: Player):
        return player.participant.vars.get(
            'group_connection_after_shock', 'none') != 'none'


class HelpDecision(Page):
    form_model  = 'player'
    form_fields = ['help_given']

    @staticmethod
    def is_displayed(player: Player):
        return player.participant.vars.get(
            'group_connection_after_shock', 'none') != 'none'

    @staticmethod
    def vars_for_template(player: Player):
        other = player.get_others_in_group()[0]
        points_p1 = int(player.participant.vars.get('points_p1', 0) or 0)
        points_p2 = int(player.participant.vars.get('points_p2', 0) or 0)
        max_help = points_p1 + points_p2
        affordable_options = [h for h in C.HELP_OPTIONS if h <= max_help]
        return dict(
            help_options=affordable_options,
            my_decision=player.task_decision,
            partner_decision=other.task_decision,
            connection=player.participant.vars.get('group_connection_after_shock', 'none'),
            max_help=max_help,
            points_p1=points_p1,
            points_p2=points_p2,
        )

    @staticmethod
    def error_message(player: Player, values):
        points_p1 = int(player.participant.vars.get('points_p1', 0) or 0)
        points_p2 = int(player.participant.vars.get('points_p2', 0) or 0)
        max_help = points_p1 + points_p2
        if (values.get('help_given') or 0) > max_help:
            if getattr(player.participant, 'language', 'de') == 'en':
                return f'You cannot give more than {max_help} points (your total from Phases 1 and 2).'
            return f'Sie können nicht mehr als {max_help} Punkte geben (Ihr Gesamtverdienst aus Phase 1 und 2).'


class HelpWaitPage(WaitPage):
    after_all_players_arrive = process_help

    def title_text(self):
        return ('Please wait'
                if self.player.participant.language == 'en'
                else 'Bitte warten')

    def body_text(self):
        return ('Please wait until the other participant is ready.'
                if self.player.participant.language == 'en'
                else 'Bitte warten Sie bis der*die andere Teilnehmer*in bereit ist.')


class BeliefHelpReceived(Page):
    form_model  = 'player'
    form_fields = ['belief_tokens_0', 'belief_tokens_15', 'belief_tokens_30',
        'belief_tokens_45', 'belief_tokens_60', 'belief_tokens_75',
        'belief_tokens_90',]

    @staticmethod
    def is_displayed(player: Player):
        return player.participant.vars.get(
            'group_connection_after_shock', 'none') != 'none'

    @staticmethod
    def vars_for_template(player: Player):
        return dict(help_options=C.HELP_OPTIONS)

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        score_belief_help(player)


class Results(Page):
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            help_received=player.help_received,
            help_given=player.help_given,
            belief_help_bonus=player.belief_help_bonus,
            connection=player.participant.vars.get(
                'group_connection_after_shock', 'none'),
        )

    @staticmethod
    def before_next_page(player, timeout_happened):
        connection_after = player.participant.vars.get('group_connection_after_shock', 'none')
        connection_bonus = 36 if connection_after == 'strong' else 0
        help_cost = player.help_given
        belief_bonus = player.belief_help_bonus

        player.participant.vars['points_p3'] = connection_bonus - help_cost + belief_bonus

        # these are needed for routing/help logic but not for LastPage
        player.participant.vars['help_received_p3'] = player.help_received
        player.participant.vars['help_given_p3'] = player.help_given

    @staticmethod
    def app_after_this_page(player: Player, upcoming_apps):
        if player.task_decision == 'A':
            return 'P4_Option_A'
        else:
            return 'P4_Option_B'


page_sequence = [
    BeliefDestruction,
    DestructionWaitPage,
    ConnectionShock,
    SWBII,
    Decision,
    WaitForBothDecisions,
    HelpDecision,
    HelpWaitPage,
    BeliefHelpReceived,
    Results
]