from otree.api import *
import json

doc = ''

class C(BaseConstants):
    NAME_IN_URL = 'prediction'
    PLAYERS_PER_GROUP = None

    INTERFACE_QUESTIONS = [
        dict(
            q="Interface",
            min=18,
            step=15,
            nb_bins=4,
            unit="",
            ages=1,
            x_axis_title = "Años",
        ),
    ]
    ATTENTION_QUESTIONS = [
        dict(
            q="Attention",
            min=200,
            step=50,
            nb_bins=6,
            unit="",
            ages=0,
            x_axis_title = "Calorías",
        ),
    ]
    BELIEF1_QUESTIONS = [
        dict(
            q="Belief1",
            min=200,
            step=50,
            nb_bins=6,
            unit="",
            ages=0,
            x_axis_title = "Calorías",
        ),
    ]
    BELIEF2_QUESTIONS = [
        dict(
            q="Belief2",
            min=200,
            step=50,
            nb_bins=6,
            unit="",
            ages=0,
            x_axis_title = "Calorías",
        ),
    ]
    BELIEF3_QUESTIONS = [
        dict(
            q="Belief3",
            min=200,
            step=50,
            nb_bins=6,
            unit="",
            ages=0,
            x_axis_title = "Calorías",
        ),
    ]
    NUM_ROUNDS = 1
    TOL = 1e-3

    AGE_GENDER_QUOTAS = {
        ('18–24', 'Female'): 22,
        ('18–24', 'Male'): 23,
        ('25–34', 'Female'): 38,
        ('25–34', 'Male'): 38,
        ('35–44', 'Female'): 36,
        ('35–44', 'Male'): 34,
        ('45–54', 'Female'): 32,
        ('45–54', 'Male'): 29,
        ('55 or older', 'Female'): 68,
        ('55 or older', 'Male'): 56,
    }

    OTHER_GENDER_QUOTA = 10

class Subsession(BaseSubsession):
    pass

class Group(BaseGroup):
    pass

class Player(BasePlayer):
    data1 = models.StringField(blank=True)
    history1 = models.LongStringField(initial="[]")
    skipped1 = models.StringField(initial="0", blank=True)
    data2 = models.StringField(blank=True)
    history2 = models.LongStringField(initial="[]")
    skipped2 = models.StringField(initial="0", blank=True)
    data3 = models.StringField(blank=True)
    history3 = models.LongStringField(initial="[]")
    skipped3 = models.StringField(initial="0", blank=True)
    data4 = models.StringField(blank=True)
    history4 = models.LongStringField(initial="[]")
    skipped4 = models.StringField(initial="0", blank=True)
    data5 = models.StringField(blank=True)
    history5 = models.LongStringField(initial="[]")
    skipped5 = models.StringField(initial="0", blank=True)
    min = models.IntegerField(label="Min value of the x axis")
    max = models.IntegerField(label="Max value of the x axis")
    nb_bins = models.IntegerField(label="Number of bins", min=2)
    consent = models.BooleanField(choices=[[True, 'Acepto'], [False, 'No acepto']], label='¿Acepta participar en esta encuesta?', widget=widgets.RadioSelect)
    age = models.StringField(choices=[['18–24', '18–24 años'], ['25–34', '25–34 años'], ['35–44', '35–44 años'], ['45–54', '45–54 años'], ['55 or older', '55 años o más']], label='¿Cuál es su edad?', widget=widgets.RadioSelect)
    gender = models.StringField(choices=[['Male', 'Hombre'], ['Female', 'Mujer'], ['Other / prefer not to say', 'Otro / Prefiero no responder']], label='¿Cuál es su género?', widget=widgets.RadioSelect)
    education = models.StringField(blank=True, choices=[['Primary school or less', 'Educación básica o menos'], ['Secondary school', 'Educación media'], ['Technical / vocational education', 'Educación técnica o profesional'], ['University degree', 'Título universitario'], ['Postgraduate degree', 'Postgrado']], label='¿Cuál es su nivel educacional?', widget=widgets.RadioSelect)
    email = models.StringField(blank=True, label='Su correo electrónico:')

