"""Contains the python model representations and their APIs."""

from copy import deepcopy
from dataclasses import dataclass
from enum import Enum
from fractions import Fraction
from typing import Tuple, cast
from deprecated import deprecated

from stormvogel import parametric
import copy
import math
import uuid
import warnings

Number = int | float | Fraction

DEPRECATED_NAMES = [("Choice", "Choices"), ("Branch", "Branches")]


def __getattr__(name):
    for old_name, new_name in DEPRECATED_NAMES:
        if name == old_name:
            warnings.warn(
                f"The '{old_name}' class or function is renamed '{new_name}'",
                DeprecationWarning,
                stacklevel=2,
            )
            return globals()[new_name]
    raise AttributeError(f"module {__name__} has no attribute {name}")


@dataclass
class Interval:
    """represents an interval value for interval models

    Args:
        bottom: the bottom (left) element of the interval
        top: the top (right) element of the interval
    """

    bottom: Number
    top: Number

    def __init__(self, bottom: Number, top: Number):
        self.bottom = bottom
        self.top = top

    def __getitem__(self, idx):
        if idx == 0:
            return self.bottom
        elif idx == 1:
            return self.top
        else:
            raise IndexError(
                "Intervals only have two elements (the bottom and top element)"
            )

    def __lt__(self, other):
        if (self.bottom, self.top) < (other.bottom, other.top):
            return True
        return False

    def __str__(self):
        return f"[{self.bottom},{self.top}]"


Value = Number | parametric.Parametric | Interval


def value_to_string(
    n: Value, use_fractions: bool = True, round_digits: int = 4, denom_limit: int = 1000
) -> str:
    """Convert a Value to a string."""
    if isinstance(n, (int, float)):
        if math.isinf(float(n)):
            return "inf"
        if use_fractions:
            return str(Fraction(n).limit_denominator(denom_limit))
        else:
            return str(round(float(n), round_digits))
    elif isinstance(
        n, Fraction
    ):  # In the case of Fraction, a denominator of zero would have caused an error before.
        if use_fractions:
            return str(n.limit_denominator(denom_limit))
        else:
            return str(round(float(n), round_digits))
    elif isinstance(n, parametric.Parametric):
        return str(n)
    elif isinstance(n, Interval):
        return f"[{value_to_string(n.bottom, use_fractions, round_digits, denom_limit)},{value_to_string(n.top, use_fractions, round_digits, denom_limit)}]"
    else:
        return str(n)


class ModelType(Enum):
    """The type of the model."""

    # implemented
    DTMC = 1
    MDP = 2
    CTMC = 3
    POMDP = 4
    MA = 5
    HMM = 6


@dataclass
class Observation:
    """Represents an observation of a state (for POMDPs and HMMs)

    Args:
        observation: the observation of a state as an integer
    """

    observation: int
    valuations: dict[str, int | float | bool] | None = None

    def format(self):
        # TODO: add a good valuations format for visualizations
        return str(self.observation)

    def __str__(self):
        return f"Obs: {self.valuations if self.valuations else ''} ({self.observation})"


def _to_state_id(state_or_id: "State | int") -> int:
    if isinstance(state_or_id, State):
        return state_or_id.id
    else:
        return state_or_id


@dataclass(frozen=True)
class Action:
    """Represents an action, e.g., in MDPs.
        Note that this action object is completely independent of its corresponding branch.
        Their relation is managed by Choices.
        Two actions with the same labels are considered equal.

    Args:
        label: The label of this action. Corresponds to Storm label.
    """

    label: str | None

    def __post_init__(self):
        if self.label == "":
            object.__setattr__(self, "label", None)

    def __lt__(self, other):
        if not isinstance(other, Action):
            return NotImplemented
        return str(self.label) < str(other.label)

    def __str__(self):
        return f"Action with label {self.label}"


# The empty action. Used for DTMCs and empty action transitions in mdps.
EmptyAction = Action(None)


