from otree.api import (
    models,
    widgets,
    BaseConstants,
    BaseSubsession,
    BaseGroup,
    BasePlayer,
    Currency as c,
    currency_range,
)
import random
import itertools
import numpy as np
from django import forms
from django.forms import widgets as django_widgets
import math


author = 'Zheng Li'

doc = """
Intertemporal Choice
"""


class Constants(BaseConstants):
    name_in_url = 'inter_main_v4'

    players_per_group = None

    # desired number of subjects for sociodemographic representativeness
    # 1000000 to relax restriction
    num_subjects = 1000000
    # stop at 95% for a category
    cap = 0.95

    # pairs of tau2 and tau1#
    decisions = [
        # [0:6] choose 2 and repeat 1 from the chosen 2
        ['In 1 week', 'Today'],
        ['In 1 month', 'Today'],
        ['In 3 months', 'Today'],
        ['In 3 years', 'Today'],
        ['In 4 years', 'Today'],
        ['In 7 years', 'Today'],
        # [6:13] choose 2
        ['In 1 month and 1 week', 'In 1 month'],
        ['In 2 months', 'In 1 month'],
        ['In 6 months', 'In 1 month'],
        ['In 1 year', 'In 1 month'],
        ['In 2 years', 'In 1 month'],
        ['In 4 years', 'In 1 month'],
        ['In 7 years', 'In 1 month'],
        # [13:16] sub-additivity set 1
        ['In 6 months', 'Today'],
        ['In 1 year', 'In 6 months'],
        ['In 1 year', 'Today'],
        # [16:19] sub-additivity set 2
        ['In 2 years', 'In 1 year'],
        ['In 3 years', 'In 2 years'],
        ['In 3 years', 'In 1 year'],
        # [19:22] sub-additivity set 3
        ['In 1 year', 'Today'],
        ['In 2 years', 'In 1 year'],
        ['In 2 years', 'Today'],
    ]

    # possible payment amounts
    amounts = range(50, 62, 2)

    # step for price list
    steps = 2

    num_rounds = 12

    # generate the certainty-uncertainty spectrum bar
    list1 = np.arange(1 * 20, - 1, -1)
    list2 = []
    list2.append(0)
    for x in range(20):
        list2.append((x+1)*5)

    # example choice list
    example_set = [('In 30 days', 'Today', 20)]
    example_steps = 2