# This I might want to remove later
def creating_session(subsession: Subsession):
    session = subsession.session

    if subsession.round_number == 1:
        if 'quota_counts' not in session.vars:
            session.vars['quota_counts'] = {
                cell: 0 for cell in C.AGE_GENDER_QUOTAS
            }

        if 'other_gender_count' not in session.vars:
            session.vars['other_gender_count'] = 0

    if 'interface' in session.config:
        for player in subsession.get_players():
            player.participant.vars['interface'] = "bins"
    else:
        import itertools

        interfaces = itertools.cycle(["click-drag", "number", "bins", "metaculus"])
        for player in subsession.get_players():
            player.interface = "bins"
            player.participant.interface = player.interface



### PAGES ###
    

# Consent page
class Consent(Page):
    form_model = 'player'
    form_fields = ['consent']


    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        player.participant.vars['rid'] = player.participant.label or ''

        if not player.consent:
            player.participant.vars['terminated'] = True
            player.participant.vars['cint_status'] = 'RIS20'

# Introduction1 page (to do: add image)
class Introduction1(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)
    
# Introduction2 page (to do: add image)    
class Introduction2(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)

# Introduction3 page
class Introduction3(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)

# Introduction4 page    
class Introduction4(Page):
    form_model = 'player'
    form_fields = ['data1', 'skipped1']

    timer_text = 'Tiempo restante para realizar su predicción:'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)

    @staticmethod
    def get_timeout_seconds(player):
        if player.session.config.get("timeout", False):
            return 90

    @staticmethod
    def live_method(player, data):
        history1 = json.loads(player.history1)
        history1.append(data["history"])
        player.history1 = json.dumps(history1)
        player.data1 = json.dumps(data["history"]["data"])

        return {player.id_in_group: dict(data1=player.data1)}
    
    @staticmethod
    def error_message(player: Player, values):
        if values.get('skipped1') == "1":
            return None
        tol = C.TOL
        raw = player.data1 or values.get('data1')
        if not raw:
            return "Por favor, ingrese una distribución antes de continuar."
        try:
            arr = json.loads(raw)
            s = sum(float(v or 0) for v in arr)
        except Exception:
            return "El formato de la distribución no es válido."
        if abs(s - 1.0) > tol:
            return "Las probabilidades deben sumar 100%% antes de continuar."
    


    @staticmethod
    def js_vars(player: Player):
        session = player.session
        participant = player.participant

        pq = C.INTERFACE_QUESTIONS[0]
        if not session.config.get("self", False):
            player.min = pq["min"]
            player.nb_bins = pq["nb_bins"]
            player.max = pq["min"] + pq["step"] * (pq["nb_bins"] - 1)

        return dict(
            interface="bins",
            prediction=True,
            yMax=1,
            min=player.min,
            step=(player.max - player.min) / (player.nb_bins - 1),
            nb_bins=player.nb_bins,
            xUnit=pq["unit"],
            xAxisTitle=pq.get("x_axis_title", ""),
            min_timeout=30,
        )
    
    @staticmethod
    def vars_for_template(player: Player):
        participant = player.participant
        pq = C.INTERFACE_QUESTIONS[0]
        return dict(
            interface_questions=pq["q"],
            interface="bins",
        )
    
