from otree.api import *
import csv
import random
from pathlib import Path


class C(BaseConstants):
    NAME_IN_URL = 'age_task'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 161 # 160正式 + 1练习
    PRACTICE_ROUNDS = 1
    TASK_TIMER = 60 # 设定 60 秒答题限制


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # 新增：用于记录被试是否同意
    consent = models.BooleanField(
        label="Do you consent to participate in this study?",
        choices=[
            [True, 'Yes, I consent to participate.'],
            [False, 'No, I do not consent.']
        ],
        widget=widgets.RadioSelect
    )
    # 改用 SliderInput 可以强制在没写自定义HTML时也有个滑块，但我们会用自定义HTML做得更漂亮
    initial_belief = models.IntegerField(
        min=0, max=100, 
        blank=True, # 允许超时为空
        label="What is the probability (0-100%) that this person is OVER 21 years old?"
    )
    final_belief = models.IntegerField(
        min=0, max=100,
        blank=True, # 允许超时为空
        label="What is your FINAL probability (0-100%) that this person is OVER 21 years old?"
    )

    # 记录该轮是否因为超时而被强制跳过
    timeout_occurred = models.BooleanField(initial=False)
    # 记录抽奖结果
    timer_draw = models.IntegerField(blank=True)
    bonus_earned = models.BooleanField(initial=False)
    folder = models.StringField()  
    image_id = models.StringField()
    truth = models.BooleanField()
    ai_prob = models.FloatField()


# ----------------------------------
# SESSION CREATION (分配图片、随机打乱)
# ----------------------------------
def creating_session(subsession):
    if subsession.round_number == 1:
        # ⚠️ 注意：这里的文件名请确保和你文件夹里的一致
        main_path = Path(__file__).parent / "image_data.csv" 
        practice_path = Path(__file__).parent / "practice_data.csv"

        # 【核心修复 1】：读取 CSV 时，直接过滤掉所有 truth 为空的无效行（比如文件末尾的回车空行）
        with open(main_path, newline='', encoding='utf-8-sig') as csvfile:
            reader = csv.DictReader(csvfile)
            subsession.session.vars['main_images'] = [row for row in reader if row.get('truth', '').strip() != '']

        with open(practice_path, newline='', encoding='utf-8-sig') as csvfile:
            reader = csv.DictReader(csvfile)
            subsession.session.vars['practice_images'] = [row for row in reader if row.get('truth', '').strip() != '']

        # 为每个被试生成独立的、随机打乱的图片读取顺序
        for p in subsession.get_players():
            participant = p.participant
            indices = list(range(len(subsession.session.vars['main_images'])))
            random.shuffle(indices) # 打乱顺序
            participant.vars['task_order'] = indices

    for p in subsession.get_players():
        participant = p.participant
        if subsession.round_number <= C.PRACTICE_ROUNDS:
            row = subsession.session.vars['practice_images'][subsession.round_number - 1]
        else:
            # 根据打乱后的索引来抽取图片
            current_task_idx = subsession.round_number - C.PRACTICE_ROUNDS - 1
            shuffled_idx = participant.vars['task_order'][current_task_idx]
            row = subsession.session.vars['main_images'][shuffled_idx]

        # 【核心修复 2】：使用保险的读取方式，防止崩溃
        p.folder = str(row.get('folder', '')).zfill(2)
        p.image_id = str(row.get('image_id', ''))
        
        truth_val = str(row.get('truth', '')).strip()
        if truth_val != '':
            p.truth = bool(int(truth_val))
            
        ai_prob_val = str(row.get('ai_prob', '')).strip()
        if ai_prob_val != '':
            p.ai_prob = float(ai_prob_val)
# ----------------------------------
# PAGES (网页渲染逻辑)
# ----------------------------------

class Consent(Page):
    form_model = 'player'
    form_fields = ['consent']

    @staticmethod
    def is_displayed(player):
        return player.round_number == 1
    @staticmethod
    def error_message(player, values):
        if values['consent'] is False:
            return "If you do not consent to participate, you cannot proceed. Please close this browser window and return the study on Prolific."
