from otree.api import *
import random
from django import forms

doc = """
In our experiment participants play the role of firms choosing whether or not to form a cartel.
Cartels are illegal, and are fined if discovered by a competition authority.
Treatments vary in the design of these fines. 
<br>
<br>

This is the app which contains the main game taking place in the experiment.
<br>
<br>

18 participants take part in each session, which consists of three supergames of predetermined length.
Ahead of each supergame, participants are matched in groups of three within two matching-pools of 9.
<br>
<br>

Here I have combined three apps to make repetition easier.
First, we allow firms the choice whether or not to join a cartel.
Then, we allow firms to independently decide on prices.
Finally, the cartels are told whether or not they were discovered and fined.
<br>
<br>

In the voting app: <br>
We allow firms the choice whether or not to join a cartel.
If a cartel is formed, a chat window is formed allowing for free chat during the decision period.
<br>
<br>

In the price-setting app: <br>
Firms simultaneously decide what price to set.
Revenue, costs, and profits are determined.
<br>
<br>

In the antitrust app: <br>
Cartels are informed whether they were discovered.
Fines are calculated and subtracted from earnings.
"""

def cumsum(lst):
    total = 0
    new = []
    for ele in lst:
        total += ele
        new.append(total)
    return new

class C(BaseConstants):
    NAME_IN_URL = 'vote_bertrand'
    
    # First supergame lasts x rounds, second supergame lasts y, etc...
    ROUNDS_PER_SG = [2, 2, 2]
    SG_ENDS = cumsum(ROUNDS_PER_SG)
    NUM_ROUNDS = sum(ROUNDS_PER_SG)
    
    # Informational texts
    INFORMATION_TEMPLATE = __name__ + '\instructions_comm.html'
    CARTEL_INFO = __name__ + '\cartel_info.html'
    INSTRUCTIONS_TEMPLATE = __name__ + '\instructions_price.html'
    
    # Production capacity to guarantee positive demand
    MAX_DEMAND = 100
    PLAYERS_PER_GROUP = 3
    
    # Constant marginal cost of production across all treatments
    MARGINAL_COST = 20
    
    # Define the Nash equilibrium for the Overcharge treatment
    NASH_PRICE = MARGINAL_COST
    NASH_QUANTITY = MAX_DEMAND - NASH_PRICE
    
    # Penalty rates for each treatment
    PROFIT_RATE = 1.2
    REVENUE_RATE = 0.8
    OVERCHARGE_RATE = 0.5
    
class Subsession(BaseSubsession):
    sg = models.IntegerField()
    period = models.IntegerField()
    is_last_period = models.BooleanField()

# This is from the Supergames app
def creating_session(subsession: Subsession):
    if subsession.round_number == 1:
        sg = 1
        period = 1
        # loop over all subsessions
        for ss in subsession.in_rounds(1, C.NUM_ROUNDS):
            ss.sg = sg
            ss.period = period
            # 'in' gives you a bool. for example: 5 in [1, 5, 6] # => True
            is_last_period = ss.round_number in C.SG_ENDS
            ss.is_last_period = is_last_period
            if is_last_period:
                sg += 1
                period = 1
            else:
                period += 1

