from pathlib import Path
from policyengine_core.taxbenefitsystems import TaxBenefitSystem
from fiscalsim_us.entities import *
from fiscalsim_us.parameters.gov.irs.uprating import (
    set_irs_uprating_parameter,
)
from policyengine_core.simulations import (
    Simulation as CoreSimulation,
    Microsimulation as CoreMicrosimulation,
    IndividualSim as CoreIndividualSim,
)
from fiscalsim_us.data import DATASETS, CPS_2023, EnhancedCPS_2023
from fiscalsim_us.variables.household.demographic.geographic.state.in_state import (
    create_50_state_variables,
)
from fiscalsim_us.tools.parameters import backdate_parameters
from fiscalsim_us.reforms import create_structural_reforms_from_parameters

from fiscalsim_us.reforms import create_structural_reforms_from_parameters

COUNTRY_DIR = Path(__file__).parent


class CountryTaxBenefitSystem(TaxBenefitSystem):
    parameters_dir = COUNTRY_DIR / "parameters"
    variables_dir = COUNTRY_DIR / "variables"
    auto_carry_over_input_variables = True
    basic_inputs = [
        "state_name",
        "employment_income",
        "age",
    ]
    modelled_policies = COUNTRY_DIR / "modelled_policies.yaml"

    def __init__(self):
        # We initialize our tax and benefit system with the general constructor
        super().__init__(entities)

        reform = create_structural_reforms_from_parameters(
            self.parameters, "2023-01-01"
        )
        if reform is not None:
            reform.apply(self)

        self.add_variables(*create_50_state_variables())

        self.parameters = set_irs_uprating_parameter(self.parameters)
        self.parameters = backdate_parameters(self.parameters)


system = CountryTaxBenefitSystem()


class Simulation(CoreSimulation):
    default_tax_benefit_system = CountryTaxBenefitSystem
    default_tax_benefit_system_instance = system
    default_role = "member"
    default_calculation_period = 2022
    default_input_period = 2022
    datasets = DATASETS

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        reform = create_structural_reforms_from_parameters(
            self.tax_benefit_system.parameters, "2023-01-01"
        )
        if reform is not None:
            self.apply_reform(reform)


class Microsimulation(CoreMicrosimulation):
    default_tax_benefit_system = CountryTaxBenefitSystem
    default_tax_benefit_system_instance = system
    default_dataset = CPS_2023
    default_dataset_year = 2023
    default_role = "member"
    default_calculation_period = 2023
    default_input_period = 2023
    datasets = DATASETS

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        reform = create_structural_reforms_from_parameters(
            self.tax_benefit_system.parameters, "2023-01-01"
        )
        if reform is not None:
            self.apply_reform(reform)


class IndividualSim(CoreIndividualSim):  # Deprecated
    tax_benefit_system = CountryTaxBenefitSystem
    entities = {entity.key: entity for entity in entities}
    default_dataset = CPS_2023

    default_roles = dict(
        tax_unit="member",
        spm_unit="member",
        household="member",
        family="member",
    )
    required_entities = [
        "tax_unit",
        "spm_unit",
        "household",
        "family",
    ]
