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 sel_measures import iat_blocks
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>Creative Writing Test</li>
    <li>Implicit Association Test</li>
    <li>Trolley Problem</li>
    <li>Ambiguity Task</li>
    <li>Big Five</li>
    <li>Grit Task</li>
    <li>Wason Logic Test</li>
    <li>Go No Go Test</li>
</ul>
"""


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

    solo_tasks_elective = ['moral', 'lying']
    solo_tasks_with_treatment = ['iat', 'redistribute', 'qv', 'ambiguity', 'trolley', 'self_affirmation',
                                 'student_adviser']

    # creativity
    creativity_min_for_task = 10
    creativity_bonus = c(5)

    # IAT
    LEFT, RIGHT = iat_blocks.LEFT, iat_blocks.RIGHT
    FIRST, SECOND = iat_blocks.LEFT, iat_blocks.RIGHT
    LEFT_KEYCODE = 69
    LEFT_KEY_NAME = '"E"'
    RIGHT_KEYCODE = 73
    RIGHT_KEY_NAME = '"I"'
    META_KEYCODE = 32
    META_KEY_NAME = "Space Bar"
    OR = "or"

    # trolley problem
    trolley_scenarios = ['bystander', 'bystander', 'footbridge', 'footbridge']
    trolley_skins = ['white', 'black']
    trolley_label_bystander = 'Is it morally permissible to divert the train and thus prevent five deaths for the price' \
                               'of a lifetime?'
    trolley_label_footbridge = 'Is it morally permissible to push man over the catwalk in order to save' \
                                'the other five people?'

    # ambiguity
    ambiguity_num_balls = 60
    ambiguity_balls_red = 20
    ambiguity_balls_other = ambiguity_num_balls - ambiguity_balls_red
    ambiguity_bals_colors = ['red', 'black', 'white']
    ambiguity_bonus = c(6)
    ambiguity_bonus_version2_high = c(10)
    ambiguity_bonus_version2_low = c(0.5)

    # big five
    bigfive_items = ['imagination', 'job', 'trust', 'relaxed', 'critic', 'nervous', 'reserved', 'lazy', 'sociable',
                     'artistic']
    bigfive_label_general = 'Je me vois comme quelqu’un...'
    bigfive_labels = ['Who has an active imagination', 'Who is rigorous in his work',
                       'Who can be trusted', 'Who is relaxed and has good stress management',
                       'Who tends to notice the faults of others', 'Who gets nervous quickly',
                       'Who is reserved', 'Who tends to be lazy',
                       'Who is outgoing, sociable', 'Who has little artistic interest']
    bigfive_choices = [(n, str(n)) for n in range(1, 6)]

    # 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

    # risky behaviour
    risky_options = ['Never', 'Once or twice', 'Once a month or less', '2 or 3 days a month',
                      'Once or twice a week', '3 or 5 days a week', 'Almost every day']
    risky_smoke_choices = [(n, l) for n, l in enumerate(risky_options)]
    risky_drink_choices = [(n, l) for n, l in enumerate(risky_options)]

    # self-confidence
    self_confidence_label = 'How do you think your level is compared to the students in your class in the' \
                            'following disciplines:'

    # wason
    wason_cards = ['d', 'k', '3', '7']
    wason_corrects = [True, False, False, True]

    # gonogo
    gonogo_num_go_trials = 20
    gonogo_num_no_go_trials = 5
    gonogo_time_limit = 2000

    # empathy
    empathy_labels = [
       "It pains me when someone is treated unfairly",
        "I have compassion for people who are unlucky",
        "When someone has been mistreated, I want to protect that person",
        "I have compassion for people who have problems",
        "I am a soft hearted person (with a good heart / compassion)",
        "I try to understand how other people feel",
        "I can put myself in someone else's shoes and understand that person's grief",
        "My friends tell me about their problems", "I know immediately when a friend of mine is sad"]

    # curiosity
    curiosity_labels = [
        "It's fun to take something apart and see what's inside",
        "When I don't know the answer to a question, it annoys me. So I stubbornly try to get it,"
        "Mysterious and unknown things arouse my curiosity",
        "I occupy my leisure time with interesting activities",
        "I ask questions about things I don't understand",
        "Knowing that I'm going to learn new things turns me on",
        "I like discovering new places", "I am very enthusiastic about everything I do"
    ]


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

        # grit
        numbers_easy = [i for pair in [(n, 100 - n) for n in [randint(1, 100 - 1) for _ in range(5)]] for i in pair]
        numbers_difficult = [i for pair in [(n, 100 - n) for n in [randint(1, 100 - 1) for _ in range(10)]] for i in
                             pair]
        numbers_easy_example = [i for pair in [(n, 100 - n) for n in [randint(1, 100 - 1) for _ in range(5)]] for i in
                                pair]
        numbers_difficult_example = [i for pair in [(n, 100 - n) for n in [randint(1, 100 - 1) for _ in range(10)]] for
                                     i in pair]
        self.session.vars['table_easy_example'] = [numbers_easy_example[n:n + 2] for n in
                                                   range(0, len(numbers_easy_example), 2)]
        self.session.vars['table_difficult_example'] = [numbers_difficult_example[n:n + 4] for n in
                                                        range(0, len(numbers_difficult_example), 4)]

        # gonogo
        gonogo_trials = ['go' for n in range(Constants.gonogo_num_go_trials)] + ['nogo' for n in range(
            Constants.gonogo_num_no_go_trials)]

        for p in self.get_players():

            p.participant.vars['year'] = 1

            # grit
            shuffle(numbers_easy)
            shuffle(numbers_difficult)
            table_easy = [numbers_easy[n:n + 2] for n in range(0, len(numbers_easy), 2)]
            table_difficult = [numbers_difficult[n:n + 4] for n in range(0, len(numbers_difficult), 4)]
            p.participant.vars['table_easy'] = table_easy
            p.participant.vars['table_difficult'] = table_difficult

            # gonogo
            shuffle(gonogo_trials)
            p.participant.vars['gonogo_trials'] = gonogo_trials

            # trolley
            p.trolley_scenario = next(treatments['trolley_scenarios'])
            p.trolley_skin = next(treatments['trolley_skins'])

            # ambiguity
            p.ambiguity_version = next(treatments['ambiguity'])

            # iat
            iat_version = next(treatments['iat'])
            if p.participant.vars.get('self_affirmation_version') and iat_version == 'gender_stem':
                p.iat_version = next((treatments['iat']))
            else:
                p.iat_version = iat_version

            # creativity rating
            p.creativity_rating_see_names = next(treatments['creativity_rating_with_name'])
            p.creativity_rating_with_incentive = next(treatments['creativity_rating_with_incentive'])


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # creativity
    creativity = models.LongStringField(label='')
    creativity_rating_see_names = models.BooleanField()
    creativity_rating_with_incentive = models.BooleanField()
    rating_text_left = models.IntegerField()
    rating_text_right = models.IntegerField()
    creativity_rating = models.StringField(
        choices=[('L', 'Text on the left'), ('R', 'Text on the right')], widget=widgets.RadioSelect())
    creativity_num_evaluations = models.IntegerField(initial=0)
    creativity_rating_estimate = models.StringField(
        choices=[('L', 'Text on the left'), ('R', 'Text on the right')], widget=widgets.RadioSelect())

    # IAT
    iat_version = models.StringField()
    iat_tables = models.LongStringField()

    # trolley problem
    trolley_scenario = models.StringField()
    trolley_skin = models.StringField()
    trolley_decision = models.BooleanField(label='', choices=[(0, "No, it is not morally allowed"),
                                                              (1, "Yes, it is morally permitted")],
                                           widget=widgets.RadioSelect)

    # ambiguity
    ambiguity_version = models.IntegerField()
    ambiguity_decision = models.StringField(label='Which lottery do you choose?',
                                            choices=[('A', 'Lottery A'), ('B', 'Lottery B')])
    ambiguity_payoff = models.CurrencyField()

    def set_ambiguity_payoff(self):
        black_balls = randint(0, Constants.ambiguity_balls_other)
        balls = ['red' for _ in range(Constants.ambiguity_balls_red)] + ['black' for _ in range(black_balls)] + \
                ['white' for _ in range(Constants.ambiguity_balls_other - black_balls)]
        ball = choice(balls)
        if self.ambiguity_decision == 'A':
            if ball == 'red':
                if self.ambiguity_version == 2:
                    self.ambiguity_payoff = Constants.ambiguity_bonus_version2_high
                else:
                    self.ambiguity_payoff = Constants.ambiguity_bonus
            elif ball == 'black':
                self.ambiguity_payoff = 0
            elif ball == 'white':
                if self.ambiguity_version == 1:
                    self.ambiguity_payoff = choice([0, Constants.ambiguity_bonus_version2_low])
                else:
                    self.ambiguity_payoff = choice([0, Constants.ambiguity_bonus])
        else:
            if ball == 'red':
                self.ambiguity_payoff = 0
            elif ball == 'black':
                if self.ambiguity_version == 2:
                    self.ambiguity_payoff = Constants.ambiguity_bonus_version2_high
                else:
                    self.ambiguity_payoff = choice([0, Constants.ambiguity_bonus])
            elif ball == 'white':
                if self.ambiguity_version == 2:
                    self.ambiguity_payoff = choice([0, Constants.ambiguity_bonus_version2_high])
                else:
                    self.ambiguity_payoff = Constants.ambiguity_bonus

    # big five
    for t, l in zip(Constants.bigfive_items, Constants.bigfive_labels):
        locals()[t] = models.IntegerField(label=l)
    del t, l

    # grit
    grit_decision = models.StringField(label='Now please choose which type of grid you want to use.',
                                       choices=[('easy ',' An easy grid'),
                                                ('difficult ',' A difficult grid')], widget=widgets.RadioSelect())
    grit_num_correct = models.IntegerField()
    grit_payoff = models.CurrencyField()

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

    # risky behaviour
    risky_smoke = models.IntegerField(label='During the past 12 months, how often have you smoked?',
                                      choices=Constants.risky_smoke_choices, widget=widgets.RadioSelect)
    risky_drink = models.IntegerField(label='During the past 12 months, how often have you been drunk?',
                                      choices=Constants.risky_drink_choices, widget=widgets.RadioSelect)

    # self-confidence
    self_confidence = models.StringField()
    ses_no = models.BooleanField(label="I didn't take this class.", widget=widgets.CheckboxInput, blank=True)
    math_no = models.BooleanField(label="I didn't take this class.", widget=widgets.CheckboxInput, blank=True)

    # wason
    for c in Constants.wason_cards:
        locals()[f'wason_card_{c}'] = models.BooleanField(initial=False)
    del c
    wason_is_correct = models.BooleanField()

    def set_wason_is_correct(self):
        if [getattr(self, f'wason_card_{c}') for c in Constants.wason_cards] == Constants.wason_corrects:
            self.wason_is_correct = True
        else:
            self.wason_is_correct = False

    # gonogo
    gonogo_trials = models.LongStringField()

    # empathy
    for n, l in enumerate(Constants.empathy_labels):
        locals()[f'empathy_{n}'] = models.IntegerField(label=l)
    del n, l

    # curiosity
    for n, l in enumerate(Constants.curiosity_labels):
        locals()[f'curiosity_{n}'] = models.IntegerField(label=l)
    del n, l

    # technical issues
    technical_issues = models.BooleanField(
        label="Did you encounter any technical difficulties while answering these questions?",
        choices=[(True, "Yes, there were some technical issues."),
                 (False, "No, there were no technical issues.")], widget=widgets.RadioSelect()
    )
    distractions = models.BooleanField(
        label="Were you distracted by other students while answering these questions?",
        choices=[(True, "Yes"), (False, "No")], widget=widgets.RadioSelect()
    )

    # email
    email = EmailField(max_length=70, blank=True)

    comments = models.LongStringField(label="If you have any comments about the games or questions, "
                                            "please let us know below.")