# GROUPS 
class Group(BaseGroup):
    # Cartel formation
    yes_votes_cast = models.IntegerField(initial=0)
    votes_needed = models.IntegerField(default=C.PLAYERS_PER_GROUP, field_maybe_none=False)
    cartel_formed = models.BooleanField(initial=True)

    # Bertrand-game outcomes
    unit_price = models.CurrencyField(initial=0, doc="""Lowest price this round""")
    total_units = models.IntegerField(initial=0, doc="""Total units sold this round""")
    
    # Competition authority
    cartel_discovered = models.BooleanField(initial=False)
       
    @property
    def prev_round_group(group):
        return  group.in_round(group.round_number - 1)
    
    def unit_price(group):
        
        players = group.get_players()
        return min([p.price for p in players])
            
    def total_units(group):  
              
        return C.MAX_DEMAND - group.unit_price()

    def set_units(group):
        
        players = group.get_players()
        lowest_price = group.unit_price()
        num_players_with_lowest_price = sum(p.price == lowest_price for p in players)

        if num_players_with_lowest_price > 0:
            units_per_player = group.total_units() / num_players_with_lowest_price
        else:
            units_per_player = 0

        for p in players:
            if p.price == lowest_price:
                p.units = units_per_player
            else:
                p.units = 0
    
    def set_revenue(group):
        
        players = group.get_players()
        
        for p in players:
            p.round_revenue = p.units * p.price
            
    def set_cost(group):
        
        players = group.get_players()
        
        for p in players:
            p.round_cost = C.MARGINAL_COST * p.units
    
    def set_no_fine_payoffs(group):
        
        players = group.get_players()
        
        for p in players:
            p.no_fine_payoff = (p.price - C.MARGINAL_COST) * p.units
            
    def check_if_cartel_discovered(group):
        if not group.cartel_formed:
            return False
                
        group.cartel_discovered =  random.random() < 0.05
   
    ########### TREATMENTS ################
    # Define the fine for being in a cartel
    def set_cartel_fines(group):        
        for p in group.get_players():
            mode = group.session.config['mode']

            if mode == "revenue":
                p.cartel_fine = C.REVENUE_RATE * p.round_revenue
            elif mode == "profit":
                p.cartel_fine = C.PROFIT_RATE * p.no_fine_payoff
            elif mode == "overcharge":
                p.cartel_fine = C.OVERCHARGE_RATE * max(0, (p.price - C.NASH_PRICE) * C.NASH_QUANTITY)
            else:
                raise ValueError(f"Mode not recognized {mode}")

    def set_group_split(group):
        
        group_split = group % 2
        
        return group_split 
       
    def set_payoffs(group):
        
        players = group.get_players()
        
        for p in players:
            p.payoff = (p.price - C.MARGINAL_COST) * p.units - p.cartel_fine
    
# PLAYERS
class Player(BasePlayer):
    # Cartel decision
    join_cartel = models.BooleanField(widget=widgets.RadioSelect, choices=[
        (True, 'Yes, I want to join the cartel'),
        (False, 'No, I do not want to join the cartel')
    ])
    voted_cartel = models.BooleanField(initial=True)
    
    # Bertrand-game outcomes
    units = models.FloatField(
        initial=0,
        doc="""Units sold by this player in this round""",
        widget=forms.TextInput(attrs={'step': 'any'}),
    )

    # Remove decimal point and trailing zeroes
    @property
    def formatted_units(player):
        return '{:.1f}'.format(player.units).rstrip('0').rstrip('.')
    
    round_revenue = models.CurrencyField(initial=0)
    round_cost = models.CurrencyField(initial=0)
    no_fine_payoff = models.CurrencyField(initial=0)
    
    # Competition authority
    cartel_fine = models.CurrencyField(initial=0)
    
    # Calculator
    calculator_chat_timeout_seconds = models.IntegerField(initial=0)
    calculator_me = models.IntegerField(
        min = 0,
        label = 'Me',
        )
    calculator_r1 = models.IntegerField(
        min = 0,
        label = 'Rival 1',
        )
    calculator_r2 = models.IntegerField(
        min = 0,
        label = 'Rival 2',
        )
    
    @property
    def is_cartel_formed(player):
        return player.group.cartel_formed
    
    # Define the price decision
    price = models.IntegerField(
        min=0,
        doc="""Price set by this player""",
        label="What price will you set?",
    )
        
    @property
    def prev_player(player):
        return  player.in_round(player.round_number - 1)
    
    @property
    def other_players(player):
        return player.get_others_in_group()

    @property
    def running_payoff(player) -> Currency:        
        return sum(( prev_self.payoff for prev_self in player.in_all_rounds() ))

### CHAT
class Message(ExtraModel):
    group = models.Link(Group)
    sender = models.Link(Player)
    text = models.StringField()
    

def to_dict(msg: Message):
    return dict(sender=msg.sender.id_in_group, text=msg.text)

def custom_export(players):
    yield ['Session', 'Group', 'Participant Code', 'Participant ID' 'Round_number', 'Chat']
    # 'filter' without any args returns everything
    msgs = Message.filter()
    for msg in msgs:
        player = msg.sender
        participant = player.participant
        group = player.group
        session = player.session
        yield [session.code, group.id, participant.code, participant.id_in_session, player.round_number, msg.text]


