from otree.api import *
import random
import json

doc = """
Payment page: draw one random prediction per machine, compare with
researcher-defined correct answers, and pay a bonus for each correct one.
"""


class C(BaseConstants):
    NAME_IN_URL = 'Pay'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1
    BONUS_PER_CORRECT = cu(1.50)  # bonus per correctly-predicted drawn guess
    CURRENCY = '£'


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # JSON of per-machine draw details (for researcher export)
    payment_details = models.LongStringField(blank=True, field_maybe_none=True)
    bonus_total = models.CurrencyField(initial=0)
    n_machines = models.IntegerField(initial=0)
    n_correct_drawn = models.IntegerField(initial=0)


class Feedback(Page):
    """Payment mode A – one random draw per machine, bonus for each correct."""

    @staticmethod
    def vars_for_template(player: Player):
        # If payment was already computed, reuse it (prevents refresh re-rolling)
        if player.field_maybe_none('payment_details') is not None:
            details = json.loads(player.payment_details)
            bonus = player.bonus_total
            show_up_fee = cu(player.session.config.get('participation_fee', 0) or 0)
            return dict(
                show_up_fee=show_up_fee,
                bonus=bonus,
                n_correct=player.n_correct_drawn,
                n_machines=player.n_machines,
                total_payment=show_up_fee + bonus,
            )

        participant = player.participant

        order_names = participant.order_names          # e.g. ['2L', '1L']
        all_guess_values = participant.guess_values    # flat list length = n_machines * gpg
        all_guess_configs = participant.guess_configs  # flat list, same length
        correct_preds = participant.correct_predictions  # dict of dicts

        n_machines = len(order_names)
        # Compute guesses-per-machine from actual data (works for both 12 and 16)
        gpg = len(all_guess_values) // n_machines if n_machines > 0 else 0

        details = []
        n_correct = 0

        for m_idx in range(n_machines):
            machine_name = order_names[m_idx]
            start = m_idx * gpg
            end = start + gpg

            m_guesses = all_guess_values[start:end]
            m_configs = all_guess_configs[start:end]

            # Guard against fewer guesses than expected
            actual_count = len(m_guesses)
            if actual_count == 0:
                continue
            draw_idx = random.randint(0, actual_count - 1)
            drawn_guess = m_guesses[draw_idx]
            drawn_config = m_configs[draw_idx]

            # Look up the researcher-defined correct answer
            correct_answer = correct_preds.get(machine_name, {}).get(drawn_config, None)
            is_correct = (drawn_guess == correct_answer)
            if is_correct:
                n_correct += 1

            details.append({
                'machine': machine_name,
                'draw_index': draw_idx + 1,
                'config': drawn_config,
                'guess': drawn_guess,
                'correct_answer': correct_answer,
                'is_correct': is_correct,
            })

        bonus = n_correct * C.BONUS_PER_CORRECT
        player.payoff = bonus
        player.bonus_total = bonus
        player.n_machines = n_machines
        player.n_correct_drawn = n_correct
        player.payment_details = json.dumps(details)

        show_up_fee = cu(player.session.config.get('participation_fee', 0) or 0)
        total_payment = show_up_fee + bonus

        return dict(
            show_up_fee=show_up_fee,
            bonus=bonus,
            n_correct=n_correct,
            n_machines=n_machines,
            total_payment=total_payment,
        )


