from .transition import Transition
from ..utils.utils import format_simulation_output, create_definitions, apply_overrides, evaluate, compute_simulation_dates, apply_initial_conditions, multinomial
from .simulation_output import Trajectory
from .simulation_results import SimulationResults
import numpy as np 
import pandas as pd
from ..population.population import Population, load_epydemix_population
from typing import List, Dict, Optional, Union, Any, Callable, Tuple
import copy
import inspect


class EpiModel:
    """
    EpiModel: A compartmental epidemic model simulator
    
    Example:
        >>> model = EpiModel(compartments=["S", "I", "R"], parameters={"beta": 0.3, "gamma": 0.1})
        >>> model.add_transition(source="S", target="I", params=("beta", "I"))
        >>> model.add_transition(source="I", target="R", params="gamma")
        >>> results = model.run_simulations(
        ...     start_date="2020-01-01",
        ...     end_date="2020-12-31",
        ...     Nsim=100
        ... )
    """

    def __init__(self, 
                 name : str = "EpiModel",
                 compartments: Optional[List] = None, 
                 parameters: Optional[Dict] = None,
                 population_name: str = "epydemix_population", 
                 population_data_path: Optional[str] = None, 
                 contact_layers: Optional[List] = None, 
                 contacts_source: Optional[str] = None, 
                 age_group_mapping: Optional[Dict] = None, 
                 supported_contacts_sources: Optional[List] = None, 
                 use_default_population: bool = True) -> None:
            """
            Initializes the EpiModel instance.

            Args:
                name (str, optional): The name of the epidemic model. Defaults to "EpiModel".
                compartments (list, optional): A list of compartments for the model. Defaults to an empty list if None.
                parameters (dict, optional): A dictionary of parameters for the model. Defaults to an empty dictionary if None.
                population_name (str, optional): The name of the population to load or create. Defaults to 'epydemix_population'.
                population_data_path (str or None, optional): The path to the population data.
                contact_layers (list or None, optional): A list of contact layers to load for the population. Defaults to 
                    ["school", "work", "home", "community"] if None.
                contacts_source (str or None, optional): The source of contact data. If None, the function will automatically detect the source.
                age_group_mapping (dict or None, optional): A mapping of age groups for the population.
                supported_contacts_sources (list or None, optional): A list of supported contact data sources. Defaults to 
                    ["prem_2017", "prem_2021", "mistry_2021"] if None.
                use_default_population (bool, optional): If True, creates a default population; if False, tries to load the population from the provided path and population name. Defaults to True.
            Returns:
                None
            """
            # Initialize core attributes
            self.name = name
            self.transitions = {}
            self.transitions_list = []
            self.interventions = []
            self.compartments = []
            self.compartments_idx = {}
            self.transitions_idx = {}
            self.transition_functions = {}
            self.parameters = {}
            self.definitions = {}
            self.overrides = {}
            self.Cs = {}

            # Handle default empty lists for compartments and contact layers
            if compartments is None:
                compartments = []
            self.add_compartments(compartments)
            
            if contact_layers is None:
                contact_layers = ["school", "work", "home", "community"]

            # Add parameters if provided
            if parameters:
                self.add_parameter(parameters_dict=parameters)

            # Handle default contact sources if not provided
            if supported_contacts_sources is None:
                supported_contacts_sources = ["prem_2017", "prem_2021", "mistry_2021"]

            # Load or create population based on use_default_population flag
            self.population = self._load_or_create_population(
                population_name,
                population_data_path,
                contact_layers,
                contacts_source,
                age_group_mapping,
                supported_contacts_sources,
                use_default_population
            )

            # Initalize functions to compute transition rates
            self.register_transition_kind(kind="spontaneous", function=compute_spontaneous_transition_rate)
            self.register_transition_kind(kind="mediated", function=compute_mediated_transition_rate)


    def __repr__(self) -> str:
        """
        Returns a string representation of the EpiModel object,
        summarizing its name, compartments, transitions, parameters, and population.

        Returns:
            str: String representation of the EpiModel object.
        """
        repr_str = f"EpiModel(name='{self.name}')\n"

        # Compartments summary
        repr_str += f"Compartments: {len(self.compartments)}\n"
        if len(self.compartments) > 0:
            repr_str += "  " + ", ".join(self.compartments) + "\n"
        else:
            repr_str += "  No compartments defined\n"

        # Transitions summary
        repr_str += f"Transitions: {len(self.transitions_list)}\n"
        if len(self.transitions_list) > 0:
            repr_str += "  Transitions between compartments:\n"
            for tr in self.transitions_list:
                repr_str += f"    {tr.source} -> {tr.target}, params: {tr.params} (kind: {tr.kind})\n"
        else:
            repr_str += "  No transitions defined\n"

        # Parameters summary
        repr_str += f"Parameters: {len(self.parameters)}\n"
        if len(self.parameters) > 0:
            repr_str += "  Model parameters:\n"
            for param_name, param_value in self.parameters.items():
                repr_str += f"    {param_name}: {param_value}\n"
        else:
            repr_str += "  No parameters defined\n"

        # Population summary
        repr_str += f"Population: {self.population.name}\n"
        repr_str += f"  Population size: {sum(self.population.Nk)} individuals\n"
        repr_str += f"  Demographic groups: {len(self.population.Nk_names)}\n"
        if len(self.population.Nk_names) > 0:
            repr_str += "    " + ", ".join([str(name) for name in self.population.Nk_names]) + "\n"

        return repr_str


    def _load_or_create_population(self, 
                                population_name: str, 
                                population_data_path: Optional[str], 
                                contact_layers: List, 
                                contacts_source: Optional[str], 
                                age_group_mapping: Optional[Dict], 
                                supported_contacts_sources: List, 
                                use_default_population: bool) -> Population:
        """
        Loads a population from a file or creates a default population if specified.

        Args:
            population_name (str): The name of the population.
            population_data_path (str or None): The path to the population data file. If None, data may be fetched online.
            contact_layers (list): A list of contact layers to load for the population (e.g., "home", "work", "school").
            contacts_source (str or None): The source of contact data. If None, load_population will auto-determine the source.
            age_group_mapping (dict or None): A mapping for age groups. If None, a default mapping is used.
            supported_contacts_sources (list): A list of supported contact data sources (e.g., "prem_2017", "prem_2021").
            use_default_population (bool): If True, creates a default population with a single contact matrix and population size of 100000. 
                If False, attempts to load population data from a local or online source.

        Returns:
            Population: A population object, either created manually or loaded from data.
        """
        if use_default_population:
            # Create a default population manually
            population = Population(name=population_name)
            population.add_contact_matrix(np.array([[1.]]))  # Default contact matrix
            population.add_population(np.array([100000]))    # Default population size
            return population
        else:
            # Load population using load_population, which handles both online and local sources
            return load_epydemix_population(
                population_name, 
                contacts_source=contacts_source, 
                path_to_data=population_data_path, 
                layers=contact_layers, 
                age_group_mapping=age_group_mapping, 
                supported_contacts_sources=supported_contacts_sources
            )


    def import_epydemix_population(self, 
                            population_name: str, 
                            population_data_path: Optional[str] = None, 
                            contact_layers: Optional[List] = None, 
                            contacts_source: Optional[str] = None, 
                            age_group_mapping: Optional[Dict] = None, 
                            supported_contacts_sources: Optional[List] = None) -> None:
        """
        Sets or updates the population for the model.

        Args:
            population_name (str): The name of the population to set.
            population_data_path (str or None, optional): The path to the population data file. If None, data may be fetched online.
            contact_layers (list or None, optional): A list of contact layers to load for the population. Defaults to 
                ["school", "work", "home", "community"] if None.
            contacts_source (str or None, optional): The source of contact data. If None, the function will attempt to determine the source automatically.
            age_group_mapping (dict or None, optional): A mapping for age groups. If None, a default mapping is used.
            supported_contacts_sources (list or None, optional): A list of supported contact data sources (e.g., "prem_2017", "prem_2021"). 
                Defaults to ["prem_2017", "prem_2021", "mistry_2021"] if None.

        Returns:
            None
        """
        # Handle default contact layers if not provided
        if contact_layers is None:
            contact_layers = ["school", "work", "home", "community"]

        # Handle default contact sources if not provided
        if supported_contacts_sources is None:
            supported_contacts_sources = ["prem_2017", "prem_2021", "mistry_2021"]

        # Load a new population using the same logic as in the constructor
        self.population = load_epydemix_population(
            population_name, 
            contacts_source=contacts_source, 
            path_to_data=population_data_path, 
            layers=contact_layers, 
            age_group_mapping=age_group_mapping, 
            supported_contacts_sources=supported_contacts_sources
        )


    def add_compartments(self, compartments: Union[List, object]) -> None:
        """
        Adds compartments to the epidemic model.

        Args:
            compartments (list or object): A list of compartments or a single compartment to be added.

        Returns:
            None
        """
        # Ensure we are always working with a list, even if a single compartment is passed
        if not isinstance(compartments, (list, tuple)):
            compartments = [compartments]

        # Add compartments to the model
        self.compartments.extend(compartments)

        # Determine the current maximum index in compartments_idx or set to -1 if empty
        max_idx = max(self.compartments_idx.values(), default=-1)

        # Add each compartment with corresponding transitions and index
        for i, compartment in enumerate(compartments, start=1):
            self.transitions[compartment] = []  # Initialize empty transitions for each compartment
            self.compartments_idx[compartment] = max_idx + i  # Assign a unique index to each compartment

    @property
    def n_compartments(self) -> int:
        """Get total number of compartments."""
        return len(self.compartments)

    
    def clear_compartments(self) -> None: 
        """
        This method resets the `compartments` list and clears the `compartments_idx` dictionary.

        Returns:
            None
        """
        self.compartments = []
        self.compartments_idx = {}


    def add_parameter(self, 
                    parameter_name: Optional[str] = None, 
                    value: Any = None, 
                    parameters_dict: Optional[Dict] = None) -> None:
        """ 
        Adds new parameters to the model.

        Args:
            parameter_name (str, optional): The name of the parameter to add. Ignored if `parameters_dict` is provided.
            value (any, optional): The value of the parameter to add. Ignored if `parameters_dict` is provided.
            parameters_dict (dict, optional): A dictionary of parameter names and values. If provided, 
                `parameter_name` and `value` are ignored and the parameters are updated using this dictionary.

        Raises:
            ValueError: If neither `parameter_name`/`value` nor `parameters_dict` is provided.
        
        Returns:
            None
        """
        if parameters_dict:
            self.parameters.update(parameters_dict)
        elif parameter_name is not None and value is not None:
            self.parameters.update({parameter_name: value})
        else:
            raise ValueError("Either parameter_name and value or parameters_dict must be provided.")


    def get_parameter(self, parameter_name: str) -> Any:
        """
        Retrieves the value of a specified parameter.

        Args:
            parameter_name (str): The name of the parameter to retrieve.

        Returns:
            any: The value of the specified parameter.

        Raises:
            KeyError: If the parameter with the given name is not found.
        """
        return self.parameters[parameter_name]


    def delete_parameter(self, parameter_name: str) -> Any:
        """
        Deletes a parameter from the dictionary of the model's parameters.

        Args:
            parameter_name (str): The name of the parameter to delete.

        Returns:
            any: The value of the deleted parameter.

        Raises:
            KeyError: If the parameter with the given name is not found.
        """
        return self.parameters.pop(parameter_name)


    def clear_parameters(self) -> None:
        """
        Clears all parameters from the model's parameters dictionary.

        Returns:
            None
        """
        self.parameters = {}


    def override_parameter(self, start_date: str, end_date: str, parameter_name: str, value: Any) -> None:
        """
        Adds an override for a parameter with the specified name.

        Args:
            start_date (str): The start date for the override period.
            end_date (str): The end date for the override period.
            parameter_name (str): The name of the parameter to override.
            value (any): The value to override the parameter with.

        Returns:
            None
        """
        override_dict = {"start_date": start_date, "end_date": end_date, "value": value}

        if parameter_name in self.overrides:
            self.overrides[parameter_name].append(override_dict)
        else:
            self.overrides[parameter_name] = [override_dict]


    def delete_override(self, parameter_name: str) -> None:
        """
        Deletes overrides for a specific parameter.

        Args:
            parameter_name (str): The name of the parameter whose overrides are to be deleted.

        Returns:
            None
        """
        self.overrides.pop(parameter_name, None)


    def clear_overrides(self) -> None:
        """
        Clears all parameter overrides from the model.

        Returns:
            None
        """
        self.overrides = {}


    def add_transition(self, source: str, target: str, kind: str, params: Any) -> None:
        """
        Adds a transition to the epidemic model.

        Args:
            source (str): The source compartment of the transition.
            target (str): The target compartment of the transition.
            kind (str): The kind of transition (e.g., spontaneous or mediated).
            params (Any): The parameters involved in the transition.

        Raises:
            ValueError: If the source or target is not in the compartments list.

        Returns:
            None
        """
        missing_compartments = [comp for comp in [source, target] if comp and comp not in self.compartments]
        if missing_compartments:
            raise ValueError(f"These compartments are not in the compartments list: {', '.join(missing_compartments)}")
        
        transition = Transition(source=source, target=target, kind=kind, params=params)
        self.transitions_list.append(transition)
        self.transitions[source].append(transition)

        transition_name = f"{source}_to_{target}"
        if transition_name not in self.transitions_idx:
            self.transitions_idx[transition_name] = len(self.transitions_idx)


    def register_transition_kind(self, kind: str, function: Callable):
        """
        Registers a new transition kind providing a function to compute the rate of the transition.

        Args:
            kind (str): The kind of transition (e.g., spontaneous or mediated).
            function (Callable): The function to register.

        Returns:
            None
        """
        validate_transition_function(function)
        self.transition_functions[kind] = function

    @property
    def n_transitions(self) -> int:
        """Get total number of transitions."""
        return len(self.transitions_list)
        
    
    def clear_transitions(self) -> None:
        """
        Clears all transitions from the model.

        This method resets the `transitions_list` and reinitializes the `transitions` dictionary, 
        clearing all existing transitions while keeping the structure intact for each compartment.

        Returns:
            None
        """
        self.transitions_list = []
        self.transitions = {comp: [] for comp in self.compartments}


    def add_intervention(self, 
                        layer_name: str, 
                        start_date: Union[str, pd.Timestamp], 
                        end_date: Union[str, pd.Timestamp], 
                        reduction_factor: Optional[float] = None, 
                        new_matrix: Optional[np.ndarray] = None, 
                        name: str = "") -> None:
        """
        Adds an intervention to the epidemic model.

        Args:
            layer_name (str): The name of the layer to which the intervention applies.
            start_date (str or datetime): The start date of the intervention.
            end_date (str or datetime): The end date of the intervention.
            reduction_factor (float, optional): The factor by which to reduce the contact matrix. Default is None.
            new_matrix (np.ndarray, optional): A new contact matrix to use during the intervention. Default is None.
            name (str, optional): The name of the intervention. Default is an empty string.

        Raises:
            ValueError: If neither reduction_factor nor new_matrix is provided.

        Returns:
            None
        """
        # Convert start_date and end_date to pandas datetime
        start_date = pd.to_datetime(start_date)
        end_date = pd.to_datetime(end_date)

        # Check if both reduction_factor and new_matrix are None
        if reduction_factor is None and new_matrix is None:
            raise ValueError("Either reduction_factor or new_matrix must be provided")

        # Append the intervention to the interventions list
        self.interventions.append({
            "layer": layer_name, 
            "start_date": start_date, 
            "end_date": end_date, 
            "reduction_factor": reduction_factor, 
            "new_matrix": new_matrix,
            "name": name
        })

 
    def clear_interventions(self) -> None:
        """
        Clears all interventions from the model.

        This method resets the `interventions` list, removing all existing interventions.

        Returns:
            None
        """
        self.interventions = []


    def apply_intervention(self, intervention: Dict, simulation_dates: List[pd.Timestamp]) -> None:
        """
        Applies an intervention to the contact matrices for specified simulation dates.

        Args:
            intervention (dict): A dictionary containing intervention details with the following keys:
                - "layer" (str): The name of the layer to which the intervention applies.
                - "start_date" (pd.Timestamp): The start date of the intervention.
                - "end_date" (pd.Timestamp): The end date of the intervention.
                - "reduction_factor" (float, optional): The factor by which to reduce the contact matrix.
                - "new_matrix" (np.ndarray, optional): A new contact matrix to use during the intervention.
            simulation_dates (list of pd.Timestamp): A list of dates for which the simulation is run.

        Raises:
            ValueError: If neither reduction_factor nor new_matrix is provided in the intervention.

        Returns:
            None
        """
        # Early validation of the intervention inputs
        reduction_factor = intervention.get("reduction_factor")
        new_matrix = intervention.get("new_matrix")
        
        if reduction_factor is None and new_matrix is None:
            raise ValueError("Intervention must have either a reduction_factor or a new_matrix")

        # Extract commonly used values
        layer = intervention["layer"]
        start_date = intervention["start_date"]
        end_date = intervention["end_date"]

        # Apply the intervention to the relevant dates
        for date in filter(lambda d: start_date <= d <= end_date, simulation_dates):
            if reduction_factor is not None:
                self.Cs[date][layer] *= reduction_factor
            else:  # If reduction_factor is None, we assume new_matrix is provided
                self.Cs[date][layer] = new_matrix


    def compute_contact_reductions(self, simulation_dates: List[pd.Timestamp]) -> None:
        """
        Computes the contact reductions for a population over the given simulation dates.

        This function applies interventions to the contact matrices and computes the overall contact matrix 
        for each date in the simulation period.

        Args:
            simulation_dates (list of pd.Timestamp): A list of dates over which the simulation is run.

        Returns:
            None: The function updates the instance variable `self.Cs` with the contact matrices for each date, 
                including the overall contact matrix after applying interventions.
        """
        # Initialize contact matrices for each simulation date, copying the population's initial contact matrices
        self.Cs = {date: {layer: np.copy(matrix) for layer, matrix in self.population.contact_matrices.items()} 
                for date in simulation_dates}

        # Apply interventions to the contact matrices
        for intervention in self.interventions:
            self.apply_intervention(intervention, simulation_dates)

        # Compute the overall contact matrix for each date by summing up the contact matrices across layers
        for date in self.Cs:
            self.Cs[date]["overall"] = np.sum(np.array(list(self.Cs[date].values())), axis=0)


    def create_default_initial_conditions(self, percentage_in_agents: float = 0.0005) -> Dict[str, np.ndarray]:
        """
        Creates default initial conditions for the epidemic model. If initial conditions are not provided, 
        assigns a percentage of the population to agent compartments and the rest to source compartments, 
        considering different age groups.

        Args:
            percentage_in_agents (float): Percentage of population to place in agent compartments. Defaults to 0.05%.
            initial_conditions (dict, optional): A dictionary specifying initial conditions. If None, this function creates default conditions.

        Returns:
            dict: A dictionary with initial conditions for each compartment, with values as arrays representing different age groups.
        """

        population = self.population.Nk

        # Initialize initial conditions dictionary
        initial_conditions_dict = {}

        # Total population for each age group
        total_population_per_age_group = np.array(population)

        # Get compartments that are agents in transitions
        agent_compartments = [tr.params[1] for tr in self.transitions_list if tr.kind == "mediated"]
        
        # Get compartments that are sources in transitions with agents
        source_compartments = [tr.source for tr in self.transitions_list if tr.kind == "mediated"]

        # Total number of agent compartments
        num_agent_compartments = len(agent_compartments)

        # Distribute population into agent compartments
        if num_agent_compartments > 0:
            population_in_agents_per_age_group = (total_population_per_age_group * percentage_in_agents).astype(int)
            population_per_agent_per_age_group = population_in_agents_per_age_group // num_agent_compartments

            for comp in agent_compartments:
                initial_conditions_dict[comp] = population_per_agent_per_age_group.copy()

        # Assign the rest of the population to source compartments
        remaining_population_per_age_group = total_population_per_age_group - np.sum(list(initial_conditions_dict.values()), axis=0)
        num_source_compartments = len(source_compartments)

        if num_source_compartments > 0:
            population_per_source_per_age_group = remaining_population_per_age_group // num_source_compartments

            for comp in source_compartments:
                initial_conditions_dict[comp] = population_per_source_per_age_group.copy()

        # If some compartments are missing, ensure they get zero population
        for comp in self.compartments:
            if comp not in initial_conditions_dict:
                initial_conditions_dict[comp] = np.zeros_like(total_population_per_age_group)

        return initial_conditions_dict


    def set_population(self, population: Population) -> None:
        """
        Sets the population for the epidemic model.

        Args:
            population (Population): The population object to set.

        Returns:
            None
        """
        self.population = population
    

    def run_simulations(self, 
                       start_date: Union[str, pd.Timestamp] = "2020-01-01", 
                       end_date: Union[str, pd.Timestamp] = "2020-12-31", 
                       initial_conditions_dict: Optional[Dict[str, np.ndarray]] = None, 
                       Nsim: int = 100, 
                       percentage_in_agents: float = 0.0005,
                       dt: Optional[float] = 1.,
                       resample_frequency: Optional[str] = "D",
                       resample_aggregation_compartments: Optional[Union[str, dict]] = "last",
                       resample_aggregation_transitions: Optional[Union[str, dict]] = "sum",
                       fill_method: Optional[str] = "ffill", 
                       apply_linear_approximation: bool = False, 
                       rng: Optional[np.random.Generator] = None) -> SimulationResults:
        """
        Simulates the epidemic model multiple times over the given time period.

        Args:
            start_date (str or pd.Timestamp): The start date of the simulation. Default is "2020-01-01".
            end_date (str or pd.Timestamp): The end date of the simulation. Default is "2020-12-31".
            initial_conditions_dict (dict, optional): A dictionary of initial conditions for the simulation.
            Nsim (int, optional): The number of simulation runs to perform (default is 100).
            percentage_in_agents (float, optional): The percentage of the population to initialize in the agents compartment.
            dt (float, optional): The time step for the simulation, expressed in days. Default is 1 (day).
            resample_frequency (str, optional): The frequency at which to resample the simulation results. Default is "D" (daily).
            resample_aggregation_compartments (str, optional): The aggregation method to use when resampling the compartments. Default is "last".
            resample_aggregation_transitions (str, optional): The aggregation method to use when resampling the transitions. Default is "sum".
            fill_method (str, optional): Method to fill NaN values after resampling. Default is "ffill".
            apply_linear_approximation (bool, optional): Whether to use linear approximation to the probabilities. Default is False.
            rng (np.random.Generator, optional): Random number generator. Default is None.

        Returns:
            SimulationResults: An object containing all simulation trajectories.

        Raises:
            RuntimeError: If the simulation fails.
        """
        
        if rng is None:
            rng = np.random.default_rng()

        # Run multiple simulations and collect trajectories
        try:
            # Pre-compute simulation dates, initial conditions, and contact matrices to speed up the simulation
            if initial_conditions_dict is None:
                initial_conditions_dict = self.create_default_initial_conditions(percentage_in_agents=percentage_in_agents)
            simulation_dates = compute_simulation_dates(start_date, end_date, dt=dt)
            self.compute_contact_reductions(simulation_dates)
            contact_matrices = [self.Cs[date] for date in simulation_dates]

            trajectories = []
            for _ in range(Nsim):
                trajectory = simulate(
                    self, 
                    start_date=start_date,
                    end_date=end_date,
                    dt=dt,
                    initial_conditions_dict=initial_conditions_dict,
                    percentage_in_agents=percentage_in_agents,
                    resample_frequency=resample_frequency,
                    resample_aggregation_compartments=resample_aggregation_compartments,
                    resample_aggregation_transitions=resample_aggregation_transitions,
                    fill_method=fill_method, 
                    apply_linear_approximation=apply_linear_approximation, 
                    rng=rng, 
                    simulation_dates=simulation_dates,
                    contact_matrices=contact_matrices
                )
                trajectories.append(trajectory)
        except Exception as e:
            raise RuntimeError(f"Simulation failed: {str(e)}") from e

        # Return SimulationResults with all trajectories
        return SimulationResults(
            trajectories=trajectories,
            parameters=self.parameters
        )


