from otree.api import (
    models, widgets, BaseConstants, BaseSubsession, BaseGroup, BasePlayer,
    Currency as c, currency_range
)
from django.db.models import EmailField
from treatments import *
from random import shuffle, randint, choice

author = 'Tommaso Batistoni - t.batistoni@ucl.ac.uk'

doc = """
Set of measures to capture socio-emotional dimensions at the individual level.<br>
The set includes among others:<br>
<ul>
    <li>Grit Task</li>
</ul>
"""


class Constants(BaseConstants):
    name_in_url = 'reduced_sel'
    players_per_group = None
    num_rounds = 1

    # grit
    grit_num_correct_for_bonus = 3
    grit_time_for_task = 1
    grit_bonus_high = c(8)
    grit_bonus_low = c(3)
    grit_total = 100
    grit_num_pairs_easy = 5
    grit_num_pairs_difficult = 10
    grit_choice_label = 'Now please choose which type of grid you want to use.'
    grit_choices = [('easy', 'An easy grid '), (' difficult ',' A difficult grid ')]


class Subsession(BaseSubsession):
    def creating_session(self):

        # grit
        numbers_easy_example = self.get_easy_numbers()
        numbers_difficult_example = self.get_difficult_numbers()
        self.session.vars['table_easy_example'] = self.easy_numbers_to_table(numbers_easy_example)
        self.session.vars['table_difficult_example'] = self.difficult_numbers_to_table(numbers_difficult_example)

        numbers_easy_round1 = self.get_easy_numbers()
        numbers_difficult_round1 = self.get_difficult_numbers()
        numbers_easy_round2 = self.get_easy_numbers()
        numbers_difficult_round2 = self.get_difficult_numbers()

        for p in self.get_players():
            # grit
            shuffle(numbers_easy_round1)
            shuffle(numbers_difficult_round1)
            p.participant.vars['table_easy_round1'] = self.easy_numbers_to_table(numbers_easy_round1)
            p.participant.vars['table_difficult_round1'] = self.difficult_numbers_to_table(numbers_difficult_round1)
            p.grit_round1_hardship = next(treatments['grit_round1_hardship'])

            shuffle(numbers_easy_round2)
            shuffle(numbers_difficult_round2)
            p.participant.vars['table_easy_round2'] = self.easy_numbers_to_table(numbers_easy_round2)
            p.participant.vars['table_difficult_round2'] = self.difficult_numbers_to_table(numbers_difficult_round2)

    def get_easy_numbers(self):
        return [i for pair in [(n, 100 - n) for n in [randint(1, 100 - 1) for _ in range(5)]] for i in pair]

    def get_difficult_numbers(self):
        return [i for pair in [(n, 100 - n) for n in [randint(1, 100 - 1) for _ in range(10)]] for i in pair]

    def easy_numbers_to_table(self, numbers):
        return [numbers[n:n + 2] for n in range(0, len(numbers), 2)]

    def difficult_numbers_to_table(self, numbers):
        return [numbers[n:n + 4] for n in range(0, len(numbers), 4)]


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # grit
    grit_round1_hardship = models.BooleanField()
    grit_decision_round1 = models.StringField(
        label='',
        choices=Constants.grit_choices,
        widget=widgets.RadioSelect()
    )
    grit_decision_round2 = models.StringField(
        label='',
        choices=Constants.grit_choices,
        widget=widgets.RadioSelect()
    )
    grit_num_correct_round1 = models.IntegerField()
    grit_num_correct_round2 = models.IntegerField()
    grit_payoff_round1 = models.CurrencyField()
    grit_payoff_round2 = models.CurrencyField()

    def get_grit_round1_realized_choice(self):
        if self.grit_round1_hardship:
            return 'difficult'
        else:
            return self.grit_decision_round1

    def set_grit_payoff_round1(self):
        self.grit_payoff_round1 = self._get_grit_payoff(
            self.grit_num_correct_round1,
            self.get_grit_round1_realized_choice()
        )

    def set_grit_payoff_round2(self):
        self.grit_payoff_round2 = self._get_grit_payoff(
            self.grit_num_correct_round2,
            self.grit_decision_round2
        )

    def _get_grit_payoff(self, grit_num_correct, grit_decision):
        grit_payoff = 0
        if grit_num_correct >= Constants.grit_num_correct_for_bonus:
            if grit_decision == 'easy':
                grit_payoff = Constants.grit_bonus_low
            else:
                grit_payoff = Constants.grit_bonus_high
        return grit_payoff