<commit_before>import logging
import warnings
from itertools import product
from collections import namedtuple

import numpy as np
import pandas as pd

from pgmpy.factors.base import BaseFactor
from pgmpy.utils import StateNameMixin
from pgmpy.extern import tabulate

State = namedtuple("State", ["var", "state"])


class DiscreteFactor(BaseFactor, StateNameMixin):
    """
    Base class for DiscreteFactor.
    """

    def __init__(self, variables, cardinality, values, state_names={}):
        """
        Initialize a `DiscreteFactor` class.

        Defined above, we have the following mapping from variable
        assignments to the index of the row vector in the value field:
        +-----+-----+-----+-------------------+
        |  x1 |  x2 |  x3 |    phi(x1, x2, x3)|
        +-----+-----+-----+-------------------+
        | x1_0| x2_0| x3_0|     phi.value(0)  |
        +-----+-----+-----+-------------------+
        | x1_0| x2_0| x3_1|     phi.value(1)  |
        +-----+-----+-----+-------------------+
        | x1_0| x2_1| x3_0|     phi.value(2)  |
        +-----+-----+-----+-------------------+
        | x1_0| x2_1| x3_1|     phi.value(3)  |
        +-----+-----+-----+-------------------+
        | x1_1| x2_0| x3_0|     phi.value(4)  |
        +-----+-----+-----+-------------------+
        | x1_1| x2_0| x3_1|     phi.value(5)  |
        +-----+-----+-----+-------------------+
        | x1_1| x2_1| x3_0|     phi.value(6)  |
        +-----+-----+-----+-------------------+
        | x1_1| x2_1| x3_1|     phi.value(7)  |
        +-----+-----+-----+-------------------+

        Parameters
        ----------
        variables: list, array-like
            List of variables on which the factor is to be defined i.e. scope of the factor.

        cardinality: list, array_like
            List of cardinalities/no.of states of each variable. `cardinality`
            array must have a value corresponding to each variable in
            `variables`.

        values: list, array_like
            List of values of factor.
            A DiscreteFactor's values are stored in a row vector in the value
            using an ordering such that the left-most variables as defined in
            `variables` cycle through their values the fastest. Please refer
            to examples for usage examples.

        Examples
        --------
        >>> import numpy as np
        >>> from pgmpy.factors.discrete import DiscreteFactor
        >>> phi = DiscreteFactor(['x1', 'x2', 'x3'], [2, 2, 2], np.ones(8))
        >>> phi
        <DiscreteFactor representing phi(x1:2, x2:2, x3:2) at 0x7f8188fcaa90>
        >>> print(phi)
        +------+------+------+-----------------+
        | x1   | x2   | x3   |   phi(x1,x2,x3) |
        |------+------+------+-----------------|
        | x1_0 | x2_0 | x3_0 |          1.0000 |
        | x1_0 | x2_0 | x3_1 |          1.0000 |
        | x1_0 | x2_1 | x3_0 |          1.0000 |
        | x1_0 | x2_1 | x3_1 |          1.0000 |
        | x1_1 | x2_0 | x3_0 |          1.0000 |
        | x1_1 | x2_0 | x3_1 |          1.0000 |
        | x1_1 | x2_1 | x3_0 |          1.0000 |
        | x1_1 | x2_1 | x3_1 |          1.0000 |
        +------+------+------+-----------------+
        """
        if isinstance(variables, str):
            raise TypeError("Variables: Expected type list or array like, got string")

        values = np.array(values, dtype=float)

        if len(cardinality) != len(variables):
            raise ValueError(
                "Number of elements in cardinality must be equal to number of variables"
            )

        if values.size != np.product(cardinality):
            raise ValueError(f"Values array must be of size: {np.product(cardinality)}")

        if len(set(variables)) != len(variables):
            raise ValueError("Variable names cannot be same")

        self.variables = list(variables)
        self.cardinality = np.array(cardinality, dtype=int)
        self.values = values.reshape(self.cardinality)

        # Set the state names
        super(DiscreteFactor, self).store_state_names(
            variables, cardinality, state_names
        )

    def scope(self):
        """
        Returns the scope of the factor i.e. the variables on which the factor is defined.

        Returns
        -------
        list: List of variable names in the scope of the factor.

        Examples
        --------
        >>> from pgmpy.factors.discrete import DiscreteFactor
        >>> phi = DiscreteFactor(['x1', 'x2', 'x3'], [2, 3, 2], np.ones(12))
        >>> phi.scope()
        ['x1', 'x2', 'x3']
        """
        return self.variables

    def get_cardinality(self, variables):
        """
        Returns the cardinality/no.of states of each variable in `variables`.

        Parameters
        ----------
        variables: list, array-like
                A list of variable names.

        Returns
        -------
        dict: Dictionary of the form {variable: variable_cardinality}

        Examples
        --------
        >>> from pgmpy.factors.discrete import DiscreteFactor
        >>> phi = DiscreteFactor(['x1', 'x2', 'x3'], [2, 3, 2], range(12))
        >>> phi.get_cardinality(['x1'])
        {'x1': 2}
        >>> phi.get_cardinality(['x1', 'x2'])
        {'x1': 2, 'x2': 3}
        """
        if isinstance(variables, str):
            raise TypeError("variables: Expected type list or array-like, got type str")

        if not all([var in self.variables for var in variables]):
            raise ValueError("Variable not in scope")

        return {var: self.cardinality[self.variables.index(var)] for var in variables}

    def get_value(self, **kwargs):
        """
        Returns the value of the given variable states.

        Parameters
        ----------
        kwargs: named arguments of the form variable=state_name
            Spcifies the state of each of the variable for which to get
            the value.

        Returns
        -------
        float: The value of states.

        Examples
        --------
        >>> from pgmpy.utils import get_example_model
        >>> model = get_example_model("asia")
        >>> phi = model.get_cpds("either").to_factor()
        >>> phi.get_value(lung="yes", tub="no", either="yes")
        1.0
        """
        for variable in kwargs.keys():
            if variable not in self.variables:
                raise ValueError(f"Factor doesn't have the variable: {variable}")

        index = []
        for var in self.variables:
            if var not in kwargs.keys():
                raise ValueError(f"Variable: {var} not found in arguments")
            elif isinstance(kwargs[var], str):
                index.append(self.name_to_no[var][kwargs[var]])
            else:
                logging.info(f"Using {var} state as number instead of name.")
                index.append(kwargs[var])
        return self.values[tuple(index)]

    def set_value(self, value, **kwargs):
        """
        Sets the probability value of the given variable states.

        Parameters<commit_msg>Fixes DiscreteFactor.get_value to work in the case when variables have int state names<commit_after>import logging
import warnings
from itertools import product
from collections import namedtuple

import numpy as np
import pandas as pd

from pgmpy.factors.base import BaseFactor
from pgmpy.utils import StateNameMixin
from pgmpy.extern import tabulate

State = namedtuple("State", ["var", "state"])


class DiscreteFactor(BaseFactor, StateNameMixin):
    """
    Base class for DiscreteFactor.
    """

    def __init__(self, variables, cardinality, values, state_names={}):
        """
        Initialize a `DiscreteFactor` class.

        Defined above, we have the following mapping from variable
        assignments to the index of the row vector in the value field:
        +-----+-----+-----+-------------------+
        |  x1 |  x2 |  x3 |    phi(x1, x2, x3)|
        +-----+-----+-----+-------------------+
        | x1_0| x2_0| x3_0|     phi.value(0)  |
        +-----+-----+-----+-------------------+
        | x1_0| x2_0| x3_1|     phi.value(1)  |
        +-----+-----+-----+-------------------+
        | x1_0| x2_1| x3_0|     phi.value(2)  |
        +-----+-----+-----+-------------------+
        | x1_0| x2_1| x3_1|     phi.value(3)  |
        +-----+-----+-----+-------------------+
        | x1_1| x2_0| x3_0|     phi.value(4)  |
        +-----+-----+-----+-------------------+
        | x1_1| x2_0| x3_1|     phi.value(5)  |
        +-----+-----+-----+-------------------+
        | x1_1| x2_1| x3_0|     phi.value(6)  |
        +-----+-----+-----+-------------------+
        | x1_1| x2_1| x3_1|     phi.value(7)  |
        +-----+-----+-----+-------------------+

        Parameters
        ----------
        variables: list, array-like
            List of variables on which the factor is to be defined i.e. scope of the factor.

        cardinality: list, array_like
            List of cardinalities/no.of states of each variable. `cardinality`
            array must have a value corresponding to each variable in
            `variables`.

        values: list, array_like
            List of values of factor.
            A DiscreteFactor's values are stored in a row vector in the value
            using an ordering such that the left-most variables as defined in
            `variables` cycle through their values the fastest. Please refer
            to examples for usage examples.

        Examples
        --------
        >>> import numpy as np
        >>> from pgmpy.factors.discrete import DiscreteFactor
        >>> phi = DiscreteFactor(['x1', 'x2', 'x3'], [2, 2, 2], np.ones(8))
        >>> phi
        <DiscreteFactor representing phi(x1:2, x2:2, x3:2) at 0x7f8188fcaa90>
        >>> print(phi)
        +------+------+------+-----------------+
        | x1   | x2   | x3   |   phi(x1,x2,x3) |
        |------+------+------+-----------------|
        | x1_0 | x2_0 | x3_0 |          1.0000 |
        | x1_0 | x2_0 | x3_1 |          1.0000 |
        | x1_0 | x2_1 | x3_0 |          1.0000 |
        | x1_0 | x2_1 | x3_1 |          1.0000 |
        | x1_1 | x2_0 | x3_0 |          1.0000 |
        | x1_1 | x2_0 | x3_1 |          1.0000 |
        | x1_1 | x2_1 | x3_0 |          1.0000 |
        | x1_1 | x2_1 | x3_1 |          1.0000 |
        +------+------+------+-----------------+
        """
        if isinstance(variables, str):
            raise TypeError("Variables: Expected type list or array like, got string")

        values = np.array(values, dtype=float)

        if len(cardinality) != len(variables):
            raise ValueError(
                "Number of elements in cardinality must be equal to number of variables"
            )

        if values.size != np.product(cardinality):
            raise ValueError(f"Values array must be of size: {np.product(cardinality)}")

        if len(set(variables)) != len(variables):
            raise ValueError("Variable names cannot be same")

        self.variables = list(variables)
        self.cardinality = np.array(cardinality, dtype=int)
        self.values = values.reshape(self.cardinality)

        # Set the state names
        super(DiscreteFactor, self).store_state_names(
            variables, cardinality, state_names
        )

    def scope(self):
        """
        Returns the scope of the factor i.e. the variables on which the factor is defined.

        Returns
        -------
        list: List of variable names in the scope of the factor.

        Examples
        --------
        >>> from pgmpy.factors.discrete import DiscreteFactor
        >>> phi = DiscreteFactor(['x1', 'x2', 'x3'], [2, 3, 2], np.ones(12))
        >>> phi.scope()
        ['x1', 'x2', 'x3']
        """
        return self.variables

    def get_cardinality(self, variables):
        """
        Returns the cardinality/no.of states of each variable in `variables`.

        Parameters
        ----------
        variables: list, array-like
                A list of variable names.

        Returns
        -------
        dict: Dictionary of the form {variable: variable_cardinality}

        Examples
        --------
        >>> from pgmpy.factors.discrete import DiscreteFactor
        >>> phi = DiscreteFactor(['x1', 'x2', 'x3'], [2, 3, 2], range(12))
        >>> phi.get_cardinality(['x1'])
        {'x1': 2}
        >>> phi.get_cardinality(['x1', 'x2'])
        {'x1': 2, 'x2': 3}
        """
        if isinstance(variables, str):
            raise TypeError("variables: Expected type list or array-like, got type str")

        if not all([var in self.variables for var in variables]):
            raise ValueError("Variable not in scope")

        return {var: self.cardinality[self.variables.index(var)] for var in variables}

    def get_value(self, **kwargs):
        """
        Returns the value of the given variable states. Assumes that the arguments
        specified are state names, and falls back to considering it as state no if
        can't find the state name.

        Parameters
        ----------
        kwargs: named arguments of the form variable=state_name
            Spcifies the state of each of the variable for which to get
            the value.

        Returns
        -------
        float: The value of states.

        Examples
        --------
        >>> from pgmpy.utils import get_example_model
        >>> model = get_example_model("asia")
        >>> phi = model.get_cpds("either").to_factor()
        >>> phi.get_value(lung="yes", tub="no", either="yes")
        1.0
        """
        for variable in kwargs.keys():
            if variable not in self.variables:
                raise ValueError(f"Factor doesn't have the variable: {variable}")

        index = []
        for var in self.variables:
            if var not in kwargs.keys():
                raise ValueError(f"Variable: {var} not found in arguments")
            else:
                try:
                    index.append(self.name_to_no[var][kwargs[var]])
                except KeyError:
                    logging.info(f"Using {var} state as number instead of name.")
                    index.append(kwargs[var])
        return self.values[tuple(index)]

    def set_value(self, value, **kwargs):
        """
        Sets the probability value of the given variable states.

        Parameters
{'fraction_non_alphanumeric': 0.17359905562113742, 'fraction_numerical': 0.028748003610860356, 'mean_word_length': 1.5020847810979847, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 101, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2876184', 'n_tokens_mistral': 4303, 'n_tokens_neox': 4115, 'n_words': 1192}