def simulate(epimodel, 
             start_date: Union[str, pd.Timestamp] = "2020-01-01", 
             end_date: Union[str, pd.Timestamp] = "2020-12-31", 
             initial_conditions_dict: Optional[Dict[str, np.ndarray]] = None, 
             percentage_in_agents: float = 0.0005,
             dt: Optional[float] = 1.,
             resample_frequency: Optional[str] = "D",
             resample_aggregation_compartments: Optional[Union[str, dict]] = "last",
             resample_aggregation_transitions: Optional[Union[str, dict]] = "sum",
             fill_method: Optional[str] = "ffill",
             apply_linear_approximation: bool = False,
             rng: Optional[np.random.Generator] = None,
             contact_matrices: Optional[List[Dict[str, np.ndarray]]] = None,
             simulation_dates: Optional[List[pd.Timestamp]] = None,
             **kwargs) -> Trajectory:
    """
    Runs a simulation of the epidemic model over the specified simulation dates.

    Args:
        epimodel (EpiModel): The epidemic model instance to simulate.
        start_date (str or pd.Timestamp): The start date of the simulation. Default is "2020-01-01".
        end_date (str or pd.Timestamp): The end date of the simulation. Default is "2020-12-31".
        initial_conditions_dict (dict, optional): A dictionary of initial conditions for the simulation.
        percentage_in_agents (float, optional): The percentage of the population to initialize in the agents compartment.
        dt (float, optional): The time step for the simulation, expressed in days. Default is 1 (day).
        resample_frequency (str, optional): The frequency at which to resample the simulation results. Default is "D" (daily).
        resample_aggregation_compartments (str, optional): The aggregation method to use when resampling the compartments. Default is "last".
        resample_aggregation_transitions (str, optional): The aggregation method to use when resampling the transitions. Default is "sum".
        fill_method (str, optional): The method to use when filling NaN values after resampling. Default is "ffill".
        apply_linear_approximation (bool, optional): Whether to use linear approximation to the probabilities. Default is False.
        rng (np.random.Generator, optional): Random number generator. Default is None.
        contact_matrices (list, optional): A list of contact matrices for the simulation. Default is None.
        simulation_dates (list, optional): A list of simulation dates. Default is None.
        **kwargs: Additional parameters to overwrite model parameters during the simulation.

    Returns:
        Trajectory: The trajectory of the simulation

    Raises:
        ValueError: If the model has no transitions defined.
    """
    if rng is None:
        rng = np.random.default_rng()

    # check that the model has transitions
    if len(epimodel.transitions_list) == 0:
        raise ValueError("The model has no transitions defined. Please add transitions before running simulations.")
    
    # Compute the simulation dates if not provided
    if simulation_dates is None:
        simulation_dates = compute_simulation_dates(start_date, end_date, dt=dt)

    # Compute initial conditions if not provided
    if initial_conditions_dict is None:
        initial_conditions_dict = epimodel.create_default_initial_conditions(percentage_in_agents=percentage_in_agents)
        
    # Compute the contact reductions based on the interventions if not provided
    if contact_matrices is None:
        epimodel.compute_contact_reductions(simulation_dates)
        contact_matrices = [epimodel.Cs[date] for date in simulation_dates]

    # Update parameters if any are provided via kwargs (needed for calibration purposes)
    parameters = epimodel.parameters.copy()
    parameters.update(kwargs)

    # Compute the definitions and apply overrides
    epimodel.definitions = create_definitions(parameters, len(simulation_dates), epimodel.population.Nk.shape[0])
    epimodel.definitions = apply_overrides(epimodel.definitions, epimodel.overrides, simulation_dates)

    # Initialize population in different compartments and demographic groups
    initial_conditions = apply_initial_conditions(epimodel, initial_conditions_dict)

    # Run simulation with pre-computed contacts
    compartments_evolution, transitions_evolution = stochastic_simulation(
        T=len(simulation_dates),
        contact_matrices=contact_matrices,  
        epimodel=epimodel,
        parameters=epimodel.definitions,
        initial_conditions=initial_conditions,
        dt=dt, 
        apply_linear_approximation=apply_linear_approximation, 
        rng=rng
    )

    # Format the simulation output
    results = format_simulation_output(compartments_evolution, transitions_evolution, 
                                       epimodel.compartments_idx, epimodel.transitions_idx, 
                                       epimodel.population.Nk_names)
    trajectory = Trajectory(compartments=results["compartments"], transitions=results["transitions"], 
                            dates=simulation_dates, compartment_idx=epimodel.compartments_idx, 
                            transitions_idx=epimodel.transitions_idx, parameters=epimodel.definitions)

    # Only resample if necessary
    if resample_frequency is not None:
        # Check if resampling is needed (simulation dates frequency != requested frequency)
        sim_freq = pd.infer_freq(simulation_dates)
        if sim_freq != resample_frequency:
            trajectory.resample(resample_frequency, 
                                resample_aggregation_compartments, 
                                resample_aggregation_transitions, 
                                fill_method)
    return trajectory


