from otree.api import (
    cu,
    currency_range,
    models,
    widgets,
    BaseConstants,
    BaseSubsession,
    BaseGroup,
    BasePlayer,
    ExtraModel,
    WaitPage,
    Page,
)
import random


doc = ''


class C(BaseConstants):
    NAME_IN_URL = 'pd_game'
    PLAYERS_PER_GROUP = None  # keep None, we handle grouping manually
    NUM_ROUNDS = 16

    PRE_TREATMENT_ROUNDS = 5
    TREATMENT_START = 6
    TREATMENT_END = 10
    VOTING_ROUND = 11
    POST_VOTING_START = 12

    PAYOFF_CC = 10
    PAYOFF_CD = 0
    PAYOFF_DC = 15
    PAYOFF_DD = 5

    HARMONY_TAX = 7

    MORAL_MESSAGE = "Don't do to others what you wouldn't want done to you"

    GROUP_SIZE_PREFERRED = 6
    GROUP_SIZE_FALLBACK = 4


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    winning_game = models.StringField(initial='A')
    is_treated = models.BooleanField(initial=False)


class Player(BasePlayer):
    # ---- Sequential PD fields (used in all regular rounds) ----
    # As First Mover
    fm_choice = models.BooleanField(
        label='Your action if you play first',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )
    # Frequency belief: how many of the OTHER players in the group will do Action A as FM?
    # Stored as an integer count (0 .. group_size - 1)
    belief_count = models.IntegerField(
        label='How many of the other players in your group do you think will do Action A?',
        min=0
    )
    # As Second Mover, contingent on FM's action
    sm_choice_if_fm_coop = models.BooleanField(
        label='Your action as Second Player if the First Player does Action A',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )
    sm_choice_if_fm_defect = models.BooleanField(
        label='Your action as Second Player if the First Player does Action B',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )

    # Assigned role and resolved choices
    is_first_mover = models.BooleanField(initial=False)
    # Actual choices that were implemented (resolved after role assignment)
    actual_choice = models.BooleanField(initial=True)

    partner_id_in_group = models.IntegerField()
    round_payoff = models.CurrencyField(initial=cu(0))

    vote = models.StringField(
        label='Which game do you vote for?',
        choices=[
            ['A', 'Game 1'],
            ['B', 'Game 2'],
            ['C', 'Game 3']
        ]
    )

    # Voting round: sequential PD strategy + frequency belief per option
    # Option A
    fm_choice_a = models.BooleanField(
        label='[Option A] Your action as First Player',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )
    belief_count_a = models.IntegerField(
        label='[Option A] How many of the other players do you think will do Action A?',
        min=0
    )
    sm_coop_a = models.BooleanField(
        label='[Option A] Your action as Second Player if the First Player does Action A',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )
    sm_defect_a = models.BooleanField(
        label='[Option A] Your action as Second Player if The First Player does Action B',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )

    # Option B
    fm_choice_b = models.BooleanField(
        label='[Option B] Your action as First Player',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )
    belief_count_b = models.IntegerField(
        label='[Option B] How many of the other players do you think will do Action A?',
        min=0
    )
    sm_coop_b = models.BooleanField(
        label='[Option B] Your action as Second Player if the First Player does Action A',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )
    sm_defect_b = models.BooleanField(
        label='[Option B] Your action as Second Player if the First Player does Action B',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )

    # Option C
    fm_choice_c = models.BooleanField(
        label='[Option C] Your action as First Player',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )
    belief_count_c = models.IntegerField(
        label='[Option C] How many of the other players do you think will do Action A?',
        min=0
    )
    sm_coop_c = models.BooleanField(
        label='[Option C] Your action as Second Player if the First Player does Action A',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )
    sm_defect_c = models.BooleanField(
        label='[Option C] Your action as Second Player if the First Player does Action B',
        choices=[[True, 'Action A'], [False, 'Action B']]
    )


