import numpy as np

import rlmesh.proto.gym_pb2 as gpb
import rlmesh.proto.env_pb2 as epb

from rlmesh.gym.utils import NP2DT, DT2NP


def tensor_from_ndarray(arr: np.ndarray) -> epb.Tensor:
    arr = np.asarray(arr)
    dt = NP2DT.get(arr.dtype.type, gpb.DType.FLOAT32)
    return epb.Tensor(dtype=dt, shape=list(arr.shape), raw=arr.tobytes(order="C"))


def ndarray_from_tensor(t: epb.Tensor) -> np.ndarray:
    ty = DT2NP.get(t.dtype, np.float32)
    arr = np.frombuffer(bytes(t.raw), dtype=ty)
    shape = tuple(int(x) for x in t.shape) if t.shape else ()
    return arr.reshape(shape) if shape else arr
