import jax
import jax.numpy as jnp
import equinox as eqx
from typing import Union, Tuple
from local_coordinates.base import AbstractBatchableObject, auto_vmap


class ScalarObj(AbstractBatchableObject):
  """Unbatched object with a single scalar."""
  value: jnp.ndarray

  @property
  def batch_size(self) -> Union[Tuple[int], int, None]:
    if self.value.ndim == 0:
      return None
    return self.value.shape[0]


class VectorObj(AbstractBatchableObject):
  """Object whose batch_size is inferred from a leading dimension."""
  data: jnp.ndarray

  @property
  def batch_size(self) -> Union[Tuple[int], int, None]:
    if self.data.ndim == 1:
      return None
    return self.data.shape[0]


def test_batch_size_unbatched():
  obj = ScalarObj(value=jnp.array(1.0))
  assert obj.batch_size is None


def test_batch_size_batched():
  obj = ScalarObj(value=jnp.array([1.0, 2.0, 3.0]))
  assert obj.batch_size == 3


def test_zeros_like():
  obj = VectorObj(data=jnp.array([5.0, 6.0, 7.0]))
  z = VectorObj.zeros_like(obj)
  assert jnp.allclose(z.data, jnp.zeros(3))


def test_shape_property():
  obj = VectorObj(data=jnp.array([1.0, 2.0]))
  shapes = obj.shape
  assert shapes.data == (2,)


def test_getitem():
  obj = VectorObj(data=jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]))
  sliced = obj[1]
  assert jnp.allclose(sliced.data, jnp.array([3.0, 4.0]))


def test_getitem_slice():
  obj = VectorObj(data=jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]))
  sliced = obj[:2]
  assert sliced.data.shape == (2, 2)
  assert jnp.allclose(sliced.data[0], jnp.array([1.0, 2.0]))


class SumObj(AbstractBatchableObject):
  data: jnp.ndarray

  @property
  def batch_size(self) -> Union[Tuple[int], int, None]:
    if self.data.ndim == 1:
      return None
    return self.data.shape[0]

  @auto_vmap
  def total(self):
    return jnp.sum(self.data)


def test_auto_vmap_unbatched():
  obj = SumObj(data=jnp.array([1.0, 2.0, 3.0]))
  result = obj.total()
  assert jnp.isclose(result, 6.0)


def test_auto_vmap_batched():
  obj = SumObj(data=jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]))
  result = obj.total()
  expected = jnp.array([3.0, 7.0, 11.0])
  assert jnp.allclose(result, expected)