# ---------- Helper functions ----------
def game_label(game_code):
    labels = {
        'A': 'Game 1',
        'B': 'Game 2',
        'C': 'Game 3',
    }
    return labels.get(game_code, game_code)


def is_treatment_phase(round_number):
    return C.TREATMENT_START <= round_number <= C.TREATMENT_END


def active_game_code(group: Group):
    if group.round_number <= C.PRE_TREATMENT_ROUNDS:
        return 'A'
    if is_treatment_phase(group.round_number):
        return 'B' if group.is_treated else 'A'
    if group.round_number == C.VOTING_ROUND:
        return 'A'
    if group.round_number >= C.POST_VOTING_START:
        stored = group.session.vars.get('winning_game_by_group', {})
        return stored.get(group.id_in_subsession, 'A')
    return 'A'


def show_moral_message(group: Group):
    game_code = active_game_code(group)
    return game_code == 'B'


def payoff_matrix_rows(game_code):
    rows = [
        dict(player_action='Action A', partner_action='Action A', payoff=cu(C.PAYOFF_CC), partner_payoff=cu(C.PAYOFF_CC)),
        dict(player_action='Action A', partner_action='Action B',    payoff=cu(C.PAYOFF_CD), partner_payoff=cu(C.PAYOFF_DC)),
        dict(player_action='Action B',    partner_action='Action A', payoff=cu(C.PAYOFF_DC), partner_payoff=cu(C.PAYOFF_CD)),
        dict(player_action='Action B',    partner_action='Action B',    payoff=cu(C.PAYOFF_DD), partner_payoff=cu(C.PAYOFF_DD)),
    ]
    if game_code == 'C':
        for row in rows:
            if row['player_action'] == 'Action B':
                row['payoff'] = cu(row['payoff'] - C.HARMONY_TAX)
            if row['partner_action'] == 'Action B':
                row['partner_payoff'] = cu(row['partner_payoff'] - C.HARMONY_TAX)
    return rows


def belief_buttons(group_size):
    """
    Returns a list of dicts for the frequency-belief buttons.
    group_size is the number of OTHER players (i.e., total group - 1).
    Buttons run from 0 ("None") to group_size ("All").
    """
    buttons = []
    for k in range(0, group_size + 1):
        if k == 0:
            label = 'None'
        elif k == group_size:
            label = 'All ({k})'.format(k=k)
        else:
            label = str(k)
        buttons.append(dict(value=k, label=label))
    return buttons


