from cProfile import label
from collections import Counter
from random import choices

from otree.api import *
import itertools
import random
import json

from otree.database import IntegerField
from sqlalchemy import Integer
doc = """
Your app description
"""

class C(BaseConstants):
    NAME_IN_URL = 'spectators'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1
    PAY = 100
    BONUS = 100
    TREATMENTS = ['control', 'immig', 'redis', 'combined']
class Subsession(BaseSubsession):
    spectators_drawn = models.BooleanField(initial=False)
class Group(BaseGroup):
    pass
class Player(BasePlayer):
    p_id = models.StringField(label="<br> Please enter your unique Prolific ID: ")
    is_chosen = models.BooleanField(initial=False)
    consent = models.BooleanField(initial=False)
    treatment = models.StringField(initial='')
    earned = models.CurrencyField(initial=0)
    audit = models.IntegerField(
        label="<b>What is your audit decision?</b>",
        widget=widgets.RadioSelect)
    penalty = models.IntegerField(min=10, max=100)
    reason = models.LongStringField(label="Could you briefly explain why you made exactly these "
                                          "audit and penalty decisions?")
    appr = models.IntegerField()
    fair = models.IntegerField()
    strategic = models.IntegerField()

    ###### Demographic Questions ######
    age = models.IntegerField(label="Please enter your age: ", min=18, max=80)
    sex = models.StringField(
        label="<br>What is your sex, as recorded on legal/official documents?",
        choices=["Male", "Female", "Prefer not to say"],
        widget=widgets.RadioSelect,
    )
    ethnicity = models.StringField(
        label="<br>What ethnic group do you belong to?",
        choices=["White", "Asian", "Black", "Mixed",
                 "Other"],
        widget=widgets.RadioSelect
    )
    ethnicity_other = models.StringField(blank=True, label="If Mixed or Other, please specify:")
    education = models.StringField(
        label="<br>What is the highest level of education you have completed?",
        choices=["No formal qualifications", "Primary education", "Secondary education", "Vocational qualification",
                 "Undergraduate degree", "Postgraduate degree", "Doctorate degree"],
        widget=widgets.RadioSelect
    )
    job = models.StringField(
        label="<br>What is your current employment status?",
        choices=["Employed / Self-employed", "Out of work / Unemployed", "Student / Part-time employment", "Retired"],
        widget=widgets.RadioSelect
    )
    income = models.StringField(
        label="<br>What was your total household income before taxes in 2024 (in GBP)?",
        choices=["Under 10,000", "10,000 to 19,999", "20,000 to 29,999", "30,000 to 39,999", "40,000 to 49,999",
                 "50,000 to 59,999", "60,000 to 69,999", "70,000 to 79,999", "80,000 to 99,999", "100,000 or more"],
        widget=widgets.RadioSelect
    )
    children = models.StringField(
        label="<br>How many children do you have?",
        choices=["I do not have children", "1", "2", "3", "4", "5 or more"]
    )
    spectrum = models.StringField(
        label="<br>On economic policy matters, where do you see yourself on the liberal/conservative spectrum?",
        choices=[
            ('vlib', 'Very Liberal'),
            ('lib', 'Liberal'),
            ('mod', 'Moderate'),
            ('cons', 'Conservative'),
            ('vcons', 'Very Conservative'),
        ],
        widget=widgets.RadioSelect
    )
    political = models.StringField(
        label="<br>Which political party do you most identify with?",
        choices=["Conservative Party", "Labour Party", "Liberal Democrats", "Green Party", "Reform UK",
                 "None / Not affiliated", "Other"],
        widget=widgets.RadioSelect
    )
    political_other = models.StringField(blank=True, label="If Other, please specify:")
    comment = models.LongStringField(blank=True, label="<br>Do you have any comments or suggestions regarding the study? "
                                                   "If not, you can leave this field blank.")


    ###### Baseline Questions ######
    b1 = models.StringField(
        label="<br>How would you rate the overall quality of public services in the United Kingdom?",
        choices=[
            (1, 'Very low'),
            (2, 'Somewhat low'),
            (3, 'Neutral'),
            (4, 'Somewhat high'),
            (5, 'Very high'),
        ],
        widget=widgets.RadioSelectHorizontal)
    b2 = models.StringField(
        label="<br>How easy do you think it is for people in the UK to improve their economic "
              "situation through hard work?",
        choices=[
            (1, 'Very difficult'),
            (2, 'Somewhat difficult'),
            (3, 'Neutral'),
            (4, 'Somewhat easy'),
            (5, 'Very easy'),
        ],
        widget=widgets.RadioSelectHorizontal)
    b3 = models.IntegerField(
        label="<br>What percentage of the UK population do you think has completed university-level education?",
        min=0, max=100,
        choices=[
            (0, '0%'), (10, '10%'), (20, '20%'), (30, '30%'), (40, '40%'), (50, '50%'),
            (60, '60%'), (70, '70%'), (80, '80%'), (90, '90%'), (100, '100%'),
        ],
        widget=widgets.RadioSelectHorizontal
        )
    b4 = models.StringField(
        label="<br>How important do you think education is for economic success in the UK today?",
        choices=[
            (1, 'Not important'),
            (2, 'Slightly important'),
            (3, 'Moderately important'),
            (4, 'Somewhat important'),
            (5, 'Very important'),
        ],
        widget=widgets.RadioSelectHorizontal)
    b5 = models.IntegerField(
        label="<br>What percentage of the UK population do you think is currently employed?",
        min=0, max=100,
        choices=[
            (0, '0%'), (10, '10%'), (20, '20%'), (30, '30%'), (40, '40%'), (50, '50%'),
            (60, '60%'), (70, '70%'), (80, '80%'), (90, '90%'), (100, '100%'),
        ],
        widget=widgets.RadioSelectHorizontal
        )

    ###### Control ######
    c1 = models.StringField(
        label="<br>In what type of area do you think most people in the United Kingdom live?",
        choices=[
            (5, 'Large cities'),
            (4, 'Medium-sized towns'),
            (3, 'Small towns'),
            (2, 'Rural areas'),
            (1, 'Equally across these areas'),
        ],
        widget=widgets.RadioSelect)
    c2 = models.IntegerField(
        label="<br>What percentage of the population do you think lives in London and its surrounding metropolitan area?",
        min=0, max=100,
        choices=[
            (0, '0%'), (10, '10%'), (20, '20%'), (30, '30%'), (40, '40%'), (50, '50%'),
            (60, '60%'), (70, '70%'), (80, '80%'), (90, '90%'), (100, '100%'),
        ],
        widget=widgets.RadioSelectHorizontal
    )
    c3 = models.IntegerField(
        label="<br>How common do you think it is for people in the UK to move to a different city or "
              "region during their lifetime?",
        choices=[
            (1, 'Very uncommon'),
            (2, 'Somewhat uncommon'),
            (3, 'Neutral'),
            (4, 'Somewhat common'),
            (5, 'Very common'),
        ],
        widget=widgets.RadioSelectHorizontal)
    c4 = models.IntegerField(
        label="<br>What percentage of the population do you think works from home at least occasionally?",
        min=0, max=100,
        choices=[
            (0, '0%'), (10, '10%'), (20, '20%'), (30, '30%'), (40, '40%'), (50, '50%'),
            (60, '60%'), (70, '70%'), (80, '80%'), (90, '90%'), (100, '100%'),
        ],
        widget=widgets.RadioSelectHorizontal)
    c5 = models.StringField(
        label="<br>How common do you think it is for households in the UK to have internet access at home?",
        choices=[
            (1, 'Very uncommon'),
            (2, 'Somewhat uncommon'),
            (3, 'Neutral'),
            (4, 'Somewhat common'),
            (5, 'Very common'),
        ],
        widget=widgets.RadioSelectHorizontal)

    ###### Redistribution ######
    r1 = models.StringField(
        label="<br>What do you think about the following statement: ''<i>The government should allocate resources to "
              "support those in need.</i>''",
        choices=[
            (1, 'Fully Disagree'),
            (2, 'Somewhat Disagree'),
            (3, 'Neutral'),
            (4, 'Somewhat Agree'),
            (5, 'Fully Agree'),
        ],
        widget=widgets.RadioSelectHorizontal
    )
    r2 = models.IntegerField(
        label="<br>To what extent do you think social welfare programs in the UK reach people who genuinely need support?",
        choices=[
            (1, 'Not at all'),
            (2, 'To a small extent'),
            (3, 'Neutral'),
            (4, 'To a large extent'),
            (5, 'To a very large extent'),
        ],
        widget=widgets.RadioSelectHorizontal
    )
    r3 = models.IntegerField(
        label="<br>Roughly what percentage of total government spending do you think goes to social welfare benefits?",
        min=0, max=100,
        choices=[
            (0, '0%'), (10, '10%'), (20, '20%'), (30, '30%'), (40, '40%'), (50, '50%'),
            (60, '60%'), (70, '70%'), (80, '80%'), (90, '90%'), (100, '100%'),
        ],
        widget=widgets.RadioSelectHorizontal
    )
    r4 = models.StringField(
        label="<br>Overall, how fair do you think the UK tax system is?",
        choices=[
            (1, 'Very unfair'),
            (2, 'Somewhat unfair'),
            (3, 'Neutral'),
            (4, 'Somewhat fair'),
            (5, 'Very fair'),
        ],
        widget=widgets.RadioSelectHorizontal)
    r5 = models.IntegerField(
        label="<br>What percentage of adults in the UK do you think receive social welfare benefits?",
        min=0, max=100,
        choices=[
            (0, '0%'), (10, '10%'), (20, '20%'), (30, '30%'), (40, '40%'), (50, '50%'),
            (60, '60%'), (70, '70%'), (80, '80%'), (90, '90%'), (100, '100%'),
        ],
        widget=widgets.RadioSelectHorizontal
    )


    ###### Immigration ######
    i1 = models.StringField(
        label="<br>What do you think about the following statement: <i>''The current number of immigrants coming to the "
              "UK is about right.''</i>",
        choices=[
            (1, 'Fully Disagree'),
            (2, 'Somewhat Disagree'),
            (3, 'Neutral'),
            (4, 'Somewhat Agree'),
            (5, "Fully Agree"),
        ],
        widget=widgets.RadioSelectHorizontal
    )
    i2 = models.StringField(
        label="<br>In your opinion, what are the three most common countries of origin among immigrants in the UK?"
    )
    i3 = models.IntegerField(
        label="<br>What percentage of the total population do you think are first-generation immigrants (born abroad)?",
        min=0, max=100,
        choices=[
            (0, '0%'), (10, '10%'), (20, '20%'), (30, '30%'), (40, '40%'), (50, '50%'),
            (60, '60%'), (70, '70%'), (80, '80%'), (90, '90%'), (100, '100%'),
        ],
        widget=widgets.RadioSelectHorizontal

    )
    i4 = models.IntegerField(
        label="<br>What percentage of immigrants do you think are currently employed?",
        min=0, max=100,
        choices=[
            (0, '0%'), (10, '10%'), (20, '20%'), (30, '30%'), (40, '40%'), (50, '50%'),
            (60, '60%'), (70, '70%'), (80, '80%'), (90, '90%'), (100, '100%'),
        ],
        widget=widgets.RadioSelectHorizontal

    )
    i5 = models.StringField(
        label="<br>In your view, what is the overall impact of immigration on the United Kingdom?",
        choices=[
            (1, 'Very negative'),
            (2, 'Somewhat negative'),
            (3, 'Neutral'),
            (4, 'Somewhat positive'),
            (5, 'Very positive'),
        ],
        widget=widgets.RadioSelectHorizontal
    )
    attention_check = models.StringField(
        label="<br>To prove that you are reading questions attentively, please select ''Fully Agree'' here.",
        choices=[
            (1, 'Fully Disagree'),
            (2, 'Somewhat Disagree'),
            (3, 'Neutral'),
            (4, 'Somewhat Agree'),
            (5, 'Fully Agree'),
        ],
        widget=widgets.RadioSelectHorizontal
    )


