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

doc = """
Exp2 (Player A):
Consent -> Instructions -> Comprehension -> ShowMessage -> IN/OUT -> Beliefs -> Demographics -> Thank you
Message is drawn from message_bank.csv; A-switch is matched to B-switch.
Belief elicitation uses Vanberg 5-point scale and scoring table (in points).
"""


class C(BaseConstants):
    NAME_IN_URL = 'exp2'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 8   # change to 10 later if you want 10 cases per Player A


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # Consent
    consent = models.BooleanField(
        label='I am at least 18 years old and I consent to participate in this study.'
    )

    # -----------------------------
    # Comprehension questions (Player A)
    # -----------------------------
    cq1 = models.StringField(
        label="1) Are messages binding in this experiment?",
        choices=[
            ['binding', 'Yes, messages are binding'],
            ['not_binding', 'No, messages are not binding'],
        ],
        widget=widgets.RadioSelect
    )

    cq2 = models.StringField(
        label="2) Will you be informed whether a switch occurred?",
        choices=[
            ['Yes', 'Yes, you will be informed whether a switch occurred'],
            ['No', 'No, you will not know whether a switch occurred or not'],
        ],
        widget=widgets.RadioSelect
    )

    cq3 = models.StringField(
        label="3) If you choose IN and Player B chooses ROLL, how much can you receive?",
        choices=[
            ['0_or_12', "0 (with probability 1/6) or 12 (with probability 5/6)"],
            ['5', "5"],
            ['14', "14"],
        ],
        widget=widgets.RadioSelect
    )

    cq4 = models.StringField(
        label="4) Which choice gives you a guaranteed payoff above 0?",
        choices=[
            ['in', 'Choose IN'],
            ['out', 'Choose OUT'],
        ],
        widget=widgets.RadioSelect
    )

    # -----------------------------
    # Message info from message_bank.csv
    # -----------------------------
    msg_id = models.StringField(initial='')
    a_switch = models.IntegerField(initial=0)          # from b_switch column
    assigned_message = models.LongStringField(initial='')
    b_choice = models.StringField(initial='')          # 'roll' or 'dont_roll'

    # optional, only for analysis later
    b_load = models.IntegerField(initial=0)
    message_original = models.LongStringField(initial='')

    # -----------------------------
    # Decision + outcomes
    # -----------------------------
    in_out = models.StringField(
        choices=[['in', 'IN'], ['out', 'OUT']],
        widget=widgets.RadioSelect
    )

    die_roll = models.IntegerField(initial=0)          # only if IN and b_choice='roll'

    points_main = models.FloatField(initial=0)
    points_belief = models.FloatField(initial=0)
    points_total = models.FloatField(initial=0)

    # -----------------------------
    # First-order belief (Vanberg-style 5-point scale)
    # -----------------------------
    expect_roll_cat = models.StringField(
        label="Please select one option. I think that Player B...",
        choices=[
            ['c_roll', 'Certainly chose ROLL'],
            ['p_roll', 'Probably chose ROLL'],
            ['unsure', 'Unsure'],
            ['p_dont', "Probably chose DON’T ROLL"],
            ['c_dont', "Certainly chose DON’T ROLL"],
        ],
        widget=widgets.RadioSelect
    )

    # Demographics
    age = models.IntegerField(label="Age:", min=18, max=120, blank=True)

    gender = models.StringField(
        label="Gender:",
        choices=[['female', 'Female'],
                 ['male', 'Male'],
                 ['nonbinary', 'Non-binary / third gender'],
                 ['prefer_not', 'Prefer not to say']],
        blank=True
    )

    education = models.StringField(
        label="Highest completed education:",
        choices=[['hs', 'High school or equivalent'],
                 ['ba', "Bachelor's degree"],
                 ['ma', "Master's degree"],
                 ['phd', "PhD / doctorate"],
                 ['other', 'Other'],
                 ['prefer_not', 'Prefer not to say']],
        blank=True
    )


# --------- Belief scoring (Vanberg table in points) ---------

def vanberg_first_order_points(choice: str, actual_roll: bool) -> float:
    pay_if_roll = {
        'c_roll': 6.5,
        'p_roll': 6,
        'unsure': 5,
        'p_dont': 3.5,
        'c_dont': 1.5,
    }
    pay_if_dont = {
        'c_roll': 1.5,
        'p_roll': 3.5,
        'unsure': 5,
        'p_dont': 6,
        'c_dont': 6.5,
    }

    return (pay_if_roll if actual_roll else pay_if_dont)[choice]


