from otree.api import *
import pandas as pd
import random
import logging

doc = """
Allows 4 enumerators to enter their designated room
"""


class C(BaseConstants):
    NAME_IN_URL = 'treat_assign'
    PLAYERS_PER_GROUP = 4
    NUM_ROUNDS = 1


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    gameID_final = models.StringField()
    gameID_init = models.StringField()
    game_id_same = models.BooleanField()
    go_to_next_pages = models.BooleanField()

class Player(BasePlayer):
    pass_key = models.IntegerField(
        label=""
    )
    enum_id = models.IntegerField()
    district = models.IntegerField()
    male = models.BooleanField()
    household = models.IntegerField()
    enum_name = models.StringField()
    input_id = models.IntegerField()
    game_id = models.IntegerField()
    uid = models.StringField()
    coin_tosser = models.BooleanField()
    household_consent = models.BooleanField()
    all_consent = models.BooleanField()
    consent_status = models.BooleanField()
    num_consents = models.IntegerField()
    num_players = models.IntegerField()
    one_no_consent_household = models.BooleanField()
    both_no_consent_household = models.BooleanField()
    coin_toss_result = models.BooleanField(
        choices = [
            [True, "টসে জিতেসে"],
            [False, "টসে হেরেছে"]
        ],
        label = ""
    )
    consent_resp = models.BooleanField(
        choices = [
            [True, "হ্যাঁ"],
            [False, "না"]
        ],
        label = "Are you willing to participate in the survey?"
    )
    consent_other_req = models.BooleanField(
        choices = [
            [True, "হ্যাঁ"],
            [False, "না"]
        ],
        label = "Before beginning the survey, should we take permission from any other members of the household?"
    )

    consent_o_hh_mem_1 = models.BooleanField(blank = True, initial = False)
    consent_o_hh_mem_2 = models.BooleanField(blank = True, initial=False)
    consent_o_hh_mem_3 = models.BooleanField(blank = True, initial=False)
    consent_o_hh_mem_4 = models.BooleanField(blank = True, initial = False)
    consent_o_hh_mem_5 = models.BooleanField(blank = True, initial=False)
    consent_o_hh_mem_oth = models.BooleanField(blank = True, initial=False)
    consent_o_hh_mem_oth_text = models.LongStringField(blank = True)

    consent_o_hh = models.BooleanField(
        choices = [
            [True, "হ্যাঁ"],
            [False, "না"]
        ],
        label = "Did all required household members consent?"
    )
    lottery_draw = models.IntegerField(
        choices = [
            [10,"WH"],
            [21,"BB"],
            [15,"BH"],
            [14,"WB"]
        ],
        label="Which paper did the respondent pull out of the bag?"
        #lottery values set as color times role for ease of calculation
        #White = 2, Black = 3, Headset = 5, Bowl = 7. e.g White-headset = 2*5 = 10
        #White*role % 2 = 0 => same team
    )
    player_role = models.BooleanField(
        choices=[
            [True, "Punisher"],
            [False, "Punished"]
        ]
    )
    player_color_match = models.BooleanField(
        choices=[
            [True, "White"],
            [False, "Black"]
        ]
    )
    opponent_id = models.IntegerField()
    type = models.IntegerField()
    treatment = models.IntegerField()
    subtreatment = models.IntegerField()
    treat = models.StringField()
    matchID = models.IntegerField()
    hh1_punisher = models.IntegerField(initial=2)
    matchID_pair = models.IntegerField(
        initial = 1
    )
    player_final_gameID = models.StringField()


# PAGES

class FirstWaitPage(WaitPage):
    @staticmethod
    def after_all_players_arrive(group: Group):
        players = group.get_players()
        for player in players:
            part_lab = player.participant.label.split("_")
            player.enum_id = int(part_lab[0])
            player.district = int(part_lab[2])
            player.male = not (player.enum_id % 2)

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

class ConsentOther(Page):
    @staticmethod
    def is_displayed(player):
        return player.consent_resp == True

    form_model = 'player'
    form_fields = ['consent_other_req']

class ConsentOtherWhom(Page):
    form_model = 'player'
    form_fields = ['consent_o_hh_mem_1', 'consent_o_hh_mem_2', 'consent_o_hh_mem_3', 'consent_o_hh_mem_4', 'consent_o_hh_mem_5', 'consent_o_hh_mem_oth', 'consent_o_hh_mem_oth_text']
    @staticmethod
    def is_displayed(player):
        return player.consent_resp and player.consent_other_req == True

class ConsentOtherText(Page):
    form_model = 'player'
    form_fields = ['consent_o_hh']
    @staticmethod
    def is_displayed(player):
        return player.consent_resp and player.consent_other_req == True

class AllConsentWaitPage(WaitPage):
    def after_all_players_arrive(self):
        for player in self.group.get_players():
            all_consent = (player.consent_resp and not player.consent_other_req) or \
                                 (player.consent_resp and player.consent_other_req and player.consent_o_hh)
            player.all_consent = all_consent
            player.participant.vars['all_consent'] = all_consent
            print("all_consent:", all_consent)