class FeedbackSingleCase(Page):
    """Payment mode B – one random machine is chosen, then one random
    observation from that machine is drawn.  The subject is paid only
    if that single draw is correct."""

    @staticmethod
    def vars_for_template(player: Player):
        # If payment was already computed, reuse it (prevents refresh re-rolling)
        if player.field_maybe_none('payment_details') is not None:
            details = json.loads(player.payment_details)
            bonus = player.bonus_total
            show_up_fee = cu(player.session.config.get('participation_fee', 0) or 0)
            return dict(
                show_up_fee=show_up_fee,
                bonus=bonus,
                n_correct=player.n_correct_drawn,
                n_machines=player.n_machines,
                chosen_machine=details[0]['machine'] if details else '',
                total_payment=show_up_fee + bonus,
            )

        participant = player.participant

        order_names = participant.order_names
        all_guess_values = participant.guess_values
        all_guess_configs = participant.guess_configs
        correct_preds = participant.correct_predictions

        n_machines = len(order_names)
        # Compute guesses-per-machine from actual data
        gpg = len(all_guess_values) // n_machines if n_machines > 0 else 0

        # 1) Pick one random machine
        chosen_m_idx = random.randint(0, n_machines - 1)
        machine_name = order_names[chosen_m_idx]

        # 2) Pick one random guess from that machine
        start = chosen_m_idx * gpg
        end = start + gpg
        m_guesses = all_guess_values[start:end]
        m_configs = all_guess_configs[start:end]

        actual_count = len(m_guesses)
        if actual_count == 0:
            draw_idx = 0
            drawn_guess = None
            drawn_config = None
        else:
            draw_idx = random.randint(0, actual_count - 1)
            drawn_guess = m_guesses[draw_idx]
            drawn_config = m_configs[draw_idx]

        # 3) Check correctness
        correct_answer = correct_preds.get(machine_name, {}).get(drawn_config, None)
        is_correct = (drawn_guess == correct_answer)
        n_correct = 1 if is_correct else 0

        details = [{
            'machine': machine_name,
            'draw_index': draw_idx + 1,
            'config': drawn_config,
            'guess': drawn_guess,
            'correct_answer': correct_answer,
            'is_correct': is_correct,
        }]

        bonus = n_correct * C.BONUS_PER_CORRECT
        player.payoff = bonus
        player.bonus_total = bonus
        player.n_machines = n_machines
        player.n_correct_drawn = n_correct
        player.payment_details = json.dumps(details)

        show_up_fee = cu(player.session.config.get('participation_fee', 0) or 0)
        total_payment = show_up_fee + bonus

        return dict(
            show_up_fee=show_up_fee,
            bonus=bonus,
            n_correct=n_correct,
            n_machines=n_machines,
            chosen_machine=machine_name,
            total_payment=total_payment,
        )


# ── Helper: draw a random observation matching a light config ────────
def _draw_observation_sound(case_def, freq, config_tuple, rng=None):
    """
    Given a machine's case definition (list of rows) and frequency vector,
    find all rows whose light-config matches *config_tuple*, weight them
    by frequency, draw one at random, and return its sound value (last col).
    Returns None if no matching row exists.
    rng: optional random.Random instance; falls back to global random if None.
    """
    n_cols = len(case_def[0])
    config_cols = n_cols - 2          # exclude Others and Sound

    matching_rows = []
    matching_weights = []
    for i, row in enumerate(case_def):
        row_config = tuple(int(x) for x in row[:config_cols])
        if row_config == config_tuple:
            matching_rows.append(row)
            matching_weights.append(freq[i] if freq[i] > 0 else 1)

    if not matching_rows:
        return None

    _rng = rng if rng is not None else random
    drawn_row = _rng.choices(matching_rows, weights=matching_weights, k=1)[0]
    return int(drawn_row[-1])         # sound column


class FeedbackDrawAll(Page):
    """Payment mode C – for EVERY machine, pick one random prediction,
    draw a random observation with the same light config from the
    original case data, and pay a bonus if the sounds match."""

    @staticmethod
    def vars_for_template(player: Player):
        # If payment was already computed, reuse it (prevents refresh re-rolling)
        if player.field_maybe_none('payment_details') is not None:
            details = json.loads(player.payment_details)
            bonus = player.bonus_total
            show_up_fee = cu(player.session.config.get('participation_fee', 0) or 0)
            return dict(
                show_up_fee=show_up_fee,
                bonus=bonus,
                n_correct=player.n_correct_drawn,
                n_machines=player.n_machines,
                total_payment=show_up_fee + bonus,
            )

        participant = player.participant

        order_names   = participant.order_names
        all_guess_values  = participant.guess_values
        all_guess_configs = participant.guess_configs
        cases_ordered = participant.cases_ordered   # [[case_def, freq], …]

        n_machines = len(order_names)
        gpg = len(all_guess_values) // n_machines if n_machines > 0 else 0

        details = []
        n_correct = 0

        for m_idx in range(n_machines):
            machine_name = order_names[m_idx]
            case_def, freq = cases_ordered[m_idx]

            start = m_idx * gpg
            end   = start + gpg
            m_guesses = all_guess_values[start:end]
            m_configs = all_guess_configs[start:end]

            actual_count = len(m_guesses)
            if actual_count == 0:
                continue

            draw_idx     = random.randint(0, actual_count - 1)
            drawn_guess  = m_guesses[draw_idx]
            drawn_config = m_configs[draw_idx]

            # Parse the config string "(0, 1, 0)" back into a tuple
            config_tuple = tuple(int(x) for x in drawn_config.strip('()').split(','))

            # Draw a random observation with the same lights from the case data
            drawn_sound = _draw_observation_sound(case_def, freq, config_tuple)

            is_correct = (drawn_guess == drawn_sound) if drawn_sound is not None else False
            if is_correct:
                n_correct += 1

            details.append({
                'machine': machine_name,
                'draw_index': draw_idx + 1,
                'config': drawn_config,
                'guess': drawn_guess,
                'drawn_sound': drawn_sound,
                'is_correct': is_correct,
            })

        bonus = n_correct * C.BONUS_PER_CORRECT
        player.payoff = bonus
        player.bonus_total = bonus
        player.n_machines = n_machines
        player.n_correct_drawn = n_correct
        player.payment_details = json.dumps(details)

        show_up_fee   = cu(player.session.config.get('participation_fee', 0) or 0)
        total_payment = show_up_fee + bonus

        return dict(
            show_up_fee=show_up_fee,
            bonus=bonus,
            n_correct=n_correct,
            n_machines=n_machines,
            total_payment=total_payment,
        )