@dataclass()
class State[ValueType: Value]:
    """Represents a state in a Model.

    Args:
        labels: The labels of this state. Corresponds to Storm labels.
        valuations: The valuations of this state. Corresponds to Storm valuations/features.
        id: The id of this state.
        model: The model this state belongs to.
        observation: the observation of this state in case the model is a POMDP or HMM.
        name: the name of this state. (unique identifier in case ids change)
    """

    labels: list[str]
    valuations: dict[str, int | float | bool]
    id: int
    model: "Model[ValueType]"
    observation: Observation | list[tuple[ValueType, Observation]] | None
    name: str

    def __init__(
        self,
        labels: list[str],
        valuations: dict[str, int | float | bool],
        id: int,
        model,
        observation: Observation | list[tuple[ValueType, Observation]] | None = None,
        name: str | None = None,
    ):
        self.model = model

        if id in self.model.states.keys():
            raise RuntimeError(
                "There is already a state with this id. Make sure the id is unique."
            )

        self.labels = labels
        self.valuations = valuations
        self.id = id
        self.observation = observation

        # names must be unique, because they are the identifier of a state in case ids are changed
        if name is None:
            # if the user does not provide a name, we try to pick a random string
            while True:
                name = uuid.uuid4().hex[:8]
                if name not in self.model.used_names:
                    break
            self.model.used_names.add(name)
            self.name = name
        else:
            if name in self.model.used_names:
                raise RuntimeError(
                    "There is already a state with this name. Make sure the name is unique."
                )
            self.model.used_names.add(name)
            self.name = name

    def add_label(self, label: str):
        """adds a new label to the state"""
        if label in self.labels:
            raise RuntimeError(f"The label {label} is already present in this state.")

        self.labels.append(label)

    def set_observation(
        self, observation: Observation | list[tuple[ValueType, Observation]]
    ) -> Observation | list[tuple[ValueType, Observation]]:
        """sets the observation for this state"""
        if self.model.type in (ModelType.POMDP, ModelType.HMM):
            self.observation = observation
            return self.observation
        else:
            raise RuntimeError("The model this state belongs to is not a pomdp")

    def get_observation(self) -> Observation | list[tuple[ValueType, Observation]]:
        """gets the observation"""
        if self.model.supports_observations():
            if self.observation is not None:
                return self.observation
            else:
                raise RuntimeError(
                    "This state does not have an observation yet. Add one with the set_observation function."
                )
        else:
            raise RuntimeError(
                "The model this state belongs to does not support observations"
            )

    def set_choice(self, choice: "Choices | ChoiceShorthand"):
        """Set the choice for this state."""
        self.model.set_choice(self, choice)

    def add_choice(self, choices: "Choices | ChoiceShorthand"):
        """Add choices to this state."""
        self.model.add_choice(self, choices)

    def has_choices(self) -> bool:
        return self.get_choices() is not None

    def nr_choices(self) -> int:
        """The number of choices in this state."""
        choices = self.get_choices()
        return len(choices) if choices is not None else 1

    @deprecated(version="0.10.0", reason="Use get_choices() instead")
    def get_choice(self) -> "Choices | None":
        return self.get_choices()

    def get_choices(self) -> "Choices | None":
        """Gets the choices of this state."""
        if self.id in self.model.choices:
            return self.model.get_choices(self)
        else:
            return None

    def add_valuation(self, variable: str, value: int | bool | float):
        """Adds a valuation to the state."""
        self.valuations[variable] = value

    def get_valuation(self, variable: str) -> int | bool | float:
        return self.valuations[variable]

    def available_actions(self) -> list[Action]:
        """returns the list of all available actions in this state"""
        if self.model.supports_actions():
            choices = self.get_choices()
            if choices is not None:
                return choices.actions()
        return [EmptyAction]

    def get_outgoing_transitions(
        self, action: Action | None = None
    ) -> list[tuple[ValueType, "State[ValueType]"]] | None:
        """gets the outgoing transitions of this state (after a specific action)"""

        choice = self.get_choices()
        assert choice is not None

        # if the model supports actions we need to provide an action
        if action and self.model.supports_actions():
            if self.id in self.model.choices:
                return choice.choice[action].branch
        elif not action and self.model.supports_actions():
            raise RuntimeError("You need to provide a specific action")
        else:
            if self.id in self.model.choices:
                return choice.choice[EmptyAction].branch

    def is_absorbing(self) -> bool:
        """returns whether the state has a nonzero transition going to another state or not"""

        # if the state has no choice it is trivially true
        choice = self.get_choices()
        if choice is None:
            return True

        # for all actions we check if the state has outgoing transitions to a different state with value != 0
        for _, branch in choice:
            for transition in branch:
                if transition[1] != self:
                    return False
        return True

    def is_initial(self):
        """Returns whether this state is initial."""
        return "init" in self.labels

    def __str__(self):
        res = f"id: {self.id}, labels: {self.labels}, valuations: {self.valuations}"
        if self.model.supports_observations() and self.observation is not None:
            res += f", observation: {self.observation}"
        return res

    def __eq__(self, other):
        if not isinstance(other, State):
            return False

        if self.id != other.id:
            return False

        if (
            self.model.supports_observations()
            and self.observation is not None
            and other.observation is not None
            and self.observation != other.observation
        ):
            return False

        return (
            sorted(self.labels) == sorted(other.labels)
            and self.valuations == other.valuations
        )

    def __lt__(self, other):
        if not isinstance(other, State):
            return NotImplemented
        return self.id < other.id


@dataclass(order=True)
class Branches[ValueType: Value]:
    """Represents branches, which is a distribution over states.

    Args:
        branch: The branch as a list of tuples representing the transitions.
            The first element is the probability value and the second element is the target state.
    """

    branch: list[tuple[ValueType, State[ValueType]]]

    def __init__(self, *args):
        if len(args) == 1 and isinstance(args[0], list):
            self.branch = args[0]
        elif len(args) == 2:
            self.branch = [(args[0], args[1])]
        else:
            raise TypeError(
                "expects either (list of (value,state) tuples) or (value, state)"
            )

    def get_successors(self) -> set[int]:
        return set([b[1].id for b in self.branch])

    def sort_states(self):
        """sorts the branch list by states"""
        self.branch.sort(key=lambda x: x[1])

    def __str__(self):
        parts = []
        for prob, state in self.branch:
            parts.append(f"{prob} -> {state}")
        return ", ".join(parts)

    def __eq__(self, other):
        if not isinstance(other, Branches):
            return False

        return sorted(self.branch) == sorted(other.branch)

    def __add__(self, other):
        return Branches(self.branch + other.branch)

    def sum_probabilities(self) -> Number:
        sum = 0
        for prob, _ in self.branch:
            if not isinstance(prob, Number):
                raise ValueError(
                    "Summing probabilities is only defined for numerical values."
                )
            sum += prob
        return sum

    def __iter__(self):
        return iter(self.branch)


class Choices[ValueType: Value]:
    """Represents a choice, which map actions to branches.
        Note that an EmptyAction may be used if we want a non-action choice.
        Note that a single Choices might correspond to multiple 'arrows'.

    Args:
        choice: The choice dictionary. For each available action, we have a branch containing the transitions.
    """

    choice: dict[Action, Branches[ValueType]]

    def __init__(self, choice: dict[Action, Branches[ValueType]]):
        # Input validation, see RuntimeError.
        if len(choice) > 1 and EmptyAction in choice:
            raise RuntimeError(
                "It is impossible to create a choice that contains more than one action, and an emtpy action"
            )
        self.choice = choice

    def actions(self) -> list[Action]:
        """Returns the actions for the choices"""
        return list(self.choice.keys())

    def __str__(self):
        parts = []
        for action, branch in self:
            if action == EmptyAction:
                parts.append(f"{branch}")
            else:
                parts.append(f"{action} => {branch}")
        return "; ".join(parts + [])

    def has_empty_action(self) -> bool:
        # Note that we don't have to deal with the corner case where there are both empty and non-empty choices. This is dealt with at __init__.
        return self.choice.keys() == {EmptyAction}

    def __eq__(self, other):
        if not isinstance(other, Choices):
            return False

        if len(self.choice) != len(other.choice):
            return False

        for action, other_action in zip(
            sorted(self.choice.keys()), sorted(other.choice.keys())
        ):
            if not (
                action == other_action and self.choice[action] == other.choice[action]
            ):
                return False

        return True

    def sum_probabilities(self, action: Action) -> Number:
        return self.choice[action].sum_probabilities()

    def is_stochastic(self, epsilon: Number) -> bool:
        """returns whether the probabilities in the branches sum to 1"""
        return all([abs(self.sum_probabilities(a) - 1) <= epsilon for a in self.choice])

    def has_zero_transition(self) -> bool:
        for _, branch in self:
            for transition in branch:
                if isinstance(transition[0], Number) and transition[0] == 0:
                    return True
        return False

    def __getitem__(self, item):
        return self.choice[item]

    def __iter__(self):
        return iter(self.choice.items())

    def __len__(self) -> int:
        return len(self.choice)


