from otree.api import *
# importing dependency
import math
import numpy as np
import itertools
import ast


doc = """
Your app description
"""


class Constants(BaseConstants):
    name_in_url = 'collaboration_under_complexity'
    players_per_group = None  # Changes the number of players per group None means every group has 1 player
    num_rounds = 4
    max_trials = 20


    def dependency_structure(N,K):
        dep = {}
        for i in np.arange(N):
            indi = list(range(N))
            indi.remove(i)
            dep[i] = np.random.choice(indi, K, replace=False).tolist()
        return dep

    def nk_landscape(N,K):
        nk = np.random.rand(N, 2 ** (K + 1), N)
        return nk

    def max_score_func(N,K,dep,nk):

        def cal_fit_raw(N, K, dep, nk, bit_str):
            if len(bit_str) == N:

                score = 0
                ind_dict = list(itertools.product(range(2), repeat=K + 1))
                for i in range(len(bit_str)):
                    dep_bit = dep[i].copy()
                    dep_bit.append(i)
                    inter_bits = tuple([bit_str[x] for x in dep_bit])
                    ind = ind_dict.index(inter_bits)
                    score += nk[i][ind][i]
                return score
            else:
                return print('The length of bit string should be equal to the length of N')

        sol_list = list(itertools.product(range(2), repeat=N))
        max_score = 0
        for sol in sol_list:
            if cal_fit_raw(N, K, dep, nk,sol) > max_score:
                max_score = cal_fit_raw(N, K, dep, nk,sol)
                max_sol = sol
        return max_score, max_sol

    def cal_fit(max_score,a, N, K, dep, nk, bit_str):

        def cal_fit_raw(N, K, dep, nk, bit_str):
            if len(bit_str) == N:

                score = 0
                ind_dict = list(itertools.product(range(2), repeat=K + 1))
                for i in range(len(bit_str)):
                    dep_bit = dep[i].copy()
                    dep_bit.append(i)
                    inter_bits = tuple([bit_str[x] for x in dep_bit])
                    ind = ind_dict.index(inter_bits)
                    score += nk[i][ind][i]
                return score
            else:
                return print('The length of bit string should be equal to the length of N')


        raw_over_max = cal_fit_raw(N, K, dep, nk, bit_str) / max_score
        return math.pow(raw_over_max, a)

C = Constants

class Subsession(BaseSubsession):
    N = models.IntegerField()
    K = models.IntegerField()
    a = models.IntegerField()
    nk = models.LongStringField()
    dep = models.LongStringField()
    max = models.FloatField()
    min = models.FloatField()
    max_score = models.FloatField()
    max_sol = models.StringField()
    l_shape = models.StringField()
    cut_off = models.FloatField()


def creating_session(subsession: Subsession):
    S = subsession
    # soum = 7
    S.N = 8
    S.a = 8
    if S.session.config['complexity']=="Low":
        S.K = 2
    if S.session.config['complexity'] == "Medium":
        S.K = 3
    if S.session.config['complexity']=="High":
        S.K = 4
    landscape = C.nk_landscape(S.N,S.K)
    S.l_shape = str(landscape.shape)
    S.nk = str(dict(enumerate(landscape.flatten(), 1)))
    S.dep = str(C.dependency_structure(S.N,S.K))
    S.max = np.max(landscape)
    S.min = np.min(landscape)
    S.cut_off = round(np.median(landscape)/3,5)
    S.max_score, shp = C.max_score_func(S.N,S.K,
                                              ast.literal_eval(S.dep),
                                              np.fromiter(ast.literal_eval(S.nk).values(),
                                                          dtype=float).reshape(ast.literal_eval(S.l_shape)) )
    S.max_sol = str(shp)


class Group(BaseGroup):
    pass

