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

    tasks_baseline = {
        {
            'name': 'risk_base_25_1_10_21',
            'type': 'risk',
            'specification': ['base', 25, '1%', 10, '21%', 0]
        },
        {
            'name': 'risk_base_25_40_10_60',
            'type': 'risk',
            'specification': ['base', 25, '40%', 10, '60%', 0]
        },
        {
            'name': 'risk_base_25_80_10_100',
            'type': 'risk',
            'specification': ['base', 25, '80%', 10, '100%', 1]
        },
    }
    tasks_treatment = {
        {
            'name': 'risk_treat_25_60_10_40_y61',
            'type': 'risk',
            'specification': ['treat', 25, 'Y reduced by 60%', 10, 'Y reduced by 40%', 61]
        },
        {
            'name': 'risk_treat_25_21_10_1_y61',
            'type': 'risk',
            'specification': ['treat', 25, 'Y reduced by 21%', 10, 'Y reduced by 1%', 61]
        },
        {
            'name': 'risk_treat_25_19_10_39_y61',
            'type': 'risk',
            'specification': ['treat', 25, 'Y plus 19%', 10, 'Y plus 39%', 61]
        },
    }

    num_rounds = len(tasks_baseline) + len(tasks_treatment)
    intertemporal_rounds = len(tasks_intertemporal['baseline'])
    cons_choice_rounds = len(tasks_cons_choice['baseline'])
    risk_rounds = len(tasks_risk['baseline'])


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_risk'] = random.choice(['baseline', 'treatment'])
                p.participant.vars['condition_time'] = random.choice(['baseline', 'treatment'])
                p.participant.vars['condition_wine'] = random.choice(['baseline'])
                blocks_and_pages_shuffled_risk = random.sample([random.choice([[0, 1], [0, 1]]), random.choice([[2, 3], [3, 2]])], 2)
                blocks_and_pages_shuffled_time = random.sample([random.choice([[0, 1], [0, 1]]), random.choice([[2, 3], [3, 2]])], 2)
                p.participant.vars['page_sequence_risk'] = [item for sublist in blocks_and_pages_shuffled_risk for item in sublist]
                p.participant.vars['page_sequence_time'] = [item for sublist in blocks_and_pages_shuffled_time for item in sublist]
                p.participant.vars['page_sequence_wine'] = random.sample(range(0, 3), 3)
        for p in self.get_players():
            p.condition_risk = p.participant.vars['condition_risk']
            p.condition_time = p.participant.vars['condition_time']
            p.condition_cons_choice = p.participant.vars['condition_wine']
            p.set_current_task_number()


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_risk = models.StringField()
    condition_time = models.StringField()
    condition_cons_choice = models.StringField()

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

    choice_time = models.CharField()
    choice_risk = models.CharField()
    choice_wine = models.CharField()

    def set_current_task_number(self):
        if self.round_number <= Constants.intertemporal_rounds:
            index = self.participant.vars['page_sequence_time'][self.round_number - 1]
            setattr(self, 'task_identifier', Constants.tasks_intertemporal[self.condition_time][index]['name'])
        elif self.round_number > Constants.intertemporal_rounds + Constants.cons_choice_rounds:
            index = self.participant.vars['page_sequence_risk'][self.round_number - Constants.intertemporal_rounds - Constants.cons_choice_rounds - 1]
            setattr(self, 'task_identifier', Constants.tasks_risk[self.condition_risk][index]['name'])
        else:
            index = self.participant.vars['page_sequence_wine'][self.round_number - Constants.intertemporal_rounds - 1]
            setattr(self, 'task_identifier', Constants.tasks_cons_choice[self.condition_cons_choice][index]['name'])

    def specification(self):
        if self.round_number <= Constants.intertemporal_rounds:
            index = self.participant.vars['page_sequence_time'][self.round_number - 1]
            return Constants.tasks_intertemporal[self.condition_time][index]['specification']
        elif self.round_number > Constants.intertemporal_rounds + Constants.cons_choice_rounds:
            index = self.participant.vars['page_sequence_risk'][self.round_number - Constants.intertemporal_rounds - Constants.cons_choice_rounds - 1]
            return Constants.tasks_risk[self.condition_risk][index]['specification']
        else:
            index = self.participant.vars['page_sequence_wine'][self.round_number - Constants.intertemporal_rounds - 1]
            return Constants.tasks_cons_choice[self.condition_cons_choice][index]['specification']

    # def set_type(self):
    #     index = self.participant.vars['page_sequence_vign'][self.round_number - 1]
    #     return Constants.task_sets[int(self.round_number-1)][index]['type']