from otree.api import *


author = 'Your name here'
doc = """
Results are presented here
"""


class Constants(BaseConstants):
    name_in_url = 'Results'
    players_per_group = None
    num_rounds = 1


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    Results_warnings = models.IntegerField()
    Results_time_spent = models.FloatField(blank=True)


def custom_export(players):
    yield ['Session', 'Prolific ID', 'Payoff']
    for p in players:
        p_vars = p.participant.vars
        if p.participant._index_in_pages == p.participant._max_page_index:
            try:
                prolific_id = p_vars['prolific_id']
                payoff = float(p.participant.payoff.to_real_world_currency(p.session))
                yield [p.session.code, prolific_id, payoff]
            except KeyError:
                pass
        else:
            pass


# FUNCTIONS
# PAGES
class Results(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == Constants.num_rounds

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

    @staticmethod
    def js_vars(player: Player):
        return dict(page_name='Results')

    @staticmethod
    def vars_for_template(player: Player):
        # make code less bulky
        p_vars = 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 * cu(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 * cu(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 = player.participant.vars['treatment']
        rel_decision = player.participant.vars['WTP_rel_decision']
        try:
            gen4_wtp_cost = 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 = player.participant.payoff.to_real_world_currency(player.session)
        money = player.participant.payoff_plus_participation_fee
        participation_fee = player.session.config['participation_fee']
        reward = player.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):
    @staticmethod
    def js_vars(player: Player):
        return dict(completionlink=player.subsession.session.config['completionlink'])


page_sequence = [Results, LastPage]