from otree.api import *
import random
from constants import *

doc = """Generation 1 Endowment Creation through coin flip and token grab tasks"""


class C(BaseConstants):
    NAME_IN_URL = 'A_Endowment_Creation_Gen1_D'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = Round_FirstTask  # Import from shared constants


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    def assign_edmt(self):
        """Assign endowments based on task performance and reason"""
        same_gen_players = [
            p for p in self.subsession.get_players()
            if p.participant.vars.get('Ana_generation') == 1 and p.participant.vars.get('state') == "active"
        ]

        # Assign GroupID
        for idx, p in enumerate(same_gen_players):
            p.Ana_GroupID = (idx // 3) + 1
            p.participant.vars['Ana_GroupID'] = p.Ana_GroupID

        reasons = self.session.config.get('GroupReason', DEFAULT_GROUP_REASONS).copy()
        random.shuffle(reasons)

        group_ids = sorted(set(p.Ana_GroupID for p in same_gen_players))

        for group_idx, gid in enumerate(group_ids):
            grp = [p for p in same_gen_players if p.Ana_GroupID == gid]
            if len(grp) != 3:
                continue

            reason = reasons[group_idx % len(reasons)]

            for p in grp:
                p.Ana_EndowmentReason = reason
                p.participant.vars['Ana_EndowmentReason'] = reason

                ranking_map = {"Moral": 1, "Immoral": 2, "Lucky": 3, "Unlucky": 4}
                p.participant.vars['Ana_RankingType'] = ranking_map.get(reason, 1)

            if reason == "Immoral":
                sorted_players = sorted(grp, key=lambda x: x.participant.vars.get('N_TokenTake', 0), reverse=True)
            elif reason == "Moral":
                sorted_players = sorted(grp, key=lambda x: x.participant.vars.get('N_TokenTake', 0))
            elif reason == "Lucky":
                sorted_players = sorted(grp, key=lambda x: x.participant.vars.get('N_CoinFlip', 0), reverse=True)
            elif reason == "Unlucky":
                sorted_players = sorted(grp, key=lambda x: x.participant.vars.get('N_CoinFlip', 0))
            else:
                sorted_players = grp

            for i, p in enumerate(sorted_players):
                p.Ana_EndowmentType = 2 if i == 0 else 1
                p.Ana_Endowment = EndowmentH if i == 0 else EndowmentL
                p.participant.vars['Ana_EndowmentType'] = p.Ana_EndowmentType
                p.participant.vars['Ana_Endowment'] = p.Ana_Endowment

            for i, p in enumerate(sorted_players):
                p.Ana_Mmb1_id = sorted_players[0].participant.vars.get('Ana_own_id')
                p.Ana_Mmb2_id = sorted_players[1].participant.vars.get('Ana_own_id')
                p.Ana_Mmb3_id = sorted_players[2].participant.vars.get('Ana_own_id')
                p.participant.vars['Ana_Mmb1_id'] = p.Ana_Mmb1_id
                p.participant.vars['Ana_Mmb2_id'] = p.Ana_Mmb2_id
                p.participant.vars['Ana_Mmb3_id'] = p.Ana_Mmb3_id

                if p.participant.vars.get('Ana_own_id') == p.Ana_Mmb1_id:
                    p.my_order = 1
                elif p.participant.vars.get('Ana_own_id') == p.Ana_Mmb2_id:
                    p.my_order = 2
                else:
                    p.my_order = 3
                p.participant.vars['my_order'] = p.my_order

                p.Ana_Edmt_Mmb1 = sorted_players[0].Ana_Endowment
                p.Ana_Edmt_Mmb2 = sorted_players[1].Ana_Endowment
                p.Ana_Edmt_Mmb3 = sorted_players[2].Ana_Endowment
                p.Ana_fname_Mmb1 = sorted_players[0].participant.vars.get('Ana_family_assignment')
                p.Ana_fname_Mmb2 = sorted_players[1].participant.vars.get('Ana_family_assignment')
                p.Ana_fname_Mmb3 = sorted_players[2].participant.vars.get('Ana_family_assignment')

                p.participant.vars['Ana_Edmt_Mmb1'] = p.Ana_Edmt_Mmb1
                p.participant.vars['Ana_Edmt_Mmb2'] = p.Ana_Edmt_Mmb2
                p.participant.vars['Ana_Edmt_Mmb3'] = p.Ana_Edmt_Mmb3
                p.participant.vars['Ana_fname_Mmb1'] = p.Ana_fname_Mmb1
                p.participant.vars['Ana_fname_Mmb2'] = p.Ana_fname_Mmb2
                p.participant.vars['Ana_fname_Mmb3'] = p.Ana_fname_Mmb3


class Player(BasePlayer):
    flip_result = models.StringField()
    take_token = models.BooleanField(
        choices=[[True, 'Take'], [False, 'Do not take']],
        widget=widgets.RadioSelectHorizontal
    )
    token_decision = models.IntegerField()

    Ana_GroupID = models.IntegerField()
    my_order = models.IntegerField()
    Ana_EndowmentReason = models.StringField()
    Ana_EndowmentType = models.IntegerField()
    Ana_Endowment = models.IntegerField()

    Ana_Mmb1_id = models.IntegerField()
    Ana_Mmb2_id = models.IntegerField()
    Ana_Mmb3_id = models.IntegerField()
    Ana_fname_Mmb1 = models.StringField()
    Ana_fname_Mmb2 = models.StringField()
    Ana_fname_Mmb3 = models.StringField()
    Ana_Edmt_Mmb1 = models.IntegerField()
    Ana_Edmt_Mmb2 = models.IntegerField()
    Ana_Edmt_Mmb3 = models.IntegerField()


class FirstTask(Page):
    @staticmethod
    def is_displayed(player):
        return (player.round_number == 1 and
                player.participant.vars.get('Ana_generation') == 1 and
                player.participant.vars.get('state') == "active")

    @staticmethod
    def vars_for_template(player):
        return {
            'Ana_generation': player.participant.vars.get('Ana_generation'),
            'Ana_own_id': player.participant.vars.get('Ana_own_id'),
            'Ana_family_assignment': player.participant.vars.get('Ana_family_assignment'),
        }


class FirstTask_CoinFlip_Intro(Page):
    @staticmethod
    def is_displayed(player):
        first_order = player.participant.vars.get('first_task_order')
        return (player.round_number == 1 and player.participant.vars.get('Ana_generation') == 1 and
                player.participant.vars.get('state') == "active" and first_order == 0) or \
            (player.round_number == 1 + Round_FirstTask // 2 and player.participant.vars.get('Ana_generation') == 1 and
             player.participant.vars.get('state') == "active" and first_order == 1)


class FirstTask_CoinFlip(Page):
    timeout_seconds = TIME_PER_ROUND
    form_model = 'player'
    form_fields = ['flip_result']

    @staticmethod
    def is_displayed(player):
        first_order = player.participant.vars.get('first_task_order')
        return (player.round_number <= Round_FirstTask // 2 and player.participant.vars.get('Ana_generation') == 1 and
                player.participant.vars.get('state') == "active" and first_order == 0) or \
            (1 + Round_FirstTask // 2 <= player.round_number <= Round_FirstTask and
             player.participant.vars.get('Ana_generation') == 1 and player.participant.vars.get(
                        'state') == "active" and first_order == 1)

    @staticmethod
    def vars_for_template(player):
        first_order = player.participant.vars.get('first_task_order')
        if first_order == 0:
            round_idx = player.round_number - 1
            round_number = player.round_number
        else:
            round_idx = player.round_number - (Round_FirstTask // 2 + 1)
            round_number = player.round_number - (Round_FirstTask // 2)

        flip_result = player.participant.vars['flip_result_list'][round_idx]
        player.flip_result = flip_result

        return {
            'flip_result': flip_result,
            'N_CoinFlip': player.participant.vars['flip_result_list'][:round_idx].count('Heads'),
            'round_number': round_number,
        }

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        first_order = player.participant.vars.get('first_task_order')
        if (first_order == 0 and player.round_number == Round_FirstTask // 2) or \
                (first_order == 1 and player.round_number == Round_FirstTask):
            player.participant.vars['N_CoinFlip'] = player.participant.vars['flip_result_list'].count('Heads')


class FirstTask_TokenGrab_Intro(Page):
    @staticmethod
    def is_displayed(player):
        first_order = player.participant.vars.get('first_task_order')
        return (player.round_number == 1 + Round_FirstTask // 2 and player.participant.vars.get(
            'Ana_generation') == 1 and
                player.participant.vars.get('state') == "active" and first_order == 0) or \
            (player.round_number == 1 and player.participant.vars.get('Ana_generation') == 1 and
             player.participant.vars.get('state') == "active" and first_order == 1)


class FirstTask_TokenGrab(Page):
    timeout_seconds = TIME_PER_ROUND
    form_model = 'player'
    form_fields = ['take_token']

    @staticmethod
    def is_displayed(player):
        first_order = player.participant.vars.get('first_task_order')
        return (1 + Round_FirstTask // 2 <= player.round_number <= Round_FirstTask and
                player.participant.vars.get('Ana_generation') == 1 and player.participant.vars.get(
                    'state') == "active" and first_order == 0) or \
            (player.round_number <= Round_FirstTask // 2 and player.participant.vars.get('Ana_generation') == 1 and
             player.participant.vars.get('state') == "active" and first_order == 1)

    @staticmethod
    def vars_for_template(player):
        ordinal_mapping = {1: '1st', 2: '2nd', 3: '3rd', 4: '4th', 5: '5th',
                           6: '6th', 7: '7th', 8: '8th', 9: '9th', 10: '10th'}

        first_order = player.participant.vars.get('first_task_order')
        if first_order == 0:
            adjusted_round = player.round_number - (Round_FirstTask // 2)
        else:
            adjusted_round = player.round_number

        round_number_th = ordinal_mapping.get(adjusted_round, f"{adjusted_round}th")

        return {
            'round_number_th': round_number_th,
            'round_number': adjusted_round,
            'N_TokenTake': player.participant.vars.get('N_TokenTake', 0),
            'TIME_PER_ROUND': TIME_PER_ROUND,
        }

    @staticmethod
    def before_next_page(player: Player, timeout_happened=False):
        if timeout_happened:
            player.token_decision = 0
        else:
            player.token_decision = 1 if player.take_token else 0
            if player.take_token:
                player.participant.vars['N_TokenTake'] = player.participant.vars.get('N_TokenTake', 0) + 1


class EndofPart1(Page):
    @staticmethod
    def is_displayed(player):
        return (player.round_number == C.NUM_ROUNDS and
                player.participant.vars.get('Ana_generation') == 1 and
                player.participant.vars.get('state') == "active")

    @staticmethod
    def before_next_page(player: Player, timeout_happened=False):
        if player.round_number == C.NUM_ROUNDS:
            player.group.assign_edmt()


page_sequence = [
    FirstTask,
    FirstTask_CoinFlip_Intro,
    FirstTask_CoinFlip,
    FirstTask_TokenGrab_Intro,
    FirstTask_TokenGrab,
    EndofPart1
]