from otree.api import *
import json
import random

doc = """
Plinko Pilot: Confirmation Bias or Joint Inference?
8 trials (4 T1 known precision, 4 T2 unknown precision).
Within-subject: T1 block then T2 block, randomized within each block.
"""

R = 500  # rows on the board for trials


def simulate_plinko(theta, n, R):
    """Simulate n balls on an R-row board from position theta.
    Each row the ball moves +1 or -1 with equal probability.
    Returns list of landing positions."""
    landings = []
    for _ in range(n):
        pos = theta
        for _ in range(R):
            pos += random.choice([-1, 1])
        landings.append(pos)
    return landings


def generate_trial(n_i, n_j, treatment):
    """Generate one trial. Redraws if signals too close (d < 3) or out of bounds."""
    while True:
        theta = random.randint(150, 850)
        landings_i = simulate_plinko(theta, n_i, R)
        landings_j = simulate_plinko(theta, n_j, R)
        x_i = round(sum(landings_i) / len(landings_i))
        x_j = round(sum(landings_j) / len(landings_j))
        if 0 <= x_i <= 1000 and 0 <= x_j <= 1000 and abs(x_j - x_i) >= 3:
            return dict(
                theta=theta,
                n_i=n_i,
                n_j=n_j,
                x_i=x_i,
                x_j=x_j,
                d=abs(x_j - x_i),
                treatment=treatment,
                landings_i=landings_i,
                landings_j=landings_j,
            )


def generate_trials_for_participant():
    """Generate 8 trials: 4 T1 (known precision), 4 T2 (unknown precision).
    One trial per (n_i, n_j) cell in each treatment block.
    T1 block first, T2 second, randomized within each."""
    N_LO, N_HI = 1, 10
    cells = [(N_LO, N_LO), (N_LO, N_HI), (N_HI, N_LO), (N_HI, N_HI)]

    t1 = [generate_trial(ni, nj, 'T1') for ni, nj in cells]
    t2 = [generate_trial(ni, nj, 'T2') for ni, nj in cells]

    random.shuffle(t1)
    random.shuffle(t2)

    return t1 + t2


class C(BaseConstants):
    NAME_IN_URL = 'plinko'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 8
    R = 500
    R_DEMO = 50
    N_J_LO = 1
    N_J_HI = 10
    PER_TRIAL_ENDOWMENT = 100
    SE_SCALING = 4


class Subsession(BaseSubsession):
    pass


def creating_session(subsession):
    if subsession.round_number == 1:
        for player in subsession.get_players():
            trials = generate_trials_for_participant()
            player.participant.vars['trial_order'] = list(range(8))
            player.participant.vars['trials'] = trials


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # Registration (round 1 only)
    student_name = models.StringField(label="Your name")
    student_id = models.StringField(label="Your student ID")

    # Trial parameters
    theta = models.IntegerField()
    n_i = models.IntegerField()
    n_j = models.IntegerField()
    x_i = models.IntegerField()
    x_j = models.IntegerField()
    d = models.IntegerField()
    treatment = models.StringField()

    # Response
    theta_hat = models.FloatField(
        label="What is your best estimate of the drop position?",
        min=0,
        max=1000,
    )

    # Computed
    squared_error = models.FloatField()
    trial_payment = models.FloatField()

    # Post-survey (last round only)
    survey_unclear = models.LongStringField(
        label="Was anything in the experiment unclear or confusing?"
    )
    survey_strategy = models.LongStringField(
        label="What strategy did you use to solve these problems?"
    )

    # Comprehension (round 1 only)
    comp_check_1 = models.IntegerField(
        label=(
            "If a ball is dropped from position 500 on a 500-row board, "
            "what is the MOST LIKELY landing position?"
        ),
        choices=[0, 250, 500, 750, 1000],
    )
    comp_check_2 = models.IntegerField(
        label=(
            "If you see 10 balls land tightly clustered around position 300, "
            "and 1 ball lands at position 320, "
            "your best estimate of the drop position should be:"
        ),
        choices=[
            [1, "Closer to 300 (the 10-ball signal)"],
            [2, "Exactly 310 (the midpoint)"],
            [3, "Closer to 320 (the 1-ball signal)"],
        ],
    )