class StartPage(WaitPage):
    @staticmethod
    def after_all_players_arrive(group: Group):
        players = group.get_players()
        if all(player.all_consent for player in players):
            group.go_to_next_pages = True
        else:
            group.go_to_next_pages = False
        print("go_to_next_pages:", group.go_to_next_pages)

        game_ids = []
        household_1_otree_players = []
        household_2_otree_players = []
        for player in players:
            part_lab = player.participant.label.split("_")
            player.enum_id = int(part_lab[0])
            player.male = not (player.enum_id % 2)
            player.game_id = int(part_lab[1])
            enum_names = pd.read_csv('./_static/enums/enum_names.csv')
            player.enum_name = enum_names.loc[enum_names['enum_id'] == player.enum_id, 'enum_name'].values[0]
            game_ids.append(player.game_id)

            # set to teams (HHID), randomize coin tosser
            if int(player.enum_id)%10<3:
                player.household = 1
                household_1_otree_players.append(player)
            else:
                player.household = 2
                household_2_otree_players.append(player)

        # Store household consent status
        for player in players:
            if player.household == 1:
                consent_status = [p.all_consent for p in household_1_otree_players]
            else:
                consent_status = [p.all_consent for p in household_2_otree_players]

            num_consents = sum(consent_status)
            num_players = len(consent_status)
            print("num_consents:", num_consents)
            print("num_players:", num_players)

            if num_consents == 0:
                player.household_consent = False  # Neither player in the household gave consent
                player.one_no_consent_household = False  # No player in the household refused consent
                player.both_no_consent_household = True  # Both players in the household refused consent
            elif num_consents == num_players:
                player.household_consent = True  # All players in the household gave consent
                player.one_no_consent_household = False  # No player in the household refused consent
                player.both_no_consent_household = False  # Both players in the household gave consent
            else:
                player.household_consent = False  # Not all players in the household gave consent
                player.one_no_consent_household = True  # At least one player in the household refused consent
                player.both_no_consent_household = False  # Both players in the household gave consent

            print("household_consent:", player.household_consent)
            print("one_no_consent_household:", player.one_no_consent_household)
            print("both_no_consent_household:", player.both_no_consent_household)

        if len(set(game_ids)) == 1:
            group.game_id_same = True
            group.gameID_init = str(group.get_player_by_id(1).game_id)


            #assign coin tosser
            temp_rand_hh1 = random.randint(0,1)
            if temp_rand_hh1 == 1:
                household_1_otree_players[0].coin_tosser = 0
                household_1_otree_players[1].coin_tosser = 1
            else:
                household_1_otree_players[0].coin_tosser = 1
                household_1_otree_players[1].coin_tosser = 0

            temp_rand_hh2 = random.randint(0,1)
            if temp_rand_hh2 == 1:
                household_2_otree_players[0].coin_tosser = 0
                household_2_otree_players[1].coin_tosser = 1
            else:
                household_2_otree_players[0].coin_tosser = 1
                household_2_otree_players[1].coin_tosser = 0

        else:
            group.game_id_same = False

class ResultsFailure(Page):
    @staticmethod
    def is_displayed(player):
        return not player.group.game_id_same

class ConsentFailure1(Page):
    @staticmethod
    def is_displayed(player):
        return player.one_no_consent_household == True and player.all_consent == True

class ConsentFailure2(Page):
    @staticmethod
    def is_displayed(player):
        return player.one_no_consent_household == True and player.all_consent == False

class ConsentFailure3(Page):
    @staticmethod
    def is_displayed(player):
        return player.household_consent and not player.group.go_to_next_pages

class ConsentFailure4(Page):
    @staticmethod
    def is_displayed(player):
        return player.both_no_consent_household == True