def stochastic_simulation(T: int,
                         contact_matrices: List[Dict[str, np.ndarray]],
                         epimodel,
                         parameters: Dict,
                         initial_conditions: np.ndarray,
                         dt: float, 
                         apply_linear_approximation: bool = False, 
                         rng: Optional[np.random.Generator] = None) -> np.ndarray:
    """
    Run a stochastic simulation of the epidemic model.
    
    Args:
        T: Number of time steps
        contact_matrices: Pre-computed list of contact matrices dictionaries (key is the layer, value is the contact matrix)
        epimodel: The epidemic model
        parameters: Model parameters
        initial_conditions: Initial population distribution
        dt: Time step size
        apply_linear_approximation (bool, optional): Whether to use linear approximation to the probabilities. Default is False.
        rng (np.random.Generator, optional): Random number generator. Default is None.
    """
    if rng is None:
        rng = np.random.default_rng()

    # Pre-allocate arrays
    N = len(epimodel.population.Nk)
    C = len(epimodel.compartments)
    
    compartments_evolution = np.zeros((T + 1, C, N), dtype=np.float64)
    transitions_evolution = np.zeros((T, epimodel.n_transitions, N), dtype=np.float64)
    compartments_evolution[0] = initial_conditions
    
    # Pre-compute population sizes and create views for better performance
    pop_sizes = epimodel.population.Nk
    comp_indices = epimodel.compartments_idx
    
    # Pre-allocate arrays for rates and transitions
    rates = np.zeros((C, N), dtype=np.float64)
    new_pop = np.zeros((C, N), dtype=np.float64)
    probs_out = np.empty(C, dtype=np.float64)

    # create a dictionary to store the data needed for the transitions
    system_data = {
        "parameters": parameters, 
        "t": 0,
        "comp_indices": comp_indices,
        "contact_matrix": None,
        "pop": None,
        "pop_sizes": pop_sizes,
        "dt": dt
        }
    
    # Simulate each time step
    for t in range(T):
        # Update system data with current state
        system_data.update({
            "t": t,
            "contact_matrix": contact_matrices[t],
            "pop": compartments_evolution[t]
        })

        new_pop[:] = compartments_evolution[t]
        
        for comp in epimodel.compartments:
            transitions = epimodel.transitions[comp]
            if not transitions: 
                continue
                
            rates.fill(0)
            source_idx = comp_indices[comp]
            current_pop = compartments_evolution[t, source_idx]
            mask = np.arange(rates.shape[0]) != source_idx

            if not np.any(current_pop):
                continue

            for tr in transitions:
                target_idx = comp_indices[tr.target]
                trans_rate = epimodel.transition_functions[tr.kind](
                    tr.params, system_data
                )
                rates[target_idx] += trans_rate
            
            delta = np.array([
                multinomial(n, r, source_idx, mask, dt, apply_linear_approximation=apply_linear_approximation, rng=rng, probs_out=probs_out) if n > 0 else np.zeros(C)
                for n, r in zip(current_pop, rates.T)
            ])

            # Store transition counts
            for tr in transitions:
                tr_name = f"{tr.source}_to_{tr.target}"
                tr_idx = epimodel.transitions_idx[tr_name]
                transitions_evolution[t, tr_idx] += delta[:, epimodel.compartments_idx[tr.target]]
            
            # Update populations
            np.subtract(new_pop[source_idx], np.sum(delta, axis=1), out=new_pop[source_idx])
            new_pop += delta.T
        
        compartments_evolution[t + 1] = new_pop
    
    return compartments_evolution[1:], transitions_evolution