def _build_lights_html(config_tuple, original_color_val):
    """Build inline HTML showing lights as the participant saw them.

    Display is always Red | Blue [| Green] circles in that order.
    original_color_val determines which data column feeds each display position:
      2-light cases: 1 = original (col0→Red, col1→Blue), 0 = flipped (col1→Red, col0→Blue)
      3-light cases: 0–5 index into the six permutations of (Red, Blue, Green)
    """
    n = len(config_tuple)
    if n == 2:
        # 2 lights: original_color is 0 or 1
        srcs = [0, 1] if original_color_val == 1 else [1, 0]
        colors = ['red', 'blue']
    else:
        # 3 lights: original_color is 0–5 permutation index
        _perms = [(0,1,2),(1,0,2),(2,1,0),(0,2,1),(1,2,0),(2,0,1)]
        perm = _perms[original_color_val] if 0 <= original_color_val < len(_perms) else (0,1,2)
        srcs = list(perm)
        colors = ['red', 'blue', 'green']

    html = '<span style="display:inline-flex;gap:8px;align-items:center;">'
    for color, src in zip(colors, srcs):
        val = config_tuple[src] if src < len(config_tuple) else 0
        cls = 'circle_{}'.format(color) if val == 1 else 'circle_{}_off'.format(color)
        html += '<div class="{}"></div>'.format(cls)
    html += '</span>'
    return html


