from otree.api import *
import random
import time


class C(BaseConstants):
    NAME_IN_URL = 'game_v1'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 5
    START_BALANCE = cu(1000000)
    PROFIT = cu(500000)
    PENALTY = cu(700000)
    RPENALTY = cu(300000)

# Moved to module level so functions/pages can read it directly
RISK_PAYOUTS = {
    1: (1400, 1400),
    2: (1200, 1800),
    3: (1000, 2200),
    4: (800, 2600),
    5: (600, 3000),
    6: (100, 3500),
}


class Subsession(BaseSubsession):
    pass


def creating_session(subsession):
    if subsession.round_number == 1:
        for player in subsession.get_players():
            player.balance = C.START_BALANCE
            treatment = random.randint(0, 1)
            player.participant.vars['groupselected'] = treatment
            player.groupselected = treatment
    else:
        for player in subsession.get_players():
            player.groupselected = player.participant.vars['groupselected']


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    balance = models.CurrencyField(initial=cu(1000000))
    oldbalance = models.CurrencyField(initial=0)
    groupselected = models.IntegerField()
    total_seconds = models.FloatField()
    start_timestamp = models.FloatField()
    end_timestamp = models.FloatField()

    spend_choice1 = models.CurrencyField(
        choices=[cu(0), cu(35000), cu(70000), cu(105000)],
        widget=widgets.RadioSelect,
        label="How much do you want to invest? "
    )
    audit_probability = models.IntegerField(initial=50)

    # Audit draw variables
    aselected = models.IntegerField(initial=0)
    mselected = models.IntegerField(initial=0)
    imselected = models.IntegerField(initial=0)
    penalty = models.CurrencyField(initial=0)
    impenalty = models.CurrencyField(initial=0)
    bel_payoff = models.CurrencyField(initial=0)
    bal_bel_payoff = models.CurrencyField(initial=0)

    # Beliefs
    beliefA = models.FloatField(min=0, max=100)
    scoreA = models.FloatField(initial=0)
    random_drawA = models.FloatField(initial=0)
    audited_outcome1 = models.IntegerField(initial=0)
    won_high_prizeA = models.BooleanField()
    beliefM = models.FloatField(min=0, max=100)
    scoreM = models.FloatField(initial=0)
    random_drawM = models.FloatField(initial=0)
    material_outcome1 = models.IntegerField(initial=0)
    won_high_prizeM = models.BooleanField()
    beliefIM = models.FloatField(min=0, max=100)
    scoreIM = models.FloatField(initial=0)
    random_drawIM = models.FloatField(initial=0)
    immaterial_outcome1 = models.IntegerField(initial=0)
    won_high_prizeIM = models.BooleanField()
    material_risk = models.FloatField(initial=0.0)

    # Demographics & Survey Fields
    gender = models.IntegerField(
         label="What is your gender?",
         choices=[[1, 'Male'], [2, 'Female'], [3, 'Non-binary'], [4, 'Other'], [5, 'Prefer not to say']],
         widget=widgets.RadioSelect
    )

    regulatorfund = models.IntegerField(
        label="Do you believe the South African audit regulator is well funded?",
        choices=[[1, 'Yes'], [2, 'No'], [3, 'Not sure']],
        widget=widgets.RadioSelect
    )

    race = models.IntegerField(
         label="What is your race?",
         choices=[[1, 'Black'], [2, 'Coloured'], [3, 'Indian'], [4, 'Asian'], [5, 'White'], [6, 'Other'], [7, 'Prefer not to say']],
         widget=widgets.RadioSelect
    )

    income = models.IntegerField(
        label="How do you think your family income compares to that of other UP students ?",
        choices=[[1, 'Below average'], [2, 'About average'], [3, 'Above average']],
        widget=widgets.RadioSelect
    )

    position = models.IntegerField(
         label="What is your current role? (pick the option which best describes your position)",
         choices=[
            [1, 'CTA student'], [2, 'Audit manager'], [3, 'Audit partner (equity)'],
            [4, 'Audit director (equity)'], [5, 'Audit associate director (not equity)'],
            [6, 'Junior partner (not equity)'], [7, 'Trainee (first year)'],
            [8, 'Trainee (second year)'], [9, 'Trainee (third year)'],
            [10, 'Consultant'], [11, 'Other']
         ],
         widget=widgets.RadioSelect
    )

    joblevel = models.IntegerField(
        label="What is your current job level? (pick the option which best describes your position)",
        choices=[[1, 'Student'], [2, 'Trainee'], [3, 'Consultant'], [4, 'Manager'], [5, 'Director'], [6, 'Other']],
        widget=widgets.RadioSelect
    )

    risk2 = models.IntegerField(
        label="You are given the following payment options for a coin flip (that is, there is 50% chance of each of the payments given in each option). Please just choose the option you would prefer.",
        choices=[
            [1, '1400 LabRands if the coin lands on heads, 1400 LabRands if the coin lands on tails'],
            [2, '1200 LabRands if the coin lands on heads, 1800 LabRands if the coin lands on tails'],
            [3, '1000 LabRands if the coin lands on heads, 2200 LabRands if the coin lands on tails'],
            [4, '800 LabRands if the coin lands on heads, 2600 LabRands if the coin lands on tails'],
            [5, '600 LabRands if the coin lands on heads, 3000 LabRands if the coin lands on tails'],
            [6, '100 LabRands if the coin lands on heads, 3500 LabRands if the coin lands on tails']
        ],
        widget=widgets.RadioSelect
    )

    risk_outcome = models.StringField()

    attentioncheck = models.IntegerField(
        label="In this experiment, what can the regulator do with money collected through fines?",
        choices=[
            [1, 'The money can ONLY be used to cover regulator operational expenses'],
            [2, 'The money can ONLY be used for the development of the accounting profession'],
            [3, 'The regulator has discretion, but has the option to use funds for operational expenses']
        ],
        widget=widgets.RadioSelect
    )

    fairness = models.IntegerField(
        label="On a scale of 0 to 10, how fairly do you think the regulator behaved in regulating penalties in this experiment? 0 would indicate very unfair; and 10 very fair.",
        choices=range(11),
        widget=widgets.RadioSelectHorizontal
    )

    Openended_1 = models.StringField(label="What made you think that the penalties were/were not levied fairly?")
    Openended_2 = models.StringField(label="Why do you believe the South African audit regulator is or is not well funded?")
    Phone = models.IntegerField(label="As noted at the start of the experiment, we will be randomly selecting 10 respondents to be paid based on their decisions. If you would like to be included for this payment, please provide your cell phone number (you can just enter 0 if you do not want to be included for possible payment). Payments will be sent by mobile money, and the pin will be 2648. Please make a note of the pin in case you are selected for payment. You will receive an SMS about the mobile money transfer if you are selected.")


