name_and_value_fun.py

from typing import Union
import numpy as np
import casadi as cas
from typing import Iterable

from base import Variable, VariableList, VariableDict
from data_type import VariableFullCodeCollectionType, VariableCodeCollectionType, VariableCodeType
from base import Model


def get_variable_collection_values(vars_all: dict[str, Union[VariableDict, VariableList]],
                                   variable_codes: VariableFullCodeCollectionType) -> list[Union[float, cas.MX]]:
    """Get the values of a collection of variables
    """
    value = []
    for name_code in variable_codes:
        value.append(vars_all[name_code[0]].val[name_code[1:]])
    return value


def get_full_var_codes(vars_all: dict[str, Variable],
                       var_code_list: VariableCodeCollectionType) -> list[VariableCodeType]:
    """ Get the full variable codes from a list possibly including string
    :vars_all, is the variable dictionary from some model
    :var_code_list, possibly has both tuples of code and string
    :full_code_list, include only full code type
    """
    full_code_list = []
    for name in var_code_list:
        if isinstance(name, str):
            # name is like: str
            index_list = vars_all[name].index_list
            m = vars_all[name].m
            if index_list is None:
                full_code_list += [(name, i) for i in range(m)]
            else:
                for k, index in enumerate(index_list):
                    full_code_list += [(name, index, i) for i in range(m)]
        elif type(name[1]) in (list, tuple) and isinstance(name[1][0], str):
            # name is like (str, [str1, str2], [0, 1, 2])
            full_code_list += [(name[0], index, i) for i in name[2] for index in name[1]]
        elif type(name[1]) in (list, tuple) and isinstance(name[1][0], int):
            # name is like (str, [0, 1, 2])
            full_code_list += [(name[0], i) for i in name[1]]
        else:
            # name has been desired format
            full_code_list.append(name)

    return full_code_list


def assign_values_to_model(model, var_code_list, var_values):
    for i, var_code in enumerate(var_code_list):
        model.set_var_component_val(var_code, var_values[i])


def assign_symbols_to_model(model: Model, var_code_list: VariableFullCodeCollectionType, var_values):
    for i, var_code in enumerate(var_code_list):
        model.set_var_component_sym(var_code, var_values[i])


def revise_data_name(var_name_code: Iterable) -> Iterable:
    if isinstance(var_name_code, str):
        raise TypeError(f"You're having {var_name_code}. var_name_code cannot be string type.")

    if len(var_name_code) == 3 and np.isnan(var_name_code[2]):
        var_name_code = (var_name_code[0], var_name_code[1])
    return var_name_code