class Subsession(BaseSubsession):
    def creating_session(self):
        # set sociodemographic counters for representative filter
        # counters: age
        self.session.vars['age_1'] = 0
        self.session.vars['age_2'] = 0
        self.session.vars['age_3'] = 0
        self.session.vars['age_4'] = 0
        self.session.vars['age_5'] = 0
        self.session.vars['age_6'] = 0
        # counters: gender
        self.session.vars['gender_1'] = 0
        self.session.vars['gender_2'] = 0
        # counters: education
        self.session.vars['education_1'] = 0
        self.session.vars['education_2'] = 0
        self.session.vars['education_3'] = 0
        self.session.vars['education_4'] = 0
        # counters: income
        self.session.vars['income_1'] = 0
        self.session.vars['income_2'] = 0
        self.session.vars['income_3'] = 0
        self.session.vars['income_4'] = 0
        self.session.vars['income_5'] = 0
        self.session.vars['income_6'] = 0
        self.session.vars['income_7'] = 0
        self.session.vars['income_8'] = 0
        self.session.vars['income_9'] = 0
        # counters: race
        self.session.vars['race_1'] = 0
        self.session.vars['race_2'] = 0
        self.session.vars['race_3'] = 0
        self.session.vars['race_4'] = 0
        self.session.vars['race_5'] = 0

        if self.round_number == 1:
            for p in self.get_players():
                # set default participant-level variables
                p.participant.vars['tests_passed'] = True
                p.participant.vars['failed_attention'] = False
                p.participant.vars['failed_demographics'] = False
                p.participant.vars['failed_comprehension'] = False

                # define participant-level sociodemo indicators
                p.participant.vars['counter_age'] = p.counter_age
                p.participant.vars['counter_gender'] = p.counter_gender
                p.participant.vars['counter_education'] = p.counter_education
                p.participant.vars['counter_income'] = p.counter_income
                p.participant.vars['counter_race'] = p.counter_race

            # randomly assign ((tau2, tau1),pay_amount)'s
                # determine which 2 sub-additivity sets to use
                p.participant.vars['sub_'] = random.sample([1, 2, 3], 2)

                # generate decisions from sub-additivity set 1
                datessub1 = Constants.decisions[13:16]
                dates_sub1 = list()
                amount_tmp = random.choice(Constants.amounts)
                for date in datessub1:
                    date = (date, amount_tmp, 1)
                    dates_sub1.append(date)

                # generate decisions from sub-additivity set 2
                datessub2 = Constants.decisions[16:19]
                dates_sub2 = list()
                amount_tmp = random.choice(Constants.amounts)
                for date in datessub2:
                    date = (date, amount_tmp, 2)
                    dates_sub2.append(date)

                # generate decisions from sub-additivity set 3
                datessub3 = Constants.decisions[19:22]
                dates_sub3 = list()
                amount_tmp = random.choice(Constants.amounts)
                for date in datessub3:
                    date = (date, amount_tmp, 3)
                    dates_sub3.append(date)

                # keep decisions from the 2 selected sub-additivity sets
                if 1 not in p.participant.vars['sub_']:
                    dates_sub = dates_sub2 + dates_sub3
                    p.participant.vars['set_identifier'] = '(2,3)'
                elif 2 not in p.participant.vars['sub_']:
                    dates_sub = dates_sub1 + dates_sub3
                    p.participant.vars['set_identifier'] = '(1,3)'
                elif 3 not in p.participant.vars['sub_']:
                    dates_sub = dates_sub1 + dates_sub2
                    p.participant.vars['set_identifier'] = '(1,2)'

                # repeat 1 sub-additivity decision
                dates_sub_extra = random.choice(dates_sub)
                dates_sub.append(dates_sub_extra)

                # generate 2 decisions with earlier date = 'Today'
                dates_1 = list()
                dates1 = random.sample(Constants.decisions[0:6], 2)
                for date in dates1:
                    amount_tmp = random.choice(Constants.amounts)
                    date = (date, amount_tmp, 0)
                    dates_1.append(date)
                # repeat 1 decision with earlier date = 'Today'
                dates_1_extra = random.choice(dates_1)
                dates_1.append(dates_1_extra)

                # generate 2 decisions with earlier date = 'In 1 month'
                dates_2 = list()
                dates2 = random.sample(Constants.decisions[6:13], 2)
                for date in dates2:
                    amount_tmp = random.choice(Constants.amounts)
                    date = (date, amount_tmp, 0)
                    dates_2.append(date)

                # randomly order all decisions and make a decision list
                p.participant.vars['dates'] = random.sample(dates_sub + dates_1 + dates_2, 12)

        for p in self.get_players():
            # assign decisions from the decision list to each round
            p.set_current_task_number()

            # pass set_identifier from participant level to player level
            p.set_identifier = p.participant.vars['set_identifier']


class Group(BaseGroup):
    pass


class Player(BasePlayer):

    psid = models.StringField(blank=True)

    # whether the participant has been passing all tests so far
    tests_passed = models.BooleanField(initial=True)

    # whether this participant should receive base payment
    completed = models.IntegerField(initial=0)

# indicators for attention, sociodemographic, voucher, and comprehension checks
    failed_attention = models.BooleanField(initial=False)
    failed_sociodemo = models.BooleanField(initial=False)
    failed_voucher = models.BooleanField(initial=False)
    failed_comprehension = models.BooleanField(initial=False)

# attention check: daylight saving question
    attention_check = models.StringField(label="")
    attention_check_num = models.IntegerField(initial=0)

# sociodemographic information
    age = models.IntegerField(
        max=100,
        min=18,
        label="Your age:",
        blank=False,
    )

    gender = models.IntegerField(label="Your sex:")

    def gender_choices(self):
        choices=[
            [1, 'Male'],
            [2, 'Female']
        ]
        random.shuffle(choices)
        return choices

    education = models.IntegerField(
        choices=[
            [1, 'No high school graduation'],
            [2, 'Currently in High school or high school graduate'],
            [3, 'Some college, but no degree'],
            [4, "Bachelor's degree"],
            [5, "Graduate or professional degree"],
        ],
        label="Your highest educational attainment:",
        blank=False
    )

    income = models.IntegerField(
        choices=[
            [1, 'Below $15,000'],
            [2, '$15,000 - $24,999'],
            [3, '$25,000 - $34,999'],
            [4, '$35,000 - $49,999'],
            [5, '$50,000 - $74,999'],
            [6, '$75,000 - $99,999'],
            [7, '$100,000 - $149,999'],
            [8, '$150,000 - $199,999'],
            [9, '$200,000 or more'],
        ],
        label="Your approximate annual household income:",
        blank=False
    )

    race = models.IntegerField(label="Your race:")

    def race_choices(self):
        choices = [
            [1, 'White'],
            [2, 'Black'],
            [3, 'Hispanic'],
            [4, 'Asian'],
            [5, 'Other']
        ]
        random.shuffle(choices)
        return choices

    # indicators for types under each sociodemographic category
    counter_age = models.IntegerField(initial=0)
    counter_gender = models.IntegerField(initial=0)
    counter_education = models.IntegerField(initial=0)
    counter_income = models.IntegerField(initial=0)
    counter_race = models.IntegerField(initial=0)

    # which 2 of the 3 sub-additivity sets
    set_identifier = models.StringField()

    # current set: 0=non-SA, 1,2,3 correspond to SA
    sub_identifier = models.IntegerField()

