variable_map

@dataclass
class IntegerBound:
    """IntegerBound

    Attributes:
        lower (int | float): lower bound of integer decision variable
        upper (int | float): upper bound of integer decision variable

    """

    lower: int | float
    upper: int | float

@dataclass
class VariableMap:
    """VarialeMap

    This class is used to manage the mapping between the label and the index of the decision variable.

    Attributes:
        var_map (dict[str, dict[tuple[int, ...], int]]): label, subscripts -> index (int)
        var_num (int): number of variables
        integer_bound (dict[str, dict[tuple[int, ...], IntegerBound]]): bounds of integer decision variables.

    """

    var_map: dict[str, dict[tuple[int, ...], int]]
    var_num: int
    integer_bound: dict[str, dict[tuple[int, ...], IntegerBound]]

    def add_variable(self, label: str, subscripts: tuple[int, ...]) -> int:
        """Add a variable to the variable map.

        Args:
            label (str): label of the variable
            subscripts (tuple[int, ...]): subscripts of the variable

        Returns:
            int: index of the variable in the variable map
        """
        if label not in self.var_map:
            self.var_map[label] = {}

        if subscripts not in self.var_map[label]:
            self.var_map[label][subscripts] = self.var_num
            self.var_num += 1
            return self.var_num - 1
        else:
            return self.var_map[label][subscripts]

    def add_int_variable(self, label, subscripts, lower, upper) -> int:
        """Add an integer variable to the variable map."""
        if label not in self.integer_bound:
            self.integer_bound[label] = {}
        self.integer_bound[label][subscripts] = IntegerBound(lower=lower, upper=upper)
        var_index = self.add_variable(label, subscripts)
        return var_index

    def to_serializable(self) -> dict:
        """Convert to a serializable object.

        Returns:
            dict: serializable object

        Examples:
            ```python
            varmap = VariableMap()
            varmap.add_variable("x", (0, 0))
            varmap.add_variable("x", (0, 1))
            varmap.to_serializable()
            # {
            #    "class": "VariableMap",
            #    "var_map": {
            #        "x": {
            #            "subscripts": [
            #                [0, 0],
            #                [0, 1]
            #            ],
            #            "index": [0, 1]
            #        }
            #    },
            #    "var_num": 2,
            #    "integer_bound": {}
            # }
            ```
        """
        varmap_serializable = {}
        for label, subscripts_map in self.var_map.items():
            varmap_serializable[label] = {"subscripts": [], "index": []}
            for subscripts, index in subscripts_map.items():
                varmap_serializable[label]["subscripts"].append(subscripts)
                varmap_serializable[label]["index"].append(index)
        integer_bound_serializable = {}
        for label, subscripts_map in self.integer_bound.items():
            integer_bound_serializable[label] = {"subscripts": [], "bound": []}
            for subscripts, bound in subscripts_map.items():
                integer_bound_serializable[label]["subscripts"].append(subscripts)
                integer_bound_serializable[label]["bound"].append(
                    {"lower": bound.lower, "upper": bound.upper}
                )
        return {
            "class": "VariableMap",
            "var_map": varmap_serializable,
            "var_num": self.var_num,
            "integer_bound": integer_bound_serializable,
        }

    @classmethod
    def from_serializable(cls, data: dict) -> "VariableMap":
        if data["class"] != "VariableMap":
            raise ValueError(f"Invalid class: {data['class']}")
        var_map = {}
        for label, subscripts_map in data["var_map"].items():
            var_map[label] = {
                tuple(subscripts): index
                for subscripts, index in zip(
                    subscripts_map["subscripts"], subscripts_map["index"]
                )
            }
        integer_bound = {}
        for label, subscripts_map in data["integer_bound"].items():
            integer_bound[label] = {
                tuple(subscripts): IntegerBound(bound["lower"], bound["upper"])
                for subscripts, bound in zip(
                    subscripts_map["subscripts"], subscripts_map["bound"]
                )
            }
        return VariableMap(var_map, data["var_num"], integer_bound)