def set_trial_params(player):
    """Load this round's trial parameters into player fields."""
    idx = player.participant.vars['trial_order'][player.round_number - 1]
    trial = player.participant.vars['trials'][idx]
    player.theta = trial['theta']
    player.n_i = trial['n_i']
    player.n_j = trial['n_j']
    player.x_i = trial['x_i']
    player.x_j = trial['x_j']
    player.d = trial['d']
    player.treatment = trial['treatment']
    player.participant.vars['current_landings_i'] = trial['landings_i']


# --- Pages ---

class Registration(Page):
    form_model = 'player'
    form_fields = ['student_name', 'student_id']

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


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


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

    @staticmethod
    def vars_for_template(player):
        return dict(R=C.R_DEMO, R_TRIAL=C.R)


class ComprehensionCheck(Page):
    form_model = 'player'
    form_fields = ['comp_check_1', 'comp_check_2']

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

    @staticmethod
    def error_message(player, values):
        errors = {}
        if values['comp_check_1'] != 500:
            errors['comp_check_1'] = (
                'Not quite. The ball bounces left and right equally, '
                'so on average it lands at the drop position.'
            )
        if values['comp_check_2'] != 1:
            errors['comp_check_2'] = (
                'Not quite. The 10-ball signal is much more precise, '
                'so your estimate should lean toward that signal.'
            )
        return errors


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

    @staticmethod
    def vars_for_template(player):
        return dict(
            R=C.R, R_DEMO=C.R_DEMO,
            N_J_LO=C.N_J_LO, N_J_HI=C.N_J_HI,
            NUM_ROUNDS=C.NUM_ROUNDS,
            PER_TRIAL=C.PER_TRIAL_ENDOWMENT,
        )


class T2Transition(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 5

    @staticmethod
    def vars_for_template(player):
        return dict(N_J_LO=C.N_J_LO, N_J_HI=C.N_J_HI, PER_TRIAL=C.PER_TRIAL_ENDOWMENT)


class Trial(Page):
    form_model = 'player'
    form_fields = ['theta_hat']

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.squared_error = (player.theta_hat - player.theta) ** 2
        player.trial_payment = round(
            max(0, C.PER_TRIAL_ENDOWMENT - player.squared_error / C.SE_SCALING), 2
        )

    @staticmethod
    def vars_for_template(player):
        set_trial_params(player)

        landings_i = player.participant.vars['current_landings_i']

        if player.round_number <= 4:
            block_trial = player.round_number
        else:
            block_trial = player.round_number - 4
        block_total = 4

        return dict(
            round_number=player.round_number,
            total_rounds=C.NUM_ROUNDS,
            block_trial=block_trial,
            block_total=block_total,
            treatment=player.treatment,
            n_i=player.n_i,
            n_j=player.n_j,
            x_i=player.x_i,
            x_j=player.x_j,
            d=player.d,
            R=C.R,
            N_J_LO=C.N_J_LO,
            N_J_HI=C.N_J_HI,
            landings_i_json=json.dumps(landings_i),
        )


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

    @staticmethod
    def vars_for_template(player):
        all_players = player.in_all_rounds()
        total_payment = sum(
            p.trial_payment for p in all_players if p.trial_payment is not None
        )
        return dict(
            total_payment=int(total_payment),
            per_trial_endowment=int(C.PER_TRIAL_ENDOWMENT),
            rounds=all_players,
        )


class PostSurvey(Page):
    form_model = 'player'
    form_fields = ['survey_unclear', 'survey_strategy']

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


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


page_sequence = [
    Registration,
    Introduction,
    PlinkoDemo,
    ComprehensionCheck,
    TaskInstructions,
    T2Transition,
    Trial,
    Results,
    PostSurvey,
    ThankYou,
]