from otree.api import *
import random
import itertools
from math import comb


doc = """
Price Information and Cost Uncertainty in Oligopoly by Ruolong Xiao and Cesar Martinelli"
"""


class C(BaseConstants):
    PLAYERS_PER_GROUP = 3
    NAME_IN_URL = 'oligopoly'
    NUM_ROUNDS = 2
    max_price = 100
    min_price = 50
    min_cost = 50
    max_cost = 70
    cap_more_info = False
    ncap_more_info = False
    baseline = True
    num_buyer = 60
    unit_demand = 2
    num_buyer_1 = 30
    num_buyer_2 = 30
    num_buyer_3 = 30

    @staticmethod
    def assign_player_cost():
        return random.randint(int(C.min_cost), int(C.max_cost))


class Subsession(BaseSubsession):
    pass


def creating_session(subsession: Subsession):
    if subsession.round_number == 1:
        subsession.session.vars['pay_round'] = random.randint(1, C.NUM_ROUNDS)
    subsession.group_randomly(fixed_id_in_group=True)


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    price = models.CurrencyField(
        min=C.min_price,
        max=C.max_price,
        doc="""Price player offers to sell product for""",
        label=f"Please enter an amount from {C.min_price} to {C.max_price} as your price:",
    )
    num_winner = models.IntegerField(initial=0)
    num_winner_b1 = models.IntegerField(initial=0)
    num_winner_b2 = models.IntegerField(initial=0)
    num_winner_b3 = models.IntegerField(initial=0)
    cost = models.CurrencyField()
    n1 = models.IntegerField(initial=0)
    n2 = models.IntegerField(initial=0)
    n3 = models.IntegerField(initial=0)
    final_pay = models.CurrencyField(initial=0)
    pay_round = models.IntegerField(initial=0)


# FUNCTIONS
def set_payoffs(group: Group):
    players = group.get_players()

    if C.baseline:
        for p in players:
            p.payoff = (p.price - p.cost) * C.num_buyer_1 * C.unit_demand / C.PLAYERS_PER_GROUP
            p.num_winner = int(C.num_buyer_1 * C.unit_demand / C.PLAYERS_PER_GROUP)
            p.num_winner_b1 += p.num_winner
        sampled_firms = [list(k) for k in itertools.combinations(players, 2)]
        for s in sampled_firms:
            winning_price = min(p.price for p in s)
            winners = [p for p in s if p.price == winning_price]
            num_buyers_per_comb = comb(int(C.PLAYERS_PER_GROUP), len(s))
            for winner in winners:
                units_sold = C.unit_demand * C.num_buyer_2 / (num_buyers_per_comb * len(winners))
                winner.payoff += (winner.price - winner.cost) * units_sold
                winner.num_winner_b2 += int(units_sold)
                winner.num_winner += int(units_sold)

    if C.cap_more_info:
        sampled_firms = [list(k) for k in itertools.combinations(players, 2)]
        for s in sampled_firms:
            winning_price = min(p.price for p in s)
            winners = [p for p in s if p.price == winning_price]
            num_buyers_per_comb = comb(int(C.PLAYERS_PER_GROUP), len(s))
            for winner in winners:
                units_sold = C.unit_demand * C.num_buyer_2 / (num_buyers_per_comb * len(winners))
                winner.payoff += (winner.price - winner.cost) * units_sold
                winner.num_winner_b2 += int(units_sold)
                winner.num_winner += int(units_sold)

        sampled_firms = [list(k) for k in itertools.combinations(players, 3)]
        for s in sampled_firms:
            winning_price = min(p.price for p in s)
            winners = [p for p in s if p.price == winning_price]
            num_buyers_per_comb = comb(int(C.PLAYERS_PER_GROUP), len(s))
            for winner in winners:
                units_sold = C.unit_demand * C.num_buyer_3 / (num_buyers_per_comb * len(winners))
                winner.payoff += (winner.price - winner.cost) * units_sold
                winner.num_winner_b3 += int(units_sold)
                winner.num_winner += int(units_sold)

    if C.ncap_more_info:
        for p in players:
            p.payoff = (p.price - p.cost) * C.num_buyer_1 * C.unit_demand / C.PLAYERS_PER_GROUP
            p.num_winner = int(C.num_buyer_1 * C.unit_demand / C.PLAYERS_PER_GROUP)
            p.num_winner_b1 += p.num_winner

        sampled_firms = [list(k) for k in itertools.combinations(players, 3)]
        for s in sampled_firms:
            winning_price = min(p.price for p in s)
            winners = [p for p in s if p.price == winning_price]
            num_buyers_per_comb = comb(int(C.PLAYERS_PER_GROUP), len(s))
            for winner in winners:
                units_sold = C.unit_demand * C.num_buyer_3 / (num_buyers_per_comb * len(winners))
                winner.payoff += (winner.price - winner.cost) * units_sold
                winner.num_winner_b3 += int(units_sold)
                winner.num_winner += int(units_sold)


def _payment_info(player: Player):
    pr = player.session.vars.get('pay_round', 1)
    payoff_val = player.in_round(pr).payoff or cu(0)
    bonus_pay = payoff_val / 200
    fee_amount = cu(10)
    total_pay = bonus_pay + fee_amount
    return pr, bonus_pay, fee_amount, total_pay


class PrePricing(Page):
    @staticmethod
    def before_next_page(player: Player):
        player.n1 = int(C.num_buyer_1 / C.PLAYERS_PER_GROUP)
        player.n2 = int(C.num_buyer_2 * comb(int(C.PLAYERS_PER_GROUP - 1), 1) / comb(C.PLAYERS_PER_GROUP, 2))
        player.n3 = 0
        if player.round_number == 1:
            player.cost = C.assign_player_cost()
        else:
            player.cost = player.in_round(1).cost


class Pricing(Page):
    form_model = 'player'
    form_fields = ['price']

    @staticmethod
    def vars_for_template(player: Player):
        start = max(1, player.round_number - 5)
        last5 = player.in_rounds(start, player.round_number - 1)
        return dict(last5=last5)


class PreResults(WaitPage):
    after_all_players_arrive = set_payoffs


class Results(Page):
    @staticmethod
    def vars_for_template(player: Player):
        others = sorted(player.get_others_in_group(), key=lambda q: q.id_in_group)
        start = max(1, player.round_number - 5)
        last5 = player.in_rounds(start, player.round_number - 1)
        me_now = dict(
            price=player.price,
            cost=player.cost,
            sales=player.num_winner,
            sales_b1=player.num_winner_b1,
            sales_b2=player.num_winner_b2,
            sales_b3=player.num_winner_b3,
            profit=player.payoff,
        )
        return dict(others=others, last5=last5, me_now=me_now)


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

    @staticmethod
    def before_next_page(player: Player):
        pr, _, _, total_pay = _payment_info(player)
        player.pay_round = pr
        player.final_pay = total_pay
        player.participant.payoff = total_pay

    @staticmethod
    def vars_for_template(player: Player):
        pr, bonus_pay, fee_amount, total_pay = _payment_info(player)
        return dict(pr=pr, final=dict(bonus_pay=bonus_pay, fee=fee_amount, total_pay=total_pay))


page_sequence = [PrePricing, Pricing, PreResults, Results, Payment]