from otree.api import *
from extdef import (
    N_MEMBERS, N_BLOCKS, N_ROUNDS_PER_BLOCK, TOTAL_ROUNDS,
    CTR_POINTS, NOISE_VALS, POLICIES, VENMO_ACCOUNT, NOISE_RANGE,
    POINTS_PER_DOLLAR, PARTICIPATION_FEE,
)

doc = 'Individual Choice: each voter independently chooses between two policy options.'


class C(BaseConstants):
    NAME_IN_URL = 'individual_choice'
    PLAYERS_PER_GROUP = N_MEMBERS   # 9 — kept for session structure; no group interaction
    NUM_ROUNDS = TOTAL_ROUNDS       # 180


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    policy_A = models.IntegerField()
    policy_B = models.IntegerField()


class Player(BasePlayer):
    ctr_point    = models.IntegerField()
    ideal_point  = models.IntegerField()
    vote         = models.StringField()   # 'A' or 'B'
    chosen_policy = models.IntegerField() # value of the chosen option
    round_payoff = models.IntegerField()


# ── Helpers ────────────────────────────────────────────────────────

def _get_indices(player):
    """Return (block, round_in_block, group_idx, player_pos) — all 0-based."""
    r = player.round_number - 1
    block = r // N_ROUNDS_PER_BLOCK
    round_in_block = r % N_ROUNDS_PER_BLOCK
    group_idx = (player.group.id_in_subsession - 1) % len(CTR_POINTS)
    player_pos = player.id_in_group - 1
    return block, round_in_block, group_idx, player_pos


def _round_values(player):
    """Compute all round-specific values directly from extdef — never reads DB fields."""
    block, rib, gidx, ppos = _get_indices(player)
    ctr_point   = CTR_POINTS[gidx][block][ppos]
    ideal_point = ctr_point + NOISE_VALS[gidx][block][rib][ppos]
    policy_A, policy_B = POLICIES[gidx][block][rib]
    return block, rib, ctr_point, ideal_point, policy_A, policy_B



# ── Pages ──────────────────────────────────────────────────────────

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

    @staticmethod
    def vars_for_template(player):
        return {
            'n_blocks':        N_BLOCKS,
            'n_rounds':        N_ROUNDS_PER_BLOCK,
            'total_rounds':    TOTAL_ROUNDS,
            'n_members':       N_MEMBERS,
            'points_per_dollar': POINTS_PER_DOLLAR,
            'participation_fee': f'{PARTICIPATION_FEE:.2f}',
        }


class BlockIntro(Page):
    @staticmethod
    def is_displayed(player):
        return (player.round_number - 1) % N_ROUNDS_PER_BLOCK == 0

    @staticmethod
    def vars_for_template(player):
        block, _, ctr_point, _, _, _ = _round_values(player)
        return {
            'block_num': block + 1,
            'n_blocks':  N_BLOCKS,
            'n_rounds':  N_ROUNDS_PER_BLOCK,
            'ctr_point': ctr_point,
            'range_lo':  ctr_point - NOISE_RANGE,
            'range_hi':  ctr_point + NOISE_RANGE,
        }


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

    @staticmethod
    def vars_for_template(player):
        block, rib, ctr_point, ideal_point, policy_A, policy_B = _round_values(player)
        return {
            'block_num':      block + 1,
            'round_in_block': rib + 1,
            'ctr_point':      ctr_point,
            'range_lo':       ctr_point - NOISE_RANGE,
            'range_hi':       ctr_point + NOISE_RANGE,
            'ideal_point':    ideal_point,
            'policy_A':       policy_A,
            'policy_B':       policy_B,
        }

    @staticmethod
    def before_next_page(player, timeout_happened):
        block, rib, ctr_point, ideal_point, policy_A, policy_B = _round_values(player)
        player.ctr_point      = ctr_point
        player.ideal_point    = ideal_point
        player.group.policy_A = policy_A
        player.group.policy_B = policy_B
        player.chosen_policy  = policy_A if player.vote == 'A' else policy_B
        player.round_payoff   = 100 - abs(player.chosen_policy - player.ideal_point)
        player.payoff         = cu(player.round_payoff)


class Results(Page):
    @staticmethod
    def vars_for_template(player):
        block, rib, ctr_point, ideal_point, policy_A, policy_B = _round_values(player)
        return {
            'block_num':      block + 1,
            'round_in_block': rib + 1,
            'ctr_point':      ctr_point,
            'range_lo':       ctr_point - NOISE_RANGE,
            'range_hi':       ctr_point + NOISE_RANGE,
            'ideal_point':    ideal_point,
            'policy_A':       policy_A,
            'policy_B':       policy_B,
            'your_vote':      player.vote,
            'chosen_policy':  player.chosen_policy,
            'round_payoff':   player.round_payoff,
        }


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

    @staticmethod
    def vars_for_template(player):
        total_pts = sum(
            p.round_payoff for p in player.in_all_rounds()
            if p.round_payoff is not None
        )
        earnings_usd = total_pts / POINTS_PER_DOLLAR
        total_usd    = earnings_usd + PARTICIPATION_FEE
        venmo_url = (
            f'https://venmo.com/{VENMO_ACCOUNT}'
            f'?txn=charge&amount={total_usd:.2f}&note=Experiment+payment'
        )
        return {
            'total_pts':        total_pts,
            'earnings_usd':     f'{earnings_usd:.2f}',
            'participation_fee': f'{PARTICIPATION_FEE:.2f}',
            'total_usd':        f'{total_usd:.2f}',
            'venmo_url':        venmo_url,
        }


page_sequence = [Instructions, BlockIntro, Decision, Results, FinalResults]