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 = 'time_framing_both_conditions'
    players_per_group = None

    tasks = [
        {
            'name': 'time_base_26_0_27_7',
            'specification': ['base', 26, 'in 0 days', 27, 'in 7 days', 0]
        },
        {
            'name': 'time_base_26_92_27_99',
            'specification': ['base', 26, 'in 92 days', 27, 'in 99 days', 0]
        },
        {
            'name': 'time_base_26_0_40_99',
            'specification': ['base', 26, 'in 0 days', 40, 'in 99 days', 0]
        },
        {
            'name': 'time_base_26_99_27_92',
            'specification': ['treat', 26, '99 days earlier than X', 27, '92 days earlier than X', 99]
        },
        {
            'name': 'time_base_26_7_27_0',
            'specification': ['treat', 26, '7 days earlier than X', 27, '0 days earlier than X', 99]
        },
        {
            'name': 'time_base_26_99_40_0',
            'specification': ['treat', 26, '99 days earlier than X', 40, '0 days earlier than X', 99]
        },
    ]
    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['condition_time'] = random.choice(['baseline', 'treatment'])
                if p.participant.vars['condition_time'] == 'baseline':
                    p.participant.vars['page_sequence_time'] = random.sample(range(0, 3), 3) + random.sample(range(3, 6), 3)
                else:
                    p.participant.vars['page_sequence_time'] = random.sample(range(3, 6), 3) + random.sample(range(0, 3), 3)
        for p in self.get_players():
            p.condition_time = p.participant.vars['condition_time']
            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()
    condition_time = models.StringField()
    task_number = models.IntegerField()
    task_type = models.StringField()
    control_time = 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_time = models.CharField()

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

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