# ---------- Session creation ----------
def creating_session(subsession: Subsession):
    session = subsession.session

    if subsession.round_number == 1:
        players = subsession.get_players()
        n = len(players)

        group_sizes = []
        remainder = n
        while remainder > 0:
            if remainder % C.GROUP_SIZE_PREFERRED == 0:
                group_sizes.extend(
                    [C.GROUP_SIZE_PREFERRED] * (remainder // C.GROUP_SIZE_PREFERRED)
                )
                remainder = 0
            elif remainder >= C.GROUP_SIZE_PREFERRED:
                if remainder - C.GROUP_SIZE_PREFERRED >= C.GROUP_SIZE_FALLBACK:
                    group_sizes.append(C.GROUP_SIZE_PREFERRED)
                    last_group = C.GROUP_SIZE_PREFERRED
                else:
                    group_sizes.append(C.GROUP_SIZE_FALLBACK)
                    last_group = C.GROUP_SIZE_FALLBACK
                remainder -= last_group
            else:
                group_sizes.append(remainder)
                remainder = 0

        random.shuffle(players)
        matrix = []
        idx = 0
        for size in group_sizes:
            matrix.append(players[idx:idx + size])
            idx += size
        subsession.set_group_matrix(matrix)

        groups = subsession.get_groups()

        session.vars['treated_groups'] = {
            group.id_in_subsession: (random.random() < 0.5)
            for group in groups
        }

        session.vars['group_matrix'] = [
            [p.id_in_subsession for p in group.get_players()]
            for group in groups
        ]

    else:
        subsession.group_like_round(1)

    # Every round: apply treatment flag, reshuffle pairs, assign FM/SM roles
    for group in subsession.get_groups():
        group.is_treated = session.vars['treated_groups'].get(
            group.id_in_subsession, False
        )

        players = group.get_players()
        random.shuffle(players)
        for i in range(0, len(players) - 1, 2):
            p1 = players[i]
            p2 = players[i + 1]
            p1.partner_id_in_group = p2.id_in_group
            p2.partner_id_in_group = p1.id_in_group
            # Randomly assign FM / SM within each pair
            if random.random() < 0.5:
                p1.is_first_mover = True
                p2.is_first_mover = False
            else:
                p1.is_first_mover = False
                p2.is_first_mover = True


# ---------- Payoffs ----------
def get_payoff(fm_choice, sm_choice, is_first_mover, is_harmony):
    if is_first_mover:
        player_choice = fm_choice
        partner_choice = sm_choice
    else:
        player_choice = sm_choice
        partner_choice = fm_choice

    if player_choice and partner_choice:
        p = C.PAYOFF_CC
    elif player_choice and not partner_choice:
        p = C.PAYOFF_CD
    elif not player_choice and partner_choice:
        p = C.PAYOFF_DC
    else:
        p = C.PAYOFF_DD

    if is_harmony and not player_choice:
        p -= C.HARMONY_TAX

    return cu(p)


def resolve_sequential_choices(fm: Player, sm: Player):
    fm_actual = fm.fm_choice
    sm_actual = sm.sm_choice_if_fm_coop if fm_actual else sm.sm_choice_if_fm_defect
    return fm_actual, sm_actual


def resolve_sequential_choices_voting(fm: Player, sm: Player, option: str):
    if option == 'A':
        fm_actual = fm.fm_choice_a
        sm_actual = sm.sm_coop_a if fm_actual else sm.sm_defect_a
    elif option == 'B':
        fm_actual = fm.fm_choice_b
        sm_actual = sm.sm_coop_b if fm_actual else sm.sm_defect_b
    else:  # C
        fm_actual = fm.fm_choice_c
        sm_actual = sm.sm_coop_c if fm_actual else sm.sm_defect_c
    return fm_actual, sm_actual


def resolve_vote(votes):
    counts = {game: votes.count(game) for game in ['A', 'B', 'C']}
    max_count = max(counts.values())
    tied_games = [game for game, count in counts.items() if count == max_count]
    return random.choice(tied_games)


def set_payoffs(group: Group):
    if group.round_number == C.VOTING_ROUND:
        votes = [p.vote for p in group.get_players()]
        winner = resolve_vote(votes)
        if 'winning_game_by_group' not in group.session.vars:
            group.session.vars['winning_game_by_group'] = {}
        group.session.vars['winning_game_by_group'][group.id_in_subsession] = winner
        group.winning_game = winner

    if group.round_number >= C.POST_VOTING_START:
        stored = group.session.vars.get('winning_game_by_group', {})
        group.winning_game = stored.get(group.id_in_subsession, 'A')

    current_game = active_game_code(group)
    is_harmony = current_game == 'C'

    for p in group.get_players():
        partner = group.get_player_by_id(p.partner_id_in_group)

        if group.round_number == C.VOTING_ROUND:
            winning_game = group.winning_game
            fm = p if p.is_first_mover else partner
            sm = partner if p.is_first_mover else p
            fm_choice, sm_choice = resolve_sequential_choices_voting(fm, sm, winning_game)
            is_harmony_vote = winning_game == 'C'
            p.round_payoff = get_payoff(fm_choice, sm_choice, p.is_first_mover, is_harmony_vote)
            p.actual_choice = fm_choice if p.is_first_mover else sm_choice
        else:
            fm = p if p.is_first_mover else partner
            sm = partner if p.is_first_mover else p
            fm_choice, sm_choice = resolve_sequential_choices(fm, sm)
            p.round_payoff = get_payoff(fm_choice, sm_choice, p.is_first_mover, is_harmony)
            p.actual_choice = fm_choice if p.is_first_mover else sm_choice

        p.participant.payoff += p.round_payoff


# ---------- Pages ----------

class Introduction(Page):
    form_model = 'player'

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

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            payoff_matrix=payoff_matrix_rows('A'),
            cumulative_payoff=player.participant.payoff,
        )