ChoiceShorthand = (
    list[tuple[Value, State]]
    | list[tuple[Value, int]]
    | list[tuple[Action, State]]
    | list[tuple[Action, int]]
    | dict[Action, list[tuple[Value, State]]]
    | dict[Action, list[tuple[Value, int]]]
)


def choice_from_shorthand[ValueType: Value](
    shorthand: ChoiceShorthand, model: "Model[ValueType]"
) -> Choices[ValueType]:
    """Get a Choice object from a ChoiceShorthand. Use for all choices in DTMCs and for empty actions in MDPs.

    There are two possible ways to define a ChoiceShorthand.
    - using only the probability and the target state (implies default action when in an MDP).
    - using only the action and the target state (implies probability=1)."""

    if isinstance(shorthand, dict):
        # We convert the shorthand so that we have states instead of ids
        converted_shorthand = dict()
        for action, branch in shorthand.items():
            converted_shorthand[action] = []
            for value, state in branch:
                if isinstance(state, int):
                    if state not in model.states:
                        raise RuntimeError(
                            "This id is not known yet. Use states instead."
                        )
                    state = model.get_state_by_id(state)
                converted_shorthand[action].append((value, state))
        shorthand = converted_shorthand

        transition_content = dict()
        for action, branch in shorthand.items():
            assert isinstance(action, Action)
            transition_content[action] = Branches(branch)
        return Choices(transition_content)
    else:
        # We convert the shorthand so that we have states instead of ids
        converted_shorthand = []
        for value_or_action, state in shorthand:
            if isinstance(state, int):
                if state not in model.states:
                    raise RuntimeError("This id is not known yet. Use states instead.")
                state = model.get_state_by_id(state)
            converted_shorthand.append((value_or_action, state))
        shorthand = converted_shorthand

        # Check the type of the first element
        first_element = shorthand[0][0]
        if isinstance(first_element, Action):
            transition_content = dict()
            for action, state in shorthand:
                assert isinstance(action, Action)
                transition_content[action] = Branches(1, state)
            return Choices(transition_content)
        elif isinstance(first_element, Value):
            return Choices(
                {EmptyAction: Branches(cast(list[tuple[Value, State]], shorthand))}
            )
        raise RuntimeError(
            f"Type of {first_element} not supported in choice {shorthand}"
        )


@dataclass()
class RewardModel[ValueType: Value]:
    """Represents a state-exit reward model.
    Args:
        name: Name of the reward model.
        model: The model this rewardmodel belongs to.
        rewards: The rewards, the keys state action pairs.
    """

    name: str
    model: "Model"
    rewards: dict[Tuple[int, Action], ValueType]
    """Rewards dict. Hashed by state id and Action.
    The function update_rewards can be called to update rewards. After this, rewards will correspond to intermediate_rewards.
    Note that in models without actions, EmptyAction will be used here."""

    def __init__(
        self, name: str, model: "Model", rewards: dict[Tuple[int, Action], ValueType]
    ):
        self.name = name
        self.rewards = rewards
        self.model = model

    def set_from_rewards_vector(self, vector: list[ValueType]) -> None:
        """Set the rewards of this model according to a (stormpy) rewards vector."""
        combined_id = 0
        self.rewards = dict()
        for id, s in self.model:
            for a in s.available_actions():
                self.rewards[s.id, a] = vector[combined_id]
                combined_id += 1

    def get_state_reward(self, state: State) -> ValueType | None:
        """Gets the reward at said state or state action pair. Return None if no reward is present."""
        if self.model.supports_actions():
            raise RuntimeError(
                "This is a model with actions. Please call the get_state_action_reward(state, action) function instead"
            )
        if (state.id, EmptyAction) in self.rewards:
            return self.rewards[state.id, EmptyAction]
        else:
            return None

    def get_state_action_reward(self, state: State, action: Action) -> ValueType | None:
        """Gets the reward at said state or state action pair. Returns None if no reward was found."""
        if self.model.supports_actions():
            if action in state.available_actions():
                if (state.id, action) in self.rewards:
                    return self.rewards[state.id, action]
                else:
                    return None
            else:
                raise RuntimeError("This action is not available in this state")
        else:
            raise RuntimeError(
                "The model this rewardmodel belongs to does not support actions"
            )

    def set_state_reward(self, state: State, value: ValueType):
        """Sets the reward at said state. If the model has actions, try to use the empty state."""
        if self.model.supports_actions():
            self.set_state_action_reward(state, EmptyAction, value)
        else:
            self.rewards[state.id, EmptyAction] = value

    def set_state_action_reward(
        self,
        state: State,
        action: Action,
        value: ValueType,
    ):
        """sets the reward at said state action pair (in case of models with actions).
        If you disable auto_update_rewards, you will need to call update_intermediate_to
        """
        if self.model.supports_actions():
            if action in state.available_actions():
                self.rewards[state.id, action] = value
            else:
                raise RuntimeError("This action is not available in this state")
        else:
            raise RuntimeError(
                "The model this rewardmodel belongs to does not support actions"
            )

    def get_reward_vector(self) -> list[ValueType]:
        """Return the rewards in a (stormpy) vector format."""
        vector = []
        for id, s in self.model:
            for a in s.available_actions():
                reward = self.rewards[s.id, a]
                if reward is None:
                    raise RuntimeError(
                        "A reward was not set. You might want to call set_unset_rewards."
                    )
                vector.append(reward)
        return vector

    def set_unset_rewards(self, value: ValueType):
        """Fills up rewards that were not set yet with the specified value.
        Use this if converting (to stormpy) doesn't work because the reward vector does not have the expected length.
        """
        for id, s in self.model:
            for a in s.available_actions():
                if (s.id, a) not in self.rewards:
                    self.rewards[s.id, a] = value

    def __lt__(self, other) -> bool:
        if not isinstance(other, RewardModel):
            return NotImplemented
        return self.name < other.name

    def __eq__(self, other) -> bool:
        if not isinstance(other, RewardModel):
            return False

        return self.name == other.name and self.rewards == other.rewards

    def __iter__(self):
        return iter(self.rewards.items())