# Attention page    
class Attention(Page):
    form_model = 'player'
    form_fields = ['data2', 'skipped2']

    timer_text = 'Tiempo restante para realizar su predicción:'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)

    @staticmethod
    def get_timeout_seconds(player):
        if player.session.config.get("timeout", False):
            return 90

    @staticmethod
    def live_method(player, data):
        history2 = json.loads(player.history2)
        history2.append(data["history"])
        player.history2 = json.dumps(history2)
        player.data2 = json.dumps(data["history"]["data"])

        return {player.id_in_group: dict(data2=player.data2)}
    
    @staticmethod
    def error_message(player: Player, values):
        if values.get('skipped2') == "1":
            return None
        tol = C.TOL
        raw = player.data2 or values.get('data2')
        if not raw:
            return "Por favor, ingrese una distribución antes de continuar."
        try:
            arr = json.loads(raw)
            s = sum(float(v or 0) for v in arr)
        except Exception:
            return "El formato de la distribución no es válido."
        if abs(s - 1.0) > tol:
            return "Las probabilidades deben sumar 100%% antes de continuar."

    @staticmethod
    def js_vars(player: Player):
        session = player.session
        participant = player.participant

        pq = C.ATTENTION_QUESTIONS[0]
        if not session.config.get("self", False):
            player.min = pq["min"]
            player.nb_bins = pq["nb_bins"]
            player.max = pq["min"] + pq["step"] * (pq["nb_bins"] - 1)

        return dict(
            interface="bins",
            prediction=True,
            yMax=1,
            min=player.min,
            step=(player.max - player.min) / (player.nb_bins - 1),
            nb_bins=player.nb_bins,
            xUnit=pq["unit"],
            xAxisTitle=pq.get("x_axis_title", ""),
            min_timeout=30,
        )
    
    @staticmethod
    def vars_for_template(player: Player):
        participant = player.participant
        pq = C.ATTENTION_QUESTIONS[0]
        return dict(
            attention_questions=pq["q"],
            interface="bins",
        )
    
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # treat skip as fail
        if player.skipped2 == "1":
            player.participant.vars['terminated'] = True
            player.participant.vars['cint_status'] = 'RIS30'
            return

        try:
            arr = json.loads(player.data2)
            arr = [float(x or 0) for x in arr]
        except Exception:
            player.participant.vars['terminated'] = True
            player.participant.vars['cint_status'] = 'RIS20'
            return

        TARGET = [0, 0, 0, 0, 0, 1]
        TOL = 1e-6

        failed = (len(arr) != len(TARGET)) or any(abs(arr[i] - TARGET[i]) > TOL for i in range(len(TARGET)))

        if failed:
            player.participant.vars['terminated'] = True
            player.participant.vars['cint_status'] = 'RIS30'
    
# PreBelief1 page
class PreBelief1(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)
    
# PreBelief2 page
class PreBelief2(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)
    
# Belief1 page    
class Belief1(Page):
    form_model = 'player'
    form_fields = ['data3', 'skipped3']

    timer_text = 'Tiempo restante para realizar su predicción:'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)

    @staticmethod
    def get_timeout_seconds(player):
        if player.session.config.get("timeout", False):
            return 90

    @staticmethod
    def live_method(player, data):
        history3 = json.loads(player.history3)
        history3.append(data["history"])
        player.history3 = json.dumps(history3)
        player.data3 = json.dumps(data["history"]["data"])

        return {player.id_in_group: dict(data3=player.data3)}
    
    @staticmethod
    def error_message(player: Player, values):
        if values.get('skipped3') == "1":
            return None
        tol = C.TOL
        raw = player.data3 or values.get('data3')
        if not raw:
            return "Por favor, ingrese una distribución antes de continuar."
        try:
            arr = json.loads(raw)
            s = sum(float(v or 0) for v in arr)
        except Exception:
            return "El formato de la distribución no es válido."
        if abs(s - 1.0) > tol:
            return "Las probabilidades deben sumar 100%% antes de continuar."

    @staticmethod
    def js_vars(player: Player):
        session = player.session
        participant = player.participant

        pq = C.BELIEF1_QUESTIONS[0]
        if not session.config.get("self", False):
            player.min = pq["min"]
            player.nb_bins = pq["nb_bins"]
            player.max = pq["min"] + pq["step"] * (pq["nb_bins"] - 1)

        return dict(
            interface="bins",
            prediction=True,
            yMax=1,
            min=player.min,
            step=(player.max - player.min) / (player.nb_bins - 1),
            nb_bins=player.nb_bins,
            xUnit=pq["unit"],
            xAxisTitle=pq.get("x_axis_title", ""),
            min_timeout=30,
        )
    
    @staticmethod
    def vars_for_template(player: Player):
        participant = player.participant
        pq = C.BELIEF1_QUESTIONS[0]
        return dict(
            belief1_questions=pq["q"],
            interface="bins",
        )


    
