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

from Multiple_Price_List.config import *
import random
from random import randrange


# ---------------------------------------------------------------------------------------------------------------- #
# *** CLASS CONSTANTS *** #
# ---------------------------------------------------------------------------------------------------------------- #

class Constants(BaseConstants):
    # ------------------------------------------------------------------------------------------------------------ #
    # --- oTree Settings (Don't Modify) --- #
    # ------------------------------------------------------------------------------------------------------------ #
    name_in_url = 'Risk'
    players_per_group = None
    num_rounds = 12 # this number must not be larger than 12 !!

    # ---------------------------------------------------------------------------------------------------------------- #
    # --- Task-specific Settings --- #
    # ---------------------------------------------------------------------------------------------------------------- #

    # number of choices between "lottery A" and "lottery B"
    # note that the number of choices determines the probabilities of high and low outcomes of lotteries "A" and "B"
    # for <num_choices = X>, the probability of outcome "high" is 1/X for the first choice, 2/X for the second, etc.
    num_choices = 16  # make sure this number is not larger than 20!!!

    # set enforce consistency to True to ensure that subjects cannot switch mutliple times
    enforce_consistency = True

    #define indices
    index_rounds = [j for j in range(1, num_rounds + 1)]
    index_choices = [j for j in range(1, num_choices + 1)]

    #define lists - 1st entry indicates high outcome of lottery, 2nd entry indicates low outcome of lottery and 3rd
    #entry indicates prob corresponding to high outcome
    list_1 = [12, 0, .17, 0.3, 0.3] # the fourth entry in the list is the starting value for the option b choice and the the fifth value is the increment increase
    list_2 = [12, 0, 0.5, 0.5, 0.5] # that's the first lottery we use for the time preferences task
    list_3 = [12, 0, .33, 0.2, 0.5]
    list_4 = [8, 0, .33, 0.4, 0.3]
    list_5 = [14, 7, .33, 7.0, 0.2]
    list_6 = [10, 5, .33, 4.9, 0.2]
    list_7 = [12.8, 4.8, .33, 5, 0.4]
    list_8 = [12, 0, 0.67, 0.6, 0.6]
    list_9 = [18, 6, .50, 6, 0.6] # that's the second lottery we use for the time preferences task
    list_10 = [12, 0, .83, 1.8, 0.6]
    list_11 = [10, 5, .33, 0.4, 0.2] # that's the fake list to capture noise
    list_12 = [10.5, 3.5, 0.5, 5, 0.20]  # that's the third lottery we use for the time preferences task

    # tie all lists together in a list called all_lists
    all_lists = [list_1, list_2, list_3, list_4, list_5, list_6, list_7, list_8, list_9, list_10,
                 list_11, list_12,]

    #create variable for low probability
    for l in all_lists:
        prob_lo = 1 - l[2]
        l.append(prob_lo)

    #create option_b choices depending on how many lines we want to have
    for l in all_lists:
        option_b = []
        for i in index_choices:
            option = l[3] + ((num_choices - i) * l[4])
            option = c(option)
            option_b.append(option)
        l.append(option_b)

    # delete elements 3 and 4 (starting value for option b and increment increase) in lists because you don't need them anymore
    for l in all_lists:
        del l[4]
        del l[3]

    # multiply all constant options for each line
    for l in all_lists:
        l[0] = c(l[0])
        lot_hi = [l[0]]*num_choices
        l[0] = lot_hi

        l[1]= c(l[1])
        lot_lo = [l[1]]*num_choices
        l[1] = lot_lo

        prob_hi = [l[2]]*num_choices
        l[2] = prob_hi

        prob_lo = [l[3]]*num_choices
        l[3] = prob_lo

    # add index to lists
    index_mpl = [j for j in range(1, len(all_lists) + 1)]
    for (a, b) in zip(all_lists, index_mpl):
        index = b
        a.append(index)



# ******************************************************************************************************************** #
# *** CLASS SUBSESSION
# ******************************************************************************************************************** #

