from otree.api import *
import random, time

relTreatments = [0,1,3] # Same walking task for basic and path-dependent decoloring
import static_data.generateAllPaths as gPaths

doc = """
Your app description
"""


class Constants(BaseConstants):
    name_in_url = 'bWalk'
    players_per_group = None
    num_rounds = 10


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # Save stopping decisions
    coinTossOutcome = models.LongStringField()
    pathDrawn = models.LongStringField()
    nodeChoice = models.LongStringField()

    # Attention Check
    attention_check_done = models.BooleanField(initial=False)
    attention_check = models.StringField(initial="-9999")

# PAGES
class atnQuest(Page):
    @staticmethod
    def is_displayed(player):
        return (player.round_number == 1 and (player.participant.treatment in relTreatments)) and not player.participant.timeoutHappened and not player.participant.demo

    @staticmethod
    def live_method(player, data):
        player.attention_check_done = True
        player.attention_check = str(data)
        return {player.id_in_group: data} # Send data back
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            testing = player.session.config["testing"],
            upTick = player.session.config["upTick"],
        ) 
    
    @staticmethod
    def js_vars(player: Player):
        return dict(
            stopEntry = player.attention_check_done,
            submittedAnswer = player.attention_check,
        ) 
    
    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant
        return participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        if player.attention_check == str(False):
            player.session.dropouts.append(player.participant.treatment)
            player.session.actives[player.participant.treatment] = player.session.actives[player.participant.treatment] - 1

        if timeout_happened:
            player.participant.timeoutHappened = True
            player.session.dropouts.append(player.participant.treatment)
            player.session.actives[player.participant.treatment] = player.session.actives[player.participant.treatment] - 1


class exitPage(Page):
    @staticmethod
    def is_displayed(player):
        return player.attention_check == str(False) and not player.participant.timeoutHappened and not player.participant.demo # Need to keep string here
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            testing = player.session.config["testing"],
        )

class taskExplain(Page):
    @staticmethod
    def is_displayed(player):
        return ((player.round_number == 1) and 
                (player.participant.treatment in relTreatments) and 
                player.in_round(1).attention_check == str(True) and 
                not player.participant.timeoutHappened) and not player.participant.demo
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            user_chosen_decision    = player.participant.bf,
            testing                 = player.session.config["testing"],
            upTick                  = player.session.config["upTick"],
        )
    
    @staticmethod
    def js_vars(player):
        return dict(
            upTick                  = player.session.config["upTick"],
            testing                 = player.session.config["testing"],
            user_chosen_decision    = player.participant.bf,
            endowment               = '£%.2f' % player.session.config["endowment"], 
            initialColor            = player.participant.plotColor
        ) 

    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant
        return participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        # Note: We need the participant variables here, need to use a dictionary to access them later
        player.participant.resolutionRound = 0
        ## Initialize for Walking Tasks
        player.participant.walkingPath = {(idx+1): [0]*6 for idx in range(Constants.num_rounds)}
        player.participant.coinTossOutcome = {(idx+1): [0]*6 for idx in range(Constants.num_rounds)}
        player.participant.continuationDec = {(idx+1): [3000]*6 for idx in range(Constants.num_rounds)}
        ## Initialize for practice round
        player.participant.walkingPathTrial = [0]*6
        player.participant.coinTossOutcomeTrial = [0]*6
        player.participant.continuationDecTrial = [3000]*6

        if timeout_happened:
            player.participant.timeoutHappened = True
            player.session.dropouts.append(player.participant.treatment)
            player.session.actives[player.participant.treatment] = player.session.actives[player.participant.treatment] - 1

class attentionScreen(Page):
    @staticmethod
    def is_displayed(player):
        return ((player.participant.treatment in relTreatments) and 
                player.in_round(1).attention_check == str(True) and 
                not player.participant.timeoutHappened) and not player.participant.demo

    @staticmethod
    def vars_for_template(player):
        return dict(
            taskNumber = player.round_number,
        )
    
    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant
        return participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        if timeout_happened:
            player.participant.timeoutHappened = True
            player.session.dropouts.append(player.participant.treatment)
            player.session.actives[player.participant.treatment] = player.session.actives[player.participant.treatment] - 1

