from typing import TYPE_CHECKING, Any, Generic, Type, TypeVar, Union, cast

import numpy as np
import orjson

from docarray.base_doc.base_node import BaseNode
from docarray.typing.proto_register import _register_proto
from docarray.typing.tensor.abstract_tensor import AbstractTensor
from docarray.utils._internal.misc import (
    import_library,
    is_jax_available,
    is_torch_available,
)

if TYPE_CHECKING:
    import tensorflow as tf  # type: ignore

    from docarray.computation.tensorflow_backend import TensorFlowCompBackend
    from docarray.proto import NdArrayProto
else:
    tf = import_library('tensorflow', raise_error=True)

torch_available = is_torch_available()
if torch_available:
    import torch

jax_available = is_jax_available()
if jax_available:
    import jax.numpy as jnp

T = TypeVar('T', bound='TensorFlowTensor')
ShapeT = TypeVar('ShapeT')

tf_base: type = type(tf.Tensor)
node_base: type = type(BaseNode)


# the mypy error suppression below should not be necessary anymore once the following
# is released in mypy: https://github.com/python/mypy/pull/14135
class metaTensorFlow(
    AbstractTensor.__parametrized_meta__,  # type: ignore
    node_base,  # type: ignore
    tf_base,  # type: ignore
):  # type: ignore
    pass