class Subsession(BaseSubsession):

    def creating_session(self):
        for p in self.get_players():
            pv = p.participant.vars

            # ------------------------------------------------------------------
            # 0) discount_seq：整体跑时不要覆盖 Welcome 分配；单测 Risk 时才给兜底
            # ------------------------------------------------------------------
            if pv.get('discount_seq') is None:
                # ✅ 单测 Risk 才会走到这里；整体流程 Welcome 已经写好了就不会改
                pv['discount_seq'] = 4   # 你想单测固定成 1 就留 1；不想固定可 random.randint(1,4)

            seq = pv['discount_seq']

            # ------------------------------------------------------------------
            # 1) 只在 round 1 设定 path/order，并清掉跨 app 残留标记
            #    （避免出现你看到的：seq=1 但 order=84 -> round1 直接进 8w）
            # ------------------------------------------------------------------
            if self.round_number == 1:
                # 清理可能从别的 app 残留的强制切段标记
                pv.pop('force_8w', None)
                pv.pop('force_4w', None)

                # 清理“只显示一次”的标记（可选，但对反复测试很有用）
                pv.pop('ins8w_shown', None)
                pv.pop('quiz8w_shown', None)
                pv.pop('mid8w_shown', None)
                pv.pop('ins4w_shown', None)
                pv.pop('quiz4w_shown', None)
                pv.pop('mid4w_shown', None)

                # seq=1：Lot1→CE + 48（先4w后8w）
                # seq=2：CE→Lot2 + 48
                # seq=3：Lot1→CE + 84（先8w后4w）
                # seq=4：CE→Lot2 + 84
                pv['path']  = 'Lot1→CE' if seq in (1, 3) else 'CE→Lot2'
                pv['order'] = '48' if seq in (1, 2) else '84'

            # --- your existing MPL randomization & exports (unchanged) ---
            list_2 = Constants.all_lists[1]
            p.participant.vars['lot_time_task'] = [list_2]

            rand_sequence = random.sample([i for i in Constants.all_lists if i != list_2], Constants.num_rounds - 1)
            rand_sequence.insert(random.randrange(0, Constants.num_rounds), list_2)

            mpl_sequence = []
            for l in rand_sequence:
                element = l[:-1]
                mpl_sequence.append(element)

            for l in mpl_sequence:
                new_prob = []
                for i in l[2]:
                    prob = "{0:.0f}".format(i * 100) + "%"
                    new_prob.append(prob)
                l[2]=new_prob

            for l in mpl_sequence:
                new_prob = []
                for i in l[3]:
                    prob = "{0:.0f}".format(i * 100) + "%"
                    new_prob.append(prob)
                l[3]=new_prob

            rand_mpl_index = []
            for l in rand_sequence:
                index = l[-1]
                rand_mpl_index.append(index)

            form_fields = ['choice_' + str(k) for k in Constants.index_choices]
            for l in mpl_sequence:
                l.append(form_fields)

            choice_index = [j for j in range(1, Constants.num_choices + 1)]
            mpl_sequence.append(choice_index)

            p.participant.vars['rand_sequence'] = rand_sequence
            p.participant.vars['mpl_sequence'] = mpl_sequence
            p.participant.vars['rand_mpl_index'] = rand_mpl_index
            p.participant.vars['form_fields'] = form_fields

            p.participant.vars['mpl_choices_made'] = [None for j in range(1, Constants.num_choices + 1)]

            p.participant.vars['cert_equi'] = []
            p.participant.vars['cert_equi_list2'] = []
            p.participant.vars['cert_equi_list9'] = []
            p.participant.vars['cert_equi_list12'] = []

            # === NEW: map discount_seq -> path/order for time tasks (TA_Lot1 / TA_CE / TA_Lot_2) ===
            # 如果你已经在别处设置了 discount_seq，这里直接读取；否则就随机一个并存起来。
            #seq = p.participant.vars.get('discount_seq')
            #if seq is None:
            #    seq = random.randint(1, 4)
            #    p.participant.vars['discount_seq'] = seq

            # seq=1：Lot1_48, CE_48
            # seq=2：CE_48,  Lot2_48
            # seq=3：Lot1_84, CE_84
            # seq=4：CE_84,  Lot2_84
            #if seq in (1, 3):
            #    p.participant.vars['path'] = 'Lot1→CE'   # 先跑 Lot1，再跑 CE（Lot2 整个被 is_displayed 跳过）
            #else:
            #    p.participant.vars['path'] = 'CE→Lot2'   # 先跑 CE，再跑 Lot2（Lot1 整个被 is_displayed 跳过）

            #p.participant.vars['order'] = '48' if seq in (1, 2) else '84'  # 段内顺序：48 或 84
            # === /NEW ===

            if 'risk_round_to_pay' not in pv:
                pv['risk_round_to_pay'] = random.randint(1, Constants.num_rounds)

            if 'risk_row_to_pay' not in pv:
                pv['risk_row_to_pay'] = random.randint(1, Constants.num_choices)

            if 'risk_choice_field' not in pv:
                pv['risk_choice_field'] = f'choice_{pv["risk_row_to_pay"]}'

            if 'risk_random_draw' not in pv:
                pv['risk_random_draw'] = random.random()


class Group(BaseGroup):
    pass