# --- FUNCTIONS ---
def calculate_belief_payoffA(player: Player):
    r = player.beliefA / 100
    player.audited_outcome1 = 1 if player.aselected >= 60 else 0
    player.scoreA = (r - 0.4) ** 2
    player.random_drawA = random.random()
    if player.scoreA < player.random_drawA:
        player.bel_payoff += cu(5000)
        player.won_high_prizeA = True
    else:
        player.won_high_prizeA = False

def calculate_belief_payoffM(player: Player):
    r = player.beliefM / 100
    player.material_outcome1 = 1 if player.aselected >= 60 and player.mselected <= player.audit_probability else 0
    player.scoreM = (r - player.audit_probability / 100) ** 2
    player.random_drawM = random.random()
    if player.scoreM < player.random_drawM:
        player.bel_payoff += cu(5000)
        player.won_high_prizeM = True
    else:
        player.won_high_prizeM = False


def calculate_belief_payoffIM(player: Player):
    r = player.beliefIM / 100
    player.immaterial_outcome1 = 1 if player.aselected >= 60 and player.imselected >= 50 else 0
    player.scoreIM = (r - 0.5) ** 2
    player.random_drawIM = random.random()
    if player.scoreIM < player.random_drawIM:
        player.bel_payoff += cu(5000)
        player.won_high_prizeIM = True
    else:
        player.won_high_prizeIM = False