# parameters of decisions
    # later payment amount at tau2
    pay_amount = models.IntegerField()
    # later date tau2
    time_date_lhs = models.StringField()
    # earlier date tau1
    time_date_rhs = models.StringField()

# whether player switched or not
    indicator_never_always_switcher = models.IntegerField()

    # indifferent payment amount
    switching_point = models.FloatField()
    # 0:0% -> 20: 100% confident
    confidence = models.FloatField()

# indicators of whether correctly answered each question
    #1
    qn_hypo_got_wrong = models.BooleanField(initial=False)
    #2
    qn_confidence_got_wrong = models.BooleanField(initial=False)
    #3
    qn_uncertain_got_wrong = models.BooleanField(initial=False)

# comprehension questions
    qn_hypo = models.IntegerField(
        choices=[
            [1, 'In making my decisions, I am asked to assume that I will actually receive all payments indicated, regardless of whether they take place now or in the future.'],
            [0, 'In making my decisions, I am asked to assume that it is less likely that I will actually receive payments that are meant to take place in the future.'],
            [0, 'In making my decisions, I am asked to assume that it is less likely that I will actually receive payments that are meant to take place now.'],
        ],
        widget=widgets.RadioSelect,
        blank=False,
        label=""
    )
    qn_confidence = models.FloatField(
        blank=False,
        label=""
    )
    qn_uncertain = models.IntegerField(
        choices=[
            [0, 'Uncertainty about whether I would actually receive the payments.'],
            [1, 'Uncertainty about how much I value the payments, assuming that I know I would receive them with certainty.'],
        ],
        widget=widgets.RadioSelect,
        blank=False,
        label=""
    )

# determine parameters in the current task
    def set_current_task_number(self):
        setattr(self, 'time_date_lhs', self.participant.vars['dates'][self.round_number - 1][0][0])
        setattr(self, 'time_date_rhs', self.participant.vars['dates'][self.round_number - 1][0][1])
        setattr(self, 'pay_amount', self.participant.vars['dates'][self.round_number - 1][1])
        setattr(self, 'sub_identifier', self.participant.vars['dates'][self.round_number - 1][2])

# set player.current_choice
    def current_choice(self):
        iet_choice = [(self.time_date_lhs, self.time_date_rhs, self.pay_amount)]
        return iet_choice

# determine number of table rows in price list page
    def table_length1(self):
        return abs(self.pay_amount) / Constants.steps + 1

# generate the list of earlier payments for option B in price list page
    def frange(self, start, stop, step):
        i = start + step
        while i < stop:
            if isinstance(i, int):
                yield i
            else:
                yield round(i, 2)
            i += step

# display the list of earlier payments for option B in price list page
    def right_side_amounts1(self):
        if self.pay_amount > 0:
            lst = self.frange(0, self.pay_amount + 1, Constants.steps)
        # creating a list of tuples from looping "lst" with the starting index = 1
        return list(enumerate(lst, 1))

# display the list of earlier payments for option B in example list page
    def right_side_amounts2(self):
        lst = self.frange(0, Constants.example_set[0][2] + 1, Constants.example_steps)
        return list(enumerate(lst, 1))

# determine if and how player switches
    def set_switching_point_and_indicator(self):
        if self.switching_point == 9999:
            # the template sets it to be 9999 when the player always chooses option A
            self.switching_point = self.pay_amount + Constants.steps
            self.indicator_never_always_switcher = 2
        elif self.switching_point == Constants.steps:
            # the template sets it to be 0 when the player always chooses option B
            self.indicator_never_always_switcher = 0
        else:
            # when a middle point is chosen (what is desired)
            self.indicator_never_always_switcher = 1