utils

def ext_vars_from(problem: jm.Problem) -> list[Variable]:
    """
    Extract variables from jijmodeling.Problem.

    Args:
        problem (jm.Problem): target problem

    Returns:
        list[Variable]: list of variables

    Examples:
        ```python
        import jijmodeling as jm
        import jijmodeling_transpiler.utils import ext_vars_from
        d = jm.Placeholder("d", dim=1)
        n = d.shape[0]
        k = jm.Placeholder("k")
        i, j = jm.Element("i", n), jm.Element("j", k)
        x, y = jm.Binary("x", shape=n), jm.Binary("y", shape=(n, k))
        problem = jm.Problem("test")
        problem += jm.Sum(i, x[i])
        problem += jm.Constraint("sample", jm.Sum(j, y[i, j]) == 1, forall=i)
        variables = ext_vars_from(problem)
        print(variables)
        # [x, y, d, n, k, i, j]
        ```
    """
    variables = jm.extract_variables(problem)
    return variables

def is_roughly_correct_type(obj: typ.Any, type_: typ.Any) -> bool:
    """Check if obj is roughly correct type.

    The roughly correct type means that the first element of the sequence and mapping only checked.

    Args:
        obj (Any): target object
        type_ (Any): target type

    Returns:
        bool: True if obj is roughly correct type, otherwise False
    """

    origin_type = typ.get_origin(type_)

    if origin_type is typ.Union:
        union_types = typ.get_args(type_)
        return any(
            is_roughly_correct_type(obj, union_type) for union_type in union_types
        )

    if origin_type is None:
        origin_type = type_
    if not isinstance(obj, origin_type):
        return False

    if isinstance(obj, Mapping):
        if len(obj) == 0:
            return True
        value = next(iter(obj.values()))
        key_type, value_type = typ.get_args(type_)

        value_check = is_roughly_correct_type(value, value_type)
        key_check = is_roughly_correct_type(next(iter(obj.keys())), key_type)
        return value_check and key_check

    if isinstance(obj, Sequence) and not isinstance(obj, str):
        if not obj:
            return True

        item_type = typ.get_args(type_)
        return is_roughly_correct_type(obj[0], item_type)

    return True

def ph_value_type_convert(
    ph_value: PH_VALUES_INTERFACE, variables: list[Variable]
) -> tuple[dict[str, np.ndarray], dict[jm.Placeholder, list]]:
    """
    Separate the input to ph_value into two types of data: data that can be.

    converted to numpy associated with a Placeholder and data that should be
    managed in a list associated with a JaggedArray.

    Args:
        ph_value (jijmodeling.type_annotations.PH_VALUES_INTERFACE): input ph_value data
        variables (list[Variable]): variable object list for the reference.

    Raises:
        TypeError: When the key of ph_value is not str or Variable.
        ValueError: When the data corresponding to the Placeholder or JaggedArray contained in the variable is not in ph_value.

    Returns:
        tuple[PLACEHOLDER_VALUES, dict[jm.Placeholder, list]]: placeholder's value and jagged array's value.

    Examples:
        ```python
        import jijmodeling as jm
        d = jm.Placeholder("d", dim=1)
        n = d.shape[0]
        k = jm.Placeholder("k")
        i = jm.Element("i", n)
        J = jm.JaggedArray("J", dim=2)
        ph_value = {"d": [1, 2, 3], "k": 3, "J": [[1, 2], [3, 4, 5]]}
        _ph_value, _jagged = ph_value_type_convert(ph_value, [J, d, i, k])
        print(_ph_value)
        # {d: array([1, 2, 3]), k: array(3)}
        print(_jagged)
        # {J: [[1, 2], [3, 4, 5]]}
        ```
    """

    ph_names: list[str] = [v.name for v in variables if isinstance(v, jm.Placeholder)]
    jagged_value = {}
    ph_array_value = {}
    for name in ph_names:
        if name not in ph_value:
            raise ValueError(f"'{name}' is needed in ph_value.")

    for name in ph_value.keys():
        try:
            ph_array_value[name] = np.array(ph_value[name])
        except ValueError:
            jagged_value[name] = ph_value[name]

    return ph_array_value, jagged_value

def raise_type_error(var_name: str, obj: typ.Any, type_: typ.Any) -> None:
    """Raise TypeError.

    Args:
        obj (Any): target object
        type_ (Any): target type
    """
    if is_roughly_correct_type(obj, type_):
        return
    raise TypeError(
        f"Type of `{var_name}` is {_type_name(obj)}, but {type_} is expected."
    )