# --------- CSV helpers ---------

def _load_message_bank():
    """
    message_bank.csv must be semicolon-separated with header:
    msg_id;b_load;b_switch;message_original;message_for_a;b_choice
    """
    csv_path = Path(__file__).resolve().parent / 'message_bank.csv'
    rows = []

    with csv_path.open('r', encoding='utf-8-sig', newline='') as f:
        reader = csv.DictReader(f, delimiter=';')
        expected = ['msg_id', 'b_load', 'b_switch', 'message_original', 'message_for_a', 'b_choice']

        if reader.fieldnames != expected:
            raise Exception(
                "message_bank.csv headers are wrong. "
                f"Expected: {';'.join(expected)} but got: {reader.fieldnames}"
            )

        for r in reader:
            msg_id = (r.get('msg_id') or '').strip()
            b_load = (r.get('b_load') or '0').strip()
            b_switch = (r.get('b_switch') or '0').strip()
            message_original = (r.get('message_original') or '').strip()
            message_for_a = (r.get('message_for_a') or '').strip()
            b_choice = (r.get('b_choice') or '').strip().lower()

            if not msg_id or not message_for_a:
                continue
            if b_choice not in ['roll', 'dont_roll']:
                continue

            rows.append({
                'msg_id': msg_id,
                'b_load': int(b_load),
                'b_switch': int(b_switch),
                'message_original': message_original,
                'message_for_a': message_for_a,
                'b_choice': b_choice,
            })

    return rows

def creating_session(subsession: Subsession):
    if subsession.round_number != 1:
        return

    players = subsession.get_players()
    rows = _load_message_bank()

    needed = len(players) * C.NUM_ROUNDS
    if len(rows) < needed:
        raise Exception(
            f"Not enough rows in message_bank.csv. "
            f"You need at least {needed}, but only have {len(rows)}."
        )

    random.shuffle(rows)

    for i, p in enumerate(players):
        start = i * C.NUM_ROUNDS
        end = start + C.NUM_ROUNDS
        assigned_rows = rows[start:end]

        for round_number, row in enumerate(assigned_rows, start=1):
            pr = p.in_round(round_number)
            pr.msg_id = row['msg_id']
            pr.b_load = row['b_load']
            pr.a_switch = row['b_switch']
            pr.message_original = row['message_original']
            pr.b_choice = row['b_choice']

            if row['b_switch'] == 1:
                pr.assigned_message = row['message_for_a']
            else:
                pr.assigned_message = row['message_original']
# --------- PAGES ---------

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

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

    @staticmethod
    def error_message(player: Player, values):
        if not values.get('consent'):
            return "You must provide consent to participate."


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


class Comprehension(Page):
    form_model = 'player'
    form_fields = ['cq1', 'cq2', 'cq3', 'cq4']

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


class ShowMessage(Page):
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            assigned_message=player.assigned_message,
            msg_id=player.msg_id,
            a_switch=player.a_switch,
            current_round=player.round_number,
            total_rounds=C.NUM_ROUNDS,
        )


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

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        if player.in_out == 'out':
            player.points_main = 5
            player.die_roll = 0
        else:
            if player.b_choice == 'dont_roll':
                player.points_main = 0
                player.die_roll = 0
            else:
                player.die_roll = random.randint(1, 6)
                player.points_main = 0 if player.die_roll == 1 else 12


class Beliefs(Page):
    form_model = 'player'
    form_fields = ['expect_roll_cat']

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        actual_roll = (player.b_choice == 'roll')
        player.points_belief = vanberg_first_order_points(
            choice=player.expect_roll_cat,
            actual_roll=actual_roll
        )
        player.points_total = player.points_main + player.points_belief


class Demographics(Page):
    form_model = 'player'
    form_fields = ['age', 'gender', 'education']

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


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

    @staticmethod
    def vars_for_template(player: Player):
        all_rounds = player.in_all_rounds()
        return dict(
            points_main=sum(p.points_main for p in all_rounds),
            points_belief=sum(p.points_belief for p in all_rounds),
            points_total=sum(p.points_total for p in all_rounds),
            die_roll=player.die_roll,
        )


page_sequence = [
    Consent, Instructions, Comprehension, ShowMessage,
    Decision, Beliefs, Demographics, ThankYou
]