def compute_spontaneous_transition_rate(params, data): 
    """
    Compute the rate of a spontaneous transition.

    Args:
        params: The parameters of the transition provided by the user.
        data: A dictionary containing the data needed for the transition. 
            - parameters: The model parameters
            - t: The current time step
    Returns:
        The rate of the transition
    """
    if isinstance(params, str):
        env_copy = copy.deepcopy(data["parameters"])
        rate_eval = evaluate(expr=params, env=env_copy)[data["t"]]
        return rate_eval
    else:
        return params


def compute_mediated_transition_rate(params, data): 
    """
    Compute the rate of a mediated transition.

    Args:
        params: The parameters of the transition provided by the user.
        data: A dictionary containing the data needed for the transition. 
            - parameters: The model parameters
            - t: The current time step
            - comp_indices: The indices of the compartments
            - contact_matrix: The contact matrix
            - pop: The population in different compartments
            - pop_sizes: The population sizes
    """
    if isinstance(params[0], str):
        env_copy = copy.deepcopy(data["parameters"])
        rate_eval = evaluate(expr=params[0], env=env_copy)[data["t"]]
    else: 
        rate_eval = params[0]
    agent_idx = data["comp_indices"][params[1]]
    interaction = np.sum(
            data["contact_matrix"]["overall"] * data["pop"][agent_idx] / data["pop_sizes"], 
            axis=1
        )
    return rate_eval * interaction


def validate_transition_function(func: Callable) -> None:
    """
    Validates that a transition rate function has the correct signature and parameters.
    
    Args:
        func: The transition rate function to validate
        
    Raises:
        ValueError: If the function signature doesn't match requirements
    """
    # Get function signature
    sig = inspect.signature(func)
    params = sig.parameters
    
    # Check number of parameters
    if len(params) != 2:
        raise ValueError(
            f"Transition function must take exactly 2 parameters. Got {len(params)}: {list(params.keys())}"
        )
    
    # Check parameter names
    expected_params = ['params', 'data']
    actual_params = list(params.keys())
    
    if actual_params != expected_params:
        raise ValueError(
            f"Transition function must have parameters named {expected_params}. Got {actual_params}"
        )
    