from otree.api import (
    models,
    widgets,
    BaseConstants,
    BaseSubsession,
    BaseGroup,
    BasePlayer,
    Currency as c,
    currency_range,
)

import numpy as np
import math

np.random.seed(1)

author = 'Jason Ralston'

doc = """
Savings/Consumption task with cognitive load
"""

def count_matches(num1, num2):
    list_num1 = [int(d) for d in str(num1)]
    list_num2 = [int(d) for d in str(num2)]

    num_correct = 0
    for i in range(len(list_num1)):
        if list_num1[i] == list_num2[i]:
            num_correct = num_correct + 1

    return num_correct



class Constants(BaseConstants):
    name_in_url = 'SavingsConsumption_poorANDlowload'
    players_per_group = None
    num_rounds = 20

    a_income_poor = -20
    a_income_rich = 0

    mode_income_poor = 5
    mode_income_rich = 25

    b_income_poor = 30
    b_income_rich = 50

    sd_income = 20

    alpha = 0.02
    multiplier = 250



class Subsession(BaseSubsession):

    def creating_session(self):
        if self.round_number == 1:
            self.session.vars['income_poor'] = np.around(np.random.triangular(Constants.a_income_poor,
                                                                              Constants.mode_income_poor,
                                                                              Constants.b_income_poor,
                                                                              Constants.num_rounds), 0).tolist()
            self.session.vars['income_rich'] = np.around(np.random.triangular(Constants.a_income_rich,
                                                                              Constants.mode_income_rich,
                                                                              Constants.b_income_rich,
                                                                              Constants.num_rounds), 0).tolist()
            self.session.vars['low_draws'] = np.random.randint(10, 98, Constants.num_rounds).tolist()
            self.session.vars['high_draws'] = np.random.randint(10000000, 99999998, Constants.num_rounds).tolist()



class Group(BaseGroup):
    pass


class Player(BasePlayer):

    # Fields the user has control over
    consumption = models.FloatField(label="How many tokens would you like to convert to points?")
    low_memorization_number = models.IntegerField(min=10, max=99, label="What number were you asked to memorize?")
    high_memorization_number = models.IntegerField(min=10000000, max=99999999, label="What number were you asked to memorize?")

    # Player attributes
    current_utility = models.FloatField(initial=0)
    total_utility = models.FloatField(initial=0, min=0)
    current_savings = models.FloatField(initial=0)
    total_savings = models.FloatField(initial=0)
    mem_num_correct = models.FloatField()
    income = models.FloatField()
    temp_utility = models.FloatField()

    def set_income(self):
        self.income = self.session.vars['income_poor'][self.subsession.round_number - 1]

    def utility_calc(self):
        self.current_utility = np.around(Constants.multiplier*(1 - math.exp(-Constants.alpha * self.consumption)), 2)

    def savings_calc(self):
        self.current_savings = np.around(self.income - self.consumption, 2)

    def update_savings(self):
        if self.subsession.round_number == 1:
            self.total_savings = self.current_savings
        else:
            self.total_savings = self.in_round(self.subsession.round_number - 1).total_savings + self.current_savings

    def update_utility(self):
        if self.subsession.round_number == 1:
            self.total_utility = self.current_utility
        else:
            self.total_utility = self.in_round(self.subsession.round_number - 1).total_utility + self.current_utility

    def mem_frac_correct(self):
        if self.subsession.round_number != Constants.num_rounds:
            self.mem_num_correct = count_matches(self.low_memorization_number,
                                                 self.session.vars['low_draws'][self.subsession.round_number - 1])
        else:
            self.mem_num_correct = 2

    def frac_utility(self):
        self.temp_utility = self.current_utility
        self.current_utility = self.current_utility * (self.mem_num_correct / len([int(d) for d in str(self.session.vars['low_draws'][self.subsession.round_number - 1])]))

    def set_consumption(self):
        self.consumption = self.income + self.in_round(self.subsession.round_number - 1).total_savings

    def get_total_savings(self):
        if self.subsession.round_number > 1:
            return self.in_round(self.subsession.round_number - 1).total_savings
        else:
            return 0

    def task_points(self):
        if self.session.config['app_sequence'].index('SavingsConsumption_poorANDlowload') == 1:
            self.participant.vars['payoff_1'] = self.total_utility
        else:
            self.participant.vars['payoff_2'] = self.total_utility

    def last_period_calcs(self):
        if self.subsession.round_number == Constants.num_rounds:
            self.set_consumption()
            self.utility_calc()
            self.savings_calc()
            self.task_points()
        else:
            pass