####### PAGES #######
class Introduction(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

class MatchingWaitPage(WaitPage):
    wait_for_all_groups = True

    @staticmethod
    def after_all_players_arrive(subsession):
                # This is the rematching mechanism   
        if subsession.period == 1:
            players = subsession.get_players()
            matrix = []
            
            if len(players) == 3:
                matrix = [[1,2,3]]

            elif len(players) == 9:
                if subsession.sg == 1:
                    matrix = [
                        [player.participant.id_in_session for player in players if player.participant.id_in_session <= 3],
                        [player.participant.id_in_session for player in players if 4 <= player.participant.id_in_session <= 6],
                        [player.participant.id_in_session for player in players if 7 <= player.participant.id_in_session <= 9],
                    ]
                elif subsession.sg == 2:
                    matrix = [
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [1, 4, 7]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [2, 5, 8]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [3, 6, 9]],
                    ]
                elif subsession.sg == 3:
                    matrix = [
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [1, 6, 8]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [2, 4, 9]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [3, 5, 7]],
                    ]
            elif len(players) == 18:
                if subsession.sg == 1:
                    matrix = [
                        [player.participant.id_in_session for player in players if player.participant.id_in_session <= 3],
                        [player.participant.id_in_session for player in players if 4 <= player.participant.id_in_session <= 6],
                        [player.participant.id_in_session for player in players if 7 <= player.participant.id_in_session <= 9],
                        [player.participant.id_in_session for player in players if 10 <= player.participant.id_in_session <= 12],
                        [player.participant.id_in_session for player in players if 13 <= player.participant.id_in_session <= 15],
                        [player.participant.id_in_session for player in players if 16 <= player.participant.id_in_session <= 18],
                    ]
                elif subsession.sg == 2:
                    matrix = [
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [1, 4, 7]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [2, 5, 8]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [3, 6, 9]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [10, 13, 16]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [11, 14, 17]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [12, 15, 18]],
                    ]
                elif subsession.sg == 3:
                    matrix = [
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [1, 6, 8]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [2, 4, 9]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [3, 5, 7]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [10, 15, 17]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [11, 13, 18]],
                        [player.participant.id_in_session for player in players if player.participant.id_in_session in [12, 14, 16]],
                    ]

            subsession.set_group_matrix(matrix)

class New_Supergame(Page):
    @staticmethod
    def is_displayed(player: Player):
        subsession = player.subsession
        return subsession.period == 1


class Cartel_Vote(Page):
    form_model = 'player'
    form_fields = ['join_cartel']
    
    def before_next_page(player, timeout_happened):

        if player.join_cartel:
            player.voted_cartel = True
            player.group.yes_votes_cast += 1
        else:
            player.voted_cartel = False
            player.group.cartel_formed = False
            
        if player.group.yes_votes_cast == len(player.group.get_players()):
            player.group.cartel_formed = True

    def is_displayed(player: Player):
        # Display the cartel vote page at the start of a new supergame
        if player.round_number == 1 or player.subsession.period == 1:
            return True
        
        # Don't display the cartel vote page if the player was in an undetected cartel in the previous round
        prev_group = player.group.prev_round_group
        did_not_form_cartel = not prev_group.cartel_formed
        cartel_was_discovered = prev_group.cartel_discovered and prev_group.cartel_formed
        
        if did_not_form_cartel or cartel_was_discovered:
            return True
        
        return False


class ResultsWaitVote(WaitPage):
    body_text = "Waiting for the other participants to vote."
    
    def is_displayed(player: Player):
        if player.round_number == 1 or player.subsession.period == 1:
            return True

        prev_group = player.group.prev_round_group
        did_not_form_cartel = not prev_group.cartel_formed
        cartel_was_discovered = prev_group.cartel_discovered and prev_group.cartel_formed
        
        if did_not_form_cartel or cartel_was_discovered:
            return True
        
        return False

class Cartel_Vote_Results(Page):
    @staticmethod
    def vars_for_template(player: Player):
        group_id = player.id_in_group
        return {
            'group_id': group_id,
            }
        
    @staticmethod
    def js_vars(player: Player):
        return dict(mode=player.group.session.config['mode'], cartel_formed = player.group.cartel_formed)
        
    def is_displayed(player: Player):
        if player.round_number == 1:
            return True

        prev_group = player.group.prev_round_group
        did_not_form_cartel = not prev_group.cartel_formed
        cartel_was_discovered = prev_group.cartel_discovered and prev_group.cartel_formed
        
        if did_not_form_cartel or cartel_was_discovered:
            return True
        
        return False

class WaitForChat(WaitPage):
    body_text = "Waiting for the other participants."

