from otree.api import *

class C(BaseConstants):
    import numpy as NP
    import RANDOM

    NAME_IN_URL = 'centipede'
    PLAYERS_PER_GROUP = 2

    NUM_NODES = 4
    NUM_GAMES = 3
    NUM_ROUNDS = NUM_NODES * NUM_GAMES

    FIRST_ROUNDS = NP.arange(1,NUM_ROUNDS,NUM_NODES)
    LAST_ROUNDS = NP.arange(NUM_NODES,NUM_ROUNDS+1,NUM_NODES)

    LARGE_PILE = 0.40
    SMALL_PILE = 0.10
    BASE = 2

    LARGE_PILES = []
    SMALL_PILES = []

    for NODE in range(NUM_NODES + 1):
        LARGE_PILES.append(LARGE_PILE * BASE ** NODE)
        SMALL_PILES.append(SMALL_PILE * BASE ** NODE)


        # Vars for Instructions/Practice
        LARGE_PILES_PRACTICE = LARGE_PILES.copy()
        SMALL_PILES_PRACTICE = SMALL_PILES.copy()

        COLUMNS_RANGE = range(4 + NUM_NODES)
        ROUNDS_RANGE = range(1, NUM_NODES + 1)

        # Merge the two Piles alternating elements
        PAYOFF_RED_PRACTICE = LARGE_PILES_PRACTICE.copy()
        for I in range(len(LARGE_PILES_PRACTICE)):
            if I % 2 != 0:
                PAYOFF_RED_PRACTICE[I] = SMALL_PILES_PRACTICE[I]

        # Merge the two Piles alternating elements
        PAYOFF_BLUE_PRACTICE = SMALL_PILES_PRACTICE.copy()
        for I in range(len(PAYOFF_BLUE_PRACTICE)):
            if I % 2 != 0:
                PAYOFF_BLUE_PRACTICE[I] = LARGE_PILES_PRACTICE[I]

        # Matrix of moves and payoffs
        # --------------------------------------------------------------------------------
        # Table of moves for each game NODE
        MOVESLIST = ['P' for n in range(NUM_NODES)]
        MOVESMATRIX = [MOVESLIST.copy()]

        for I in range(NUM_NODES):
            MC = MOVESLIST.copy()
            for J in range(NUM_NODES):
                if J <= I:
                    MC[I - J] = ''
                    MC[I] = 'T'
            MC.reverse()
            MOVESMATRIX.append(MC)

        MOVESMATRIX.reverse()

        # Zip variables needed for Instructions
        INSTRUCTIONSMATRIX = list(zip(MOVESMATRIX, LARGE_PILES_PRACTICE, SMALL_PILES_PRACTICE,
                                      PAYOFF_RED_PRACTICE, PAYOFF_BLUE_PRACTICE))

class Subsession(BaseSubsession):

    game = models.IntegerField(initial=1)
    game_node = models.IntegerField(initial=1)

    def creating_session(self):

        for x in range(1,C.NUM_ROUNDS+1):
            self.in_round(x).game = np.ceil(x/C.NUM_NODES)
            self.in_round(x).game_node = x - (np.ceil(x/C.NUM_NODES)-1)*C.NUM_NODES

        # shuffles players in groups randomly
        if self.round_number in C.FIRST_ROUNDS:
            self.group_randomly()
        else:
            x = C.FIRST_ROUNDS[int(np.ceil(self.round_number / C.NUM_NODES)) - 1]
            self.group_like_round(x)

        # shuffles players in groups randomly but keep players in fixed roles
        #if self.round_number in C.FIRST_ROUNDS:
        #    self.group_randomly(fixed_id_in_group=True)
        #else:
        #    x = C.FIRST_ROUNDS[int(np.ceil(self.round_number / C.NUM_NODES)) - 1]
        #    self.group_like_round(x)

        # shuffles players in groups as specified in matrix
        #if self.round_number == C.FIRST_ROUNDS[0]:
        #    new_structure = [[1, 2], [3, 4]]
        #    self.set_group_matrix(new_structure)
        #elif self.round_number == C.FIRST_ROUNDS[1]:
        #    new_structure = [[1, 3], [2, 4]]
        #    self.set_group_matrix(new_structure)
        #elif self.round_number == C.FIRST_ROUNDS[2]:
        #    new_structure = [[1, 4], [3, 2]]
        #    self.set_group_matrix(new_structure)
        #else:
        #    x = C.FIRST_ROUNDS[int(np.ceil(self.round_number / C.NUM_NODES)) - 1]
        #    self.group_like_round(x)

def advance_game(self):
    for g in self.get_groups():
        if g.subsession.game < C.NUM_GAMES:
            g.in_round(g.round_number + 1).game_on = True

class Group(BaseGroup):

    game_on = models.BooleanField(initial=True)
    game_outcome = models.IntegerField(initial=0)
    last_node = models.IntegerField(initial=1)

    def stop_game(self):
        players = self.get_players()
        takes = [p.take for p in players]
        if takes[0]:
            self.game_on = False
            self.game_outcome = 1
            self.last_node = self.subsession.game_node
        elif takes[1]:
            self.game_on = False
            self.game_outcome = 2
            self.last_node = self.subsession.game_node
        elif self.subsession.game_node == C.NUM_NODES and not any(takes):
            self.game_on = False
            self.last_node= self.subsession.game_node

        for group in self.in_rounds(self.round_number + 1, C.LAST_ROUNDS[self.subsession.game-1]):
            group.game_on = self.in_round(self.round_number).game_on
            group.game_outcome = self.in_round(self.round_number).game_outcome
            group.last_node = self.in_round(self.round_number).last_node

    def set_payoffs(self):
        players = self.get_players()
        takes = [p.take for p in players]
        for p in self.get_players():
            if self.subsession.game_node == C.NUM_NODES and not any(takes):
                if p.id_in_group == 1:
                    p.payoff = c(C.LARGE_PILES[-1])
                else:
                    p.payoff = C.SMALL_PILES[-1]
            elif self.subsession.game_node < C.NUM_NODES and any(takes):
                if p.take:
                    p.payoff = c(C.LARGE_PILES[self.last_node - 1])
                else:
                    p.payoff = c(C.SMALL_PILES[self.last_node - 1])

