from otree.api import *
import random 
# from numpy import arange as arange

doc = """
Lemonade Stand Task (Ederer 2013)
otree implementation
"""


class Constants(BaseConstants):
    name_in_url = 'lemonade_stand'
    players_per_group = None
    num_rounds = 3
    payment_per_correct=cu(2)

    #Format of variables: business, stadium, school    
    demand = [100.0, 60.0, 200.0]
    optimal_sugar = [1.55, 5.55, 9.55]
    optimal_lemon = [7.55, 5.55, 1.55]
    optimal_price = [7.55, 7.55, 2.55]
    optimal_flavor = [1, 1, 0]
     
    penalty_factor_sugar=[3,0.5,6]
    penalty_factor_lemon = [3, 0.5, 6]
    penalty_factor_flavor = [20.0, 0.5, 60.0]
    penalty_factor_price = [3, 0.5, 6]


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    profit_round = models.FloatField()
    profit_total = models.FloatField(initial=0)
    location=models.IntegerField(
        choices=[
            [0, 'Commercial District'],
            [1, 'Railway Station'],
            [2, 'School'],
        ],
        widget=widgets.RadioSelectHorizontal
    )
    flavor=models.IntegerField(
        choices=[
            [0, 'Plain'],
            [1, 'Masala'],
        ],
        widget=widgets.RadioSelectHorizontal
    )
    sugar = models.FloatField(min=0, max=20)
    # lemon_array = arange(0, 20, 0.1)
    # lemon = models.FloatField(
    #     choices=getattr(lemon_array, "tolist", lambda: lemon_array)(),
    #     widget=widgets.RadioSelect
    # )
    lemon = models.FloatField(min=0, max=20)
    price = models.FloatField(min=0, max=10, initial=5)

    feedback = models.StringField()

# PAGES
class MakeChoices(Page):
    form_model="player"
    form_fields = ["location", "flavor", "sugar", "lemon", "price"]
    
    @staticmethod
    def vars_for_template(player: Player):
        
        if player.round_number==1:
             return dict(
                prev_price=8.1,
                prev_lemon=7,
                prev_sugar=3,
                )
        prev_player = player.in_round(player.round_number - 1)
        return dict (
            prev_price=prev_player.price,
            prev_lemon=prev_player.lemon,
            prev_sugar=prev_player.sugar,
            )
        
    @staticmethod
    def before_next_page(player: Player, timeout_happened):

        #Calculate total penalties
        SugarP = abs(Constants.optimal_sugar[player.location] -
                     player.sugar)*Constants.penalty_factor_sugar[player.location]
        LemonP = abs(Constants.optimal_lemon[player.location] -
                     player.lemon)*Constants.penalty_factor_lemon[player.location]
        PriceP = abs(Constants.optimal_price[player.location] -
                     player.price)*Constants.penalty_factor_price[player.location]
        if (player.flavor == Constants.optimal_flavor[player.location]):
            flavorP = 0.0
        else:
            flavorP = Constants.penalty_factor_flavor[player.location]
        Penalty = SugarP + LemonP + PriceP + flavorP

        #Calculate profits
        player.profit_round = Constants.demand[player.location] - Penalty
        ##Adding profits for all rounds
        all_players = player.in_all_rounds()
        player.profit_total = 0
        for temp_player in all_players:
            player.profit_total += temp_player.profit_round

        #Feedback
        if (player.sugar < Constants.optimal_sugar[player.location]):
            sugar_feedback = "Some of your customers told you that the lemonade is not sweet enough"
        else:
            sugar_feedback = "Some of your customers told you that the lemonade is too sweet."

        if (player.lemon < Constants.optimal_lemon[player.location]):
            lemon_feedback = "Some of your customers told you that the lemonade is not sour enough."
        else:
            lemon_feedback = "Some of your customers told you that the lemonade is too sour."

        if (player.price < Constants.optimal_price[player.location]):
            price_feedback = "You have too many customers demanding lemonade. The price may be too low."
        else:
            price_feedback = "You have too few customers demanding lemonade. The price may be too high."
        
        player.feedback = random.choice(
            [sugar_feedback, lemon_feedback, price_feedback])
        
        

class Results(Page):
    pass


class FinalPage(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number==Constants.num_rounds
    

page_sequence = [MakeChoices, Results,FinalPage]