# Belief2 page    
class Belief2(Page):
    form_model = 'player'
    form_fields = ['data4', 'skipped4']

    timer_text = 'Tiempo restante para realizar su predicción:'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)

    @staticmethod
    def get_timeout_seconds(player):
        if player.session.config.get("timeout", False):
            return 90

    @staticmethod
    def live_method(player, data):
        history4 = json.loads(player.history4)
        history4.append(data["history"])
        player.history4 = json.dumps(history4)
        player.data4 = json.dumps(data["history"]["data"])

        return {player.id_in_group: dict(data4=player.data4)}
    
    @staticmethod
    def error_message(player: Player, values):
        if values.get('skipped4') == "1":
            return None
        tol = C.TOL
        raw = player.data4 or values.get('data4')
        if not raw:
            return "Por favor, ingrese una distribución antes de continuar."
        try:
            arr = json.loads(raw)
            s = sum(float(v or 0) for v in arr)
        except Exception:
            return "El formato de la distribución no es válido."
        if abs(s - 1.0) > tol:
            return "Las probabilidades deben sumar 100%% antes de continuar."

    @staticmethod
    def js_vars(player: Player):
        session = player.session
        participant = player.participant

        pq = C.BELIEF2_QUESTIONS[0]
        if not session.config.get("self", False):
            player.min = pq["min"]
            player.nb_bins = pq["nb_bins"]
            player.max = pq["min"] + pq["step"] * (pq["nb_bins"] - 1)

        return dict(
            interface="bins",
            prediction=True,
            yMax=1,
            min=player.min,
            step=(player.max - player.min) / (player.nb_bins - 1),
            nb_bins=player.nb_bins,
            xUnit=pq["unit"],
            xAxisTitle=pq.get("x_axis_title", ""),
            min_timeout=30,
        )
    
    @staticmethod
    def vars_for_template(player: Player):
        participant = player.participant
        pq = C.BELIEF2_QUESTIONS[0]
        return dict(
            belief2_questions=pq["q"],
            interface="bins",
        )
    
# Belief3 page    
class Belief3(Page):
    form_model = 'player'
    form_fields = ['data5', 'skipped5']

    timer_text = 'Tiempo restante para realizar su predicción:'

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)

    @staticmethod
    def get_timeout_seconds(player):
        if player.session.config.get("timeout", False):
            return 90

    @staticmethod
    def live_method(player, data):
        history5 = json.loads(player.history5)
        history5.append(data["history"])
        player.history5 = json.dumps(history5)
        player.data5 = json.dumps(data["history"]["data"])

        return {player.id_in_group: dict(data5=player.data5)}
    
    @staticmethod
    def error_message(player: Player, values):
        if values.get('skipped5') == "1":
            return None
        tol = C.TOL
        raw = player.data5 or values.get('data5')
        if not raw:
            return "Por favor, ingrese una distribución antes de continuar."
        try:
            arr = json.loads(raw)
            s = sum(float(v or 0) for v in arr)
        except Exception:
            return "El formato de la distribución no es válido."
        if abs(s - 1.0) > tol:
            return "Las probabilidades deben sumar 100%% antes de continuar."

    @staticmethod
    def js_vars(player: Player):
        session = player.session
        participant = player.participant

        pq = C.BELIEF3_QUESTIONS[0]
        if not session.config.get("self", False):
            player.min = pq["min"]
            player.nb_bins = pq["nb_bins"]
            player.max = pq["min"] + pq["step"] * (pq["nb_bins"] - 1)

        return dict(
            interface="bins",
            prediction=True,
            yMax=1,
            min=player.min,
            step=(player.max - player.min) / (player.nb_bins - 1),
            nb_bins=player.nb_bins,
            xUnit=pq["unit"],
            xAxisTitle=pq.get("x_axis_title", ""),
            min_timeout=30,
        )
    
    @staticmethod
    def vars_for_template(player: Player):
        participant = player.participant
        pq = C.BELIEF3_QUESTIONS[0]
        return dict(
            belief3_questions=pq["q"],
            interface="bins",
        )
    