class Instructions(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.group.go_to_next_pages

class CoinToss(Page):
    def is_displayed(player: Player):
        return player.group.go_to_next_pages and player.group.game_id_same and player.coin_tosser == 1

    form_model = 'player'
    form_fields = ['coin_toss_result']

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        other_players = player.get_others_in_group()
        for p in other_players:
            if p.household == player.household:
                if player.coin_toss_result == False:
                    p.coin_toss_result = True
                else:
                    p.coin_toss_result = False

class CoinToss_wait(WaitPage):
    def is_displayed(player):
        return player.group.go_to_next_pages and player.group.game_id_same

class Lottery(Page):
    def is_displayed(player: Player):
        return player.group.go_to_next_pages and player.group.game_id_same == True and player.coin_toss_result == True

    form_model = 'player'
    form_fields = ['lottery_draw']

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        other_players = player.get_others_in_group()
        for p in other_players:
            if p.household == player.household:
                if player.lottery_draw == 10:
                    p.lottery_draw = 21
                elif player.lottery_draw == 21:
                    p.lottery_draw = 10
                elif player.lottery_draw == 15:
                    p.lottery_draw = 14
                else:
                    p.lottery_draw = 15


class Lottery_Results_Wait(WaitPage):
    def is_displayed(player):
        return player.group.go_to_next_pages and player.group.game_id_same

    @staticmethod
    def after_all_players_arrive(group: Group):
        players = group.get_players()
        #set roles
        for player in players:
            if player.lottery_draw % 5 == 0:
                player.player_role = 1
            else:
                player.player_role = 0
            #set team
            if player.lottery_draw % 3 == 0:
                player.player_color_match = 1
            else:
                player.player_color_match = 0


class Lottery_Results_Wait_2(WaitPage):
    def is_displayed(player):
        return player.group.go_to_next_pages and player.group.game_id_same

    @staticmethod
    def after_all_players_arrive(group: Group):

        players = group.get_players()
        for player in players:
            other_players = player.get_others_in_group()
            for p in other_players:
                if player.player_color_match == p.player_color_match:
                    player.opponent_id = p.id_in_group
            opponent = player.group.get_player_by_id(player.opponent_id)
            '''
            #randomly assign treatment to only one player within each pair through role
            if player.player_role == 1:
                type = str(random.randint(0, 2))
                player.type = int(type)
                treatment = str(random.randint(0, 2))
                player.treatment = int(treatment)
                subtreatment = str(random.randint(0, 1))
                player.subtreatment = int(subtreatment)
                player.treat = type + treatment + subtreatment
                opponent.type = player.type
                opponent.treatment = player.treatment
                opponent.subtreatment = player.subtreatment
                opponent.treat = player.treat
            '''

            if opponent.male == player.male:
                player.matchID = 0
            else:
                player.matchID = 1
            if (player.matchID == 0 and player.male == 1):
                player.matchID_pair = 0
            if (player.matchID == 1 and player.male == 1 and player.household == 1):
                player.matchID_pair = 0
            if (player.matchID == 1 and player.male == 0 and player.household == 2):
                player.matchID_pair = 0


class Lottery_Results_Wait_3(WaitPage):
    def is_displayed(player: Player):
        return player.group.go_to_next_pages and player.group.game_id_same

    @staticmethod
    def after_all_players_arrive(group: Group):
        treatment_df = pd.read_csv('./_static/treatments/treatment_randomization_otree_final.csv')
        for player in group.get_players():
            opponent = player.group.get_player_by_id(player.opponent_id)
            #if opponent punisher/chooser and higher/lower HH
            if opponent.player_role == 1 and opponent.household == 1:
                player.hh1_punisher = 1
            elif opponent.player_role == 1 and opponent.household == 2:
                player.hh1_punisher = 0
            elif opponent.player_role == 0 and opponent.household == 1:
                player.hh1_punisher = 0
            else:
                player.hh1_punisher = 1

            player.player_final_gameID = group.gameID_init + str(player.matchID) + str(player.matchID_pair) + str(player.hh1_punisher)
            id_enums = treatment_df[treatment_df['surveyIDtreat'] == int(player.player_final_gameID)].loc[:,
                       ["uid","type", "treatment", "subTreatment", "enumID"]]
            player.type = int(id_enums.loc[id_enums['enumID'] == player.enum_id, 'type'].item())
            player.treatment = int(id_enums.loc[id_enums['enumID'] == player.enum_id, 'treatment'].item())
            player.subtreatment = int(id_enums.loc[id_enums['enumID'] == player.enum_id, 'subTreatment'].item())
            opponent.type = player.type
            opponent.treatment = player.treatment
            opponent.subtreatment = player.subtreatment
            player.uid = str(id_enums.loc[id_enums['enumID'] == player.enum_id, 'uid'].item())
            '''
            player.player_final_gameID = group.gameID_init + player.treat + str(player.matchID) + str(
                player.matchID_pair) + str(hh1_punisher)
            id_enums = treatment_df[treatment_df['gameIDtreat'] == int(player.player_final_gameID)].loc[:,
                       ["uid", "enumID"]]
            player.uid = str(id_enums.loc[id_enums['enumID'] == player.enum_id, 'uid'].item())
            '''

class Lottery_Results_one(Page):
    def is_displayed(player):
        return player.group.go_to_next_pages
class Lottery_Results_two(Page):
    def is_displayed(player):
        return player.group.go_to_next_pages
class ResultsSuccess(Page):
    def is_displayed(player):
        return player.group.go_to_next_pages and player.group.game_id_same

    form_model = 'player'


page_sequence = [
    FirstWaitPage,
    Consent,
    ConsentOther,
    ConsentOtherWhom,
    ConsentOtherText,
    AllConsentWaitPage,
    StartPage,
    ResultsFailure,
    ConsentFailure1,
    ConsentFailure2,
    ConsentFailure3,
    ConsentFailure4,
    Instructions,
    CoinToss,
    CoinToss_wait,
    Lottery,
    Lottery_Results_Wait,
    Lottery_Results_Wait_2,
    Lottery_Results_Wait_3,
    Lottery_Results_one,
    Lottery_Results_two,
    ResultsSuccess
]