from otree.api import Currency as c, currency_range
from ._builtin import Page, WaitPage
from .models import Constants
import random


# TODO: plays all rounds in the background - generate lots of unnecessary data
# TODO: should they see number of attempts so far?
# TODO: decision 1 fixed, 2-3 in random order

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

    def vars_for_template(self):
        return{
            'total_required': Constants.required_tasks * 2
        }

    def before_next_page(self):
        def set_correct_answers():
            self.participant.vars['correct_answers'] = 0

        def set_finished():
            self.participant.vars['finished'] = False

        set_correct_answers()
        set_finished()

        print('Decision rounds are', self.participant.vars['decision_rounds'])


class Task(Page):
    form_model = 'player'
    form_fields = ['answer']

    def is_displayed(self):
        return self.participant.vars['correct_answers'] < Constants.num_correct_needed

    def vars_for_template(self):
        x = self.player.matrix_id_in_round
        self.participant.vars['correct_answers'] = sum([p.answer_correct for p in self.player.in_previous_rounds()])
        # what is written on template is number of correct answers in previous rounds

        if self.round_number > 1:
            return{
                'matrix_id': x,
                'image_path': 'elicitation/matrix_{}.gif'.format(x),
                'solution': Constants.Matrix[x],
                'matrix_number': self.participant.vars['correct_answers'] + 1,
                'correct_answers': self.participant.vars['correct_answers'],
                'previous_correct': self.player.in_round(self.round_number - 1).answer_correct
            }
        else:
            return{
                'matrix_id': x,
                'image_path': 'elicitation/matrix_{}.gif'.format(x),
                'solution': Constants.Matrix[x],
                'matrix_number': self.participant.vars['correct_answers'] + 1,
                'correct_answers': self.participant.vars['correct_answers']
            }

    def before_next_page(self):
        x = self.player.matrix_id_in_round

        self.player.answer_correct = (self.player.answer == Constants.Matrix[x])
        self.participant.vars['correct_answers'] = sum([p.answer_correct for p in self.player.in_all_rounds()])
        # here it calculates number of correct answers in previous rounds plus current round


class Instructions(Page):
    def is_displayed(self):
        return self.participant.vars['correct_answers'] == Constants.num_correct_needed \
               and not self.participant.vars['finished']

    def vars_for_template(self):
        return{
            'total_required': Constants.required_tasks * 2
        }

    def before_next_page(self):
        self.participant.vars['finished_tasks_in_round'] = self.round_number
        self.participant.vars['finished'] = True
        print('Round number is', self.round_number)


class AllocationDecision1(Page):
    def is_displayed(self):
        if self.participant.vars['correct_answers'] < Constants.num_correct_needed:
            return False
        else:
            return self.round_number == self.participant.vars['finished_tasks_in_round'] \
                   + self.participant.vars['decision_rounds'][1]

    form_model = 'player'
    form_fields = ['w1_certain']

    def vars_for_template(self):
        return {
            'decision_number': self.participant.vars['decision_rounds'][1] + 1
        }

    def before_next_page(self):
        self.participant.vars['w1_certain'] = self.player.w1_certain


class AllocationDecision2(Page):
    def is_displayed(self):
        if self.participant.vars['correct_answers'] < Constants.num_correct_needed:
            return False
        else:
            return self.round_number == self.participant.vars['finished_tasks_in_round'] \
                   + self.participant.vars['decision_rounds'][2]

    form_model = 'player'
    form_fields = ['w1_risky_today']

    def vars_for_template(self):
        return {
            'decision_number': self.participant.vars['decision_rounds'][2] + 1
        }

    def before_next_page(self):
        self.participant.vars['w1_risky_today'] = self.player.w1_risky_today


class AllocationDecision3(Page):
    def is_displayed(self):
        if self.participant.vars['correct_answers'] < Constants.num_correct_needed:
            return False
        else:
            return self.round_number == self.participant.vars['finished_tasks_in_round'] \
                   + self.participant.vars['decision_rounds'][3]

    form_model = 'player'
    form_fields = ['w1_risky_tomorrow']

    def vars_for_template(self):
        return {
            'decision_number': self.participant.vars['decision_rounds'][3] + 1
        }

    def before_next_page(self):
        self.participant.vars['w1_risky_tomorrow'] = self.player.w1_risky_tomorrow