# Demographics page
class Demographics(Page):
    form_model = 'player'
    form_fields = ['age', 'gender', 'education']

    @staticmethod
    def is_displayed(player):
        return player.consent and not player.participant.vars.get('terminated', False)

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        key = (player.age, player.gender)

        # separate pooled quota for "Other / prefer not to say"
        if player.gender == 'Other / prefer not to say':
            other_count = player.session.vars['other_gender_count']

            if other_count >= C.OTHER_GENDER_QUOTA:
                player.participant.vars['terminated'] = True
                player.participant.vars['cint_status'] = 'RIS40'
                return

            player.participant.vars['quota_cell'] = 'OTHER_GENDER'
            return

        counts = player.session.vars['quota_counts']
        quotas = C.AGE_GENDER_QUOTAS

        if key not in quotas:
            player.participant.vars['terminated'] = True
            player.participant.vars['cint_status'] = 'RIS40'
            return

        if counts[key] >= quotas[key]:
            player.participant.vars['terminated'] = True
            player.participant.vars['cint_status'] = 'RIS40'
            return

        player.participant.vars['quota_cell'] = key
# PII page
class PII(Page):
    form_model = 'player'
    form_fields = ['email']

    @staticmethod
    def is_displayed(player):
        return player.consent == 1 and not player.participant.vars.get('terminated', False)

class End(Page):

    @staticmethod
    def vars_for_template(player: Player):
        rid = player.participant.vars.get('rid', '')
        status = player.participant.vars.get('cint_status', 'RIS20')

        if status == 'COMPLETE':
            url = f"https://notch.insights.supply/cb?token=eb54b913-40e9-4125-a10d-677313d30906&RID={rid}"
        elif status == 'RIS30':
            url = f"https://samplicio.us/s/ClientCallBack.aspx?RIS=30&RID={rid}"
        elif status == 'RIS40':
            url = f"https://samplicio.us/s/ClientCallBack.aspx?RIS=40&RID={rid}"
        else:
            url = f"https://samplicio.us/s/ClientCallBack.aspx?RIS=20&RID={rid}"


        return dict(redirect_url=url)

    @staticmethod
    def before_next_page(player, timeout_happened):
        if not player.participant.vars.get('terminated', False):
            cell = player.participant.vars.get('quota_cell')

            if cell == 'OTHER_GENDER':
                player.session.vars['other_gender_count'] += 1

            elif cell:
                counts = player.session.vars['quota_counts']
                counts[cell] += 1
                player.session.vars['quota_counts'] = counts

            player.participant.vars['cint_status'] = 'COMPLETE'

class Redirect(Page):
    @staticmethod
    def vars_for_template(player: Player):
        rid = player.participant.vars.get('rid', '')
        status = player.participant.vars.get('cint_status', 'RIS20')

        if status == 'COMPLETE':
            url = f"https://notch.insights.supply/cb?token=eb54b913-40e9-4125-a10d-677313d30906&RID={rid}"
        elif status == 'RIS30':
            url = f"https://samplicio.us/s/ClientCallBack.aspx?RIS=30&RID={rid}"
        elif status == 'RIS40':
            url = f"https://samplicio.us/s/ClientCallBack.aspx?RIS=40&RID={rid}"
        else:
            url = f"https://samplicio.us/s/ClientCallBack.aspx?RIS=20&RID={rid}"


        return dict(redirect_url=url)

# Sequencing pages
page_sequence = [Consent, Demographics, Introduction1, Introduction2, Introduction3, Introduction4, Attention, PreBelief1, PreBelief2, Belief1, Belief2, Belief3, PII, End, Redirect]