#####################################################
#####################################################
# 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 os
import pickle
import all_constants as _K
from shared_main_lib import ISLAND_COLORS, generate_candidate_graph, get_treatment, get_num_cap, _recycle_uid
c = cu

doc = ''

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

_MARKETS_CACHE = {}

def _load_markets(num_players):
    if num_players not in _MARKETS_CACHE:
        pkl_path = os.path.join(os.path.dirname(__file__), f'payoff_priority_stb_aligned_{num_players}.pkl')
        with open(pkl_path, 'rb') as f:
            raw = pickle.load(f)
        _MARKETS_CACHE[num_players] = {d['market']: d for d in raw}
    return _MARKETS_CACHE[num_players]

# Constants
class C(BaseConstants):
    NAME_IN_URL = 'stb_ali_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
    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 (positional: belief1 = prob. for rank-1 choice, etc.)
    belief0 = models.IntegerField(min=0, max=100)  # prob. of not being sent to any ranked island
    belief1 = models.IntegerField(min=0, max=100)
    belief2 = models.IntegerField(min=0, max=100, blank = True)
    belief3 = models.IntegerField(min=0, max=100, blank = True)
    belief4 = models.IntegerField(min=0, max=100, blank = True)
    belief5 = models.IntegerField(min=0, max=100, blank = True)
    belief6 = models.IntegerField(min=0, max=100, blank = True)

    ## 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_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 = _load_markets(player.session.config['num_players'])[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


###### 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

        n = player.session.config['num_players']
        table = generate_candidate_graph(payoffs=payoffs, priorities=priorities, treatment='aligned', num_players=n)
        return {
            'pf_graph': table,
            'rank_fields': [f'rank{i}' for i in range(1, get_num_cap(player) + 1)],
            'num_players': n,
            'num_quota': n // 6,
        }


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]
        n = player.session.config['num_players']
        table = generate_candidate_graph(payoffs=payoffs, priorities=priorities, treatment='aligned', num_players=n)

        ranking_rows = []
        for i in range(1, get_num_cap(player) + 1):
            val = getattr(player, f'rank{i}')
            letter = val.upper() if val else None
            bg, txt = ISLAND_COLORS[letter] if letter else (None, None)
            ranking_rows.append({'rank': i, 'island': letter, 'bg': bg, 'txt': txt})

        return {
            'pf_graph': table,
            'slider_min': 0,
            'slider_max': max(payoffs),
            'ranking_rows': ranking_rows,
            'num_players': n,
            'num_quota': n // 6,
        }


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

    @staticmethod
    def get_form_fields(player):
        n = sum(1 for i in range(1, 7) if player.field_maybe_none(f'rank{i}'))
        return ['belief0'] + [f'belief{i}' for i in range(1, n + 1)] + ['belief_smt_ct']

    @staticmethod
    def error_message(player, values):
        n = sum(1 for i in range(1, 7) if player.field_maybe_none(f'rank{i}'))
        missing_ranked = [i for i in range(1, n + 1) if values.get(f'belief{i}') is None]
        if missing_ranked or values.get('belief0') is None:
            return 'Please enter a probability for every island in your ranking and for "Not sent to any island".'
        total = sum(values[f'belief{i}'] for i in range(1, n + 1)) + values['belief0']
        if total != 100:
            return f'Your reported probabilities sum to {total}%, but must equal exactly 100%. Please adjust them.'

    @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]
        n = player.session.config['num_players']
        table = generate_candidate_graph(payoffs=payoffs, priorities=priorities, treatment='aligned', num_players=n)

        num_cap = get_num_cap(player)
        belief_rows = []
        for i in range(1, num_cap + 1):
            letter = player.field_maybe_none(f'rank{i}') or ''
            filled = bool(letter)
            bg, txt = ISLAND_COLORS.get(letter, ('#eee', '#000'))
            val = player.field_maybe_none(f'belief{i}') if filled else None
            belief_rows.append({
                'rank': i,
                'island': letter,
                'bg': bg,
                'txt': txt,
                'field_name': f'belief{i}',
                'value': '' if val is None else val,
                'filled': filled,
            })

        belief0_val = player.field_maybe_none('belief0')
        return {
            'pf_graph': table,
            'belief_rows': belief_rows,
            'belief0_value': '' if belief0_val is None else belief0_val,
            'num_players': n,
            'num_quota': n // 6,
        }


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):
        num_cap = get_num_cap(player)
        ranks = []
        beliefs = []
        for i in range(1, num_cap + 1):
            letter = player.field_maybe_none(f'rank{i}') or ''
            filled = bool(letter)
            bg, txt = ISLAND_COLORS.get(letter, ('#eee', '#000'))
            ranks.append({'rank': i, 'island': letter, 'bg': bg, 'txt': txt, 'filled': filled})
            val = player.field_maybe_none(f'belief{i}') if filled else None
            beliefs.append({'rank': i, 'island': letter, 'bg': bg, 'txt': txt,
                            'belief': val, 'filled': filled})

        return {'ranks': ranks, 'beliefs': beliefs,
                'expected_payoff': player.expected_payoff}


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]