from openfisca_tunisia.variables.base import *


# Pensions


class pension_de_retraite(Variable):
    value_type = float
    label = 'Pension'
    entity = Individu
    definition_period = MONTH
    set_input = set_input_divide_by_period


class pension_orphelin(Variable):
    value_type = float
    label = 'Pension'
    entity = Individu
    definition_period = MONTH
    set_input = set_input_divide_by_period


class pension_d_invalidite(Variable):
    value_type = float
    label = 'Pension'
    entity = Individu
    definition_period = MONTH
    set_input = set_input_divide_by_period


class pension(Variable):
    value_type = float
    label = 'Pension'
    entity = Individu
    definition_period = YEAR
    set_input = set_input_divide_by_period

    def formula(individu, period):
        return (
            individu('pension_de_retraite', period, options = [ADD])
            + individu('pension_d_invalidite', period, options = [ADD])
            + individu('pension_orphelin', period, options = [ADD])
            )


class revenu_assimile_pension(Variable):
    value_type = float
    label = 'Revenus assimilés à des pensions (pensions et rentes viagères)'
    entity = Individu
    definition_period = YEAR

    def formula_2025_01(individu, period):
        return individu('pension_de_retraite', period, options = [ADD])

    def formula(individu, period):
        return individu('pension', period, options = [ADD])


class avantages_nature_assimile_pension(Variable):
    value_type = float
    label = 'Avantages en nature assimilables à des pensions'
    entity = Individu
    definition_period = YEAR
