from otree.api import (
    models, widgets, BaseConstants, BaseSubsession, BaseGroup, BasePlayer,
    Currency as c, currency_range
)
import random

author = 'Franziska Heinicke'

doc = """
Bargaining with a nuclear option
"""


class Constants(BaseConstants):
    name_in_url = 'bargain_b'
    players_per_group = 2
    num_rounds = 11

    endowmentsA = [[30, 50, 60, 50, 30, 40, 60, 60, 40, 60, 30],
                   [30, 60, 40, 60, 60, 40, 30, 50, 60, 50, 30],
                   [40, 60, 60, 40, 60, 30, 30, 50, 60, 50, 30],
                   [30, 60, 30, 60, 40, 30, 50, 50, 40, 60, 60],
                   [30, 60, 40, 60, 30, 40, 60, 50, 30, 60, 50],
                   [30, 50, 60, 50, 40, 60, 30, 60, 40, 60, 30],
                   [60, 50, 40, 30, 60, 50, 60, 30, 60, 40, 30],
                   [60, 60, 60, 60, 50, 50, 40, 40, 30, 30, 30],
                   [30, 40, 60, 30, 60, 50, 60, 30, 40, 50, 60],
                   [30, 30, 30, 40, 40, 50, 50, 60, 60, 60, 60],
                   [60, 60, 50, 30, 40, 50, 60, 40, 60, 30, 30],
                   [60, 40, 50, 30, 30, 30, 50, 60, 60, 40, 60],
                   [50, 60, 60, 60, 40, 50, 30, 40, 30, 60, 30],
                   [50, 40, 30, 50, 60, 60, 30, 60, 30, 40, 60],
                   [60, 60, 60, 30, 50, 50, 30, 40, 30, 40, 60],
                   [60, 60, 30, 30, 60, 50, 30, 60, 40, 50, 40],
                   [30, 60, 60, 30, 30, 40, 60, 60, 50, 40, 30],
                   [30, 50, 60, 50, 30, 40, 60, 60, 40, 60, 30],
                   [30, 60, 40, 60, 60, 40, 30, 50, 60, 50, 30],
                   [40, 60, 60, 40, 60, 30, 30, 50, 60, 50, 30],
                   [30, 60, 30, 60, 40, 30, 50, 50, 40, 60, 60],
                   [30, 60, 40, 60, 30, 40, 60, 50, 30, 60, 50],
                   [30, 50, 60, 50, 40, 60, 30, 60, 40, 60, 30],
                   [60, 50, 40, 30, 60, 50, 60, 30, 60, 40, 30],
                   [60, 60, 60, 60, 50, 50, 40, 40, 30, 30, 30],
                   [30, 40, 60, 30, 60, 50, 60, 30, 40, 50, 60],
                   [30, 30, 30, 40, 40, 50, 50, 60, 60, 60, 60],
                   [60, 60, 50, 30, 40, 50, 60, 40, 60, 30, 30],
                   [60, 40, 50, 30, 30, 30, 50, 60, 60, 40, 60],
                   [50, 60, 60, 60, 40, 50, 30, 40, 30, 60, 30],
                   [50, 40, 30, 50, 60, 60, 30, 60, 30, 40, 60],
                   [60, 60, 60, 30, 50, 50, 30, 40, 30, 40, 60],
                   [60, 60, 30, 30, 60, 50, 30, 60, 40, 50, 40],
                   [30, 60, 60, 30, 30, 40, 60, 60, 50, 40, 30]
                   ]
    endowmentsB = [[30, 60, 30, 50, 60, 30, 40, 60, 60, 50, 40],
                   [40, 50, 60, 60, 40, 30, 60, 50, 30, 60, 30],
                   [30, 40, 60, 60, 50, 40, 30, 60, 30, 50, 60],
                   [30, 30, 60, 40, 60, 40, 60, 50, 30, 60, 50],
                   [60, 30, 60, 40, 40, 30, 50, 60, 30, 60, 50],
                   [30, 50, 60, 60, 30, 50, 40, 40, 60, 30, 60],
                   [60, 60, 60, 60, 50, 50, 40, 40, 30, 30, 30],
                   [60, 50, 40, 30, 60, 50, 60, 30, 60, 40, 30],
                   [30, 30, 30, 40, 40, 50, 50, 60, 60, 60, 60],
                   [30, 40, 60, 30, 60, 50, 60, 30, 40, 50, 60],
                   [60, 40, 50, 30, 60, 60, 30, 30, 50, 60, 40],
                   [50, 30, 60, 40, 60, 30, 50, 60, 30, 60, 40],
                   [60, 40, 30, 50, 60, 50, 30, 30, 40, 60, 60],
                   [50, 60, 40, 60, 60, 40, 60, 50, 30, 30, 30],
                   [50, 40, 30, 30, 60, 50, 60, 30, 40, 60, 60],
                   [30, 60, 60, 30, 40, 50, 40, 50, 30, 60, 60],
                   [30, 30, 40, 60, 40, 60, 60, 50, 50, 30, 60],
                   [30, 60, 30, 50, 60, 30, 40, 60, 60, 50, 40],
                   [40, 50, 60, 60, 40, 30, 60, 50, 30, 60, 30],
                   [30, 40, 60, 60, 50, 40, 30, 60, 30, 50, 60],
                   [30, 30, 60, 40, 60, 40, 60, 50, 30, 60, 50],
                   [60, 30, 60, 40, 40, 30, 50, 60, 30, 60, 50],
                   [30, 50, 60, 60, 30, 50, 40, 40, 60, 30, 60],
                   [60, 60, 60, 60, 50, 50, 40, 40, 30, 30, 30],
                   [60, 50, 40, 30, 60, 50, 60, 30, 60, 40, 30],
                   [30, 30, 30, 40, 40, 50, 50, 60, 60, 60, 60],
                   [30, 40, 60, 30, 60, 50, 60, 30, 40, 50, 60],
                   [60, 40, 50, 30, 60, 60, 30, 30, 50, 60, 40],
                   [50, 30, 60, 40, 60, 30, 50, 60, 30, 60, 40],
                   [60, 40, 30, 50, 60, 50, 30, 30, 40, 60, 60],
                   [50, 60, 40, 60, 60, 40, 60, 50, 30, 30, 30],
                   [50, 40, 30, 30, 60, 50, 60, 30, 40, 60, 60],
                   [30, 60, 60, 30, 40, 50, 40, 50, 30, 60, 60],
                   [30, 30, 40, 60, 40, 60, 60, 50, 50, 30, 60]
                   ]

    matching_size = 6

    max_waiting = 90

    final_waiting = 150

    first_waiting = 599

    instructions_template = 'BaNu_b/instructions.html'