# FUNCTIONS
def creating_session(subsession: Subsession):
    for p in subsession.get_players():
        p.consent = False  # default
        p.treatment = ""  # placeholder
        p.is_chosen = False
def age_error_message(player, value):
    if value < 18:
        return "You need to be at least 18 years old."
    elif value > 80:
        return "Please enter your age."
def maybe_draw_spectators(subsession, N=1):
    for p in subsession.get_players():
        p.is_chosen = False

    eligible = [p for p in subsession.get_players() if p.consent and p.treatment]

    for treatment in C.TREATMENTS:
        group_players = [p for p in eligible if p.treatment == treatment]

        if not group_players:
            continue

        chosen = random.sample(group_players, k=min(N, len(group_players)))

        for p in chosen:
            p.is_chosen = True
def audit_choices(player):
    choices = [(1, "Audit 30% (Cost £0.10)"),
                 (2, "Audit 70% (Cost £0.40)")]
    random.shuffle(choices)
    return choices


# PAGES
class S1Consent(Page):
    form_model = "player"
    form_fields = ["p_id"]

    @staticmethod
    def before_next_page(player, timeout_happened):
        # Only assign treatment if they consented
        player.consent = True

        # Get already-consented players
        subsession = player.subsession
        consented = [p for p in subsession.get_players() if p.consent and p.treatment]

        treatments = C.TREATMENTS

        # Simple round-robin based on number of consenters so far
        player.treatment = treatments[len(consented) % len(treatments)]
