from otree.api import *
import json, random


class Constants(BaseConstants):
    name_in_url = 'trust_row_experiment'
    players_per_group = None  # one per role, but pairing happens at the session level
    num_rounds = 1
    slider_payment = 0.25
    endowment = cu(10)
    multiplier = 3


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    role_choice = models.IntegerField(
        choices=[(1, 'First Row'), (2, 'Second Row'), (3, 'Third Row')],
        label="",
        widget=widgets.RadioSelect
    )

    # cross-row pairing
    partners_row3_id = models.IntegerField(blank=True, null=True)
    partner_row2_id = models.IntegerField(blank=True, null=True)

    # slider task stats
    num_correct = models.IntegerField(initial=0)
    num_incorrect = models.IntegerField(initial=0)
    attempted = models.IntegerField(initial=0)

    # ranking and pairing
    slider_rank = models.IntegerField(blank=True)   # 1..4 for row3
    is_top = models.BooleanField(initial=False)
    pair_id = models.IntegerField(blank=True)       # 1 or 2 for the two pairs
    knows_status = models.BooleanField(initial=False)  # only meaningful for row3
    tie_note = models.BooleanField(initial=False)   # show a message if their tie impacted rank

    # trust game
    sender_decision = models.CurrencyField(min=0, max=Constants.endowment, blank=True, label=" ")
    amount_received = models.CurrencyField(initial=0)     # tripled
    responder_decision = models.CurrencyField(blank=True) # what 3rd row returns

    # predictor entries per pair
    pred_pair1 = models.CurrencyField(blank=True)
    pred_pair2 = models.CurrencyField(blank=True)
    pred_pair3 = models.CurrencyField(blank=True)
    pred_pair4 = models.CurrencyField(blank=True)

    # mapping for randomized display order
    pred_order_json = models.LongStringField(blank=True)   # list of pair indexes 1..4
    pair_meta_json = models.LongStringField(blank=True)    # per pair info for templates

    # earnings
    slider_earnings = models.CurrencyField(blank=True, null=True)
    amount_kept = models.CurrencyField(blank=True, null=True)
    prediction_bonus = models.CurrencyField(blank=True, null=True)
    prediction_error = models.FloatField(blank=True, null=True)
    payout = models.CurrencyField()

    # First/Second/Third row comprehension (kept)
    comp3_q1 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])
    comp3_q2 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])
    comp3_q3 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])
    comp2_q1 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])
    comp2_q2 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])
    comp2_q3 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])
    comp1_q1 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])
    comp1_q2 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])
    comp1_q3 = models.BooleanField(label="", choices=[(1, 'True'), (0, 'False')])

    # store which displayed card (1..4) was randomly selected for payment
    chosen_item_index = models.IntegerField(blank=True, null=True)
    # store the paid prediction and the corresponding actual return
    predicted_return = models.CurrencyField(blank=True, null=True)
    actual_return = models.CurrencyField(blank=True, null=True)

    # demographics
    age = models.IntegerField(
        label="What is your age?",
        min=18, max=100)
    gender = models.IntegerField(
        choices=[(1, 'Male'), (2, 'Female'), (3, 'Other'), (4, 'Prefer not to answer')],
        label="What is your gender identity?",
        widget=widgets.RadioSelect
    )

    # choices for PEQs
    def make_field1(label):
        return models.IntegerField(
            choices=[1, 2, 3, 4, 5, 6, 7],
            label=label,
            widget=widgets.RadioSelect,
        )

    # initial PEQs
    q1 = make_field1('People should always be rewarded for their hard work.')
    q2 = make_field1('My estimates were based on what I felt was fair to return.')
    q3 = make_field1('High performers deserved to keep more money from the sender than low performers.')
    q4 = make_field1('The sender was able to signal trust through the amount they sent.')
    q5 = make_field1('I considered each pairing separately.')