class Player(BasePlayer):

    # define variables for control questions
    q1 = models.CurrencyField(
        label='How much money would you earn? '
              '(Please use "." to separate Euro and Cent values)',
        min=0
    )
    q2 = models.CurrencyField(
        label='How much money would you earn in the <b>worst case scenario</b>? '
              '(Please use "." to separate Euro and Cent values)',
        min=0
    )
    incorrect_attempts = models.IntegerField(initial=0)

    # list number to indicate which list occurred.
    list_number = models.IntegerField()
    def set_list_number(self):
        self.list_number = self.participant.vars['rand_mpl_index'][self.round_number - 1]

    # add model fields to class player
    for j in range(1, Constants.num_choices + 1):
        locals()['choice_' + str(j)] = models.StringField()
    del j

    # determine consistency
    inconsistent = models.IntegerField()

    def set_consistency(self):

        n = Constants.num_choices

        # replace A's by 1's and B's by 0's
        choices_coded = [None for j in range(1, n + 1)]
        for j in Constants.index_choices:
            if self.participant.vars['mpl_choices_made'][j-1] == "A":
                choice = 0
                choices_coded[j - 1] = choice
            else:
                choice = 1
                choices_coded[j - 1] = choice

        self.participant.vars['mpl_choices_coded'] = choices_coded

        # check for multiple switching behavior
        for j in range(1, n):
            self.inconsistent = 1 if choices_coded[j] > choices_coded[j - 1] else 0
            if self.inconsistent == 1:
                break

    #determine switching row
    switching_row = models.IntegerField()

    def set_switching_row(self):

        # set switching point to row number of first 'B' choice
        if self.inconsistent == 0:
            self.switching_row = sum(self.participant.vars['mpl_choices_coded']) + 1


    # set player's payoff
    round_to_pay = models.IntegerField()
    choice_to_pay = models.StringField()
    index_to_pay = models.IntegerField() #choice and index to pay contain the same information but in different formats. We need both!
    option_to_pay = models.StringField()
    random_draw = models.FloatField()

    selected_table = models.IntegerField(blank=True, null=True)
    selected_row = models.IntegerField(blank=True, null=True)

    selected_choice = models.StringField(blank=True)  # 'A' or 'B'

    selected_certainty = models.CurrencyField(blank=True, null=True)

    selected_lot_hi = models.CurrencyField(blank=True, null=True)
    selected_lot_lo = models.CurrencyField(blank=True, null=True)
    selected_prob_hi = models.FloatField(blank=True, null=True)
    selected_prob_lo = models.FloatField(blank=True, null=True)

    lottery_realization = models.StringField(blank=True)  # 'high', 'low', or 'not_applicable'
    final_bonus_risk = models.CurrencyField(blank=True, null=True)




    def set_payoffs(self):
        pv = self.participant.vars

        # payment draw comes from Welcome
        self.round_to_pay = pv['risk_round_to_pay']
        self.choice_to_pay = pv['risk_choice_field']
        self.index_to_pay = pv['risk_row_to_pay']

        if self.round_to_pay != self.round_number:
            self.option_to_pay = None
            self.random_draw = None
            self.payoff = c(0)
            return

        # selected page / row
        self.selected_table = self.round_to_pay
        self.selected_row = self.index_to_pay

        # participant's choice in the selected row
        self.option_to_pay = getattr(self, self.choice_to_pay)
        self.selected_choice = self.option_to_pay

        # lottery draw comes from Welcome
        self.random_draw = pv['risk_random_draw']

        # current selected MPL
        selected_mpl_raw = pv['rand_sequence'][self.round_number - 1]
        selected_mpl_display = pv['mpl_sequence'][self.round_number - 1]

        # raw values (not percentage strings)
        lot_hi = selected_mpl_raw[0][0]
        lot_lo = selected_mpl_raw[1][0]
        prob_hi = selected_mpl_raw[2][0]
        prob_lo = selected_mpl_raw[3][0]

        cert = selected_mpl_display[4][self.index_to_pay - 1]

        # store decision details
        self.selected_lot_hi = lot_hi
        self.selected_lot_lo = lot_lo
        self.selected_prob_hi = prob_hi
        self.selected_prob_lo = prob_lo
        self.selected_certainty = cert

        if self.option_to_pay == 'A':
            if self.random_draw <= prob_hi:
                self.payoff = lot_hi
                self.lottery_realization = 'high'
            else:
                self.payoff = lot_lo
                self.lottery_realization = 'low'
        else:
            self.payoff = cert
            self.lottery_realization = 'not_applicable'

        self.final_bonus_risk = self.payoff

        # store in participant.vars for final results page
        pv['payoff_risk'] = self.payoff
        pv['risk_selected_table'] = self.selected_table
        pv['risk_selected_row'] = self.selected_row
        pv['risk_selected_choice'] = self.selected_choice
        pv['risk_selected_certainty'] = self.selected_certainty
        pv['risk_selected_lot_hi'] = self.selected_lot_hi
        pv['risk_selected_lot_lo'] = self.selected_lot_lo
        pv['risk_selected_prob_hi'] = self.selected_prob_hi
        pv['risk_selected_prob_lo'] = self.selected_prob_lo
        pv['risk_lottery_realization'] = self.lottery_realization
        pv['risk_final_bonus'] = self.final_bonus_risk


    #set certainty equivalents
    cert_equi = models.CurrencyField()
    cert_equi_rounded = models.CurrencyField()

    def set_certainty_equivalents(self):
        self.cert_equi = self.participant.vars['cert_equi'][self.round_number-1]
        self.cert_equi_rounded = round(self.cert_equi + 0.01,1) # we have to add this very small number since otherwise the rounding inpython is strange


    def set_settings_for_next_task(self):
        self.participant.vars['task_to_pay'] = random.randint(0,1)
        self.participant.vars['round_to_pay'] = random.randint(1, 18)