class wTpractice(Page):
    @staticmethod
    def is_displayed(player):
        return ((player.participant.treatment in relTreatments) 
                and player.in_round(1).attention_check == str(True)
                and player.round_number == 1
                and not player.participant.timeoutHappened
                ) and not player.participant.demo
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            user_chosen_decision = player.participant.bf,
            testing = player.session.config["testing"],
            remainingRounds = 5 - player.participant.resolutionRound,
            endowment = '£%.2f' % player.session.config["endowment"], 
        )
    
    @staticmethod
    def js_vars(player):
        return dict(
            winProb                 = player.session.config["winProb"],
            upTick                  = player.session.config["upTick"],
            min_time_in_modal       = player.session.config["min_time_in_modal"],
            min_time_on_page        = player.session.config["min_time_on_page"],
            stopStrategy            = player.participant.decoloring,    
            user_chosen_decision    = player.participant.bf,
            roundNumber             = player.participant.resolutionRound,
            path                    = player.participant.walkingPathTrial,
            continuationDec         = player.participant.continuationDecTrial,
            testing                 = player.session.config["testing"],
            endowment               = '£%.2f' % player.session.config["endowment"], 
            initialColor            = player.participant.plotColor,
            allPaths                = gPaths.getAllPaths(),
        )
    
    @staticmethod
    def live_method(player,data):
        decision = data["type"]
        value = data["value"]

        imagePathGuess      = '/static/img/coin.png'
        imagePathOutcome    = '/static/img/coin_questionmark.png'
        resolutionRound     = player.participant.resolutionRound


        if decision == "nodeChoice":
            # Check if stop node --> write stop
            if data["ballContinue"] == 0 and not "Stop" in player.participant.continuationDecTrial:
                player.participant.continuationDecTrial[resolutionRound] = "Stop"
            elif data["ballContinue"] > 0 and not "Stop" in player.participant.continuationDecTrial and resolutionRound == 4:
                player.participant.continuationDecTrial[resolutionRound] = data["ballContinue"]
                player.participant.continuationDecTrial[5] = "Stop"
            else:
                player.participant.continuationDecTrial[resolutionRound] = data["ballContinue"]
        elif decision == "coin":
            # Get coin outcome
            rnd = random.uniform(0,1)
            winningSide = "H" if rnd < 0.5 else "T"

            # Compare coin toss outcome with winning side
            player.participant.resolutionRound += 1
            player.participant.coinTossOutcomeTrial[resolutionRound] = value

            if value == "H":
                imagePathGuess = '/static/img/coin_h.png'
            else:
                imagePathGuess = '/static/img/coin_t.png'

            if winningSide == "H":
                imagePathOutcome = '/static/img/coin_h.png'
            else:
                imagePathOutcome = '/static/img/coin_t.png'

            currentNode = player.participant.walkingPathTrial[resolutionRound]

            if winningSide == value:
                # Move Up
                nextNode = currentNode + resolutionRound + 1
            else:
                # Move Down
                nextNode = currentNode + resolutionRound + 2

            player.participant.walkingPathTrial[resolutionRound + 1] =  nextNode
            
        return {
            # Data is returned to live-receive on HTML page
            player.id_in_group: dict(
                type = decision,
                roundNumber = player.participant.resolutionRound,
                path = player.participant.walkingPathTrial,
                continuationDec = player.participant.continuationDecTrial,
                imagePathGuess = imagePathGuess,
                imagePathOutcome = imagePathOutcome,
        )}

    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant
        return participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.participant.resolutionRound = 0 # Reset
        if timeout_happened:
            player.participant.timeoutHappened = True
            player.session.dropouts.append(player.participant.treatment)
            player.session.actives[player.participant.treatment] = player.session.actives[player.participant.treatment] - 1

