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

author = 'Thomas Graeber'

doc = """
Assign treatment status.
"""


class Constants(BaseConstants):
    """Contains constants of the current experiment app."""
    name_in_url = 'real_effort_tax'
    players_per_group = None

    tasks = [
        {
            'name': 'effort_tax_3_23',
            'specification': ['3%', '23%']
        },
        {
            'name': 'effort_tax_40_60',
            'specification': ['40%', '60%']
        },
        {
            'name': 'effort_tax_77_97',
            'specification': ['77%', '97%']
        },
    ]
    growth_rates = [5, 40]
    num_rounds = len(tasks)


class Subsession(BaseSubsession):
    """Contains subsession-level objects of the current experiment app."""

    def before_session_starts(self):
        """Depending on session configs, get respective list of treatment conditions from constants.
        Then distribute treatment conditions and groups to participants.
        """
        if self.round_number == 1:
            for p in self.get_players():
                p.participant.vars['page_sequence_effort_tax'] = random.sample(range(0, 3), 3)
        for p in self.get_players():
            p.set_current_task_number()
            if self.round_number == Constants.num_rounds:
                q = random.sample(range(0, 2), 2)
                p.income_growth_condition_q1 = Constants.growth_rates[q[0]]
                p.income_growth_condition_q2 = Constants.growth_rates[q[1]]



class Group(BaseGroup):
    """Contains group-level objects of the current experiment app."""
    pass


class Player(BasePlayer):
    """Contains player-level objects of the current experiment app."""

    task_identifier = models.StringField()
    task_number = models.IntegerField()
    income_growth_condition_q1 = models.IntegerField()
    income_growth_condition_q2 = models.IntegerField()
    income_growth_guess1 = models.IntegerField()
    income_growth_guess2 = models.IntegerField()

    choice_task = models.CharField()

    control_probability = models.IntegerField()
    control_variable = models.IntegerField()

    def set_current_task_number(self):
        setattr(self, 'task_number', self.participant.vars['page_sequence_effort_tax'][self.round_number - 1])
        setattr(self, 'task_identifier', Constants.tasks[self.task_number]['name'])

    def specification(self):
        return Constants.tasks[self.task_number]['specification']