from otree.api import *
import random


doc = """
Your app description
"""


class C(BaseConstants):
    NAME_IN_URL = 'bayesian_learning'
    PLAYERS_PER_GROUP = None
    NUM_PRACTICE_ROUNDS = 3
    NUM_REAL_ROUNDS = 30
    NUM_ROUNDS = NUM_PRACTICE_ROUNDS + NUM_REAL_ROUNDS
    
    # Payment parameters
    TOKENS_PER_DOLLAR = 100  # 100 tokens = $1
    NUM_ROUNDS_FOR_PAYMENT = 3
    
    # Payoff parameters
    ENDOWMENT = 100
    RETURN_MULTIPLIER_HIGH = 2.0  
    RETURN_MULTIPLIER_LOW = 1.5
    INSURANCE_COST_LOW = 0.3
    INSURANCE_COST_HIGH = 0.5
    LOSS_AMOUNT = 100        
    
    TRUE_PROBABILITY = 0.6   
    SIGNAL_ACCURACY = 0.75   
    
    # Belief elicitation
    BELIEF_ELICITATION_ROUNDS = [5, 10, 15, 20, 25, 30]



class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    information_treatment = models.StringField(
        choices=['none', 'perfect', 'signal']
    )
    belief_elicitation_treatment = models.StringField(
        choices=['none', 'periodic']
    )
    framing_treatment = models.StringField(
        choices=['gain', 'loss']
    )
    incentive_structure = models.StringField(
        choices=['high', 'low']
    )    

    is_practice = models.BooleanField()

    choice = models.IntegerField(
        min=0, max=100,
    )
    
    outcome = models.BooleanField()  # True = success/no-loss, False = failure/loss
    signal = models.BooleanField()

    payoff_tokens = models.IntegerField()
    
    stated_belief = models.IntegerField(
        min=0, max=100,
        label="What do you believe is the probability of precipitation? (0-100)"
    )
    
    comments = models.LongStringField(
        label="Please provide us some comments about how you made choices over the course of this experiment:"
    )
    #def is_practice_round(self):
    #    return self.round_number <= C.NUM_PRACTICE_ROUNDS
    
    #def real_round_number(self):
    #    """The round number shown to participants (excluding practice)"""
    #    if self.is_practice_round():
    #        return self.round_number
    #    else:
    #        return self.round_number - C.NUM_PRACTICE_ROUNDS
    
    comp_q1_gain = models.IntegerField(
        label="You spend 40 tokens. Heavy snow occurs. How many tokens do you have at the end of the round?",
        choices=[
            [60, "60 tokens"],
            [80, "80 tokens"],
            [140, "140 tokens"],  # Correct: 100 - 40 + 80
            [180, "180 tokens"],
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )
    
    comp_q2_gain = models.IntegerField(
        label="You spend 40 tokens. Heavy snow does NOT occur. How many tokens do you have at the end of the round?",
        choices=[
            [0, "0 tokens"],
            [40, "40 tokens"],
            [60, "60 tokens"],  # Correct: 100 - 40
            [100, "100 tokens"],
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )
    
    comp_q3_gain = models.IntegerField(
        label="In Round 5, heavy snow occurs. In Round 6, heavy snow does NOT occur. How many tokens do you start with in Round 7?",
        choices=[
            [0, "0 tokens - I lost everything in Round 6"],
            [60, "60 tokens - what I had left after Round 6"],
            [100, "100 tokens - I start fresh each round"],  # Correct
            [140, "140 tokens - what I had after Round 5"],
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )
    
    comp_q4_gain = models.IntegerField(
        label="How many rounds will count toward your final payment?",
        choices=[
            [1, "1 round"],
            [3, "3 randomly selected rounds"],  # Correct
            [10, "10 rounds"],
            [30, "All 30 rounds"],
        ],
        widget=widgets.RadioSelect,
        blank=True,
        )
    
    comp_q1_loss = models.IntegerField(
        label="You spend 40 tokens on site protection. NO heavy rain occurs. How many tokens do you have at the end of the round?",
        choices=[
            [140, "140 tokens"],
            [160, "160 tokens"],  # Correct: 200 - 40
            [180, "180 tokens"],
            [200, "200 tokens"],
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )
    
    comp_q2_loss = models.IntegerField(
        label="You spend 40 tokens on site protection. Heavy rain occurs. How many tokens do you have at the end of the round?",
        choices=[
            [0, "0 tokens"],
            [20, "20 tokens"],  # Correct: 40 / 2 = 20
            [40, "40 tokens"],
            [80, "80 tokens"],
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )
    
    comp_q3_loss = models.IntegerField(
        label="In Round 5, NO heavy rain occurs. In Round 6, heavy rain occurs. How many tokens do you start with in Round 7?",
        choices=[
            [0, "0 tokens - I lost everything in Round 6"],
            [20, "20 tokens - what I had left after Round 6"],
            [80, "80 tokens - what I had after Round 5"],
            [100, "100 tokens - I start fresh each round"],  # Correct
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )
    
    comp_q4_loss = models.IntegerField(
        label="What does it mean if you spend 200 tokens of site protection?",
        choices=[
            [1, "You are protected as possible - if heavy rain occurs, you'll get back 100 tokens"],  # Correct
            [2, "You will definitely not have heavy rain"],
            [3, "You will get 100 extra tokens"],
            [4, "The insurance costs 50 tokens"],
        ],
        widget=widgets.RadioSelect,
        blank=True,
    )
    
    # Track attempts
    comp_attempts = models.IntegerField(initial=0)
    conf_num = models.IntegerField(initial=0)
    
    selected_for_payment = models.BooleanField(initial=False)

def creating_session(subsession):
    """Assign treatments and generate outcomes"""
    if subsession.round_number == 1:
        for player in subsession.get_players():
            # Assign treatments
          
            player.information_treatment = player.session.config.get('information_treatment', 'none')
            player.belief_elicitation_treatment = player.session.config.get('belief_elicitation_treatment', 'none')
            player.framing_treatment = player.session.config.get('framing_treatment', 'gain')
            player.incentive_structure = player.session.config.get('incentive_structure', 'high')  # high or low
    else:
        # Copy treatments from round 1
        for player in subsession.get_players():
            player.information_treatment = player.in_round(1).information_treatment
            player.belief_elicitation_treatment = player.in_round(1).belief_elicitation_treatment
            player.framing_treatment = player.in_round(1).framing_treatment
            player.incentive_structure = player.in_round(1).incentive_structure
    
    for player in subsession.get_players():
        player.is_practice = subsession.round_number <= C.NUM_PRACTICE_ROUNDS
    
    # Generate outcomes for this round
    for player in subsession.get_players():
        player.outcome = random.random() < C.TRUE_PROBABILITY
        
        if player.information_treatment == 'signal':
            if random.random() < C.SIGNAL_ACCURACY:
                player.signal = player.outcome
            else:
                player.signal = not player.outcome

def comprehension_correct(player: Player, values):
    """Check if all comprehension answers are correct"""
    if player.framing_treatment == 'gain':
        return (
            values.get('comp_q1_gain') == 140 and
            values.get('comp_q2_gain') == 60 and
            values.get('comp_q3_gain') == 100 #and
            #values.get('comp_q4_gain') == 3
        )
    else:
        return (
            values.get('comp_q1_loss') == 160 and
            values.get('comp_q2_loss') == 20 and
            values.get('comp_q3_loss') == 100 and
            values.get('comp_q4_loss') == 1
        )

def get_return_multiplier(player: Player):
    """Get return multiplier based on incentive structure"""
    if player.incentive_structure == 'high':
        return C.RETURN_MULTIPLIER_HIGH
    else:
        return C.RETURN_MULTIPLIER_LOW
        
def get_insurance_cost(player: Player):
    """Get insurance cost based on incentive structure"""
    if player.incentive_structure == 'high':
        return C.INSURANCE_COST_LOW  # Low cost = optimal to buy
    else:
        return C.INSURANCE_COST_HIGH  # High cost = optimal not to buy
        
def set_payoffs(player: Player):
    """Calculate payoffs based on choice and outcome"""
    x = player.choice
    
    if player.framing_treatment == 'gain':
        # Investment framing
        R = get_return_multiplier(player)
        if player.outcome:  # Success
            player.payoff_tokens = C.ENDOWMENT - x + int(x * R)
        else:  # Failure
            player.payoff_tokens = C.ENDOWMENT - x
    
    else:  # loss framing
        # Insurance framing
        cost = get_insurance_cost(player)
        insurance_cost = int(x * cost)
        
        if player.outcome:  # No loss
            player.payoff_tokens = C.ENDOWMENT - insurance_cost
        else:  # Loss occurs
            uncovered_loss = C.LOSS_AMOUNT - x
            player.payoff_tokens = C.ENDOWMENT - insurance_cost - uncovered_loss
    
    # Set oTree payoff (will be converted to real money later)
    player.payoff = cu(player.payoff_tokens)        
        
#def app_after_this_page(player: Player, upcoming_pages):
#    """Determine whether to show belief elicitation page"""
#    if player.belief_elicitation_treatment == 'none':
#        return
#    
#    if player.round_number not in C.BELIEF_ELICITATION_ROUNDS:
#        # Skip belief elicitation page
#        pass

def real_round_number(player: Player):
    """Get the round number as displayed to participant (excluding practice)"""
    if player.is_practice:
        return player.round_number
    else:
        return player.round_number - C.NUM_PRACTICE_ROUNDS
        
# PAGES
class Introduction(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1
    
    @staticmethod
    def vars_for_template(player: Player):
        insurance_cost = get_insurance_cost(player) if player.framing_treatment == 'loss' else None
        
        # Pre-calculate insurance examples
        if player.framing_treatment == 'loss':
            insurance_100_cost = int(100 * insurance_cost)
            insurance_100_remaining = 100 - insurance_100_cost
            insurance_60_cost = int(60 * insurance_cost)
            insurance_60_no_loss = 100 - insurance_60_cost
            insurance_60_loss = insurance_60_cost + (100 - 60)  # Cost + uncovered loss
            insurance_60_remaining = 100 - insurance_60_loss
        else:
            insurance_100_cost = None
            insurance_100_remaining = None
            insurance_60_cost = None
            insurance_60_no_loss = None
            insurance_60_loss = None
            insurance_60_remaining = None
        
        return {
            'information': player.information_treatment,
            'framing': player.framing_treatment,
            'insurance_cost': insurance_cost,
            'endowment': C.ENDOWMENT,
            # Pre-calculated values for loss frame examples
            'insurance_100_cost': insurance_100_cost,
            'insurance_100_remaining': insurance_100_remaining,
            'insurance_60_cost': insurance_60_cost,
            'insurance_60_no_loss': insurance_60_no_loss,
            'insurance_60_loss': insurance_60_loss,
            'insurance_60_remaining': insurance_60_remaining,
        }

class ComprehensionCheck(Page):
    form_model = 'player'
    
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1
    
    @staticmethod
    def get_form_fields(player: Player):
        if player.framing_treatment == 'gain':
            return ['comp_q1_gain', 'comp_q2_gain', 'comp_q3_gain']
            #return ['comp_q1_gain', 'comp_q2_gain', 'comp_q3_gain', 'comp_q4_gain']
        else:
            return ['comp_q1_loss', 'comp_q2_loss', 'comp_q3_loss', 'comp_q4_loss']
    
    @staticmethod
    def error_message(player: Player, values):
        player.comp_attempts += 1
        
        if not comprehension_correct(player, values):
            return 'One or more answers were incorrect. Please read the instructions again and try again.'
    
    @staticmethod
    def vars_for_template(player: Player):
        insurance_cost = get_insurance_cost(player) if player.framing_treatment == 'loss' else None
        return {
            'framing': player.framing_treatment,
            'insurance_cost': insurance_cost,
            'comp_attempts': player.comp_attempts,
        }

class Decision(Page):
    form_model = 'player'
    form_fields = ['choice']
    
    @staticmethod
    def vars_for_template(player: Player):
        # Determine round display
        is_practice = player.is_practice
        if is_practice:
            round_display = f"Practice Round {player.round_number}"
        else:
            real_round = real_round_number(player)
            round_display = f"Round {real_round} of {C.NUM_REAL_ROUNDS}"
        
        # Show signal if applicable
        show_signal = player.information_treatment == 'signal'
        if show_signal:
            if player.framing_treatment == 'gain':
                signal_text = "SUCCESS LIKELY" if player.signal else "FAILURE LIKELY"
            else:
                signal_text = "NO LOSS LIKELY" if player.signal else "LOSS LIKELY"
        else:
            signal_text = None
        
        # Show probability if perfect information
        show_probability = player.information_treatment == 'perfect'
        probability_text = f"{int(C.TRUE_PROBABILITY * 100)}%" if show_probability else None
        
        insurance_cost = get_insurance_cost(player) if player.framing_treatment == 'loss' else None
        
        return {
            'is_practice': is_practice,
            'round_display': round_display,
            'round_num': real_round_number(player) if not is_practice else player.round_number,
            'endowment': C.ENDOWMENT,
            'framing': player.framing_treatment,
            'show_signal': show_signal,
            'signal_text': signal_text,
            'show_probability': show_probability,
            'probability_text': probability_text,
            'insurance_cost': insurance_cost,
        }
    
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        set_payoffs(player)

class BeliefElicitation(Page):
    form_model = 'player'
    form_fields = ['stated_belief']
    
    @staticmethod
    def is_displayed(player: Player):
        if player.is_practice:
            return False
        if player.belief_elicitation_treatment != 'periodic':
            return False
        
        real_round = real_round_number(player)
        return real_round in C.BELIEF_ELICITATION_ROUNDS
    
    @staticmethod
    def vars_for_template(player: Player):
        return {
            'round_num': real_round_number(player),
            'framing': player.framing_treatment,
            'belief_elicitation': player.belief_elicitation_treatment,
        }

class Outcome(Page):
    @staticmethod
    def vars_for_template(player: Player):
        # Get history (only real rounds, not practice)
        history = []
        success_count = 0
        total_payoff = 0
        
        for past_round in player.in_all_rounds():
            if past_round.round_number <= player.round_number and past_round.choice is not None:
                if not past_round.is_practice or past_round.round_number <= C.NUM_PRACTICE_ROUNDS:
                    stated_belief = past_round.field_maybe_none('stated_belief')   
                    
                    history.append({
                        'round': real_round_number(past_round) if not past_round.is_practice else past_round.round_number,
                        'choice': past_round.choice,
                        'outcome': past_round.outcome,
                        'payoff': past_round.payoff_tokens,
                        'stated_belief': stated_belief if stated_belief else None,
                        'is_practice': past_round.is_practice,
                    })
                    if past_round.outcome:
                        success_count += 1
                    total_payoff += past_round.payoff_tokens
        
        avg_payoff = total_payoff / len(history) if history else 0
        
        insurance_cost = get_insurance_cost(player) if player.framing_treatment == 'loss' else 0
        return_multiplier = get_return_multiplier(player) if player.framing_treatment == 'gain' else 0
        
        # Pre-calculate values for display
        if player.framing_treatment == 'gain':
            investment_return = int(player.choice * return_multiplier)
            kept_amount = C.ENDOWMENT - player.choice
        else:
            insurance_premium = int(player.choice * insurance_cost)
            covered_amount = player.choice
            uncovered_amount = C.ENDOWMENT - player.choice
            investment_return = None
            kept_amount = None
        
        return {
            'round_num': real_round_number(player) if not player.is_practice else player.round_number,
            'is_practice': player.is_practice,
            'choice': player.choice,
            'choice_multiply_2': 2*player.choice,
            'endowment_subtract_choice' : C.ENDOWMENT - player.choice,
            'outcome': player.outcome,
            'payoff_tokens': player.payoff_tokens,
            'framing': player.framing_treatment,
            'history': history,
            'success_count': success_count,
            'average_payoff': avg_payoff,
            'insurance_cost': insurance_cost,
            'return_multiplier': return_multiplier,
            'choice_multiply_insurance_cost': insurance_cost*player.choice,
            'endowment': C.ENDOWMENT,
            'belief_elicitation': player.belief_elicitation_treatment,
            # Pre-calculated values for outcome display
            'investment_return': investment_return if player.framing_treatment == 'gain' else None,
            'kept_amount': kept_amount if player.framing_treatment == 'gain' else None,
            'insurance_premium': insurance_premium if player.framing_treatment == 'loss' else None,
            'covered_amount': covered_amount if player.framing_treatment == 'loss' else None,
            'uncovered_amount': uncovered_amount if player.framing_treatment == 'loss' else None,
        }

class PracticeComplete(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == C.NUM_PRACTICE_ROUNDS
    
    @staticmethod
    def vars_for_template(player: Player):
        # Get practice round results
        practice_history = []
        for r in range(1, C.NUM_PRACTICE_ROUNDS + 1):
            past = player.in_round(r)
            practice_history.append({
                'round': r,
                'choice': past.choice,
                'outcome': past.outcome,
                'payoff': past.payoff_tokens,
            })
        
        return {
            'practice_history': practice_history,
            'framing': player.framing_treatment,
            'num_practice_rounds': C.NUM_PRACTICE_ROUNDS,
        }


class FinalQuestion(Page):
    form_model = 'player'
    form_fields = ['comments']
    
    @staticmethod
    def is_displayed(player: Player):
        #return True
        return player.round_number == C.NUM_ROUNDS
    
    @staticmethod
    def vars_for_template(player: Player):
        # Get only real rounds (not practice)
        real_rounds = []
        for r in range(C.NUM_PRACTICE_ROUNDS + 1, C.NUM_ROUNDS + 1):
            real_rounds.append(player.in_round(r))
        
        # Select random rounds for payment
        #selected_rounds = random.sample(real_rounds, C.NUM_ROUNDS_FOR_PAYMENT)
        #selected_rounds_data = []
        #total_tokens = 0
        
        #for round_obj in selected_rounds:
        #    selected_rounds_data.append({
        #        'round': real_round_number(round_obj),
        #        'choice': round_obj.choice,
        #        'outcome': round_obj.outcome,
        #        'payoff': round_obj.payoff_tokens,
        #    })
        #    total_tokens += round_obj.payoff_tokens
        
       # final_payment = total_tokens / C.TOKENS_PER_DOLLAR
        
        # Calculate summary statistics
        #total_successes = sum(1 for r in real_rounds if r.outcome)
        #average_choice = sum(r.choice for r in real_rounds) / len(real_rounds)
        #average_tokens = sum(r.payoff_tokens for r in real_rounds) / len(real_rounds)
        #conf_num = random.randint(10**9, (10**10) - 1)
        #return {
        #    'selected_rounds_data': selected_rounds_data,
        #    'total_tokens': total_tokens,
        #    'final_payment': final_payment,
        #    'framing': player.framing_treatment,
        #    'total_successes': total_successes,
        #    'average_choice': average_choice,
        #    'average_tokens': average_tokens,
        #    'conf_num': conf_num,
        #}


class FinalPayment(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == C.NUM_ROUNDS
    
    @staticmethod
    def vars_for_template(player: Player):
        # Get only real rounds (not practice)
        real_rounds = []
        for r in range(C.NUM_PRACTICE_ROUNDS + 1, C.NUM_ROUNDS + 1):
            real_rounds.append(player.in_round(r))
        
        # Select random rounds for payment
        selected_rounds = random.sample(real_rounds, C.NUM_ROUNDS_FOR_PAYMENT)
        selected_rounds_data = []
        total_tokens = 0
        
        for round_obj in selected_rounds:
            selected_rounds_data.append({
                'round': real_round_number(round_obj),
                'choice': round_obj.choice,
                'outcome': round_obj.outcome,
                'payoff': round_obj.payoff_tokens,
            })
            total_tokens += round_obj.payoff_tokens
        
        final_payment = total_tokens / C.TOKENS_PER_DOLLAR
        
        # Calculate summary statistics
        total_successes = sum(1 for r in real_rounds if r.outcome)
        average_choice = sum(r.choice for r in real_rounds) / len(real_rounds)
        average_tokens = sum(r.payoff_tokens for r in real_rounds) / len(real_rounds)
        conf_num = random.randint(10**9, (10**10) - 1)
        return {
            'selected_rounds_data': selected_rounds_data,
            'total_tokens': total_tokens,
            'final_payment': final_payment,
            'framing': player.framing_treatment,
            'total_successes': total_successes,
            'average_choice': average_choice,
            'average_tokens': average_tokens,
            'conf_num': conf_num,
        }

page_sequence = [
    Introduction,
    ComprehensionCheck,
    Decision,
    Outcome,
    BeliefElicitation,
    PracticeComplete,
    FinalQuestion,
    FinalPayment,
]