# --- PAGES ---
class InformedConsent(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1


class Instructions(Page):
    @staticmethod
    def before_next_page(player, timeout_happened):
        player.participant.vars['start_timestamp'] = time.time()

    @staticmethod
    def is_displayed(player):
        return player.round_number == 1


class Start(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1

    @staticmethod
    def vars_for_template(player):
        msg = "The audit regulator is able to utilise the money earned from penalties to fund day to day operations." if player.groupselected == 1 else "The audit regulator is required to ringfence the money earned from penalties for the purpose of professional outreach."
        return dict(groupmessage=msg)


class NewsFlash(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 3

    @staticmethod
    def vars_for_template(player):
        msg = "The audit regulator is able to utilise the money earned from penalties to fund day to day operations." if player.groupselected == 1 else "The audit regulator is required to ringfence the money earned from penalties for the purpose of professional outreach."
        return dict(groupmessage=msg)


class Spend1(Page):
    form_model = 'player'
    form_fields = ['spend_choice1']

    @staticmethod
    def vars_for_template(player: Player):
        if player.round_number > 1:
            previous_round_player = player.in_round(player.round_number - 1)
            player.balance = previous_round_player.balance
        return dict(display_balance=player.balance)

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        player.balance -= player.spend_choice1
        mapping = {cu(0): 50, cu(35000): 40, cu(70000): 30, cu(105000): 20}
        player.audit_probability = mapping[player.spend_choice1]


class BeliefTaskA(Page):
    form_model = 'player'
    form_fields = ['beliefA']


class BeliefTaskM(Page):
    form_model = 'player'
    form_fields = ['beliefM']


class BeliefTaskIM(Page):
    form_model = 'player'
    form_fields = ['beliefIM']


class Profit(Page):
    @staticmethod
    def vars_for_template(player):
        return dict(profit=C.PROFIT, balance=player.balance + C.PROFIT)

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.balance += C.PROFIT
        player.aselected = random.randint(1, 100)
        player.mselected = random.randint(1, 100)
        player.imselected = random.randint(1, 100)
        if player.aselected >= 60:
            if player.mselected <= player.audit_probability:
                player.penalty = C.PENALTY
            if player.imselected >= 50:
                player.impenalty = C.RPENALTY
        player.balance -= (player.penalty + player.impenalty)


class CarryForward(Page):
    @staticmethod
    def vars_for_template(player):
        return dict(C=C)


class roundup(Page):
    @staticmethod
    def vars_for_template(player):
        total_fine = player.penalty + player.impenalty
        return dict(
            oldbalance=player.balance + total_fine - C.PROFIT + player.spend_choice1,
            profit=C.PROFIT,
            Spend=player.spend_choice1,
            totalfine=total_fine,
            balance=player.balance
        )

    @staticmethod
    def before_next_page(player, timeout_happened):
        calculate_belief_payoffA(player)
        calculate_belief_payoffM(player)
        calculate_belief_payoffIM(player)
        if player.round_number == 1:
            player.bal_bel_payoff = player.bel_payoff
        else:
            prev = player.in_round(player.round_number - 1)
            player.bal_bel_payoff = prev.bal_bel_payoff + player.bel_payoff


class Demographics(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS

    form_model = 'player'
    form_fields = ['attentioncheck', 'fairness', 'Openended_1', 'joblevel', 'position', 'gender', 'race', 'regulatorfund', 'Openended_2', 'Phone']


class Risk(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS

    form_model = 'player'
    form_fields = ['risk2']

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.risk_outcome = random.choice(['heads', 'tails'])
        heads_amount, tails_amount = RISK_PAYOUTS[player.risk2]

        if player.risk_outcome == 'heads':
            player.payoff = cu(heads_amount)
        else:
            player.payoff = cu(tails_amount)

        start = player.participant.vars.get('start_timestamp', time.time())
        end = time.time()

        player.end_timestamp = end
        player.total_seconds = end - start


class Final(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == C.NUM_ROUNDS

    @staticmethod
    def vars_for_template(player):
        risk_player = player.in_round(C.NUM_ROUNDS)
        # Fixed: risk_player.payoff instead of non-existent risk_player.risk_payoff
        final_payment = player.balance + player.bal_bel_payoff + risk_player.payoff
        rand_payment = round(max(0, float(final_payment) * 0.0002), 2)

        return dict(
            finalpayment=final_payment,
            balance=player.balance,
            beliefpayment=player.bal_bel_payoff,
            riskpayment=risk_player.payoff,
            randpayment=rand_payment
        )


page_sequence = [
    Instructions,
    Start,
    NewsFlash,
    Spend1,
    BeliefTaskA,
    BeliefTaskM,
    BeliefTaskIM,
    Profit,
    CarryForward,
    roundup,
    Demographics,
    Risk,
    Final
]