class S2Intro(Page):
    form_model = "player"

    @staticmethod
    def vars_for_template(player):
        subsession = player.subsession

        maybe_draw_spectators(subsession)

        return {}
class S3Survey(Page):
    form_model = "player"

    @staticmethod
    def get_form_fields(player):
        baseline = ["b1", "b2", "b3", "b4", "b5"]
        if player.treatment == "control":
            randomized = ["c1", "c2", "c3", "c4", "c5", "attention_check"]
            random.shuffle(randomized)
            return baseline + randomized
        elif player.treatment == "redis":
            randomized = ["r1", "r2", "r3", "r4", "r5", "attention_check"]
            random.shuffle(randomized)
            return baseline + randomized
        elif player.treatment == "immig":
            randomized = ["i1", "i2", "i3", "i4", "i5", "attention_check"]
            random.shuffle(randomized)
            return baseline + randomized
        elif player.treatment == "combined":
            randomized = ["r1", "r2", "r3", "r4", "r5", "i1", "i2", "i3", "i4", "i5", "attention_check"]
            random.shuffle(randomized)
            return baseline + randomized
class S4Info(Page):
    form_model = "player"
class S5Players(Page):
    form_model = "player"
class S6Audit(Page):
    form_model = "player"
class S7Audit(Page):
    form_model = "player"
    form_fields = ["audit"]
class S9Penalty(Page):
    form_model = "player"
    form_fields = ["penalty"]
class S10Views(Page):
    form_model = "player"
    form_fields = ["reason"]
class S11Demo(Page):
    form_model = "player"
    form_fields = ["age", "sex", "ethnicity", "ethnicity_other", "education", "job", "income", "spectrum",
                   "political", "political_other", "comment"]
class S12Finish(Page):
    form_model = "player"

    @staticmethod
    def vars_for_template(player):
        player.earned = C.PAY + C.BONUS
        if player.audit in [1]:
            player.earned -= 10
        elif player.audit in [2]:
            player.earned -= 40
        conversion_rate = player.session.config['real_world_currency_per_point']
        real_world_earned = player.earned * conversion_rate
        real_world_earned_formatted = '{:.2f}'.format(round(real_world_earned, 2))

        return {
            "pence": player.earned,
            "earned": real_world_earned_formatted,
        }

page_sequence = [S1Consent, S2Intro, S3Survey, S4Info, S5Players,
                 S6Audit, S7Audit, S9Penalty, S10Views, S11Demo, S12Finish]