# choices for prosocial scale
    def make_field2(label):
        return models.IntegerField(
            choices=[1, 2, 3, 4, 5],
            label=label,
            widget=widgets.RadioSelect,
        )

    # prosocial scale questions
    ps1 = make_field2('1) I am pleased to help my friends/colleagues in their activities')
    ps2 = make_field2('2) I share the things that I have with my friends')
    ps3 = make_field2('3) I try to help others')
    ps4 = make_field2('4) I am available for volunteer activities to help those who are in need')
    ps5 = make_field2('5) I am emphatic with those who are in need')
    ps6 = make_field2('6) I help immediately those who are in need')
    ps7 = make_field2('7) I do what I can to help others avoid getting into trouble')
    ps8 = make_field2('8) I intensely feel what others feel')
    ps9 = make_field2('9) I am willing to make my knowledge and abilities available to others')
    ps10 = make_field2('10) I try to console those who are sad')
    ps11 = make_field2('11) I easily lend money or other things')
    ps12 = make_field2('12) I easily put myself in the shoes of those who are in discomfort')
    ps13 = make_field2('13) I try to be close to and take care of those who are in need')
    ps14 = make_field2('14) I easily share with friends any good opportunity that comes to me')
    ps15 = make_field2('15) I spend time with those friends who feel lonely')
    ps16 = make_field2('16) I immediately sense my friends’ discomfort even when it is not directly communicated to me')
    ps17 = make_field2('17) I easily trust people.')


# ---------- Utility ----------

def safe_json_loads(raw_data, expected_type):
    try:
        parsed = json.loads(raw_data) if raw_data else expected_type()
        return parsed if isinstance(parsed, expected_type) else expected_type()
    except Exception:
        return expected_type()


def session_players_by_role(session, role):
    all_players = []
    for subs in session.get_subsessions():
        all_players += [p for p in subs.get_players() if p.role_choice == role]
    return all_players

# ---------- Pages ----------

class SetRole(Page):
    form_model = 'player'
    form_fields = ['role_choice']


class Consent(Page):
    pass