@dataclass
class Model[ValueType: Value]:
    """Represents a model.

    Args:
        type: The model type.
        states: The states of the model. The keys are the state's ids.
        choices: The choices of this model. The keys are the state ids.
        actions: The actions of the model, if this is a model that supports actions.
        rewards: The rewardsmodels of this model.
        exit_rates: The exit rates of the model, optional if this model supports rates.
        markovian_states: list of markovian states in the case of a ma.
    """

    type: ModelType
    # Both of these are hashed by the id of the state (=rowgroup number in the matrix)
    states: dict[int, State[ValueType]]
    choices: dict[int, Choices[ValueType]]
    actions: set[Action] | None
    observations: list[Observation] | None
    rewards: list[RewardModel]
    # In ctmcs we work with rate transitions but additionally we can optionally store exit rates (hashed by id of the state)
    exit_rates: dict[int, ValueType] | None
    # In ma's we keep track of markovian states
    markovian_states: list[State[ValueType]] | None
    parametric: bool | None
    interval: bool | None

    def __init__(self, model_type: ModelType, create_initial_state: bool = True):
        self.type = model_type
        self.choices = {}
        self.states = {}
        self.rewards = []
        self.parametric = None
        self.interval = None

        # Initialize actions if those are supported by the model type
        if self.supports_actions():
            self.actions = set()
        else:
            self.actions = None

        # Initialize rates if those are supported by the model type
        if self.supports_rates():
            self.exit_rates = {}
        else:
            self.exit_rates = None

        # Initialize possible observations if those are supported by the model type
        if self.type in (ModelType.POMDP, ModelType.HMM):
            self.observations = []
        else:
            self.observations = None

        # Initialize markovian states if applicable (in the case of MA's)
        if self.type == ModelType.MA:
            self.markovian_states = []
        else:
            self.markovian_states = None

        # We also keep track of used state names
        self.used_names = set()

        # We initialize the used id counter
        self.id_counter = 0

        # We initialize the observation id counter
        self.observation_id = 0

        # Add the initial state if specified to do so
        if create_initial_state:
            self.new_state(["init"])

    def summary(self):
        """Give a short summary of the model."""
        actions_bit = (
            f"{len(self.actions)} actions, " if self.actions is not None else ""
        )
        return (
            f"{self.type} model with {len(self.get_states())} states, "
            + actions_bit
            + f"and {len(self.get_labels())} distinct labels."
        )

    @deprecated(version="0.10.0", reason="use actions instead")
    def get_actions(self) -> set[Action] | None:
        """Return the actions of the model. Returns None if actions are not supported."""
        return self.actions

    def supports_actions(self) -> bool:
        """Returns whether this model supports actions."""
        return self.type in (ModelType.MDP, ModelType.POMDP, ModelType.MA)

    def supports_rates(self) -> bool:
        """Returns whether this model supports rates."""
        return self.type in (ModelType.CTMC, ModelType.MA)

    def supports_observations(self) -> bool:
        """Returns whether this model supports observations."""
        return self.type in (ModelType.POMDP, ModelType.HMM)

    def iterate_transitions(self) -> list[tuple[ValueType, State]]:
        """iterates through all transitions in all choices of the model"""
        transitions = []
        for transition in self.choices.values():
            for action, branch in transition:
                for transition in branch:
                    transitions.append(transition)
        return transitions

    def is_interval_model(self) -> bool:
        """Returns whether this model is an interval model, i.e., containts interval values)"""
        if self.interval is None:
            # Initialize
            for transition in self.iterate_transitions():
                if isinstance(transition[0], Interval):
                    self.interval = True
                    return self.interval
            self.interval = False
        return self.interval

    def is_parametric(self) -> bool:
        """Returns whether this model contains parametric transition values"""
        if self.parametric is None:
            # Initialize
            for transition in self.iterate_transitions():
                if isinstance(transition[0], parametric.Parametric):
                    self.parametric = True
                    return self.parametric
            self.parametric = False
        return self.parametric

    def is_stochastic(self, epsilon: Number = 0.000001) -> bool | None:
        """For discrete models: Checks if all sums of outgoing transition probabilities for all states equal 1, with at most epsilon rounding error.
        For continuous models: Checks if all sums of outgoing rates sum to 0
        """

        # if the model is parametric or an interval model, it should be trivially true, as the probabilities do
        # not even sum to a constant.
        if self.is_parametric() or self.is_interval_model():
            return True

        if not self.supports_rates():
            return all(
                [
                    self.get_choices(id).is_stochastic(epsilon)
                    for id, _ in self
                    if id in self.choices
                ]
            )
        else:
            for _, state in self:
                for action in state.available_actions():
                    sum_rates = 0
                    transitions = state.get_outgoing_transitions(action)
                    assert transitions is not None
                    for transition in transitions:
                        if isinstance(transition[0], Number):
                            sum_rates += transition[0]
                    if sum_rates != 0:
                        return False

        return True

    def normalize(self):
        """Normalizes a model (for states where outgoing transition probabilities don't sum to 1, we divide each probability by the sum)"""
        if self.is_parametric() or self.is_interval_model():
            raise RuntimeError(
                "normalize method undefined for parametric or interval models"
            )

        if not self.supports_rates():
            self.add_self_loops()
            for _, state in self:
                for action in state.available_actions():
                    # we first calculate the sum
                    sum_prob = 0
                    transitions = state.get_outgoing_transitions(action)
                    assert transitions is not None
                    for tuple in transitions:
                        if isinstance(tuple[0], Number):
                            sum_prob += tuple[0]

                    # then we divide each value by the sum
                    new_transitions = []
                    for tuple in transitions:
                        if isinstance(tuple[0], Number):
                            normalized_transition = (
                                tuple[0] / sum_prob,
                                tuple[1],
                            )
                            new_transitions.append(normalized_transition)
                    self.choices[state.id].choice[action].branch = new_transitions
        else:
            # for ctmcs and mas we currently only add self loops
            self.add_self_loops()

    def get_sub_model(self, states: list[State], normalize: bool = True) -> "Model":
        """Returns a submodel of the model based on a collection of states.
        The states in the collection are the states that stay in the model."""
        sub_model = copy.deepcopy(self)
        remove = []
        for _, state in sub_model:
            if state not in states:
                remove.append(state)
        for state in remove:
            sub_model.remove_state(state, normalize=False)

        if normalize:
            sub_model.normalize()
        return sub_model

    def parameter_valuation(self, values: dict[str, Number]) -> "Model":
        """evaluates all parametric transitions with the given values and returns the induced model"""
        evaluated_model = copy.deepcopy(self)
        for state, transition in evaluated_model.choices.items():
            for action, branch in transition:
                new_branch = []
                for tup in branch:
                    if isinstance(tup[0], parametric.Parametric):
                        tup = (tup[0].evaluate(values), tup[1])
                    new_branch.append(tup)
                evaluated_model.choices[state][action].branch = new_branch

        return evaluated_model

    def get_choice_id(self, state: State[ValueType], action: Action) -> int | None:
        """
        We calculate the appropriate choice id for a given state and action.
        This choice id corresponds to the row of the matrix in the underlying matrix.
        """
        choice_id = 0
        for s in self.states.values():
            if s == state:
                for a in s.available_actions():
                    if a == action:
                        return choice_id
                    choice_id += 1
                return None
            choice_id += s.nr_choices()
        return None

    @deprecated(version="0.10.0", reason="Use get_choice_id() instead.")
    def get_state_action_id(
        self, state: State[ValueType], action: Action
    ) -> int | None:
        return self.get_choice_id(state, action)

    def get_state_action_pair(self, id: int) -> tuple[State[ValueType], Action] | None:
        """Given an state action id, we return the corresponding state action pair
        Each state action pair corresponds to an id (which also corresponds to the row of the matrix).
        """
        i = 0
        for _, s in self:
            for a in s.available_actions():
                if id == i:
                    return (s, a)
                i += 1

    def add_self_loops(self):
        """adds self loops to all states that do not have an outgoing transition"""
        for _, state in self:
            if (
                state.get_choices() is None
            ):  # TODO what if the state has a choice but it is empty?
                self.set_choice(
                    state, [(float(0) if self.supports_rates() else float(1), state)]
                )

    def add_valuation_at_remaining_states(
        self, variables: list[str] | None = None, value: int | bool | float = 0
    ):
        """sets (dummy) value to variables in all states where they don't have a value yet"""

        # we either set it at all variables or just at a given subset of variables
        if variables is None:
            v = self.get_variables()
        else:
            v = variables

        # we set the values
        for _, state in self:
            for var in v:
                if var not in state.valuations.keys():
                    state.valuations[var] = value

    def unassigned_variables(self) -> list[tuple[State[ValueType], str]]:
        """we return a list of tuples of state variable pairs that are unassigned"""
        variables = self.get_variables()

        # we check all variables in all states
        unassigned = []
        for _, state in self:
            for variable in variables:
                if variable not in state.valuations.keys():
                    unassigned.append((state, variable))
        return unassigned

    def all_states_outgoing_transition(self) -> bool:
        """checks if all states have a choice"""
        for _, state in self:
            if state.get_choices() is None:  # TODO what if there is an empty choice
                return False
        return True

    def all_non_init_states_incoming_transition(self) -> bool:
        """checks if all states except the initial state have an incoming transition"""
        # Iterate over all transitions and remove all successor states as having an ingoing transition
        remaining_states = set(self.states.keys())
        for transition in self.iterate_transitions():
            remaining_states.discard(transition[1].id)
        for s in remaining_states:
            # Initial states could have no ingoing transition
            if not self.states[s].is_initial():
                return False
        return True

    def has_zero_transition(self) -> bool:
        """checks if the model has transitions with probability zero"""
        for _, choice in self.choices.items():
            if choice.has_zero_transition():
                return True
        return False

    def add_markovian_state(self, markovian_state: State):
        """adds a state to the markovian states (in case of markov automatas)"""
        if self.type == ModelType.MA and self.markovian_states is not None:
            self.markovian_states.append(markovian_state)
        else:
            raise RuntimeError("This model is not a MA")

    def set_choice(
        self, s: State, choices: Choices[ValueType] | ChoiceShorthand
    ) -> None:
        """Set the choice for a state."""
        if not isinstance(choices, Choices):
            choices = choice_from_shorthand(choices, self)

        if choices.has_zero_transition() and not self.supports_rates():
            raise RuntimeError("All transition probabilities should be nonzero.")

        if self.actions is not None and EmptyAction in choices.choice.keys():
            self.actions.add(EmptyAction)
        self.choices[s.id] = choices

    def add_choice(
        self, s: State, choices: Choices[ValueType] | ChoiceShorthand
    ) -> None:
        """Add new choices from a state to the model. If no choice currently exists, the result will be the same as set_choice."""

        if not isinstance(choices, Choices):
            choices = choice_from_shorthand(choices, self)

        if choices.has_zero_transition() and not self.supports_rates():
            raise RuntimeError("All transition probabilities should be nonzero.")

        try:
            existing_choices = self.get_choices(s)
        except KeyError:
            # Empty choices case, act like set_choice.
            self.set_choice(s, choices)
            return

        if not self.supports_actions():
            self.choices[s.id].choice[EmptyAction].branch += choices[EmptyAction].branch
        else:
            # Adding a choice is only valid if they are both empty or both non-empty.
            if not choices.has_empty_action() and existing_choices.has_empty_action():
                raise RuntimeError(
                    "You cannot add a choice with an non-empty action to a choice which has an empty action. Use set_choice instead."
                )
            if choices.has_empty_action() and not existing_choices.has_empty_action():
                raise RuntimeError(
                    "You cannot add a choice with an empty action to a choice which has no empty action. Use set_choice instead."
                )

            # Empty action case, add the branches together.
            if choices.has_empty_action():
                self.choices[s.id].choice[EmptyAction] += choices[EmptyAction]
            else:
                for action, branch in choices:
                    assert self.actions is not None
                    if action not in self.actions:
                        self.actions.add(action)
                    self.choices[s.id].choice[action] = branch

    def get_successor_states(self, state_or_id: State[ValueType] | int) -> set[int]:
        """Returns the set of successors of state_or_id."""
        result = set()
        for _, branches in self.get_choices(state_or_id):
            result |= branches.get_successors()
        return result

    def get_choices(self, state_or_id: State[ValueType] | int) -> Choices[ValueType]:
        """Get the choices at state s. Throws a KeyError if not present."""
        return self.choices[_to_state_id(state_or_id)]

    @deprecated(version="0.10.0", reason="use get_choices instead")
    def get_choice(self, state_or_id: State[ValueType] | int) -> Choices[ValueType]:
        return self.get_choices(state_or_id)

    def get_branches(self, state_or_id: State | int) -> Branches[ValueType]:
        """Get the branch at state s. Only intended for emtpy choices, otherwise a RuntimeError is thrown."""
        choice = self.get_choices(_to_state_id(state_or_id)).choice
        if EmptyAction not in choice:
            raise RuntimeError("Called get_branch on a non-empty choice.")
        return choice[EmptyAction]

    @deprecated(version="0.10.0", reason="use get_branches instead")
    def get_branch(self, state_or_id: State[ValueType] | int) -> Branches[ValueType]:
        """Get the branch at state s. Only intended for emtpy choices, otherwise a RuntimeError is thrown."""
        return self.get_branches(state_or_id)

    def get_action_with_label(self, label: str | None) -> Action | None:
        """Get the action with provided label"""
        assert self.actions is not None
        for action in self.actions:
            if action.label == label:
                return action
        return None

    def reassign_ids(self):
        """Reassigns the ids of states, choices and rates to be in order again.
        Mainly useful to keep consistent with storm."""

        warnings.warn(
            "Warning: Using this can cause problems in your code if there are existing references to states by id."
        )

        # we change the ids in the dictionaries of the model object
        self.states = {
            new_id: value
            for new_id, (old_id, value) in enumerate(sorted(self.states.items()))
        }

        self.choices = {
            new_id: value
            for new_id, (old_id, value) in enumerate(sorted(self.choices.items()))
        }

        if self.supports_rates and self.exit_rates is not None:
            self.exit_rates = {
                new_id: value
                for new_id, (old_id, value) in enumerate(
                    sorted(self.exit_rates.items())
                )
            }

        # we change the ids in the states themselves
        for index, state in enumerate(self.states.values()):
            state.id = index

        # we also reset the counter
        self.id_counter = max(self.states.keys()) + 1

    def remove_state(
        self, state: State, normalize: bool = True, reassign_ids: bool = False
    ):
        """Properly removes a state, it can optionally normalize the model and reassign ids automatically."""

        if state in self.states.values():
            # we remove the state from the transitions
            # first we remove incoming transitions of the state
            remove_actions_index = []
            for index, transition in self.choices.items():
                for action, branch in transition:
                    for index_tuple, tuple in enumerate(branch):
                        # remove the tuple if it references the state
                        if tuple[1].id == state.id:
                            self.choices[index].choice[action].branch.pop(index_tuple)

                    # if we have empty actions we need to remove those as well (later)
                    if branch.branch == []:
                        remove_actions_index.append((action, index))
            # here we remove those empty actions (this needs to happen after the other for loops)
            for action, index in remove_actions_index:
                self.choices[index].choice.pop(action)
                # if we have no actions at all anymore, delete the choice
                if self.choices[index].choice == {} and not index == state.id:
                    self.choices.pop(index)

            # we remove the choice with the corresponding id
            self.choices.pop(state.id)

            # We remove the state
            self.states.pop(state.id)

            # we remove the exit rates from the state when applicable
            if self.supports_rates and self.exit_rates is not None:
                self.exit_rates.pop(state.id)

            # we remove the state from the markovian state list when applicable
            if self.type == ModelType.MA and self.markovian_states is not None:
                if state in self.markovian_states:
                    self.markovian_states.remove(state)

            # we normalize the model if specified to do so
            if normalize:
                self.normalize()

            # we reassign the ids if specified to do so
            if reassign_ids:
                self.reassign_ids()
        else:
            raise RuntimeError("This state is not part of this model.")

    def remove_transitions_between_states(
        self, state0: State, state1: State, normalize: bool = True
    ):
        """
        Remove the transition(s) that start in state0 and go to state1.
        Only works on models that don't support actions.
        """
        if not self.supports_actions():
            for tuple in self.choices[state0.id][EmptyAction]:
                if tuple[1] == state1:
                    self.choices[state0.id].choice[EmptyAction].branch.remove(tuple)
            # if we have empty objects we need to remove those as well
            if self.choices[state0.id].choice[EmptyAction].branch == []:
                self.choices.pop(state0.id)

            if normalize:
                self.normalize()
        else:
            raise RuntimeError(
                "This method only works for models that don't support actions."
            )

    def new_action(self, label: str | None = None) -> Action:
        """Creates a new action and returns it."""
        if not self.supports_actions():
            raise RuntimeError(
                "Called new_action on a model that does not support actions"
            )
        assert self.actions is not None
        action = Action(label)
        self.actions.add(action)
        return action

    def get_action(self, name: str) -> Action:
        """Gets an existing action."""
        if not self.supports_actions():
            raise RuntimeError(
                "Called get_action on a model that does not support actions"
            )
        assert self.actions is not None
        found_action = None
        for action in self.actions:
            if action.label is not None and name in action.label:
                found_action = action
                break

        if found_action is None:
            raise RuntimeError(f"Could not find an action with name {name}")

        return found_action

    def action(self, label: str | None) -> Action:
        """New action or get action if it exists."""
        if not self.supports_actions():
            raise RuntimeError(
                "Called method action on a model that does not support actions"
            )
        assert self.actions is not None
        action = Action(label)

        if action not in self.actions:
            self.new_action(label)
        return action

    def new_observation(
        self,
        valuations: dict[str, int | float | bool] | None = None,
        observation_id: int | None = None,
    ) -> Observation:
        if not self.supports_observations():
            raise RuntimeError(
                "Called new_observation on a model that does not support observations"
            )
        assert self.observations is not None
        if observation_id is None:
            observation_id = self.observation_id
            self.observation_id += 1

        if (found_obs := self.get_observation(observation_id)) is not None:
            raise RuntimeError(
                f"An observation with id {observation_id} already exists, namely {found_obs}."
            )

        observation = Observation(observation_id, valuations)
        self.observations.append(observation)
        return observation

    def get_observation(self, observation_id: int) -> Observation | None:
        if not self.supports_observations():
            raise RuntimeError(
                "Called get_observation on a model that does not support observations"
            )
        assert self.observations is not None
        for observation in self.observations:
            if observation.observation == observation_id:
                return observation
        return None

    def observation(
        self,
        observation_id: int | None = None,
        valuations: dict[str, int | float | bool] | None = None,
    ) -> Observation:
        """New observation or get observation if it exists."""
        if not self.supports_observations():
            raise RuntimeError(
                "Called method observation on a model that does not support observations"
            )
        assert self.observations is not None

        if (
            observation_id is not None
            and (found_obs := self.get_observation(observation_id)) is not None
        ):
            if valuations is not None and found_obs.valuations != valuations:
                raise RuntimeError(
                    f"An observation with id {observation_id} already exists, namely {found_obs}, but has different valuations than the provided ones."
                )
            return found_obs

        return self.new_observation(valuations, observation_id)

    def new_state(
        self,
        labels: list[str] | str | None = None,
        valuations: dict[str, int | bool | float] | None = None,
        observation: Observation | list[tuple[ValueType, Observation]] | None = None,
        name: str | None = None,
    ) -> State[ValueType]:
        """Creates a new state and returns it."""

        # we set the id from the counter and increase it
        state_id = self.id_counter
        self.id_counter += 1

        if isinstance(labels, list):
            state = State(
                labels,
                valuations or {},
                state_id,
                self,
                observation=observation,
                name=name,
            )
        elif isinstance(labels, str):
            state = State(
                [labels],
                valuations or {},
                state_id,
                self,
                observation=observation,
                name=name,
            )
        elif labels is None:
            state = State(
                [], valuations or {}, state_id, self, observation=observation, name=name
            )
        else:
            raise RuntimeError("Unknown type for labels.")

        self.states[state_id] = state

        return state

    def has_only_number_values(self) -> bool:
        """Checks whether all transition values, observation probabilities, exit rates, markovian states in this model are numbers."""
        for choice in self.choices.values():
            for action, branch in choice:
                for value, _ in branch:
                    if not isinstance(value, Number):
                        return False

        if self.exit_rates is not None:
            for rate in self.exit_rates.values():
                if not isinstance(rate, Number):
                    return False

        if self.supports_observations() and self.observations is not None:
            for state in self.states.values():
                observation = state.observation
                if isinstance(observation, list):
                    for value, _ in observation:
                        if not isinstance(value, Number):
                            return False

        return True

    def make_observations_deterministic(self, reassign_ids: bool = False):
        """
        In case of POMDPs or HMMs, makes the observations deterministic by splitting states with
        multiple observations into multiple states with single observations.
        """
        if not self.has_only_number_values():
            raise RuntimeError("This method only works for models with Number values.")
        if not self.supports_observations():
            raise RuntimeError(
                "This method only works for models that support observations."
            )

        for state in list(self.states.values()):
            observation = state.observation
            if isinstance(observation, list):
                new_states_distribution: list[tuple[ValueType, State]] = []
                new_choice = state.get_choice()

                # Create new states for each observation possible in this state
                for prob, obs in observation:
                    new_state = self.new_state(
                        state.labels,
                        state.valuations,
                        obs,
                        f"{state.name} {{{str(obs.observation)}}}",
                    )
                    if new_choice is not None:
                        new_state.set_choice(deepcopy(new_choice))

                    new_states_distribution.append((prob, new_state))

                for src_state_idx, choice in self.choices.items():
                    for action, branch in choice:
                        for value, dest_state in branch:
                            if dest_state.id == state.id:
                                # we need to redirect this transition to the new states
                                # we first remove the old transition
                                self.choices[src_state_idx].choice[
                                    action
                                ].branch.remove((value, dest_state))
                                # then we add the new transitions
                                for prob, new_state in new_states_distribution:
                                    new_value = cast(Number, value) * cast(Number, prob)
                                    self.choices[src_state_idx].choice[
                                        action
                                    ].branch.append(
                                        (cast(ValueType, new_value), new_state)
                                    )

                self.remove_state(state, reassign_ids=reassign_ids)

    def get_states_with_label(self, label: str) -> list[State]:
        """Get all states with a given label."""
        # TODO: slow, not sure if that will become a problem though
        collected_states = []
        for _, state in self:
            if label in state.labels:
                collected_states.append(state)
        return collected_states

    def get_state_by_id(self, state_id: int) -> State:
        """Get a state by its id."""
        if state_id not in self.states:
            raise RuntimeError("Requested a non-existing state")
        return self.states[state_id]

    def get_state_by_name(self, state_name) -> State | None:
        """Get a state by its name."""
        names = {state.name: state for state in self.states.values()}
        if state_name not in names:
            raise RuntimeError("Requested a non-existing state")

        return names[state_name]

    def get_initial_state(self) -> State:
        """Gets the initial state (contains label "init", or has id 0)."""

        if len(self.states) == 0:
            raise RuntimeError(
                "This model does not have an initial state because it does not have any states."
            )

        init = []
        for _, state in self:
            if state.is_initial():
                init.append(state)

        if len(init) == 0:
            raise RuntimeError("This model does not have an initial state.")
        elif len(init) == 1:
            return init[0]
        else:
            raise RuntimeError(
                "This model has multiple initial states, which is not supported."
            )

    def get_ordered_labels(self) -> list[list[str]]:
        """Get all the labels of this model, ordered by id.
        IMPORTANT: If a state has no label, then a value '' is inserted!"""
        return [(s.labels if len(s.labels) > 0 else []) for s in self.get_states()]

    def get_labels(self) -> set[str]:
        """Get all labels in states of this Model."""
        collected_labels: set[str] = set()
        for _id, state in self:
            collected_labels = collected_labels | set(state.labels)
        return collected_labels

    def get_variables(self) -> set[str]:
        """gets the set of all variables present in this model (corresponding to valuations"""
        variables: set[str] = set()
        for _id, state in self.states.items():
            variables = variables | set(state.valuations.keys())
        return variables

    def get_default_rewards(self) -> RewardModel:
        """Gets the default reward model, throws a RuntimeError if there is none."""
        if len(self.rewards) == 0:
            raise RuntimeError("This model has no reward models.")
        return self.rewards[0]

    def get_rewards(self, name: str) -> RewardModel:
        """Gets the reward model with the specified name. Throws a RuntimeError if said model does not exist."""
        for model in self.rewards:
            if model.name == name:
                return model
        raise RuntimeError(f"Reward model {name} not present in model.")

    def get_parameters(self) -> set[str]:
        """Returns the set of parameters of this model"""
        parameters = set()
        for transition in self.iterate_transitions():
            if isinstance(transition[0], parametric.Parametric):
                parameters = parameters.union(transition[0].get_variables())
        return parameters

    def get_states(self) -> list[State]:
        return list(self.states.values())

    def new_reward_model(self, name: str) -> RewardModel:
        """Creates a reward model with the specified name, adds it and returns it."""
        for model in self.rewards:
            if model.name == name:
                raise RuntimeError(f"Reward model {name} already present in model.")
        reward_model = RewardModel(name, self, {})
        self.rewards.append(reward_model)
        return reward_model

    def get_rate(self, state: State) -> ValueType:
        """Gets the rate of a state."""
        if not self.supports_rates() or self.exit_rates is None:
            raise RuntimeError("Cannot get a rate of a discrete-time model.")
        return self.exit_rates[state.id]

    def set_rate(self, state: State, rate: ValueType):
        """Sets the rate of a state."""
        if not self.supports_rates() or self.exit_rates is None:
            raise RuntimeError("Cannot set a rate of a discrete-time model.")
        self.exit_rates[state.id] = rate

    @deprecated(version="0.10.0", reason="use type instead.")
    def get_type(self) -> ModelType:
        """Gets the type of this model"""
        return self.type

    @property
    def nr_states(self) -> int:
        """
        Returns the number of states in this model.
        Note that not all states need to be reachable.
        """
        return len(self.states)

    @property
    def nr_choices(self) -> int:
        """
        Returns the number of choices in the model (summed over all states).
        """
        return sum(state.nr_choices() for state in self.states.values())

    def to_dot(self) -> str:
        """Generates a dot representation of this model."""
        dot = "digraph model {\n"
        for state_id, state in self:
            dot += f'{state_id} [ label = "{state_id}: {", ".join(state.labels)}" ];\n'
        for state_id, transition in self.choices.items():
            for action, branch in transition:
                if action != EmptyAction:
                    dot += f'{state_id} [ label = "", shape=point ];\n'
        for state_id, transition in self.choices.items():
            for action, branch in transition:
                if action == EmptyAction:
                    # Only draw probabilities
                    for prob, target in branch:
                        dot += f'{state_id} -> {target.id} [ label = "{prob}" ];\n'
                else:
                    # Draw actions, then probabilities
                    dot += f'{state_id} -> {state_id} [ label = "{action.label}" ];\n'
                    for prob, target in branch:
                        dot += f'{state_id} -> {target.id} [ label = "{prob}" ];\n'

        dot += "}"
        return dot

    def __str__(self) -> str:
        res = [f"{self.type}"]
        res += ["", "States:"] + [f"{state}" for (_id, state) in self]
        res += ["", "Choices:"] + [
            f"{id}: {transition}" for (id, transition) in self.choices.items()
        ]

        if self.supports_rates() and self.exit_rates is not None:
            res += ["", "Exit rates:"] + [f"{self.exit_rates}"]

        if (
            self.supports_actions()
            and self.supports_rates()
            and self.markovian_states is not None
        ):
            markovian_states = [state.id for state in self.markovian_states]
            res += ["", "Markovian states:"] + [f"{markovian_states}"]

        return "\n".join(res)

    def __eq__(self, other) -> bool:
        if not isinstance(other, Model):
            return False

        return (
            self.actions == other.actions
            and self.type == other.type
            and self.states == other.states
            and self.choices == other.choices
            and sorted(self.rewards) == sorted(other.rewards)
            and self.exit_rates == other.exit_rates
            and self.markovian_states == other.markovian_states
        )

    def __getitem__(self, state_id: int):
        return self.states[state_id]

    def __iter__(self):
        return iter(self.states.items())


def new_dtmc(create_initial_state: bool = True) -> Model:
    """Creates a DTMC."""
    return Model(ModelType.DTMC, create_initial_state)


def new_mdp(create_initial_state: bool = True) -> Model:
    """Creates an MDP."""
    return Model(ModelType.MDP, create_initial_state)


def new_ctmc(create_initial_state: bool = True) -> Model:
    """Creates a CTMC."""
    return Model(ModelType.CTMC, create_initial_state)


def new_pomdp(create_initial_state: bool = True) -> Model:
    """Creates a POMDP."""
    return Model(ModelType.POMDP, create_initial_state)


def new_hmm(create_initial_state: bool = True) -> Model:
    """Creates a HMM."""
    return Model(ModelType.HMM, create_initial_state)


def new_ma(create_initial_state: bool = True) -> Model:
    """Creates a MA."""
    return Model(ModelType.MA, create_initial_state)


def new_model(modeltype: ModelType, create_initial_state: bool = True) -> Model:
    """More general model creation function"""
    return Model(modeltype, create_initial_state)
