from otree.api import *
import random
import pycountry

doc = """
Nationality allocation experiment
"""


class C(BaseConstants):
    NAME_IN_URL = 'allocation_app'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1
    TOTAL_MINUTES = 60

    NATIONALITIES = sorted([country.name for country in pycountry.countries])
    COUNTRY_POOL = [
    'India', 'Italy', 'Romania', 'Morocco', 'Ghana', 'Brazil',
    'Belgium', 'Netherlands', 'Germany', 'France', 'Spain', 'Portugal',
    'Poland', 'Turkey', 'China', 'Japan', 'South Korea', 'Pakistan',
    'Bangladesh', 'United States', 'Canada', 'Mexico', 'Nigeria',
    'Egypt', 'Ukraine'
]

class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    consent = models.StringField(
        choices=[
            ['yes', 'I agree, continue'],
            ['no', 'I do not agree'],
        ],
        label='Please indicate whether you agree to participate voluntarily.',
        widget=widgets.RadioSelect,
    )

    nationality = models.StringField(
        label='What is your country of nationality?'
    )
    
    level_of_study = models.StringField(
        label='What is your current level of study?',
        choices=[
            'Bachelor’s student',
            'Master’s student',
            'PhD / doctoral student',
            'Other',
        ],
    )

    student_status = models.StringField(
        label='Which of the following best describes your current student status?',
        choices=[
            'International / exchange student',
            'Local student',
            'Other',
        ],
    )

    time_in_country = models.StringField(
        label='How long have you been living in your current country of study?',
        choices=[
            'Less than 1 month',
            '1–6 months',
            'More than 6 months',
            'I am a local student',
        ],
    )

    prior_abroad_experience = models.StringField(
        label='Before your current studies, had you previously lived or studied abroad for more than 3 months?',
        choices=[
            'Yes',
            'No',
        ],
    )
    treatment = models.StringField(initial='')
    own_index = models.IntegerField(initial=0)

    alloc_1 = models.IntegerField(min=0, max=C.TOTAL_MINUTES, initial=0)
    alloc_2 = models.IntegerField(min=0, max=C.TOTAL_MINUTES, initial=0)
    alloc_3 = models.IntegerField(min=0, max=C.TOTAL_MINUTES, initial=0)
    alloc_4 = models.IntegerField(min=0, max=C.TOTAL_MINUTES, initial=0)
    alloc_5 = models.IntegerField(min=0, max=C.TOTAL_MINUTES, initial=0)
    alloc_6 = models.IntegerField(min=0, max=C.TOTAL_MINUTES, initial=0)

    diversity_importance = models.IntegerField(
        choices=range(1, 11),
        label='How important is it for you to build an internationally diverse social circle during your time abroad?',
        widget=widgets.RadioSelectHorizontal,
    )

    comfort_unfamiliar = models.IntegerField(
        choices=range(1, 11),
        label='How comfortable do you generally feel approaching students from unfamiliar backgrounds?',
        widget=widgets.RadioSelectHorizontal,
    )

    own_country_importance = models.IntegerField(
        choices=range(1, 11),
        label='Before coming abroad, how important was it for you to meet people from your own country?',
        widget=widgets.RadioSelectHorizontal,
    )

    join_matching = models.StringField(
        choices=[
            ['yes', 'Yes, I would like to participate'],
            ['no', 'No, I prefer not to participate'],
        ],
        label='Would you like to participate in the optional friend-matching exercise?',
        widget=widgets.RadioSelect,
    )

    match_email = models.StringField(
        label='Email address',
        blank=True
    )

    short_intro = models.LongStringField(
        label='Please write a short introduction about yourself, your interests, or the kind of people you would like to meet.',
        blank=True
    )

    match_name = models.StringField(
        label='Name (optional)',
        blank=True
    )

    match_phone = models.StringField(
        label='Phone number (optional)',
        blank=True
    )

    match_age = models.IntegerField(
        label='Age (optional)',
        blank=True
    )

class Consent(Page):
    form_model = 'player'
    form_fields = ['consent']