@_register_proto(proto_type_name='tensorflow_tensor')
class TensorFlowTensor(AbstractTensor, Generic[ShapeT], metaclass=metaTensorFlow):
    """
    TensorFlowTensor class with a `.tensor` attribute of type `tf.Tensor`,
    intended for use in a Document.

    This enables (de)serialization from/to protobuf and json, data validation,
    and coercion from compatible types like numpy.ndarray.

    This type can also be used in a parametrized way, specifying the shape of the
    tensor.

    In comparison to [`TorchTensor`][docarray.typing.TorchTensor] and
    [`NdArray`][docarray.typing.tensor.ndarray.NdArray],
    [`TensorFlowTensor`][docarray.typing.tensor.tensorflow_tensor.TensorFlowTensor]
    is not a subclass of `tf.Tensor` (or `torch.Tensor`, `np.ndarray` respectively).
    Instead, the `tf.Tensor` is stored in
    [`TensorFlowTensor.tensor`][docarray.typing.tensor.tensorflow_tensor.TensorFlowTensor].
    Therefore, to do operations on the actual tensor data you have to always access the
    [`TensorFlowTensor.tensor`][docarray.typing.tensor.tensorflow_tensor.TensorFlowTensor]
    attribute.

    ---

    ```python
    import tensorflow as tf
    from docarray.typing import TensorFlowTensor


    t = TensorFlowTensor(tensor=tf.zeros((224, 224)))

    # tensorflow functions
    broadcasted = tf.broadcast_to(t.tensor, (3, 224, 224))
    broadcasted = tf.broadcast_to(t.unwrap(), (3, 224, 224))

    # this will fail:
    # broadcasted = tf.broadcast_to(t, (3, 224, 224))

    # tensorflow.Tensor methods:
    arr = t.tensor.numpy()
    arr = t.unwrap().numpy()

    # this will fail:
    # arr = t.numpy()
    ```

    ---

    The [`TensorFlowBackend`] however, operates on our
    [`TensorFlowTensor`][docarray.typing.TensorFlowTensor] instances.
    Here, you do not have to access the `.tensor` attribute,
    but can instead just hand over your
    [`TensorFlowTensor`][docarray.typing.TensorFlowTensor] instance.

    ---

    ```python
    import tensorflow as tf
    from docarray.typing import TensorFlowTensor


    zeros = TensorFlowTensor(tensor=tf.zeros((3, 224, 224)))

    comp_be = zeros.get_comp_backend()
    reshaped = comp_be.reshape(zeros, (224, 224, 3))
    assert comp_be.shape(reshaped) == (224, 224, 3)
    ```

    ---

    You can use [`TensorFlowTensor`][docarray.typing.TensorFlowTensor] in a Document as follows:

    ---

    ```python
    from docarray import BaseDoc
    from docarray.typing import TensorFlowTensor
    import tensorflow as tf


    class MyDoc(BaseDoc):
        tensor: TensorFlowTensor
        image_tensor: TensorFlowTensor[3, 224, 224]
        square_crop: TensorFlowTensor[3, 'x', 'x']
        random_image: TensorFlowTensor[
            3, ...
        ]  # first dimension is fixed, can have arbitrary shape


    # create a document with tensors
    doc = MyDoc(
        tensor=tf.zeros((128,)),
        image_tensor=tf.zeros((3, 224, 224)),
        square_crop=tf.zeros((3, 64, 64)),
        random_image=tf.zeros((3, 128, 256)),
    )

    # automatic shape conversion
    doc = MyDoc(
        tensor=tf.zeros((128,)),
        image_tensor=tf.zeros((224, 224, 3)),  # will reshape to (3, 224, 224)
        square_crop=tf.zeros((3, 128, 128)),
        random_image=tf.zeros((3, 64, 128)),
    )

    # !! The following will raise an error due to shape mismatch !!
    from pydantic import ValidationError

    try:
        doc = MyDoc(
            tensor=tf.zeros((128,)),
            image_tensor=tf.zeros((224, 224)),  # this will fail validation
            square_crop=tf.zeros((3, 128, 64)),  # this will also fail validation
            random_image=tf.zeros(4, 64, 128),  # this will also fail validation
        )
    except ValidationError as e:
        pass
    ```

    ---
    """

    __parametrized_meta__ = metaTensorFlow

    def __init__(self, tensor: tf.Tensor):
        super().__init__()
        self.tensor = tensor

    def __getitem__(self, item):
        from docarray.computation.tensorflow_backend import TensorFlowCompBackend

        tensor = self.unwrap()
        if tensor is not None:
            tensor = tensor[item]
        return TensorFlowCompBackend._cast_output(t=tensor)

    def __setitem__(self, index, value):
        """Set a slice of this tensor's `tf.Tensor`"""
        t = self.unwrap()
        value = tf.cast(value, dtype=t.dtype)
        var = tf.Variable(t)
        var[index].assign(value)
        self.tensor = tf.constant(var)

    def __iter__(self):
        """Iterate over the elements of this tensor's `tf.Tensor`."""
        for i in range(len(self)):
            yield self[i]

    @classmethod
    def _docarray_validate(
        cls: Type[T],
        value: Union[T, np.ndarray, str, Any],
    ) -> T:
        if isinstance(value, TensorFlowTensor):
            return cast(T, value)
        elif isinstance(value, tf.Tensor):
            return cls._docarray_from_native(value)
        elif isinstance(value, np.ndarray):
            return cls._docarray_from_ndarray(value)
        elif isinstance(value, AbstractTensor):
            return cls._docarray_from_ndarray(value._docarray_to_ndarray())
        elif torch_available and isinstance(value, torch.Tensor):
            return cls._docarray_from_native(value.detach().cpu().numpy())
        elif jax_available and isinstance(value, jnp.ndarray):
            return cls._docarray_from_native(value.__array__())
        elif isinstance(value, str):
            value = orjson.loads(value)

        try:
            arr: tf.Tensor = tf.constant(value)
            return cls(tensor=arr)
        except Exception:
            pass  # handled below

        raise ValueError(
            f'Expected a tensorflow.Tensor compatible type, got {type(value)}'
        )

    @classmethod
    def _docarray_from_native(cls: Type[T], value: Union[tf.Tensor, T]) -> T:
        """
        Create a `TensorFlowTensor` from a `tf.Tensor` or `TensorFlowTensor`
        instance.

        :param value: instance of `tf.Tensor` or `TensorFlowTensor`
        :return: a `TensorFlowTensor`
        """
        if isinstance(value, TensorFlowTensor):
            if cls.__unparametrizedcls__:  # None if the tensor is parametrized
                value.__class__ = cls.__unparametrizedcls__  # type: ignore
            else:
                value.__class__ = cls
            return cast(T, value)
        else:
            if cls.__unparametrizedcls__:  # None if the tensor is parametrized
                cls_param_ = cls.__unparametrizedcls__
                cls_param = cast(Type[T], cls_param_)
            else:
                cls_param = cls

            return cls_param(tensor=value)

    @staticmethod
    def get_comp_backend() -> 'TensorFlowCompBackend':
        """Return the computational backend of the tensor"""
        from docarray.computation.tensorflow_backend import TensorFlowCompBackend

        return TensorFlowCompBackend()

    def _docarray_to_json_compatible(self) -> np.ndarray:
        """
        Convert `TensorFlowTensor` into a json compatible object
        :return: a representation of the tensor compatible with orjson
        """
        return self.unwrap().numpy()

    def to_protobuf(self) -> 'NdArrayProto':
        """
        Transform self into an NdArrayProto protobuf message.
        """
        from docarray.proto import NdArrayProto

        nd_proto = NdArrayProto()

        value_np = self.tensor.numpy()
        nd_proto.dense.buffer = value_np.tobytes()
        nd_proto.dense.ClearField('shape')
        nd_proto.dense.shape.extend(list(value_np.shape))
        nd_proto.dense.dtype = value_np.dtype.str

        return nd_proto

    @classmethod
    def from_protobuf(cls: Type[T], pb_msg: 'NdArrayProto') -> 'T':
        """
        Read ndarray from a proto msg.
        :param pb_msg:
        :return: a `TensorFlowTensor`
        """
        source = pb_msg.dense
        if source.buffer:
            x = np.frombuffer(bytearray(source.buffer), dtype=source.dtype)
            return cls.from_ndarray(x.reshape(source.shape))
        elif len(source.shape) > 0:
            return cls.from_ndarray(np.zeros(source.shape))
        else:
            raise ValueError(
                f'Proto message {pb_msg} cannot be cast to a TensorFlowTensor.'
            )

    @classmethod
    def from_ndarray(cls: Type[T], value: np.ndarray) -> T:
        """Create a `TensorFlowTensor` from a numpy array.

        :param value: the numpy array
        :return: a `TensorFlowTensor`
        """
        return cls._docarray_from_native(tf.convert_to_tensor(value))

    def unwrap(self) -> tf.Tensor:
        """
        Return the original `tf.Tensor` without any memory copy.

        The original view rest intact and is still a Document `TensorFlowTensor`
        but the return object is a pure `tf.Tensor` but both object share
        the same memory layout.

        ---

        ```python
        from docarray.typing import TensorFlowTensor
        import tensorflow as tf

        t1 = TensorFlowTensor.validate(tf.zeros((3, 224, 224)), None, None)
        # here t1 is a docarray TensorFlowTensor
        t2 = t1.unwrap()
        # here t2 is a pure tf.Tensor but t1 is still a Docarray TensorFlowTensor
        ```

        ---
        :return: a `tf.Tensor`
        """
        return self.tensor

    def __len__(self) -> int:
        return len(self.tensor)

    @classmethod
    def _docarray_from_ndarray(cls: Type[T], value: np.ndarray) -> T:
        """Create a `tensor from a numpy array
        PS: this function is different from `from_ndarray` because it is private under the docarray namesapce.
        This allows us to avoid breaking change if one day we introduce a Tensor backend with a `from_ndarray` method.
        """
        return cls.from_ndarray(value)

    def _docarray_to_ndarray(self) -> np.ndarray:
        """cast itself to a numpy array"""
        return self.tensor.numpy()

    @property
    def shape(self):
        return tf.shape(self.tensor)
