#####################################################
#####################################################
# README
#
# Program Name: __init__.py
# Purpose:  Interface Code — combined nocap + cap treatments
#####################################################
#
# Author: Andrew Olsen
# Date Created: 04.23.2026
# Last Updated: 04.23.2026
#
#####################################################

#### Modules
from otree.api import *
import numpy as np
import itertools
import pandas as pd
import random
import math
import os
import pickle
import stb_rdm_constants as _K
c = cu

doc = ''

####################
#### File Paths ####
####################

_pkl_path = os.path.join(os.path.dirname(__file__), 'payoff_priority_stb_random_150.pkl')
with open(_pkl_path, 'rb') as _f:
    _raw_markets = pickle.load(_f)
MARKETS = {d['market']: d for d in _raw_markets}

# Constants
class C(BaseConstants):
    NAME_IN_URL = 'stb_rdm_main'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = _K.NUM_ROUNDS_PER_BLOCK * 2
    NUM_ROUNDS_PER_BLOCK = _K.NUM_ROUNDS_PER_BLOCK
    NUM_CAP_NOC = _K.NUM_CAP_NOC
    NUM_CAP_CAP = _K.NUM_CAP_CAP
    NUM_PLAYERS = _K.NUM_PLAYERS
    NUM_QUOTA = _K.NUM_QUOTA
    BDM_PAY = _K.BDM_PAY
    TIMEOUT_SECONDS = _K.TIMEOUT_SECONDS


# Subsessions
class Subsession(BaseSubsession):
    pass


# Groups
class Group(BaseGroup):
    pass


# Player
class Player(BasePlayer):

    round_id = models.IntegerField()

    ## Belief of acceptance
    beliefA = models.IntegerField(min=0, max=100)
    beliefB = models.IntegerField(min=0, max=100)
    beliefC = models.IntegerField(min=0, max=100)
    beliefD = models.IntegerField(min=0, max=100)
    beliefE = models.IntegerField(min=0, max=100)
    beliefF = models.IntegerField(min=0, max=100)

    ## Rank Orders (rank1-rank6 covers nocap; cap rounds leave rank3-rank6 blank)
    rank1 = models.StringField()
    rank2 = models.StringField(blank=True)
    rank3 = models.StringField(blank=True)
    rank4 = models.StringField(blank=True)
    rank5 = models.StringField(blank=True)
    rank6 = models.StringField(blank=True)

    smt_ct = models.IntegerField(initial=0)

    expected_payoff = models.IntegerField()
    ep_smt_ct = models.IntegerField(initial=0)

    belief_smt_ct = models.IntegerField(initial=0)

    market_id = models.IntegerField()
    payoffA = models.IntegerField()
    payoffB = models.IntegerField()
    payoffC = models.IntegerField()
    payoffD = models.IntegerField()
    payoffE = models.IntegerField()
    payoffF = models.IntegerField()
    priorityA = models.IntegerField()
    priorityB = models.IntegerField()
    priorityC = models.IntegerField()
    priorityD = models.IntegerField()
    priorityE = models.IntegerField()
    priorityF = models.IntegerField()
    
    uid = models.IntegerField()

    cap = models.IntegerField()  # 1 = cap treatment (ROL capped at 2), 0 = nocap

    approach_block1 = models.LongStringField()
    approach_block2 = models.LongStringField()


_ISLANDS = ['A', 'B', 'C', 'D', 'E', 'F']


def get_treatment(player):
    """Return 'nocap' or 'cap' for this player's current round.

    Odd uid  → rounds 1-10 nocap, rounds 11-20 cap.
    Even uid → rounds 1-10 cap,   rounds 11-20 nocap.
    """
    uid = player.participant.uid
    first_block = player.round_number <= C.NUM_ROUNDS_PER_BLOCK
    is_odd = (uid % 2 == 1)
    if first_block:
        return 'nocap' if is_odd else 'cap'
    else:
        return 'cap' if is_odd else 'nocap'


def get_num_cap(player):
    return C.NUM_CAP_NOC if get_treatment(player) == 'nocap' else C.NUM_CAP_CAP


def get_market_data(player):
    treatment = get_treatment(player)
    block_round = (player.round_number - 1) % C.NUM_ROUNDS_PER_BLOCK
    if treatment == 'nocap':
        mkt_idx = player.participant.mkt_ord_nocap[block_round]
    else:
        mkt_idx = player.participant.mkt_ord_cap[block_round]
    mkt = MARKETS[mkt_idx]
    uid = player.participant.uid
    payoffs    = [int(mkt['Payoffs_Mat'][uid][i])  for i in range(6)]
    priorities = [int(mkt['Priority_Mat'][uid][i]) for i in range(6)]
    return mkt_idx, payoffs, priorities