class NoConsentExit(Page):
    @staticmethod
    def is_displayed(player):
        return player.consent == 'no'
    
class Background(Page):
    form_model = 'player'
    form_fields = [
        'nationality',
        'level_of_study',
        'student_status',
        'time_in_country',
        'prior_abroad_experience',
    ]
    allow_back_button = True
    preserve_unsubmitted_inputs = True

    @staticmethod
    def is_displayed(player):
        return player.consent == 'yes'

    @staticmethod
    def vars_for_template(player):
        return dict(nationalities=C.NATIONALITIES)

    @staticmethod
    def error_message(player, values):
        if values['nationality'] not in C.NATIONALITIES:
            return 'Please select a valid country from the autocomplete suggestions.'

    @staticmethod
    def before_next_page(player, timeout_happened):
        if not player.treatment:
            player.treatment = random.choice(['control', 'treatment'])

        if 'profiles' not in player.participant.vars:
            own_country = player.nationality.strip()

            available_countries = [
                c for c in C.COUNTRY_POOL
                if c.lower() != own_country.lower()
            ]

            selected_others = random.sample(available_countries, 5)
            profiles = selected_others + [own_country]
            random.shuffle(profiles)

            player.participant.vars['profiles'] = profiles
            player.own_index = profiles.index(own_country)

class Allocation(Page):
    form_model = 'player'
    form_fields = ['alloc_1', 'alloc_2', 'alloc_3', 'alloc_4', 'alloc_5', 'alloc_6']
    allow_back_button = True
    preserve_unsubmitted_inputs = True
    @staticmethod
    def is_displayed(player):
        return player.consent == 'yes'
    @staticmethod
    def vars_for_template(player):
        return dict(
            profiles=player.participant.vars['profiles'],
            treatment=player.treatment,
            total_minutes=C.TOTAL_MINUTES,
        )

    @staticmethod
    def error_message(player, values):
        total = sum(values.values())

        for field_name, value in values.items():
            if not float(value).is_integer():
                return 'Please enter whole numbers only.'

        if total != C.TOTAL_MINUTES:
            return f'The total must equal exactly {C.TOTAL_MINUTES} minutes. Your current total is {total}.'

class FollowUp(Page):
    form_model = 'player'
    form_fields = ['diversity_importance', 'comfort_unfamiliar', 'own_country_importance']
    allow_back_button = True
    preserve_unsubmitted_inputs = True
    @staticmethod
    def is_displayed(player):
        return player.consent == 'yes'

class MatchingInvite(Page):
    form_model = 'player'
    form_fields = ['join_matching']
    @staticmethod
    def is_displayed(player):
        return player.consent == 'yes'
    # no back button here: this is where responses become locked


class MatchingForm(Page):
    form_model = 'player'
    form_fields = ['match_email', 'short_intro', 'match_name', 'match_phone', 'match_age']
    @staticmethod
    def is_displayed(player):
        return player.consent == 'yes'
    @staticmethod
    def is_displayed(player):
        return player.join_matching == 'yes'

    @staticmethod
    def error_message(player, values):
        if player.join_matching == 'yes':
            if not values['match_email']:
                return 'Email is required if you want to participate in the friend-matching exercise.'
            if not values['short_intro']:
                return 'Please write a short introduction if you want to participate in the friend-matching exercise.'
            if '@' not in values['match_email']:
                return 'Please enter a valid email address.'

class Results(Page):
    @staticmethod
    def is_displayed(player):
        return player.consent == 'yes'
    @staticmethod
    def vars_for_template(player):
        profiles = player.participant.vars['profiles']
        allocations = [
            player.alloc_1,
            player.alloc_2,
            player.alloc_3,
            player.alloc_4,
            player.alloc_5,
            player.alloc_6,
        ]
        own_minutes = allocations[player.own_index]

        return dict(
            profiles=profiles,
            allocations=allocations,
            own_minutes=own_minutes,
            total_minutes=C.TOTAL_MINUTES,
        )


page_sequence = [
    Consent,
    NoConsentExit,
    Background,
    Allocation,
    FollowUp,
    MatchingInvite,
    MatchingForm,
    Results,
]