from otree.api import *
import itertools
import random
import time
import numpy as np

author = 'Emily Tran'
doc = """
Part 3: Lottery Outcome & Belief Elicitation
"""


class C(BaseConstants):
    NAME_IN_URL = 'Part2_beliefs'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1
    TIMEOUT_DEFAULT = 15
    ENDOWMENT = 300
    # probability of success
    PROB_HIGH = 0.75
    PROB_LOW = 0.25
    # # payment for members' beliefs
    # BELIEFPAYMENT = 200
    # # returns from lottery
    # RETURN_WIN = {
    #     1: 250,
    #     2: 200,
    #     3: 300,
    # }
    # RETURN_LOSE = {
    #     1: 50,
    #     2: 0,
    #     3: 50,
    # }


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # outcome based on performance of task
    high_outcome = models.BooleanField(initial=False)

    # outcome of lottery
    lottery_outcome = models.BooleanField(initial=False)

    # prior beliefs
    prior = models.IntegerField(
        min=0,
        max=100,
    )
    # posterior beliefs
    posterior = models.IntegerField(
        min=0,
        max=100,
    )

    # store whether subject drops out
    droppedouteffort = models.BooleanField(initial=False)
    droppedoutprior = models.BooleanField(initial=False)
    droppedoutposterior = models.BooleanField(initial=False)


# FUNCTIONS
def get_timeout_seconds(player):
    import time
    return player.participant.expiry - time.time()


# def creating_session(subsession):  # this is run before the start of every round
#     for player in subsession.get_players():
#         participant = player.participant
#         participant.num_correct_ravens = player.num_correct_ravens
#         participant.payoff_ravens = player.payoff_ravens
#         # participant.high_outcome = player.participant.high_outcome
#         # participant.high_outcome = True
#         # if player.id_in_group == 1:
#         #     player.leader = True
#         #     player.member = False
#         # else:
#         #     player.leader = False
#         #     player.member = True


# PAGES
class TaskWaitPage(Page):
    wait_for_all_groups = True

    @staticmethod
    def before_next_page(player, timeout_happened):
        participant = player.participant

        if timeout_happened:
            participant.is_dropout = True

        if player.participant.num_correct_ravens >= 2:  # if task performance >= threshold
            player.high_outcome = True

        if player.participant.num_correct_ravens < 2:  # if task performance < threshold
            player.high_outcome = False


class Performance_Outcome(Page):
    @staticmethod
    def is_displayed(player):
        # return player.round_number == 1 and player.participant.is_dropout is False
        return player.participant.is_dropout is False

    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant

        if participant.is_dropout:
            return 1  # instant timeout, 1 second
        else:
            return player.participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        participant = player.participant

        if timeout_happened:
            participant.is_dropout = True

        if player.high_outcome is True:  # if player achieves the high outcome
            player.lottery_outcome = np.random.choice([1, 0], p=[C.PROB_HIGH, C.PROB_LOW])
            # Show the result
            print(f'Win Probability: {0.75}, Result: {player.lottery_outcome}')

        if player.high_outcome is False:  # if player achieves the low outcome
            player.lottery_outcome = np.random.choice([1, 0], p=[C.PROB_LOW, C.PROB_HIGH])
            # Show the result
            print(f'Win Probability: {0.25}, Result: {player.lottery_outcome}')

    @staticmethod
    def vars_for_template(player):  # pass variables to the HTML templates
        participant = player.participant
        return dict(
            par_vars=participant,
        )


class Lottery_Outcome(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1 and player.participant.is_dropout is False

    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant

        if participant.is_dropout:
            return 1  # instant timeout, 1 second
        else:
            return player.participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        participant = player.participant

        if timeout_happened:
            participant.is_dropout = True

    @staticmethod
    def vars_for_template(player):  # pass variables to the HTML templates
        participant = player.participant
        return dict(
            par_vars=participant,
        )


class Prior(Page):
    form_model = 'player'
    form_fields = ['prior']

    # @staticmethod
    # def is_displayed(player):
    #     return player.member

    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant

        if participant.is_dropout:
            return 1  # instant timeout, 1 second
        else:
            return player.participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        participant = player.participant

        if timeout_happened:
            participant.prior = 50
            participant.droppedoutprior = True
            participant.is_dropout = True

        participant.expiry = time.time() + C.TIMEOUT_DEFAULT * 60

    @staticmethod
    def vars_for_template(player):  # pass variables to the HTML templates
        participant = player.participant
        return dict(
            par_vars=participant,
        )


class Posterior(Page):
    form_model = 'player'
    form_fields = ['posterior']

    # @staticmethod
    # def is_displayed(player):
    #     return player.member

    @staticmethod
    def get_timeout_seconds(player):
        participant = player.participant

        if participant.is_dropout:
            return 1  # instant timeout, 1 second
        else:
            return player.participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        participant = player.participant

        if timeout_happened:
            participant.posterior = 50
            participant.droppedoutposterior = True
            participant.is_dropout = True

        # participant.expiry = time.time() + C.TIMEOUT_DEFAULT * 60

    @staticmethod
    def vars_for_template(player):  # pass variables to the HTML templates
        participant = player.participant

        return dict(
            par_vars=participant,
        )


class H_WaitPage(Page):
    wait_for_all_groups = True  # not needed? each group can proceed at their own pace

    @staticmethod
    def get_timeout_seconds(player):
        if player.participant.is_dropout:
            return 1  # instant timeout, 1 second
        else:
            return player.participant.expiry - time.time()

    @staticmethod
    def before_next_page(player, timeout_happened):
        participant = player.participant

        if timeout_happened:
            participant.posterior = 50
            participant.droppedoutposterior = True
            participant.is_dropout = True

    @staticmethod
    def vars_for_template(player):
        participant = player.participant

        # Add participant variables
        participant.droppedoutprior = player.droppedoutprior
        participant.droppedoutposterior = player.droppedoutposterior
        player.participant.high_outcome = player.high_outcome
        player.participant.lottery_outcome = player.lottery_outcome

        return dict(
            par_vars=participant,
        )

    # @staticmethod
    # def before_next_page(player, timeout_happened):
    #     participant = player.participant
    #     # group = player.group
    #
    #     # Add participant variable
    #     # player.participant.group.set_round_payoffs()
    #     # participant.leaderdroppedout = group.leaderdroppedout
    #     # participant.groupreturnsuccess = group.groupreturnsuccess
    #     # participant.groupreturnfailure = group.
    #     # participant.droppedouteffort = player.droppedouteffort
    #     participant.droppedoutprior = player.droppedoutprior
    #     participant.droppedoutposterior = player.droppedoutposterior


page_sequence = [
    TaskWaitPage,
    Performance_Outcome,
    Lottery_Outcome,
    Prior,
    Posterior,
    H_WaitPage,
]