def generate_candidate_graph(payoffs=[100, 200, 300, 250, 120, 150],
                             priorities=[100, 100, 100, 100, 100, 100]):
    """Generate HTML for candidate payoff table and number lines"""

    bar_colors = ['rgb(21, 96, 130)', 'rgb(233, 113, 50)',
                  'rgb(25, 107, 36)', 'rgb(15, 158, 213)',
                  'rgb(160, 43, 147)', 'rgb(209, 209, 209)']
    islands = ['A', 'B', 'C', 'D', 'E', 'F']

    pct_lower_vals = []
    for priority_now in priorities:
        pct_lower = 100 * (priority_now - 1) / (C.NUM_PLAYERS - 1)
        if pct_lower != 100:
            pct_lower_vals.append(f'{pct_lower:3.1f}%')
        else:
            pct_lower_vals.append(f'{pct_lower:3.0f}%')

    html = '<div class="info-table-container">'
    html += '<table class="info-table">'

    html += '<tr>'
    html += '<th class="row-label">Island</th>'
    for i, island in enumerate(islands):
        txt_color = '#222' if island == 'F' else 'white'
        html += (f'<th data-island="{island}"><span class="island-badge" '
                 f'style="background-color:{bar_colors[i]};color:{txt_color}">'
                 f'{island}</span></th>')
    html += '</tr>'

    html += '<tr>'
    html += '<th class="row-label">Points</th>'
    for i, pf in enumerate(payoffs):
        html += f'<td data-island="{islands[i]}">{pf:.0f}</td>'
    html += '</tr>'

    html += '<tr>'
    html += '<th class="row-label">Priority Score</th>'
    for i, pr in enumerate(priorities):
        html += f'<td data-island="{islands[i]}">{pr}</td>'
    html += '</tr>'

    html += '<tr>'
    html += '<th class="row-label">% Lower Priority</th>'
    for i, pl in enumerate(pct_lower_vals):
        html += f'<td data-island="{islands[i]}">{pl}</td>'
    html += '</tr>'

    html += '</table>'
    html += '</div>'

    svg_width = 600
    pad = 50
    line_width = svg_width - 2 * pad
    spacing = 20
    r = 9

    def make_svg(label, values, v_min, v_max):
        items = []
        for i, val in enumerate(values):
            x = pad + (val - v_min) / (v_max - v_min) * line_width
            items.append({'x': x, 'color': bar_colors[i], 'island': islands[i]})

        groups = {}
        for item in items:
            key = round(item['x'])
            groups.setdefault(key, []).append(item)

        for key, group in groups.items():
            n = len(group)
            group.sort(key=lambda p: p['island'])
            for j, p in enumerate(group):
                p['y'] = -(n - 1) / 2 * spacing + j * spacing

        min_sep = 2 * r + 2
        for _ in range(300):
            moved = False
            for idx_a in range(len(items)):
                for idx_b in range(idx_a + 1, len(items)):
                    a, b = items[idx_a], items[idx_b]
                    dx = abs(a['x'] - b['x'])
                    if dx >= min_sep:
                        continue
                    needed_dy = math.sqrt(max(0.0, min_sep ** 2 - dx ** 2))
                    dy = b['y'] - a['y']
                    if abs(dy) < needed_dy - 0.01:
                        extra = (needed_dy - abs(dy)) / 2
                        if dy >= 0:
                            a['y'] -= extra
                            b['y'] += extra
                        else:
                            a['y'] += extra
                            b['y'] -= extra
                        moved = True
            if not moved:
                break

        min_y = min(item['y'] for item in items)
        offset = (r + 8) - min_y
        for item in items:
            item['y'] += offset
        y_axis_local = int(offset)

        max_y = max(item['y'] for item in items)
        svg_h = max(70, int(max_y) + r + 20)

        s = f'<div class="number-line-section">'
        s += f'<div class="number-line-label">{label}</div>'
        s += f'<svg width="{svg_width}" height="{svg_h}" class="number-line-svg" overflow="visible">'

        s += (f'<line x1="{pad}" y1="{y_axis_local}" x2="{pad + line_width}" '
              f'y2="{y_axis_local}" stroke="#888" stroke-width="2"/>')

        s += (f'<line x1="{pad}" y1="{y_axis_local - 5}" x2="{pad}" '
              f'y2="{y_axis_local + 5}" stroke="#888" stroke-width="2"/>')
        s += (f'<text x="{pad - 10}" y="{y_axis_local + 5}" text-anchor="end" '
              f'font-size="16" fill="#444">{v_min}</text>')
        s += (f'<text x="{pad - 10}" y="{y_axis_local + 18}" text-anchor="end" '
              f'font-size="11" fill="#bbb">(Lowest)</text>')

        x_max = pad + line_width
        s += (f'<line x1="{x_max}" y1="{y_axis_local - 5}" x2="{x_max}" '
              f'y2="{y_axis_local + 5}" stroke="#888" stroke-width="2"/>')
        s += (f'<text x="{x_max + 10}" y="{y_axis_local + 5}" text-anchor="start" '
              f'font-size="16" fill="#444">{v_max}</text>')
        s += (f'<text x="{x_max + 10}" y="{y_axis_local + 18}" text-anchor="start" '
              f'font-size="11" fill="#bbb">(Highest)</text>')

        for item in items:
            cx = item['x']
            cy = item['y']
            color = item['color']
            island = item['island']
            text_color = '#222' if island == 'F' else 'white'

            s += (f'<g class="island-dot" data-island="{island}" style="cursor:pointer">')
            s += (f'<circle cx="{cx:.1f}" cy="{cy:.1f}" r="{r}" fill="{color}" '
                  f'stroke="white" stroke-width="1"/>')
            s += (f'<text x="{cx:.1f}" y="{cy + 4:.1f}" text-anchor="middle" '
                  f'font-size="13" font-weight="bold" fill="{text_color}" '
                  f'style="pointer-events:none">{island}</text>')
            s += '</g>'

        s += '</svg></div>'
        return s

    html += make_svg('Points', payoffs, min(payoffs), max(payoffs))
    html += make_svg('Priority Scores', priorities, 1, C.NUM_PLAYERS)

    return html


