from otree.api import *
import json
import random

class Constants(BaseConstants):
    name_in_url = 'trust_game_observer'
    players_per_group = None
    num_rounds = 1
    endowment = cu(10)
    multiplier = 3


class Subsession(BaseSubsession):
    def creating_session(subsession):
        for p in subsession.get_players():
            p.participant.vars['role_choice'] = None


class Group(BaseGroup):
    def set_pairings(group):
        row2 = [p for p in group.get_players() if p.participant.vars['role_choice'] == 2]
        row3 = [p for p in group.get_players() if p.participant.vars['role_choice'] == 3]
        pairs = list(zip(row2, row3))

        for i, (sender, responder) in enumerate(pairs):
            sender.participant.partner_id = responder.id_in_subsession
            responder.participant.partner_id = sender.id_in_subsession

        for p in group.get_players():
            p.participant.all_pairs = [
                {
                    'pair_id': i + 1,
                    'sender_id': sender.id_in_subsession,
                    'responder_id': responder.id_in_subsession,
                    'responder_rank': responder.participant.vars.get('responder_rank', 'unknown')
                }
                for i, (sender, responder) in enumerate(pairs)
            ]


class Player(BasePlayer):
    performance_condition = models.StringField(choices=["performance known", "performance unknown"])
    role_choice = models.IntegerField(choices=[(1, 'First Row (Observer)'), (2, 'Second Row (Sender)'), (3, 'Third Row (Responder)')])
    send_amount = models.CurrencyField(min=0, max=Constants.endowment, blank=True)
    return_amount = models.CurrencyField(min=0, blank=True)
    predictions = models.LongStringField(blank=True)
    num_correct = models.IntegerField(initial=0)
    num_incorrect = models.IntegerField(initial=0)
    attempted = models.IntegerField(initial=0)
    top_performer = models.BooleanField()


def get_player_by_id(subsession, id_in_subsession):
    for p in subsession.get_players():
        if p.id_in_subsession == id_in_subsession:
            return p
    return None


class ConditionSelect(Page):
    form_model = 'player'
    form_fields = ['performance_condition']


class RoleSelect(Page):
    form_model = 'player'
    form_fields = ['role_choice']

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.participant.vars['role_choice'] = player.role_choice


class WaitForAllRoles(WaitPage):
    wait_for_all_groups = True

    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        group = subsession.get_groups()[0]
        group.set_pairings()


class Slider(Page):
    timeout_seconds = 180

    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 3

    @staticmethod
    def js_vars(player: Player):
        return dict(lower1=101, upper1=500, lower2=11, upper2=99)

    @staticmethod
    def live_method(player: Player, data):
        if data == 0:
            player.num_correct += 1
        else:
            player.num_incorrect += 1
        player.attempted = player.num_correct + player.num_incorrect


class RankThirdRowPerformance(WaitPage):
    wait_for_all_groups = True

    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 3

    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        row3 = [p for p in subsession.get_players() if p.role_choice == 3]
        sorted_row3 = sorted(row3, key=lambda p: p.num_correct, reverse=True)
        midpoint = len(sorted_row3) // 2

        for idx, p in enumerate(sorted_row3):
            if len(sorted_row3) == 1:
                rank = 'top'
            elif len(sorted_row3) % 2 == 0:
                rank = 'top' if idx < midpoint else 'bottom'
            else:
                rank = 'top' if idx <= midpoint else 'bottom'
            p.top_performer = rank == 'top'
            p.participant.vars['top_performer'] = p.top_performer
            p.participant.vars['responder_rank'] = rank


class SenderDecision(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.participant.vars['role_choice'] == 2

    form_model = 'player'
    form_fields = ['send_amount']


class WaitForSenderDecision(WaitPage):
    @staticmethod
    def is_displayed(player: Player):
        return player.participant.vars['role_choice'] in [1, 3]


class ResponderDecision(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.participant.vars['role_choice'] == 3

    @staticmethod
    def vars_for_template(player: Player):
        sender = get_player_by_id(player.subsession, player.participant.partner_id)
        tripled = sender.send_amount * Constants.multiplier
        player.participant.tripled_amount = tripled
        return {'tripled_amount': tripled}

    form_model = 'player'
    form_fields = ['return_amount']


class ObserverPrediction(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.participant.vars['role_choice'] == 1

    @staticmethod
    def vars_for_template(player: Player):
        show_rank = player.performance_condition == 'performance known'
        pairs_info = []
        for pair in player.participant.all_pairs:
            sender = get_player_by_id(player.subsession, pair['sender_id'])
            amount_sent = sender.send_amount
            tripled_amount = amount_sent * Constants.multiplier
            rank_info = pair['responder_rank'] if show_rank else None
            pairs_info.append({
                'pair_id': pair['pair_id'],
                'amount_sent': amount_sent,
                'tripled_amount': tripled_amount,
                'responder_rank': rank_info
            })
        return {'pairs': pairs_info, 'show_rank': show_rank}

    form_model = 'player'
    form_fields = ['predictions']


class ResultsWaitPage(WaitPage):
    wait_for_all_groups = True


class Results(Page):
    @staticmethod
    def vars_for_template(player: Player):
        role = player.participant.vars['role_choice']
        if role == 2:
            responder = get_player_by_id(player.subsession, player.participant.partner_id)
            payoff = Constants.endowment - player.send_amount + responder.return_amount
        elif role == 3:
            sender = get_player_by_id(player.subsession, player.participant.partner_id)
            received = sender.send_amount * Constants.multiplier
            payoff = received - player.return_amount
        elif role == 1:
            predictions = json.loads(player.predictions)
            chosen_pair = random.choice(player.participant.all_pairs)
            responder = get_player_by_id(player.subsession, chosen_pair['responder_id'])
            actual = responder.return_amount
            predicted = cu(predictions.get(str(chosen_pair['pair_id']), 0))
            diff = abs(actual - predicted)
            if diff == 0:
                payoff = cu(10)
            elif diff <= 1:
                payoff = cu(8)
            elif diff <= 2:
                payoff = cu(5)
            elif diff <= 3:
                payoff = cu(2)
            else:
                payoff = cu(0)
        else:
            payoff = cu(0)

        player.payoff = payoff
        return {'payoff': payoff}


page_sequence = [
    ConditionSelect,
    RoleSelect,
    WaitForAllRoles,
    Slider,
    RankThirdRowPerformance,
    SenderDecision,
    WaitForSenderDecision,
    ResponderDecision,
    ObserverPrediction,
    ResultsWaitPage,
    Results,
]