# -*- coding: utf-8 -*-
# Copyright (C) 2018-2026 Intel Corporation
# SPDX-License-Identifier: Apache-2.0


from typing import Optional
from collections.abc import Callable
from openvino import Op, Type, Shape, Tensor, PartialShape, TensorVector


class PostponedConstant(Op):
    """Postponed Constant is a way to materialize a big constant.

    This class materializes a big constant only when it is going to be serialized
    to IR and then immediately disposes of it.
    """

    def __init__(
        self, element_type: Type, shape: Shape, maker: Callable[[], Tensor], name: Optional[str] = None
    ) -> None:
        """Creates a PostponedConstant.

        :param element_type: Element type of the constant.
        :type element_type: openvino.Type
        :param shape: Shape of the constant.
        :type shape: openvino.Shape
        :param maker: A callable that returns a Tensor.
        :type maker: Callable[[], Tensor]
        :param name: Optional name for the constant.
        :type name: Optional[str]

        :Example of a maker that returns a Tensor:

        .. code-block:: python

            class Maker:
                def __call__(self) -> ov.Tensor:
                    tensor_data = np.array([2, 2, 2, 2], dtype=np.float32)
                    return ov.Tensor(tensor_data)
        """
        super().__init__(self)
        self.get_rt_info()["postponed_constant"] = True  # value doesn't matter
        self.m_element_type = element_type
        self.m_shape = shape
        self.m_maker = maker
        if name is not None:
            self.friendly_name = name
        self.constructor_validate_and_infer_types()

    def evaluate(self, outputs: TensorVector, _: list[Tensor]) -> bool:  # type: ignore
        outputs[0] = self.m_maker()
        return True

    def validate_and_infer_types(self) -> None:
        self.set_output_type(0, self.m_element_type, PartialShape(self.m_shape))

    def clone_with_new_inputs(self, new_inputs: list[Tensor]) -> Op:
        return PostponedConstant(self.m_element_type, self.m_shape, self.m_maker, self.friendly_name)

    def has_evaluate(self) -> bool:
        return True


# `maker` is a function that returns ov.Tensor that represents a target Constant
def make_postponed_constant(element_type: Type,
                            shape: Shape,
                            maker: Callable[[], Tensor],
                            name: Optional[str] = None) -> Op:
    return PostponedConstant(element_type, shape, maker, name)