class FeedbackDrawSingle(Page):
    """Payment mode D – pick ONE random machine, then one random
    prediction from it, draw a random observation with the same light
    config, and pay a bonus only if the sounds match."""

    @staticmethod
    def vars_for_template(player: Player):
        # If payment was already computed, reuse it (prevents refresh re-rolling)
        if player.field_maybe_none('payment_details') is not None:
            details = json.loads(player.payment_details)
            bonus = player.bonus_total
            show_up_fee = cu(player.session.config.get('participation_fee', 0) or 0)
            participant = player.participant
            detail0 = details[0] if details else {}
            machine_name = detail0.get('machine', '')
            order_names = participant.order_names
            chosen_m_idx = order_names.index(machine_name) if machine_name in order_names else 0
            drawn_config = detail0.get('config') or '()'
            drawn_guess = detail0.get('guess')
            drawn_sound = detail0.get('drawn_sound')
            is_correct = detail0.get('is_correct', False)
            try:
                config_tuple = tuple(int(x) for x in drawn_config.strip('()').split(',') if x.strip())
            except Exception:
                config_tuple = ()
            orig_colors = getattr(participant, 'original_color', [])
            original_color_val = orig_colors[chosen_m_idx] if chosen_m_idx < len(orig_colors) else 1
            lights_html = _build_lights_html(config_tuple, original_color_val) if config_tuple else '—'
            guess_text = 'Ding' if drawn_guess == 1 else ('No Ding' if drawn_guess == 0 else '—')
            correct_text = 'Ding' if drawn_sound == 1 else ('No Ding' if drawn_sound == 0 else '—')
            return dict(
                show_up_fee=show_up_fee,
                bonus=bonus,
                n_correct=player.n_correct_drawn,
                n_machines=player.n_machines,
                chosen_machine=machine_name,
                total_payment=show_up_fee + bonus,
                lights_html=lights_html,
                guess_text=guess_text,
                correct_text=correct_text,
                is_correct=is_correct,
                draw_index=detail0.get('draw_index', '—'),
            )

        participant = player.participant

        order_names   = participant.order_names
        all_guess_values  = participant.guess_values
        all_guess_configs = participant.guess_configs
        cases_ordered = participant.cases_ordered

        n_machines = len(order_names)
        gpg = len(all_guess_values) // n_machines if n_machines > 0 else 0

        # Use a per-participant RNG so draws are independent across subjects.
        # participant.code is oTree's unique random token per participant.
        rng = random.Random(participant.code)

        # 1) Pick one random machine
        chosen_m_idx = rng.randint(0, n_machines - 1)
        machine_name = order_names[chosen_m_idx]
        case_def, freq = cases_ordered[chosen_m_idx]

        # 2) Pick one random prediction from that machine
        start = chosen_m_idx * gpg
        end   = start + gpg
        m_guesses = all_guess_values[start:end]
        m_configs = all_guess_configs[start:end]

        actual_count = len(m_guesses)
        if actual_count == 0:
            drawn_guess  = None
            drawn_config = None
            drawn_sound  = None
            is_correct   = False
            draw_idx     = 0
            lights_html  = '—'
            guess_text   = '—'
            correct_text = '—'
        else:
            draw_idx     = rng.randint(0, actual_count - 1)
            drawn_guess  = m_guesses[draw_idx]
            drawn_config = m_configs[draw_idx]

            config_tuple = tuple(int(x) for x in drawn_config.strip('()').split(','))
            drawn_sound  = _draw_observation_sound(case_def, freq, config_tuple, rng=rng)
            is_correct   = (drawn_guess == drawn_sound) if drawn_sound is not None else False

            orig_colors = getattr(participant, 'original_color', [])
            original_color_val = orig_colors[chosen_m_idx] if chosen_m_idx < len(orig_colors) else 1
            lights_html  = _build_lights_html(config_tuple, original_color_val)
            guess_text   = 'Ding' if drawn_guess == 1 else 'No Ding'
            correct_text = 'Ding' if drawn_sound == 1 else ('No Ding' if drawn_sound is not None else '—')

        n_correct = 1 if is_correct else 0

        details = [{
            'machine': machine_name,
            'draw_index': draw_idx + 1,
            'config': drawn_config,
            'guess': drawn_guess,
            'drawn_sound': drawn_sound,
            'is_correct': is_correct,
        }]

        bonus = n_correct * C.BONUS_PER_CORRECT
        player.payoff = bonus
        player.bonus_total = bonus
        player.n_machines = n_machines
        player.n_correct_drawn = n_correct
        player.payment_details = json.dumps(details)

        show_up_fee   = cu(player.session.config.get('participation_fee', 0) or 0)
        total_payment = show_up_fee + bonus

        return dict(
            show_up_fee=show_up_fee,
            bonus=bonus,
            n_correct=n_correct,
            n_machines=n_machines,
            chosen_machine=machine_name,
            total_payment=total_payment,
            lights_html=lights_html,
            guess_text=guess_text,
            correct_text=correct_text,
            is_correct=is_correct,
            draw_index=draw_idx + 1,
        )


# ── Choose payment mode ─────────────────────────────────────────────
# Feedback             → mode A: one draw per machine, checked against
#                        researcher-defined correct_predictions
# FeedbackSingleCase   → mode B: one random machine, one draw, checked
#                        against researcher-defined correct_predictions
# FeedbackDrawAll      → mode C: one draw per machine, checked against
#                        a random observation drawn from the case data
# FeedbackDrawSingle   → mode D: one random machine, one draw, checked
#                        against a random observation drawn from the case data
# Uncomment the desired mode:
# page_sequence = [Feedback]
# page_sequence = [FeedbackSingleCase]
# page_sequence = [FeedbackDrawAll]
page_sequence = [FeedbackDrawSingle]