def _recycle_uid(player):
    uid = player.participant.uid
    recycled = player.session.recycled_uids
    if uid != 1000 and uid not in recycled:
        player.session.recycled_uids = recycled + [uid]


###### Pages

class Intro_Cap(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number in (1, C.NUM_ROUNDS_PER_BLOCK + 1) and get_treatment(player) == 'cap'

    @staticmethod
    def vars_for_template(player):
        first_block = player.round_number <= C.NUM_ROUNDS_PER_BLOCK
        return {
            'round_start':    1 if first_block else C.NUM_ROUNDS_PER_BLOCK + 1,
            'round_end':      C.NUM_ROUNDS_PER_BLOCK if first_block else C.NUM_ROUNDS,
            'cap_num':        C.NUM_CAP_CAP,
            'is_second_block': not first_block,
        }


class Intro_NoCap(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number in (1, C.NUM_ROUNDS_PER_BLOCK + 1) and get_treatment(player) == 'nocap'

    @staticmethod
    def vars_for_template(player):
        first_block = player.round_number <= C.NUM_ROUNDS_PER_BLOCK
        return {
            'round_start':    1 if first_block else C.NUM_ROUNDS_PER_BLOCK + 1,
            'round_end':      C.NUM_ROUNDS_PER_BLOCK if first_block else C.NUM_ROUNDS,
            'cap_num':        C.NUM_CAP_NOC,
            'is_second_block': not first_block,
        }


class Ranks(Page):
    form_model = 'player'
    preserve_unsubmitted_inputs = True
    timeout_seconds = C.TIMEOUT_SECONDS

    @staticmethod
    def get_form_fields(player):
        return [f'rank{i}' for i in range(1, get_num_cap(player) + 1)] + ['smt_ct']

    @staticmethod
    def before_next_page(player, timeout_happened):
        if timeout_happened:
            player.participant.timed_out = True
            _recycle_uid(player)

    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        if player.participant.timed_out:
            return 'dq_fail'

    @staticmethod
    def error_message(player, values):
        allowed = set('ABCDEFabcdef')
        field_strings = [f'rank{i}' for i in range(1, get_num_cap(player) + 1)]

        if not values.get('rank1', ''):
            return {'rank1': 'Please rank at least one island.'}

        for field_name in field_strings:
            value = values.get(field_name, '')
            if not value:
                continue
            if len(value) != 1:
                return {field_name: 'Please enter one island.'}
            if value not in allowed:
                return {field_name: 'Please enter one of: A, B, C, D, E, F'}

        for i in range(1, len(field_strings)):
            val_prev = values.get(field_strings[i - 1], '')
            val_curr = values.get(field_strings[i], '')
            if val_curr and not val_prev:
                return {field_strings[i]: 'Please fill ranks in order without gaps.'}

        non_empty = [(field_strings[i], values[field_strings[i]].upper())
                     for i in range(len(field_strings))
                     if values.get(field_strings[i], '')]
        for ii in range(len(non_empty)):
            for jj in range(ii + 1, len(non_empty)):
                if non_empty[ii][1] == non_empty[jj][1]:
                    return {non_empty[jj][0]: 'Islands can only be ranked once.'}

    @staticmethod
    def vars_for_template(player):
        mkt_id, payoffs, priorities = get_market_data(player)

        player.market_id = mkt_id
        for i, letter in enumerate(_ISLANDS):
            setattr(player, f'payoff{letter}',   payoffs[i])
            setattr(player, f'priority{letter}', priorities[i])
            
        ## Display player UID for debug
        setattr(player, 'uid', player.participant.uid)

        player.cap = 1 if get_treatment(player) == 'cap' else 0

        table = generate_candidate_graph(payoffs=payoffs, priorities=priorities)
        return {
            'pf_graph': table,
            'rank_fields': [f'rank{i}' for i in range(1, get_num_cap(player) + 1)],
        }


class ExpectedPayoff(Page):
    form_model = 'player'
    form_fields = ['expected_payoff', 'ep_smt_ct']
    preserve_unsubmitted_inputs = True
    timeout_seconds = C.TIMEOUT_SECONDS

    @staticmethod
    def before_next_page(player, timeout_happened):
        if timeout_happened:
            player.participant.timed_out = True
            _recycle_uid(player)

    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        if player.participant.timed_out:
            return 'dq_fail'

    @staticmethod
    def vars_for_template(player):
        payoffs    = [getattr(player, f'payoff{l}')   for l in _ISLANDS]
        priorities = [getattr(player, f'priority{l}') for l in _ISLANDS]
        table = generate_candidate_graph(payoffs=payoffs, priorities=priorities)
        return {
            'pf_graph': table,
            'slider_min': 0,
            'slider_max': max(payoffs),
        }


class Beliefs(Page):
    form_model = 'player'
    form_fields = ['beliefA', 'beliefB', 'beliefC', 'beliefD', 'beliefE', 'beliefF', 'belief_smt_ct']
    preserve_unsubmitted_inputs = True
    timeout_seconds = C.TIMEOUT_SECONDS

    @staticmethod
    def before_next_page(player, timeout_happened):
        if timeout_happened:
            player.participant.timed_out = True
            _recycle_uid(player)

    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        if player.participant.timed_out:
            return 'dq_fail'

    @staticmethod
    def vars_for_template(player):
        payoffs    = [getattr(player, f'payoff{l}')   for l in _ISLANDS]
        priorities = [getattr(player, f'priority{l}') for l in _ISLANDS]
        table = generate_candidate_graph(payoffs=payoffs, priorities=priorities)
        return {'pf_graph': table}


ISLAND_COLORS = {
    'A': ('rgb(21, 96, 130)',  'white'),
    'B': ('rgb(233, 113, 50)', 'white'),
    'C': ('rgb(25, 107, 36)',  'white'),
    'D': ('rgb(15, 158, 213)', 'white'),
    'E': ('rgb(160, 43, 147)', 'white'),
    'F': ('rgb(209, 209, 209)', '#222'),
}

class Summary(Page):
    timeout_seconds = C.TIMEOUT_SECONDS

    @staticmethod
    def before_next_page(player, timeout_happened):
        if timeout_happened:
            player.participant.timed_out = True
            _recycle_uid(player)

    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        if player.participant.timed_out:
            return 'dq_fail'

    @staticmethod
    def vars_for_template(player):
        ranks = []
        for i in range(1, get_num_cap(player) + 1):
            val = getattr(player, f'rank{i}')
            if val:
                letter = val.upper()
                bg, txt = ISLAND_COLORS[letter]
                ranks.append({'rank': i, 'island': letter, 'bg': bg, 'txt': txt})

        beliefs = []
        for letter in ['A', 'B', 'C', 'D', 'E', 'F']:
            bg, txt = ISLAND_COLORS[letter]
            beliefs.append({'island': letter,
                            'belief': getattr(player, f'belief{letter}'),
                            'bg': bg, 'txt': txt})

        return {'ranks': ranks, 'beliefs': beliefs}


class Approach(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player):
        return player.round_number in (C.NUM_ROUNDS_PER_BLOCK, C.NUM_ROUNDS)

    @staticmethod
    def get_form_fields(player):
        if player.round_number == C.NUM_ROUNDS_PER_BLOCK:
            return ['approach_block1']
        else:
            return ['approach_block2']

    @staticmethod
    def vars_for_template(player):
        first_block = player.round_number == C.NUM_ROUNDS_PER_BLOCK
        return {
            'round_start': 1 if first_block else C.NUM_ROUNDS_PER_BLOCK + 1,
            'round_end':   C.NUM_ROUNDS_PER_BLOCK if first_block else C.NUM_ROUNDS,
            'field_name':  'approach_block1' if first_block else 'approach_block2',
        }


page_sequence = [Intro_Cap, Intro_NoCap, Ranks, ExpectedPayoff, Beliefs, Summary, Approach]