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

class Results(Page):
    def is_displayed(self):
        return self.round_number == Constants.num_rounds

    form_model = 'player'
    form_fields = ['Results_warnings','Results_time_spent']

    def js_vars(self):
        return dict(
            page_name='Results')

    def vars_for_template(self):
        # make code less bulky
        p_vars = self.player.participant.vars

        # Ability Stages
        pwd_rounds = p_vars['num_rounds_pwd_ability']
        pwd_ability = sum([p_vars[f'pwd_ability_round_{round}'] for round in range(1,pwd_rounds+1)])
        pwd_ability_payoff = sum([p_vars[f'pwd_ability_payoff_round_{round}'] for round in range(1,pwd_rounds+1)])
        possible_pwd_task_ability  = p_vars['num_pwd_tasks_pwd_ability']*pwd_rounds
        pwd_diff = pwd_ability_payoff- pwd_ability*c(10)

        math_rounds = p_vars['num_rounds_math_ability']
        math_ability = sum([p_vars[f'math_ability_round_{round}'] for round in range(1,math_rounds+1)])
        math_ability_payoff = sum([p_vars[f'math_ability_payoff_round_{round}'] for round in range(1,math_rounds+1)])
        possible_math_task_ability  = p_vars['num_math_tasks_math_ability']*math_rounds
        math_diff = math_ability_payoff- math_ability*c(10)
        # Part 3
        gen_rounds = int(p_vars['gen_num_rounds'])
        gen3_pwd = sum([p_vars[f'base_pwd_round_{round}'] for round in range(1,gen_rounds+1)])
        gen3_pwd_payoff = sum([p_vars[f'base_pwd_payoff_round_{round}'] for round in range(1,gen_rounds+1)])
        gen3_math = sum([p_vars[f'base_math_round_{round}'] for round in range(1,gen_rounds+1)])
        gen3_math_payoff = sum([p_vars[f'base_math_payoff_round_{round}'] for round in range(1,gen_rounds+1)])
        gen3_payoff = gen3_math_payoff + gen3_pwd_payoff
        # Part 4
        gen4_pwd = sum([p_vars[f'treat_pwd_round_{round}'] for round in range(1,gen_rounds+1)])
        gen4_pwd_payoff = sum([p_vars[f'treat_pwd_payoff_round_{round}'] for round in range(1,gen_rounds+1)])
        gen4_math = sum([p_vars[f'treat_math_round_{round}'] for round in range(1,gen_rounds+1)])
        gen4_math_payoff = sum([p_vars[f'treat_math_payoff_round_{round}'] for round in range(1,gen_rounds+1)])
        gen4_payoff = gen4_math_payoff + gen4_pwd_payoff
        gen_possible_pwd_task  = int(p_vars['gen_num_pwd_tasks'])*gen_rounds
        gen_possible_math_task  = int(p_vars['gen_num_math_tasks'])*gen_rounds
        ############# NEW ####################
        treatment = self.player.participant.vars['treatment']
        rel_decision = self.player.participant.vars['WTP_rel_decision']
        try:
            gen4_wtp_cost = self.player.participant.vars['WTP_cost']*gen_rounds
            gen4_total_payoff = gen4_math_payoff + gen4_pwd_payoff - gen4_wtp_cost
        except:
            gen4_wtp_cost = None
            gen4_total_payoff = gen4_math_payoff + gen4_pwd_payoff


        # Task Treatment
        #treatment = p_vars['task_treatment']

        # Total payoff
        total_payoff = math_ability_payoff+pwd_ability_payoff+gen3_payoff+gen4_total_payoff
        money_wo_fee = self.participant.payoff.to_real_world_currency(self.session)
        money = self.participant.payoff_plus_participation_fee
        participation_fee=self.session.config['participation_fee']
        reward = self.session.config['prolific_reward']

        return dict(
            pwd_rounds=pwd_rounds,
            pwd_ability=pwd_ability,
            pwd_ability_payoff=pwd_ability_payoff,
            possible_pwd_task_ability=possible_pwd_task_ability,
            math_rounds=math_rounds,
            math_ability=math_ability,
            math_ability_payoff=math_ability_payoff,
            possible_math_task_ability=possible_math_task_ability,
            gen_rounds=gen_rounds,
            gen3_pwd=gen3_pwd,
            gen3_pwd_payoff=gen3_pwd_payoff,
            gen3_math=gen3_math,
            gen3_math_payoff=gen3_math_payoff,
            gen3_payoff=gen3_payoff,
            gen4_pwd=gen4_pwd,
            gen4_pwd_payoff=gen4_pwd_payoff,
            gen4_math=gen4_math,
            gen4_math_payoff=gen4_math_payoff,
            gen4_payoff=gen4_payoff,
            gen4_wtp_cost=gen4_wtp_cost,
            gen4_total_payoff=gen4_total_payoff,
            treatment=treatment,
            rel_decision=rel_decision,
            gen_possible_pwd_task=gen_possible_pwd_task,
            general_possible_math_task=gen_possible_math_task,
            total_payoff=total_payoff,
            money_wo_fee=money_wo_fee,
            money=money,
            participation_fee=participation_fee,
            #treatment=treatment,
            math_diff=math_diff,
            pwd_diff=pwd_diff,
            #gen3_payoff=gen3_payoff,
            #gen4_pwd=gen4_pwd,
            #gen4_pwd_payoff=gen4_pwd_payoff,
            #gen4_math=gen4_math,
            #gen4_math_payoff=gen4_math_payoff,
            #gen4_payoff=gen4_payoff,
            reward=reward,
        )

class LastPage(Page):

    def js_vars(self):
        return dict(
            completionlink=
              self.subsession.session.config['completionlink']
        )


page_sequence = [
    Results,
    LastPage
]