class walkingTask(Page):
    @staticmethod
    def is_displayed(player):
        if player.participant.demo:
            player.in_round(1).attention_check = str(True)
            
        return ((player.participant.treatment in relTreatments) and 
                player.in_round(1).attention_check == str(True) and
                not player.participant.timeoutHappened)
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            user_chosen_decision = player.participant.bf,
            testing = player.session.config["testing"],
            remainingRounds = 5 - player.participant.resolutionRound,
            endowment = '£%.2f' % player.session.config["endowment"], 
            upTick = '£%.2f' % player.session.config["upTick"],
        )
    
    @staticmethod
    def js_vars(player):
        taskNumber = player.round_number
        return dict(
            winProb                 = player.session.config["winProb"],
            upTick                  = player.session.config["upTick"],
            min_time_in_modal       = player.session.config["min_time_in_modal"],
            min_time_on_page        = player.session.config["min_time_on_page"],
            stopStrategy            = player.participant.decoloring,    
            user_chosen_decision    = player.participant.bf,
            taskNumber              = taskNumber,
            roundNumber             = player.participant.resolutionRound,
            path                    = player.participant.walkingPath[taskNumber],
            continuationDec         = player.participant.continuationDec[taskNumber],
            testing                 = player.session.config["testing"],
            endowment               = '£%.2f' % player.session.config["endowment"], 
            initialColor            = player.participant.plotColor,
            allPaths                = gPaths.getAllPaths(),
            costly                  = (player.participant.treatment == 3),
            costPerBall             = '£%.2f' % player.session.config["costPerBall"],
        )
    
    @staticmethod
    def live_method(player,data):
        decision = data["type"]
        value = data["value"]

        imagePathGuess      = '/static/img/coin.png'
        imagePathOutcome    = '/static/img/coin_questionmark.png'
        resolutionRound     = player.participant.resolutionRound
        taskNumber          = player.round_number


        if decision == "nodeChoice":
            # Check if stop node --> write stop
            if data["ballContinue"] == 0 and not "Stop" in player.participant.continuationDec[taskNumber]:
                player.participant.continuationDec[taskNumber][resolutionRound] = "Stop"
            elif data["ballContinue"] > 0 and not "Stop" in player.participant.continuationDec[taskNumber] and resolutionRound == 4:
                player.participant.continuationDec[taskNumber][resolutionRound] = data["ballContinue"]
                player.participant.continuationDec[taskNumber][5] = "Stop"
            else:
                player.participant.continuationDec[taskNumber][resolutionRound] = data["ballContinue"]
        elif decision == "coin":
            # Compare coin toss outcome with winning side
            player.participant.resolutionRound += 1
            winningSides = player.participant.winningSideDict[player.round_number]
            winningSide = winningSides[int(resolutionRound)]
            player.participant.coinTossOutcome[taskNumber][resolutionRound] = value

            if value == "H":
                imagePathGuess = '/static/img/coin_h.png'
            else:
                imagePathGuess = '/static/img/coin_t.png'

            if winningSide == "H":
                imagePathOutcome = '/static/img/coin_h.png'
            else:
                imagePathOutcome = '/static/img/coin_t.png'

            currentNode = player.participant.walkingPath[taskNumber][resolutionRound]

            if winningSide == value:
                # Move Up
                nextNode = currentNode + resolutionRound + 1
            else:
                # Move Down
                nextNode = currentNode + resolutionRound + 2

            player.participant.walkingPath[taskNumber][resolutionRound + 1] =  nextNode
        
        elif decision == "skip" and value == "skip":
            player.participant.coinTossOutcome[taskNumber] = ['H', 'T', 'H', 'T', 'H']
            player.participant.walkingPath[taskNumber] = [0,1,3,7,11,17]
            player.participant.continuationDec[taskNumber] = [100,100,100,100,'Stop',0]

            
        return {
            # Data is returned to live-receive on HTML page
            player.id_in_group: dict(
                type = decision,
                roundNumber = player.participant.resolutionRound,
                path = player.participant.walkingPath[taskNumber],
                continuationDec = player.participant.continuationDec[taskNumber],
                imagePathGuess = imagePathGuess,
                imagePathOutcome = imagePathOutcome,
        )}

    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant
        return participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        taskNumber = player.round_number
        # Save all outcomes
        player.coinTossOutcome = str(player.participant.coinTossOutcome[taskNumber])
        player.pathDrawn = str(player.participant.walkingPath[taskNumber])
        player.nodeChoice = str(player.participant.continuationDec[taskNumber])

        player.participant.resolutionRound = 0 # Reset

        if taskNumber == 10:
            # Move from active into complete
            player.session.actives[player.participant.treatment] = player.session.actives[player.participant.treatment] - 1
            player.session.completes[player.participant.treatment] = player.session.completes[player.participant.treatment] + 1
        
        if timeout_happened:
            player.participant.timeoutHappened = True
            player.session.dropouts.append(player.participant.treatment)
            player.session.actives[player.participant.treatment] = player.session.actives[player.participant.treatment] - 1

page_sequence = [
    atnQuest,
    taskExplain,
    wTpractice,
    attentionScreen,
    walkingTask,
    exitPage
    ]