class Calculator_Chat(Page):
    form_model = 'player'
    form_fields = ['calculator_me', 'calculator_r1', 'calculator_r2']

    '''@staticmethod
    def get_timeout_seconds(player):
        if player.round_number == 1 or player.subsession.period == 1:
            player.calculator_chat_timeout_seconds = 60
        else:
            player.calculator_chat_timeout_seconds = 30
        return player.calculator_chat_timeout_seconds'''

    #### CHAT and CALCULATOR ####
    @staticmethod
    def js_vars(player: Player):
        return dict(
            my_id = player.id_in_group,
            mode = player.group.session.config['mode'],
            cartel_formed = player.group.cartel_formed,
            )

    def live_method(player: Player, data):
        my_id = player.id_in_group
        group = player.group

        if 'text' in data:
            text = data['text']
            msg = Message.create(group=group, sender=player, text=text)
            return {0: [to_dict(msg)]}
        return {my_id: [to_dict(msg) for msg in Message.filter(group=group)]}
    
    @staticmethod
    def vars_for_template(player: Player):
        nickname = f"Firm {player.id_in_subsession}"
        return dict(my_nickname=nickname)


class Set_Price(Page):
    form_model = 'player'
    form_fields = ['price']
    
    @staticmethod
    def js_vars(player: Player):
        return dict(
            my_id = player.id_in_group,
            mode = player.group.session.config['mode'],
            cartel_formed = player.group.cartel_formed,
            )
    
    
class ResultsWaitBertrand(WaitPage):
    body_text = "Waiting for the other participants to decide."
    
    def after_all_players_arrive(self):
        group = self.group
        group.unit_price()
        group.total_units()
        group.set_units()
        group.set_revenue()
        group.set_cost()
        group.set_no_fine_payoffs()
        group.check_if_cartel_discovered()
        if group.cartel_discovered:
            group.set_cartel_fines()
        group.set_payoffs()

class Bertrand_Results(Page):
    @staticmethod
    def vars_for_template(player: Player):
        
        group_id = player.id_in_group
        other_group_ids =  [p.id_in_group for p in player.group.get_players() if p.id_in_group != player.id_in_group]
        
        template = {}
        
        for (j, other_player) in enumerate(player.other_players):
            template[f"other_player_{j + 1}_price"] = other_player.price
            template[f"other_player_{j + 1}_no_fine_payoff"] = other_player.no_fine_payoff
            template[f"other_player_{j + 1}_payoff"] = other_player.payoff
                    
        return {
            'group_id': group_id,
            'other_group_ids': other_group_ids,
            **template
            }
        
    
class Antitrust(Page):
    @staticmethod
    def vars_for_template(player: Player):
        result = {
            'cartel_formed': player.group.cartel_formed,
            'cartel_discovered': player.group.cartel_discovered,
            'cartel_fine': player.cartel_fine if player.group.cartel_discovered else 0.,
            'payoff_this_round': player.payoff,
            'total_payoff': player.running_payoff
        }
        
        return result
    
class History(Page):
    timeout_seconds = 60
        
    @staticmethod
    def vars_for_template(player: Player):
        subsession = player.subsession
        current_sg = subsession.sg
        
        other_group_ids = [p.id_in_group for p in player.group.get_players() if p.id_in_group != player.id_in_group]
        
        history = []
        for j, r in enumerate(player.in_all_rounds()):
            if r.subsession.sg == current_sg:
                other_prices = {f"other_player_{k + 1}_price": other_player.price for k, other_player in enumerate(r.get_others_in_group())}
                history.append((
                    r.subsession.period,
                    r.price,
                    r.formatted_units,
                    r.payoff,
                    r.running_payoff,
                    r.cartel_fine,
                    r.group.cartel_formed,
                    r.group.cartel_discovered,
                    other_prices,
                ))

        return {
            'other_group_ids': other_group_ids,
            'history': history,
        }

    
class End_of_Supergame(Page):
    timeout_seconds = 60
    
    @staticmethod
    def vars_for_template(player: Player):
        subsession = player.subsession
        current_sg = subsession.sg
        
        history = []
        for j, r in enumerate(player.in_all_rounds()):
            if r.subsession.sg == current_sg:
                history.append((
                    r.subsession.period,
                    r.price,
                    r.formatted_units,
                    r.payoff,
                    r.running_payoff,
                    r.cartel_fine,
                    r.group.cartel_formed,
                    r.group.cartel_discovered,
                ))

        return {'history': history}
    
    @staticmethod
    def is_displayed(player: Player):
        subsession = player.subsession
        return subsession.is_last_period == True

page_sequence = [
    Introduction,
    MatchingWaitPage,
    New_Supergame, 
    Cartel_Vote, 
    ResultsWaitVote, 
    Cartel_Vote_Results,
    WaitForChat,
    Calculator_Chat,
    Set_Price, 
    ResultsWaitBertrand, 
    Bertrand_Results, 
    Antitrust,
    History,
    End_of_Supergame,
    ]