class Subsession(BaseSubsession):
    # new fields for group assignment
    a_num_grouped = models.IntegerField(initial=0)
    b_num_grouped = models.IntegerField(initial=0)

    def group_by_arrival_time_method(self, wp):
        # get waiting players directly from the database
        # Why? because we can ask it to retrieve the associated participant model in one go, instead of loading it
        # only when it is used. This reduces the number of database queries
        STALE_THRESHOLD_SECONDS = 20
        waiting_players = list(self.player_set.filter(
            _gbat_arrived=True,
            _gbat_grouped=False,
            participant___last_request_timestamp__gte=time.time() - STALE_THRESHOLD_SECONDS,
        ).prefetch_related('participant'))

        # we also cache the list of all players for future use:
        all_players = self.get_players()

        if self.round_number == 1:
            d = defaultdict(list)

            # use just one loop to separate the lists of waiting players
            # a_players = [p for p in waiting_players if p.participant.vars['type'] == 'A']
            # b_players = [p for p in waiting_players if p.participant.vars['type'] == 'B']

            # I have integrated the part where the matching-group is determined.
            # I now store the numbers of A and B players that have been grouped on the session level
            # and calculate the suitable matching group number from that.
            a_players, b_players = [], []
            for wp in waiting_players:
                if wp.participant.vars['type'] == 'A':
                    a_players.append(wp)
                    if wp.matching_group == 0:
                        self.assign_matching_group(wp, type_a=True)
                else:
                    b_players.append(wp)
                    if wp.matching_group == 0:
                        self.assign_matching_group(wp, type_a=False)

            # now we are going to match people
            for p in waiting_players:
                if p.matching_group != 0:
                    if time.time() - p.participant.vars['wait_page_arrival'] > Constants.first_waiting:
                        self.set_first_wait_dropout_vars(p)
                        return [p]
                    else:
                        matching_number = p.participant.vars['matching']
                        players_in_this_matching = d[matching_number]
                        players_in_this_matching.append(p)

                        a_players_this_matching = [p for p in players_in_this_matching if
                                                   p.participant.vars['type'] == 'A']
                        b_players_this_matching = [p for p in players_in_this_matching if
                                                   p.participant.vars['type'] == 'B']

                        # only if we actually make a match, we set the new player variables
                        if len(a_players_this_matching) >= 1 and len(b_players_this_matching) >= 1:
                            # only if we make it past this first test, we actually consider all other players:
                            a_players_total = [p for p in all_players if
                                               p.matching_group == matching_number and p.participant.vars[
                                                   'type'] == 'A']
                            b_players_total = [p for p in all_players if
                                               p.matching_group == matching_number and p.participant.vars[
                                                   'type'] == 'B']

                            if len(a_players_total) == Constants.matching_size and len(b_players_total) == Constants.matching_size:
                                new_group = [a_players_this_matching[0], b_players_this_matching[0]]
                                self.set_first_round_player_vars(new_group)
                                return new_group

        else:
            # list of players in previous round
            prev_round = self.round_number - 1
            prev_round_players = self.in_round(prev_round).get_players()

            for p in waiting_players:
                # we do not need to set these variables over and over again
                if not p.paying_round:
                    self.update_player_vars(p)

                # pull up the check if the player is active. If they are not, we can return immediately, before
                # going through all the list comprehensions below
                if not p.active:
                    return [p]

                # we can only get here, if player is active, no need for the else clause
                matching_number = p.participant.vars['matching']

                total_now = [p for p in all_players if p.matching_group == matching_number and p.active]
                total_last = [p for p in prev_round_players if p.matching_group == matching_number and p.active]

                if len(total_last) == len(total_now) or time.time() - p.participant.vars[
                    'wait_page_arrival'] > Constants.max_waiting:
                    # now that we have passed checks on totals, we also separate by player type. we do not need to do it earlier
                    a_players = [p for p in waiting_players if
                                 p.matching_group == matching_number and p.participant.vars['type'] == 'A' and p.active]
                    b_players = [p for p in waiting_players if
                                 p.matching_group == matching_number and p.participant.vars['type'] == 'B' and p.active]

                    if len(a_players) >= 1 and len(b_players) >= 1:
                        # shuffle b_players
                        random.shuffle(b_players)
                        return [a_players[0], b_players[0]]

                    if time.time() - p.participant.vars['wait_page_arrival'] > Constants.final_waiting:
                        self.handle_final_wait_exceeded(p)
                        return [p]


    # some helper methods
    def assign_matching_group(self, player, type_a):
        if type_a:
            player.matching_group = self.a_num_grouped // Constants.matching_size + 1
            self.a_num_grouped += 1
        else:
            player.matching_group = self.b_num_grouped // Constants.matching_size + 1
            self.b_num_grouped += 1

        player.participant.vars['matching'] = player.matching_group


    def set_first_wait_dropout_vars(self, player):
        player.participant.vars.update({
            'first_out': True,
            'active': False,
            'drop_out': False
        })
        player.active = False
        player.matching_group = 0


    def set_first_round_player_vars(self, new_group):
        for player in new_group:
            player.paying_round = player.participant.vars['paying_round']
            player.participant.vars.update({
                'active': True,
                'drop_out': False,
                'first_out': False,
                'cumulative_payoff': 0
            })


    def update_player_vars(self, player):
        player.paying_round = player.participant.vars['paying_round']
        player.matching_group = player.participant.vars['matching']
        player.active = player.participant.vars['active']
        player.drop_out = player.participant.vars['drop_out']


    def handle_final_wait_exceeded(self, player):
        player.participant.vars['active'] = False
        player.active = False
        # if the cumulative payoff is 0, pay for the previous round
        if player.participant.vars['cumulative_payoff'] == 0:
            prev_round = self.round_number - 1
            prev_round_player = player.in_round(prev_round)

            player.paying_round = prev_round
            player.round_payoff = prev_round_player.round_payoff
            player.payoff = prev_round_player.round_payoff

            player.participant.vars.update({
                'paying_round': prev_round,
                'cumulative_payoff': prev_round_player.round_payoff,
                'disagreement': prev_round_player.group.failed,
            })