class Choice(Page):
    form_model = 'player'
    form_fields = ['fm_choice', 'sm_choice_if_fm_coop', 'sm_choice_if_fm_defect', 'belief_count']

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

    @staticmethod
    def vars_for_template(player: Player):
        current_game = active_game_code(player.group)
        group_size = len(player.group.get_players()) - 1
        return dict(
            current_game=current_game,
            current_game_label=game_label(current_game),
            show_moral_message=show_moral_message(player.group),
            moral_message=C.MORAL_MESSAGE,
            is_harmony=(current_game == 'C'),
            payoff_matrix=payoff_matrix_rows(current_game),
            cumulative_payoff=player.participant.payoff,
            group_size=group_size,
            belief_buttons=belief_buttons(group_size),
        )


class Vote(Page):
    form_model = 'player'
    form_fields = [
        'vote',
        'fm_choice_a', 'sm_coop_a', 'sm_defect_a', 'belief_count_a',
        'fm_choice_b', 'sm_coop_b', 'sm_defect_b', 'belief_count_b',
        'fm_choice_c', 'sm_coop_c', 'sm_defect_c', 'belief_count_c',
    ]

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

    @staticmethod
    def vars_for_template(player: Player):
        group_size = len(player.group.get_players()) - 1
        return dict(
            matrix_a=payoff_matrix_rows('A'),
            matrix_b=payoff_matrix_rows('B'),
            matrix_c=payoff_matrix_rows('C'),
            moral_message=C.MORAL_MESSAGE,
            cumulative_payoff=player.participant.payoff,
            group_size=group_size,
            belief_buttons=belief_buttons(group_size),
        )


class ResultsWaitPage(WaitPage):
    @staticmethod
    def after_all_players_arrive(group: Group):
        set_payoffs(group)


class Results(Page):
    form_model = 'player'

    @staticmethod
    def vars_for_template(player: Player):
        partner = player.group.get_player_by_id(player.partner_id_in_group)
        group_size = len(player.group.get_players()) - 1

        if player.round_number == C.VOTING_ROUND:
            displayed_game = player.group.winning_game
            fm = player if player.is_first_mover else partner
            sm = partner if player.is_first_mover else player
            fm_choice, sm_choice = resolve_sequential_choices_voting(fm, sm, displayed_game)
            partner_actual_choice = sm_choice if player.is_first_mover else fm_choice
        else:
            displayed_game = active_game_code(player.group)
            fm = player if player.is_first_mover else partner
            sm = partner if player.is_first_mover else player
            fm_choice, sm_choice = resolve_sequential_choices(fm, sm)
            partner_actual_choice = sm_choice if player.is_first_mover else fm_choice

        return dict(
            partner_actual_choice=partner_actual_choice,
            player_actual_choice=player.actual_choice,
            is_first_mover=player.is_first_mover,
            cumulative_payoff=player.participant.payoff,
            displayed_game=displayed_game,
            displayed_game_label=game_label(displayed_game),
            winning_game_label=game_label(player.group.winning_game),
            show_moral_message=show_moral_message(player.group) if player.round_number != C.VOTING_ROUND else False,
            moral_message=C.MORAL_MESSAGE,
            is_harmony=(displayed_game == 'C'),
            group_size=group_size,
        )


page_sequence = [Introduction, Choice, Vote, ResultsWaitPage, Results]