class ThirdRowOverview(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 3


class ThirdRowStage1(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 3


class Slider(Page):
    timeout_seconds = 180
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 3
    @staticmethod
    def js_vars(player: Player):
        return dict(lower1=101, upper1=500, lower2=11, upper2=99)
    @staticmethod
    def live_method(player: Player, data):
        if data == 0:
            player.num_correct += 1
        else:
            player.num_incorrect += 1
        player.attempted = player.num_correct + player.num_incorrect


class RankAndPairRow3(WaitPage):
    """Ranks row-3, creates two top-bottom pairs, and randomly assigns knowledge to one pair."""
    wait_for_all_groups = True

    @staticmethod
    def is_displayed(player: Player):
        # run once after sliders, show only to row3 to keep UX clean
        return player.role_choice == 3

    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        session = subsession.session
        row3 = session_players_by_role(session, 3)
        row2 = session_players_by_role(session, 2)

        if len(row3) < 4 or len(row2) < 4:
            session.vars['pairings_ready'] = False
            return

        rng = list(row3)
        random.shuffle(rng)
        rng.sort(key=lambda p: p.num_correct, reverse=True)
        scores = [p.num_correct for p in rng]
        tie_at_2 = (scores[1] == scores[2])

        for i, p in enumerate(rng):
            p.slider_rank = i + 1
            p.is_top = (i in [0, 1])
            if tie_at_2 and i in [1, 2]:
                p.tie_note = True

        top2 = [p for p in rng if p.is_top][:2]
        bottom2 = [p for p in rng if not p.is_top][:2]
        s_rng = list(row2)
        random.shuffle(s_rng)

        pairs = []
        for pid, (t, b, s_top, s_bottom) in enumerate(zip(top2, bottom2, s_rng[:2], s_rng[2:4]), start=1):
            t.pair_id = pid
            b.pair_id = pid
            t.partner_row2_id = s_top.id_in_subsession
            b.partner_row2_id = s_bottom.id_in_subsession
            s_top.partners_row3_id = t.id_in_subsession
            s_bottom.partners_row3_id = b.id_in_subsession
            pairs.append({
                'pair_id': pid,
                'top_row3_id': t.id_in_subsession,
                'bottom_row3_id': b.id_in_subsession,
                'sender_top_id': s_top.id_in_subsession,
                'sender_bottom_id': s_bottom.id_in_subsession,
                'tie': (t.tie_note or b.tie_note),
            })

        knowledge_pair = random.choice([1, 2])
        for p in row3:
            p.knows_status = (p.pair_id == knowledge_pair)

        # ✅ no p.save() needed here
        session.vars['pairs'] = pairs
        session.vars['knowledge_pair'] = knowledge_pair
        session.vars['pairings_ready'] = True


class ThirdRowStage2(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 3


class SecondRowOverview(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 2


class SecondRowTask(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 2


class SecondRowDecision(Page):
    form_model = 'player'
    form_fields = ['sender_decision']
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 2

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # compute amount received for their matched 3rd-row partner
        amt = player.sender_decision or cu(0)
        # store on sender now
        player.amount_received = amt * Constants.multiplier
        # also copy it onto their 3rd-row partner so the partner can see it
        subs = player.subsession
        partner = next((pl for pl in subs.get_players()
                        if getattr(pl, 'id_in_subsession', None) == getattr(player, 'partners_row3_id', None)), None)
        if partner:
            partner.amount_received = player.amount_received


class WaitForAllSends(WaitPage):
    template_name = 'HMX/WaitForAllSends.html'
    wait_for_all_groups = True
    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        # nothing else needed; this just ensures 3rd-row sees actual amounts
        return


class ThirdRowReturnDecision(Page):
    """Only one decision: how much to return, given their actual tripled amount."""
    form_model = 'player'
    form_fields = ['responder_decision']
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 3
    @staticmethod
    def vars_for_template(player: Player):
        # text helpers for templates
        perf_str = 'high performer' if player.is_top else 'low performer'
        if player.knows_status:
            info = f'Note: You were a {perf_str} in the math slider task.'
        else:
            info = ''
        return dict(
            amount_received=player.amount_received,
            knowledge_text=info,
            tie_note=player.tie_note,
        )
    @staticmethod
    def error_message(player: Player, values):
        x = values.get('responder_decision') or cu(0)
        if x < 0 or x > player.amount_received:
            return f'Return must be between 0 and {player.amount_received}.'


class FirstRowOverview(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 1


class FirstRowTask(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 1


class BuildPredictorView(WaitPage):
    wait_for_all_groups = True

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

    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        session = subsession.session

        if not session.vars.get('pairings_ready'):
            # No valid pairs → predictors should see nothing
            meta = []
        else:
            pairs = session.vars['pairs']
            players = subsession.get_players()

            def getp(pid):
                return next((pl for pl in players if pl.id_in_subsession == pid), None)

            meta = []
            for pair in pairs:
                t = getp(pair['top_row3_id'])
                b = getp(pair['bottom_row3_id'])
                s_t = getp(pair['sender_top_id'])
                s_b = getp(pair['sender_bottom_id'])

                def know_label(responder):
                    if responder.knows_status:
                        return f'knows they were a {"high" if responder.is_top else "low"} performer'
                    return 'does not know their performance status'

                meta.append({
                    'pair_id': pair['pair_id'],
                    'slot': 1,
                    'sender_id': s_t.id_in_subsession,
                    'responder_id': t.id_in_subsession,
                    'amount_sent': float((s_t.sender_decision or cu(0))),
                    'amount_received': float((s_t.sender_decision or cu(0)) * Constants.multiplier),
                    'knowledge_text': know_label(t),
                    'tie': pair['tie'],
                })

                meta.append({
                    'pair_id': pair['pair_id'],
                    'slot': 2,
                    'sender_id': s_b.id_in_subsession,
                    'responder_id': b.id_in_subsession,
                    'amount_sent': float((s_b.sender_decision or cu(0))),
                    'amount_received': float((s_b.sender_decision or cu(0)) * Constants.multiplier),
                    'knowledge_text': know_label(b),
                    'tie': pair['tie'],
                })

        # Now store meta + a randomized order per predictor
        for p in [pl for pl in subsession.get_players() if pl.role_choice == 1]:
            p.pair_meta_json = json.dumps(meta)

            if len(meta) == 4:
                order = [1, 2, 3, 4]
                random.shuffle(order)
                p.pred_order_json = json.dumps(order)
            else:
                # no pairs → no items
                p.pred_order_json = json.dumps([])


class FirstRowPrediction(Page):
    """Predictor sees all 4 interactions in random order and enters a predicted return for each."""
    form_model = 'player'
    form_fields = ['pred_pair1', 'pred_pair2', 'pred_pair3', 'pred_pair4']

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

    @staticmethod
    def vars_for_template(player: Player):
        order = safe_json_loads(player.pred_order_json, list)
        meta = safe_json_loads(player.pair_meta_json, list)
        items = []
        for i, idx in enumerate(order, start=1):
            item = meta[idx - 1]
            item['amount_sent_fmt'] = f"{item['amount_sent']:.2f}"
            item['amount_received_fmt'] = f"{item['amount_received']:.2f}"
            items.append({'display_index': i, **item})
        return dict(items=items)

    @staticmethod
    def error_message(player: Player, values):
        order = safe_json_loads(player.pred_order_json, list)
        meta = safe_json_loads(player.pair_meta_json, list)
        for i, idx in enumerate(order, start=1):
            field = f'pred_pair{i}'
            amt_rcv = meta[idx - 1]['amount_received']
            val = values.get(field) or cu(0)
            if val < 0 or val > amt_rcv:
                return f'Prediction {i} must be between 0 and {amt_rcv}.'

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # Compute predictor payoff here
        order = safe_json_loads(player.pred_order_json, list)
        meta = safe_json_loads(player.pair_meta_json, list)

        if len(order) != 4 or len(meta) != 4:
            player.prediction_bonus = cu(0)
            player.payout = cu(10)
            player.payoff = player.payout
            player.chosen_item_index = None
            player.predicted_return = None
            player.actual_return = None
            return

        chosen_idx = random.choice(order)  # 1..4
        chosen_meta = meta[chosen_idx - 1]

        subsession = player.subsession
        players = subsession.get_players()
        by_id = {p.id_in_subsession: p for p in players}
        responder = by_id.get(chosen_meta['responder_id'])

        actual_return = responder.responder_decision or cu(0)

        pos = order.index(chosen_idx) + 1
        field_map = {
            1: player.pred_pair1,
            2: player.pred_pair2,
            3: player.pred_pair3,
            4: player.pred_pair4,
        }
        predicted = field_map[pos] or cu(0)

        player.chosen_item_index = chosen_idx
        player.predicted_return = predicted
        player.actual_return = actual_return

        error = abs(float(predicted) - float(actual_return))
        player.prediction_error = error

        if error == 0:
            bonus = cu(10)
        elif error <= 1:
            bonus = cu(8)
        elif error <= 2:
            bonus = cu(5)
        elif error <= 3:
            bonus = cu(2)
        else:
            bonus = cu(0)

        player.prediction_bonus = bonus
        player.payout = cu(10) + bonus
        player.payoff = player.payout


class ComprehensionCheck(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice < 3
    @staticmethod
    def get_form_fields(player):
        if player.role_choice == 2:
            return ['comp2_q1', 'comp2_q2', 'comp2_q3']
        elif player.role_choice == 1:
            return ['comp1_q1', 'comp1_q2', 'comp1_q3']
        return []
    @staticmethod
    def error_message(player, values):
        wrong = 'At least one answer is incorrect. Please review and try again.'
        if player.role_choice == 2 and not (values['comp2_q1'] and not values['comp2_q2'] and values['comp2_q3']):
            return wrong
        if player.role_choice == 1 and not (values['comp1_q1'] and not values['comp1_q2'] and values['comp1_q3']):
            return wrong


class ComprehensionCheckThirdRow(Page):
    form_model = 'player'
    form_fields = ['comp3_q1', 'comp3_q2', 'comp3_q3']
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 3
    form_model = 'player'
    @staticmethod
    def error_message(player, values):
        wrong = 'At least one answer is incorrect. Please review and try again.'
        if not (values['comp3_q1'] and not values['comp3_q2'] and values['comp3_q3']):
            return wrong


class WaitRow2and3ForPay(WaitPage):
    wait_for_all_groups = True

    @staticmethod
    def is_displayed(player: Player):
        # Only rows 2 and 3 need to wait for each other
        return player.role_choice in [2, 3]

    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        # When ALL Row2 + Row3 players are done, compute their payoffs
        players = subsession.get_players()
        by_id = {p.id_in_subsession: p for p in players}

        for responder in [p for p in players if p.role_choice == 3]:
            # Row 3 stores the id of their Row 2 partner in partner_row2_id
            partner_id = responder.field_maybe_none('partner_row2_id')
            if partner_id is None:
                continue

            sender = by_id.get(partner_id)
            if not sender:
                continue

            # Slider earnings
            responder.slider_earnings = round(responder.num_correct * Constants.slider_payment, 2)
            sender.slider_earnings = round(sender.num_correct * Constants.slider_payment, 2)

            amount_sent = sender.sender_decision or cu(0)
            amount_rcv = amount_sent * Constants.multiplier
            returned = responder.responder_decision or cu(0)

            # Responder payoff
            responder.amount_received = amount_rcv
            responder.amount_kept = amount_rcv - returned
            responder.payout = cu(5 + responder.slider_earnings + responder.amount_kept)
            responder.payoff = responder.payout

            # Sender payoff
            sender.responder_decision = returned
            sender.amount_kept = Constants.endowment - amount_sent
            sender.payout = cu(10 + sender.amount_kept + returned)
            sender.payoff = sender.payout


# class FinalizeAndPay(WaitPage):
#     wait_for_all_groups = True
#
#     @staticmethod
#     def after_all_players_arrive(subsession: Subsession):
#         players = subsession.get_players()
#         # map by id_in_subsession
#         by_id = {p.id_in_subsession: p for p in players}
#
#         # Responder slider earnings and final payoff
#         for p in players:
#             p.slider_earnings = round(p.num_correct * Constants.slider_payment, 2)
#
#         # Sender and Responder payoffs
#         for s in [p for p in players if p.role_choice == 2]:
#             partner_id = s.field_maybe_none('partners_row3_id')
#             r = by_id.get(partner_id)
#             if not r:
#                 # if pairing somehow failed, skip to avoid crashes
#                 continue
#
#             amount_sent = s.sender_decision or cu(0)
#             amount_rcv = amount_sent * Constants.multiplier
#             ret = r.field_maybe_none('responder_decision') or cu(0)
#
#             r.amount_received = amount_rcv
#             r.amount_kept = amount_rcv - ret
#             r.payout = cu(5 + r.slider_earnings + r.amount_kept)
#
#             s.responder_decision = ret
#             s.amount_kept = Constants.endowment - amount_sent
#             s.payout = cu(10 + s.amount_kept + ret)
#
#         # Predictor payoffs: pick one random displayed item and score against the corresponding responder return
#         for pr in [p for p in players if p.role_choice == 1]:
#             order = safe_json_loads(pr.pred_order_json, list)
#             meta = safe_json_loads(pr.pair_meta_json, list)
#             if len(order) != 4 or len(meta) != 4:
#                 pr.prediction_bonus = cu(0)
#                 pr.payout = cu(10)
#                 # optional: clear chosen fields
#                 pr.chosen_item_index = None
#                 pr.predicted_return = None
#                 pr.actual_return = None
#                 continue
#
#             chosen_idx = random.choice(order)  # 1..4 (index into meta, not position on page)
#             chosen_meta = meta[chosen_idx - 1]
#             r = by_id.get(chosen_meta['responder_id'])
#             actual_return = r.responder_decision or cu(0)
#
#             # position on the page = where that chosen_idx appeared in the randomized order
#             pos = order.index(chosen_idx) + 1
#             field_map = {1: pr.pred_pair1, 2: pr.pred_pair2, 3: pr.pred_pair3, 4: pr.pred_pair4}
#             predicted = field_map[pos] or cu(0)
#
#             # store for template/export
#             pr.chosen_item_index = chosen_idx
#             pr.predicted_return = predicted
#             pr.actual_return = actual_return
#
#             error = abs(float(predicted) - float(actual_return))
#             pr.prediction_error = error
#
#             if error == 0:
#                 bonus = cu(10)
#             elif error <= 1:
#                 bonus = cu(8)
#             elif error <= 2:
#                 bonus = cu(5)
#             elif error <= 3:
#                 bonus = cu(2)
#             else:
#                 bonus = cu(0)
#
#             pr.prediction_bonus = bonus
#             pr.payout = cu(10) + bonus
#
#         # set oTree payoff
#         for p in players:
#             p.payoff = p.payout


class Results(Page):
    pass


class Payout(Page):
    pass


class WaitForPairings(WaitPage):
    wait_for_all_groups = True

    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 2

    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        # nothing else needed, just ensures pairings are set
        return


class PEQ(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 1

    form_model = 'player'
    form_fields = ['q1',
                   'q2',
                   'q3',
                   'q4',
                   'q5',
                   ]


class ProsocialPEQ(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.role_choice == 1

    form_model = 'player'
    form_fields = ['ps1',
                   'ps2',
                   'ps3',
                   'ps4',
                   'ps5',
                   'ps6',
                   'ps7',
                   'ps8',
                   'ps9',
                   'ps10',
                   'ps11',
                   'ps12',
                   'ps13',
                   'ps14',
                   'ps15',
                   'ps16',
                   'ps17'
                   ]


class Demographics(Page):
    form_model = 'player'
    form_fields = ['age', 'gender']


page_sequence = [
    SetRole,
    Consent,
    ThirdRowOverview,
    ThirdRowStage1,
    Slider,
    FirstRowOverview,
    FirstRowTask,
    SecondRowOverview,
    SecondRowTask,
    ComprehensionCheck,
    RankAndPairRow3,
    WaitForPairings,
    SecondRowDecision,
    ThirdRowStage2,
    ComprehensionCheckThirdRow,
    WaitForAllSends,
    ThirdRowReturnDecision,
    WaitRow2and3ForPay,
    BuildPredictorView,
    FirstRowPrediction,
    # FinalizeAndPay,
    Results,
    PEQ,
    ProsocialPEQ,
    Demographics,
    Payout,
]