class Group(BaseGroup):

    endowA = models.IntegerField()
    endowB = models.IntegerField()

    currentA = models.IntegerField(initial=100)
    historyA = models.CharField(initial="")
    currentB = models.IntegerField(initial=100)
    historyB = models.CharField(initial="")


    firstpage_done = models.BooleanField(initial=False)
    final_offer = models.IntegerField(initial=0)
    time_final = models.CharField()
    accept_player = models.CharField()
    failed = models.BooleanField(initial=False)

    arrival = models.CharField()

    def set_time(self):
        import time
        self.arrival = str(int(time.time() * 1000))

    def set_payoffs(self):
        playerA = self.get_player_by_role('A')
        playerB = self.get_player_by_role('B')
        if self.failed:
            playerA.round_payoff = self.endowA
            playerB.round_payoff = self.endowB
        else:
            playerA.round_payoff = self.final_offer
            playerB.round_payoff = 100 - self.final_offer
        players = self.get_players()
        for p in players:
            if p.paying_round == self.round_number:
                p.payoff = p.round_payoff
                p.participant.vars['cumulative_payoff'] = p.round_payoff
                if self.failed:
                    p.participant.vars['disagreement'] = True
                else:
                    p.participant.vars['disagreement'] = False
            else:
                p.payoff = 0


class Player(BasePlayer):

    matching_group = models.IntegerField(initial=0)

    active = models.BooleanField(initial=True)

    drop_out = models.BooleanField(initial=False)

    paying_round = models.IntegerField()

    round_payoff = models.IntegerField()

    def role(self):
        if self.participant.vars['type'] == 'A':
            return 'A'
        else:
            return 'B'