class Player(BasePlayer):

    take = models.BooleanField(
        label='',
        widget=widgets.RadioSelectHorizontal,
    )

class FirstPage(Page):
    def is_displayed(self):
        return self.round_number == 1


class Welcome(FirstPage):
    pass


class Instructions(FirstPage):

    def vars_for_template(self):
        return dict(
            turns=int(C.NUM_ROUNDS / 2),
            instructionsMatrix=C.INSTRUCTIONSMATRIX,
            rounds_range=C.ROUNDS_RANGE,
            large_pile_practice=C.LARGE_PILES,
            small_pile_practice=C.SMALL_PILES,
            large_pile_practice_second=C.LARGE_PILES[1],
            small_pile_practice_second=C.SMALL_PILES[1],
            large_pile_practice_third=C.LARGE_PILES[2],
            small_pile_practice_third=C.SMALL_PILES[2],
            large_pile_practice_last=C.LARGE_PILES[-2],
            small_pile_practice_last=C.SMALL_PILES[-2],
            large_pile_practice_pass=C.LARGE_PILES[-1],
            small_pile_practice_pass=C.SMALL_PILES[-1]
        )


class Practice1Page1(FirstPage):
    pass


class Practice1Page2(FirstPage):
    def vars_for_template(self):
        return dict(
            large_pile_practice=C.LARGE_PILES,
            small_pile_practice=C.SMALL_PILES,
            large_pile_practice_second=C.LARGE_PILES[1],
            small_pile_practice_second=C.SMALL_PILES[1],
        )


class Practice1Page3(FirstPage):

    def vars_for_template(self):
        return dict(
            large_pile_practice=C.LARGE_PILES,
            small_pile_practice=C.SMALL_PILES,
            large_pile_practice_second=C.LARGE_PILES[1],
            small_pile_practice_second=C.SMALL_PILES[1],
        )


class Practice1Page4(FirstPage):

    def vars_for_template(self):
        return dict(
            large_pile_practice=C.LARGE_PILES,
            small_pile_practice=C.SMALL_PILES,
            large_pile_practice_second=C.LARGE_PILES[1],
            small_pile_practice_second=C.SMALL_PILES[1],
        )


class Practice2Page1(FirstPage):

    def vars_for_template(self):
        return dict(
            large_pile_practice_second=C.LARGE_PILES[1],
            small_pile_practice_second=C.SMALL_PILES[1],
            large_pile_practice_third=C.LARGE_PILES[2],
            small_pile_practice_third=C.SMALL_PILES[2]
        )


class Practice2Page2(FirstPage):

    def vars_for_template(self):
        return dict(
            large_pile_practice_last=C.LARGE_PILES[-2],
            small_pile_practice_last=C.SMALL_PILES[-2],
            large_pile_practice_pass=C.LARGE_PILES[-1],
            small_pile_practice_pass=C.SMALL_PILES[-1],
        )


class Practice2Page3(FirstPage):
    pass


class WaitPage1(WaitPage):

    def is_displayed(self):
        return self.round_number == 1

    wait_for_all_groups = True


class Decision(Page):

    form_model = 'player'
    form_fields = ['take']

    def is_displayed(self):
        if self.id_in_group == 1 and self.round_number % 2 != 0 and self.group.game_on:
            return True
        elif self.id_in_group == 2 and self.round_number % 2 == 0 and self.group.game_on:
            return True
        else:
            return False

    def vars_for_template(self):
        return dict(
            game = self.subsession.game,
            game_node = self.subsession.game_node,
            large_pile = C.LARGE_PILES[self.subsession.game_node - 1],
            small_pile = C.SMALL_PILES[self.subsession.game_node - 1]
        )

def before_next_page(self):
    return self.group.stop_game(), self.group.set_payoffs()

class WaitPage2(WaitPage):
    def after_all_players_arrive(self):
        pass

class Results(Page):
    def is_displayed(self):
        if self.round_number in C.LAST_ROUNDS:
            return True
        else:
            return False

    def vars_for_template(self):
        return dict(
            game=self.subsession.game,
            last_node=self.group.last_node,
            large_pile=C.LARGE_PILES[self.group.last_node-1],
            small_pile=C.SMALL_PILES[self.group.last_node-1],
            large_pile_pass=C.LARGE_PILES[-1],
            small_pile_pass=C.SMALL_PILES[-1]
        )


class WaitPage3(WaitPage):
    def is_displayed(self):
        if self.round_number in C.LAST_ROUNDS:
            return True
        else:
            return False
    wait_for_all_groups = True
    after_all_players_arrive = 'advance_game'


page_sequence = [
    Welcome,
    Instructions,
    #Practice1Page1,
    #Practice1Page2,
    #Practice1Page3,
    #Practice1Page4,
    #Practice2Page1,
    #Practice2Page2,
    WaitPage1,
    Decision,
    WaitPage2,
    Results,
    WaitPage3
                ]