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 = 61

    NATIONALITIES = sorted([country.name for country in pycountry.countries])
    COUNTRY_POOL = [
    'Canada',
    'Germany',
    'Japan',
    'Brazil',
    'Morocco',
    'Italy',
    'Poland',
    'Korea, Republic of'
]

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',
        ],
    )
    
    background_time_sec = models.FloatField(blank=True)
    allocation_time_sec = models.FloatField(blank=True)
    followup_time_sec = models.FloatField(blank=True)

    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)

    socialized_with_most = models.StringField(
    label='So far this semester, who have you socialized with most?',
    choices=[
            'Mainly people from my own country',
            'Mainly other international students',
            'Mainly local students',
            'A mix of everyone',
        ],
    )

    homesick_frequency = models.IntegerField(
        choices=range(1, 6),
        label='How often do you feel homesick or miss your home country?',
        widget=widgets.RadioSelectHorizontal,
    )

    own_country_importance_before_uni = models.IntegerField(
        choices=range(1, 6),
        label='Before starting at this university, how important was it for you to meet people from your own country?',
        widget=widgets.RadioSelectHorizontal,
    )

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

    languages_spoken = models.StringField(
        label='How many languages do you speak at a conversational level?',
        choices=[
            '1',
            '2',
            '3',
            '4 or more',
        ],
    )

    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',
        'background_time_sec',
    ]
    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.'

        if values['student_status'] == 'Local student' and values['time_in_country'] != 'I am a local student':
            return 'If you selected Local student, please choose "I am a local student" for time in country.'

        if values['student_status'] == 'International / exchange student' and values['time_in_country'] == 'I am a local student':
            return 'If you selected International / exchange student, please choose a non-local option for time in country.'
    @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)

        if player.alloc_1 is None:
            player.alloc_1 = 0
            player.alloc_2 = 0
            player.alloc_3 = 0
            player.alloc_4 = 0
            player.alloc_5 = 0
            player.alloc_6 = 0

class Allocation(Page):
    form_model = 'player'
    form_fields = ['alloc_1', 'alloc_2', 'alloc_3', 'alloc_4', 'alloc_5', 'alloc_6', 'allocation_time_sec',]
    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):
        allocation_fields = [f'alloc_{i}' for i in range(1, 7)]
        total = sum(values[f] for f in allocation_fields)

        for field_name in allocation_fields:
            value = values[field_name]
            if value is None:
                return 'Please fill in all allocation fields.'
            if value < 0:
                return 'Minutes cannot be negative.'

        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 = [
        'socialized_with_most',
        'homesick_frequency',
        'own_country_importance_before_uni',
        'comfort_unfamiliar',
        'languages_spoken',
        'followup_time_sec',
    ]
    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' and player.join_matching == 'yes'

    @staticmethod
    def error_message(player, values):
        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,
        )

def custom_export(players):
    yield [
        'session_code',
        'participant_code',
        'nationality',
        'own_index',
        'profile_1',
        'profile_2',
        'profile_3',
        'profile_4',
        'profile_5',
        'profile_6',
        'alloc_1',
        'alloc_2',
        'alloc_3',
        'alloc_4',
        'alloc_5',
        'alloc_6',
    ]

    for p in players:
        profiles = p.participant.vars.get('profiles', [None, None, None, None, None, None])

        yield [
            p.session.code,
            p.participant.code,
            p.nationality,
            p.own_index,
            profiles[0] if len(profiles) > 0 else None,
            profiles[1] if len(profiles) > 1 else None,
            profiles[2] if len(profiles) > 2 else None,
            profiles[3] if len(profiles) > 3 else None,
            profiles[4] if len(profiles) > 4 else None,
            profiles[5] if len(profiles) > 5 else None,
            p.alloc_1,
            p.alloc_2,
            p.alloc_3,
            p.alloc_4,
            p.alloc_5,
            p.alloc_6,
        ]

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