class Player(BasePlayer):

    bit_1 = models.IntegerField(choices=[[0, 'Red'],[1, 'Black']],initial=0)
    bit_2 = models.IntegerField(choices=[[0, 'Red'], [1, 'Black']],initial=0)
    bit_3 = models.IntegerField(choices=[[0, 'Red'], [1, 'Black']],initial=1)
    bit_4 = models.IntegerField(choices=[[0, 'Red'], [1, 'Black']],initial=1)
    bit_5 = models.IntegerField(choices=[[0, 'Red'], [1, 'Black']],initial=0)
    bit_6 = models.IntegerField(choices=[[0, 'Red'], [1, 'Black']],initial=0)
    bit_7 = models.IntegerField(choices=[[0, 'Red'], [1, 'Black']],initial=1)
    bit_8 = models.IntegerField(choices=[[0, 'Red'], [1, 'Black']],initial=0)
    # Make sure the default value should match the default colour in the html display


    num_trials = models.IntegerField(max =25, initial=0)
    test_score = models.FloatField(initial= 0)
    own_score = models.FloatField(initial=0)
    test_bit = models.StringField(initial="r")
    own_bit = models.StringField(initial="r")
    search_cost = models.IntegerField(initial=0)
    score_choice = models.StringField(initial="Own_Score")
    rewards_points = models.FloatField()
    is_finished = models.BooleanField(initial=False)
    group_score = models.FloatField()

# PAGES
class Welcome(Page):
    pass