class Introduction(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1


class AIAdvice(Page):
    @staticmethod
    def is_displayed(player):
        return not player.timeout_occurred # 如果上一页超时了，这一页直接跳过（或者你可以删掉这句强制显示）

    @staticmethod
    def vars_for_template(player):
        return dict(
            image_path=f"wiki_crop/{player.folder}/{player.image_id}",
            initial=player.initial_belief,
            ai_prob_pct=int(player.ai_prob * 100),
            practice=(player.round_number <= C.PRACTICE_ROUNDS)
        )

class InitialBelief(Page):
    form_model = 'player'
    form_fields = ['initial_belief']
    
    @staticmethod
    def vars_for_template(player):
        return dict(
            image_path=f"wiki_crop/{player.folder}/{player.image_id}",
            practice=(player.round_number <= C.PRACTICE_ROUNDS)
        )

class FinalBelief(Page):
    form_model = 'player'
    form_fields = ['final_belief']

    @staticmethod
    def vars_for_template(player):
        return dict(
            image_path=f"wiki_crop/{player.folder}/{player.image_id}",
            initial=player.initial_belief,
            ai_prob_pct=int(player.ai_prob * 100),
            practice=(player.round_number <= C.PRACTICE_ROUNDS)
        )

class Feedback(Page):
    @staticmethod
    def vars_for_template(player):
        return dict(
            image_path=f"wiki_crop/{player.folder}/{player.image_id}",
            initial=player.initial_belief,
            final=player.final_belief,
            ai_prob_pct=int(player.ai_prob * 100),
            truth="OVER 21" if player.truth else "21 OR UNDER",
            practice=(player.round_number <= C.PRACTICE_ROUNDS)
        )

class PracticeTransition(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == C.PRACTICE_ROUNDS

# ----------------------------------
# 结算与抽奖逻辑
# ----------------------------------
def custom_export(players):
    # 这里用于未来导出完整数据，暂留空
    pass

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

    @staticmethod
    def vars_for_template(player):
        # 1. 获取该玩家过去所有正式轮次的数据
        past_players = player.in_all_rounds()[C.PRACTICE_ROUNDS:] 
        
        # 2. 随机抽取一轮作为支付轮
        import random
        selected_round_player = random.choice(past_players)
        
        # 3. 提取关键数据
        is_over_21 = selected_round_player.truth 
        # 获取最终信念（如果有极端情况没抓到值，保底设为50）
        final_belief_val = selected_round_player.final_belief if selected_round_player.final_belief is not None else 50
        reported_prob_over = final_belief_val / 100.0 
        
        # 4. 二次计分规则 (Quadratic Scoring Rule)
        if is_over_21:
            win_chance = 1.0 - (1.0 - reported_prob_over)**2
        else:
            win_chance = 1.0 - reported_prob_over**2
        
        win_chance_pct = round(win_chance * 100, 2)
        
        # 5. 拼接那张被选中的图片的路径 (这是解决不显示图片的关键)
        chosen_image_path = f"wiki_crop/{selected_round_player.folder}/{selected_round_player.image_id}"
        
        # 将结果存入 participant 变量，方便下一页抽奖用
        player.participant.vars['payment_win_chance'] = win_chance_pct

        return dict(
            chosen_image_path=chosen_image_path,
            truth_text="OVER 21" if is_over_21 else "UNDER 21",
            reported_prob_over=final_belief_val,
            reported_prob_under=100 - final_belief_val,
            win_chance_pct=win_chance_pct
        )

class BonusPayment(Page):
    form_model = 'player'
    form_fields = ['timer_draw'] 

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

    @staticmethod
    def vars_for_template(player):
        win_chance_pct = player.participant.vars.get('payment_win_chance', 0)
        return dict(win_chance_pct=win_chance_pct)

    @staticmethod
    def before_next_page(player, timeout_happened):
        win_chance_pct = player.participant.vars.get('payment_win_chance', 0)
        # 如果 timer_draw 有值且小于中奖概率，则记录赢了 bonus
        if player.timer_draw is not None and player.timer_draw < win_chance_pct:
            player.bonus_earned = True
        else:
            player.bonus_earned = False

# 新增：感谢与事后说明页面
class Debrief(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS

# 记得把 Debrief 加进 page_sequence 的最后！
page_sequence = [
    Consent, Introduction, InitialBelief, AIAdvice, 
    FinalBelief, Feedback, PracticeTransition, MainResults, BonusPayment, Debrief
]