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

author = 'Your name here'

doc = """
Your app description
"""


class Constants(BaseConstants):
    name_in_url = 'task'
    players_per_group = None
    num_rounds = 4


class Subsession(BaseSubsession):
    def creating_session(self):
        for p in self.get_players():
            if self.round_number <= self.session.vars['N_rounds']:
                p.n_firms = p.participant.vars[self.round_number]['Portfolio_1']['N_firms']
                p.n_firms_pay = p.participant.vars[self.round_number]['Portfolio_1']['N_firms_pay']
                p.winning_prob_prior = p.participant.vars[self.round_number]['Portfolio_1']['winning_prob_prior']
                p.n_firms_good_1 = p.participant.vars[self.round_number]['Portfolio_1']['N_firms_good']
                p.n_firms_good_2 = p.participant.vars[self.round_number]['Portfolio_2']['N_firms_good']
                p.firms_1 = p.participant.vars[self.round_number]['Portfolio_1']['industry_composition']
                p.firms_2 = p.participant.vars[self.round_number]['Portfolio_2']['industry_composition']

                if p.participant.vars['treatment'] == 'control':
                    portfolio_types = list(list(p.session.vars['portfolios']['types'].values())[0].keys())
                    # p.assigned = random.choice(portfolio_types)
                    p.assigned = choice(portfolio_types,
                                        size=1,
                                        p=[self.session.vars['P_good_assigned_prob'],
                                           1-self.session.vars['P_good_assigned_prob']
                                           ]
                                        )[0]
                    p.own = [x for x in p.participant.vars[self.round_number].keys() if
                             p.participant.vars[self.round_number][x]['type'] == p.assigned][0]


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    n_firms = models.IntegerField()
    n_firms_pay = models.IntegerField()
    winning_prob_prior = models.FloatField()
    n_firms_good_1 = models.IntegerField()
    n_firms_good_2 = models.IntegerField()
    firms_1 = models.StringField()
    firms_2 = models.StringField()

    assigned = models.StringField()
    choice = models.StringField()
    own = models.StringField()
    own_type = models.StringField()

    belief_1 = models.IntegerField(min=0, max=100)
    belief_2 = models.IntegerField(min=0, max=100)