class Welcome_2(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

class Welcome_3(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

class Game_instructions_l(Page):
    pass

class Game_instructions_2(Page):
    pass

class Game_instructions_31(Page):
    pass

class Game_instructions_4(Page):
    pass

class Stage_1_instructions_l(Page):
    pass


def live_method(player: Player, data):
    # this function is used to communicate with the live dynamic page in the game.
    # It takes 'data' processes it and sends response as 'resp'
    my_id = player.id_in_group
    print("inlive_method")

    sawyer = player.group.subsession

    # guard
    if player.is_finished:  # check if player finished and you need to move page
        return {my_id: dict(finished=True)}
    resp = {}

    print(data)
    if 'bit_1' in data:
        if data['bit_1']==255: player.bit_1 = 0
        else: player.bit_1 =1
    if 'bit_2' in data:
        if data['bit_2']==255: player.bit_2 = 0
        else: player.bit_2 =1
    if 'bit_3' in data:
        if data['bit_3']==255: player.bit_3 = 0
        else: player.bit_3 =1
    if 'bit_4' in data:
        if data['bit_4']==255: player.bit_4 = 0
        else: player.bit_4 =1
    if 'bit_5' in data:
        if data['bit_5']==255: player.bit_5 = 0
        else: player.bit_5 =1
    if 'bit_6' in data:
        if data['bit_6']==255: player.bit_6 = 0
        else: player.bit_6 =1
    if 'bit_7' in data:
        if data['bit_7']==255: player.bit_7 = 0
        else: player.bit_7 =1
    if 'bit_8' in data:
        if data['bit_8']==255: player.bit_8 = 0
        else: player.bit_8 =1

    if 'cost' in data:
        cost = data['cost']
        player.search_cost = cost
        player.num_trials += 1
    bit = str(player.bit_1).replace("1","B").replace("0","R") +str(" ")+ str(player.bit_2).replace("1","B").replace("0","R") +str(" ")+ \
          str(player.bit_3).replace("1","B").replace("0","R") +str(" ")+ str(player.bit_4).replace("1","B").replace("0","R") +str(" ")+ \
          str(player.bit_5).replace("1","B").replace("0","R") +str(" ")+ str(player.bit_6).replace("1","B").replace("0","R") +str(" ")+ \
          str(player.bit_7).replace("1","B").replace("0","R") +str(" ")+ str(player.bit_8).replace("1","B").replace("0","R") +str(" ")
    player.test_bit = bit
    # loc = str(0) + 'b' + bit
    # random.seed(int(loc, 2))  ## this is where score comptation ( input and output) from NK landscape happens
    bit_str = [player.bit_1,player.bit_2,player.bit_3,player.bit_4,player.bit_5,player.bit_6,player.bit_7,player.bit_8]


    player.test_score = C.cal_fit(sawyer.max_score,sawyer.a, sawyer.N,sawyer.K,ast.literal_eval(sawyer.dep), np.fromiter(ast.literal_eval(sawyer.nk).values(), dtype=float).reshape(ast.literal_eval(sawyer.l_shape)),bit_str )
    if (player.test_score >= player.own_score): ## update teh best test score
        player.own_score = player.test_score
        player.own_bit = player.test_bit
    resp.update(test_score=round(player.test_score,5), own_score = round(player.own_score,5), own_bit = player.own_bit)
    if player.num_trials == C.max_trials:
        player.is_finished = True
        resp.update(finished=True)
    resp.update(num_trials=player.num_trials)
    print(resp)
    return {my_id: resp}


class Stage_1_Demo_11(Page):

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

    live_method = live_method
    @staticmethod
    def vars_for_template(player: Player):
        return


class Stage_1_waitpage(WaitPage):

    def after_all_players_arrive(self):
        players = self.group.get_players()
        for p in players:
            # p = players[0]
            p.bit_1 = 0 ; p.bit_2 = 0; p.bit_3 = 1; p.bit_4 = 1; p.bit_5 = 0; p.bit_6 = 0; p.bit_7 = 1;  p.bit_8 = 0 # randomly matched to default values see players
            p.num_trials = 0; p.test_score = 0; p.own_score = 0; p.test_bit = "R R B B R R B R";  p.own_bit = "R R B B R R B R"
            p.search_cost = 0
            p.is_finished =False

    #median = np.median(np.fromiter(ast.literal_eval(sawyer.nk).values(), dtype=float).reshape(ast.literal_eval(sawyer.l_shape)))

class Stage_11(Page):
    live_method = live_method
    @staticmethod
    def vars_for_template(player: Player):
        return

def live_method2(player: Player, data):
    my_id = player.id_in_group
    print("inlive2")
    print(data)

    if 'Score_choice' in data:
        player.score_choice = data['Score_choice']

    return

class Stage_2(Page):
    live_method = live_method2
    @staticmethod
    def vars_for_template(player: Player):
        return


class ResultsWaitPage(WaitPage):

    def after_all_players_arrive(self):
        players = self.group.get_players()
        player = players[0]
        # group score is computed here.
        # GS = 0
        # for player in players:
        #     GS = GS + player.own_score
        # group_score = GS
        group_score = float(np.random.normal(0, 0.1, 1) + player.group.subsession.cut_off)

        # for player in players:
        if player.own_score>=group_score:
            corr_choice = "Own_Score"
        if player.own_score<group_score:
            corr_choice = "Group_Score"
        if player.own_score>= player.group.subsession.cut_off and player.score_choice == corr_choice:
            # compute player reward here.
            player.rewards_points = 10* max(player.own_score, group_score)- 0.1 * player.search_cost
        else:
            player.rewards_points = 0
        # player.payoff = player.payoff + player.rewards_points
        player.group_score = round(group_score,5)


class Round_result(Page):

    @staticmethod
    def vars_for_template(player: Player):
        if player.rewards_points==0:
            RESULT = "LOSE"
        else:
            RESULT = "WIN"
        REWARD = round(player.rewards_points,5)
        return { 'RESULT': RESULT, 'REWARD': REWARD}

class Game_ResultsWaitPage(WaitPage):
    def set_payoffs(self,players):
        for player in players:
            player.payoff = player.payoff + player.rewards_points

    def after_all_players_arrive(self):
        players = self.group.get_players()
        self.set_payoffs(players)

class Game_result(Page):

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == C.num_rounds

    @staticmethod
    def vars_for_template(player: Player):
        REWARD = player.participant.payoff + 2 # the base participation paff-off.
        return {'REWARD':REWARD}

page_sequence = [
    # Welcome,
    Welcome_2,
    Welcome_3,
    # Game_instructions_l,
    # Game_instructions_2,
    # Game_instructions_31,
    # Stage_1_instructions_l,
    Stage_1_Demo_11,
    Stage_1_waitpage,
    Stage_11,
    Stage_2,
    ResultsWaitPage,
    Round_result,
    Game_ResultsWaitPage,
    Game_result
]