from otree.api import *

c = cu

doc = ''


class C(BaseConstants):
    NAME_IN_URL = 'Experiment'
    PLAYERS_PER_GROUP = 3
    NUM_ROUNDS = 10
    MANAGER_ROLE = 'Manager'
    WORKER_A_ROLE = 'Worker_A'
    WORKER_B_ROLE = 'Worker_B'


class Subsession(BaseSubsession):
    pass


def group_by_arrival_time_method(subsession: Subsession, waiting_players):
    manager_players = [p for p in waiting_players if p.participant.is_manager]
    worker_players = [p for p in waiting_players if not p.participant.is_manager]

    if len(manager_players) >= 1 and len(worker_players) >= 2:
        manager_players.sort(key=lambda p: p.participant.wait_page_arrival)
        worker_players.sort(key=lambda p: p.participant.wait_page_arrival)
        return [manager_players[0], worker_players[0],
                worker_players[1]]  # worker_0 will be in worker_a role and worker_1 will be in worker_b role
    print("manager:", manager_players)
    print("worker: ", worker_players)
    print("\n")


class Group(BaseGroup):
    worker_a_effort = models.IntegerField(min=1, max=10, label='')
    worker_b_effort = models.IntegerField(min=1, max=10, label='')
    is_bad_dp = models.BooleanField(initial=False)
    manager_compensation = models.IntegerField(min=60, max=100, label='')


def set_player_payoffs(group: Group):
    manager = group.get_player_by_role(C.MANAGER_ROLE)
    worker_a = group.get_player_by_role(C.WORKER_A_ROLE)
    worker_b = group.get_player_by_role(C.WORKER_B_ROLE)

    manager.payoff = 100 + 4 * (group.worker_a_effort + group.worker_b_effort) - group.manager_compensation
    worker_a.payoff = group.manager_compensation / 2 - group.worker_a_effort
    worker_b.payoff = group.manager_compensation / 2 - group.worker_b_effort


def set_participant_payoffs(group: Group):
    for player in group.get_players():
        history = player.in_all_rounds()
        participant = player.participant
        import random
        if not participant.is_dropout:
            first_payoff = 13 if participant.is_high_dimra else 17
            participant.payoff = random.choice(history).payoff + first_payoff
            print(history)
        else:
            participant.payoff = 0

class Player(BasePlayer):
    pass


class Lobby(WaitPage):
    group_by_arrival_time = True


class ManagerSelection(Page):
    form_model = 'group'
    form_fields = ['manager_compensation']
    timeout_seconds = 60

    @staticmethod
    def is_displayed(player: Player):
        return player.role == C.MANAGER_ROLE

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        group = player.group
        participant = player.participant
        if timeout_happened:
            group.is_bad_dp = True
            player.participant.is_dropout = True
            player.group.manager_compensation = 80


class WorkerWait(WaitPage):
    @staticmethod
    def is_displayed(player: Player):
        return player.role == C.WORKER_A_ROLE or player.role == C.WORKER_B_ROLE


class WorkerASelection(Page):
    form_model = 'group'
    form_fields = ['worker_a_effort']
    timeout_seconds = 60

    @staticmethod
    def is_displayed(player: Player):
        return player.role == C.WORKER_A_ROLE

    @staticmethod
    def vars_for_template(player: Player):
        session = player.session
        group = player.group
        participant = player.participant
        if player.id_in_group % 2:
            other_worker = player.group.get_player_by_id(2)
        else:
            other_worker = player.group.get_player_by_id(3)
        print("AAA", player.group.manager_compensation)
        return dict(
            show_choice=session.config['show_choice'],
            is_high_dimra=player.group.get_player_by_role(C.WORKER_B_ROLE).participant.is_high_dimra,
            wage=player.group.manager_compensation / 2,
            other_id=other_worker.participant.id_in_session,
            treatment=player.participant.treatment
        )

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        group = player.group
        participant = player.participant
        if timeout_happened:
            group.is_bad_dp = True
            player.participant.is_dropout = True
            player.group.worker_a_effort = 4


class WorkerBSelection(Page):
    form_model = 'group'
    form_fields = ['worker_b_effort']
    timeout_seconds = 60

    @staticmethod
    def is_displayed(player: Player):
        return player.role == C.WORKER_B_ROLE

    @staticmethod
    def vars_for_template(player: Player):
        session = player.session
        group = player.group
        participant = player.participant
        if player.id_in_group % 2:
            other_worker = player.group.get_player_by_id(2)
        else:
            other_worker = player.group.get_player_by_id(3)
        return dict(
            show_choice=session.config['show_choice'],
            is_high_dimra=player.group.get_player_by_role(C.WORKER_A_ROLE).participant.is_high_dimra,
            wage=player.group.manager_compensation / 2,
            other_id=other_worker.participant.id_in_session,
            treatment=player.participant.treatment
        )

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        group = player.group
        participant = player.participant
        if timeout_happened:
            group.is_bad_dp = True
            player.participant.is_dropout = True
            player.group.worker_b_effort = 4


class ManagerWait(WaitPage):
    @staticmethod
    def is_displayed(player: Player):
        return player.role == C.MANAGER_ROLE


class SetPlayerPayoffs(WaitPage):
    after_all_players_arrive = set_player_payoffs


class SetWaitPageArrival(Page):
    form_model = 'player'

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        participant = player.participant
        import time
        participant.wait_page_arrival = time.time()


class SetParticipantPayoffs(WaitPage):
    after_all_players_arrive = set_participant_payoffs

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


page_sequence = [Lobby, ManagerSelection, WorkerWait, WorkerASelection, WorkerBSelection, ManagerWait, SetPlayerPayoffs,
                 SetWaitPageArrival, SetParticipantPayoffs]