class PaymentChoice(Page):
    def is_displayed(self):
        if self.participant.vars['correct_answers'] < Constants.num_correct_needed:
            return False
        else:
            return self.round_number == self.participant.vars['finished_tasks_in_round'] + len(Constants.decisions) - 1

    def before_next_page(self):
        stochastic_offset = Constants.dw*random.choice([-1, 1])
        print('stochastic offset is now', stochastic_offset)

        # Marc's suggestion
        # 1. Payoff if risky today choice matters
        def risky_today_payoff():
            self.player.rand1 = 0
            w_rt = self.participant.vars['w1_risky_today']
            self.player.w1_rewarded = w_rt + stochastic_offset
            return Constants.p1*self.player.w1_rewarded + Constants.p2*(Constants.w - w_rt) \
                   + Constants.required_tasks*(Constants.p1 + Constants.p2)

        # 2. Payoff if risky future choice matters
        def risky_future_payoff():
            self.player.rand1 = 1
            w_rf = self.participant.vars['w1_risky_tomorrow']
            self.player.w2_rewarded = Constants.w - w_rf + stochastic_offset
            return Constants.p1*w_rf + Constants.p2*self.player.w2_rewarded \
                   + Constants.required_tasks*(Constants.p1 + Constants.p2)

        # 3. Payoff if certain choice matters
        def certain_payoff():
            self.player.rand1 = 2
            w_c = self.participant.vars['w1_certain']
            return Constants.p1*w_c + Constants.p2*(Constants.w - w_c) \
                   + Constants.required_tasks*(Constants.p1 + Constants.p2)

        # 4. Choose (randomize) which matters
        def pick_choice_set_payoff():
            payoff = random.choice([risky_today_payoff, risky_future_payoff, certain_payoff])
            self.player.payoff = payoff()

        pick_choice_set_payoff()


class Results(Page):
    def is_displayed(self):
        if self.participant.vars['correct_answers'] < Constants.num_correct_needed:
            return False
        else:
            return self.round_number == self.participant.vars['finished_tasks_in_round'] + len(Constants.decisions) - 1

    def vars_for_template(self):
        p = self.player
        q = self.participant
        chosen_w = q.vars['w1_risky_today'] if p.rand1 == 0 else q.vars['w1_risky_tomorrow'] \
            if p.rand1 == 1 else q.vars['w1_certain']

        def rounds_in_next_apps():
            q.vars['w1'] = chosen_w
            q.vars['w2'] = Constants.w - chosen_w

        rounds_in_next_apps()

        if p.rand1 == 0:
            return{
                'w1': chosen_w,
                'w2': Constants.w - chosen_w,
                'total_tasks_for_today': self.participant.vars['w1'] + Constants.required_tasks,
                'total_tasks_for_tomorrow': self.participant.vars['w2'] + Constants.required_tasks,
                'rewarded_tasks_for_today': p.w1_rewarded + Constants.required_tasks,
                'decision_number': q.vars['decision_rounds'][2] + 1 if p.rand1 == 0 \
                    else q.vars['decision_rounds'][3] + 1 if p.rand1 == 1 \
                    else q.vars['decision_rounds'][1] + 1
            }
        elif p.rand1 == 1:
            return{
                'w1': chosen_w,
                'w2': Constants.w - chosen_w,
                'total_tasks_for_today': self.participant.vars['w1'] + Constants.required_tasks,
                'total_tasks_for_tomorrow': self.participant.vars['w2'] + Constants.required_tasks,
                'rewarded_tasks_for_tomorrow': p.w2_rewarded + Constants.required_tasks,
                'decision_number': q.vars['decision_rounds'][2] + 1 if p.rand1 == 0 \
                    else q.vars['decision_rounds'][3] + 1 if p.rand1 == 1 \
                    else q.vars['decision_rounds'][1] + 1
            }

        else:
            return{
                'w1': chosen_w,
                'w2': Constants.w - chosen_w,
                'decision_number': q.vars['decision_rounds'][2] + 1 if p.rand1 == 0 \
                    else q.vars['decision_rounds'][3] + 1 if p.rand1 == 1 \
                    else q.vars['decision_rounds'][1] + 1
            }

    def before_next_page(self):

        def set_participant_payoff():
            p = self.player
            q = self.participant
            q.vars['payoff'] = p.payoff
            payment_schemes = {0: 'risky_today', 1: 'risky_tomorrow', 2: 'certain'}
            q.vars['chosen_scheme'] = [p.rand1, payment_schemes[p.rand1]]
            q.vars['w_rewarded'] = [p.w1_rewarded, Constants.w - q.vars['w1_risky_today']] if p.rand1 == 0 \
                else [q.vars['w1_risky_tomorrow'], p.w2_rewarded] if p.rand1 == 1 \
                else [q.vars['w1_certain'], Constants.w - q.vars['w1_certain']]

        set_participant_payoff()

        print('Participant vars are', self.participant.vars)


page_sequence = [
    WelcomePage,
    Task,
    Instructions,
    AllocationDecision1,
    AllocationDecision2,
    AllocationDecision3,
    PaymentChoice,
    Results
]