"""Adapter: wrap a Gym-style env so it implements DiscoRL's Environment API.

This is a minimal adapter intended for evaluation / inference (batching=1
or small batches). It converts Gym observations/rewards/dones into the
`types.EnvironmentTimestep` structure expected by DiscoRL and keeps a
Python-side list of env instances for the batch.

Notes:
- This adapter does not attempt to JIT or vectorize with JAX. It simply
  converts numpy -> jax arrays before returning timesteps so the DiscoRL
  agent (Haiku/JAX) can consume them.
- For training at scale you can rework this into a true batched env that
  runs multiple envs in parallel / in subprocesses.
"""

from typing import Any, Tuple

import numpy as np
import jax
import jax.numpy as jnp

from disco_rl.environments import base
from disco_rl import types
try:
  from dm_env import StepType
except ImportError:
  # Fallback if dm_env not available
  class StepType:
    MID = 0
    LAST = 1


class GymToDiscoEnv(base.Environment):
  """Wrap a Gym-compatible environment class.

  The wrapped `gym_env_cls` must follow the Gym API (reset() -> obs, info,
  step(action) -> obs, reward, terminated, truncated, info).
  
  Args:
    gym_env_cls: factory/class that creates Gym environment instances.
    batch_size: number of parallel env instances to manage.
    env_settings: dict of kwargs to pass to gym_env_cls.
    discrete_actions: optional array of shape (num_actions, action_dim) mapping
      discrete action indices to continuous action vectors. If provided, the
      action_spec becomes discrete (int32 scalar indices) and step() will
      map indices to continuous actions before sending to underlying env.
  """

  def __init__(
      self,
      gym_env_cls: Any,
      batch_size: int = 1,
      env_settings=None,
      discrete_actions: np.ndarray | None = None,
  ):
    self.batch_size = batch_size
    env_settings = {} if env_settings is None else env_settings
    # Create multiple env instances for simple batching.
    self._envs = [gym_env_cls(**env_settings) for _ in range(batch_size)]
    self._discrete_actions = (
        np.asarray(discrete_actions) if discrete_actions is not None else None
    )

    # Build single action/observation specs in the simple form expected by
    # DiscoRL (a mapping with key 'observation'). We keep dtype/shape simple.
    obs_space = self._envs[0].observation_space
    act_space = self._envs[0].action_space

    # Use dm_env-like BoundedArray for actions if available, else a simple
    # placeholder (the agent only queries shape/dtype in most places).
    try:
      from dm_env import specs as dm_specs

      # If discrete_actions is provided, create a discrete action spec;
      # otherwise use the original continuous spec.
      if self._discrete_actions is not None:
        num_actions = len(self._discrete_actions)
        self._single_action_spec = dm_specs.BoundedArray(
            shape=(), dtype=np.int32, minimum=0, maximum=num_actions - 1
        )
      else:
        self._single_action_spec = dm_specs.BoundedArray(
            act_space.shape, act_space.dtype, act_space.low, act_space.high
        )
      self._single_observation_spec = {
          'observation': dm_specs.Array(shape=obs_space.shape, dtype=obs_space.dtype)
      }
    except Exception:
      # Fallback to simple numpy-shape descriptors.
      if self._discrete_actions is not None:
        num_actions = len(self._discrete_actions)
        self._single_action_spec = type('ActionSpec', (), {
            'shape': (),
            'dtype': np.int32,
            'low': 0,
            'high': num_actions - 1,
        })
      else:
        self._single_action_spec = act_space
      self._single_observation_spec = {'observation': obs_space}

    # Keep last observations / states for each env
    self._last_obs = [None] * batch_size
    self._dones = [True] * batch_size

  def single_action_spec(self):
    return self._single_action_spec

  def single_observation_spec(self):
    return self._single_observation_spec

  def _obs_to_timestep(self, obs, reward, done):
    # Convert a single env's raw outputs into types.EnvironmentTimestep
    # DiscoRL expects a mapping for observation (e.g. {'observation': ...}).
    obs_map = {'observation': jnp.asarray(obs, dtype=jnp.float32)}
    step_type = jnp.array(StepType.LAST if done else StepType.MID, dtype=jnp.int32)
    return types.EnvironmentTimestep(observation=obs_map, step_type=step_type, reward=jnp.array(float(reward), dtype=jnp.float32))

  def step(self, state_unused, actions) -> Tuple[Any, types.EnvironmentTimestep]:
    # actions expected to be a batched array with shape (batch_size, ...)
    # For simplicity we iterate over envs sequentially.
    # Support actions provided as numpy/jax arrays.
    actions = [np.array(a) for a in list(actions)] if hasattr(actions, '__iter__') else [np.array(actions)]

    obs_batch = []
    reward_batch = []
    done_batch = []

    for i, env in enumerate(self._envs):
      act = actions[i] if i < len(actions) else actions[0]
      # If discrete_actions is provided, map the action index to continuous action.
      if self._discrete_actions is not None:
        act = self._discrete_actions[int(act)]
      # Convert to python scalar if necessary
      obs, reward, terminated, truncated, info = env.step(act)
      done = bool(terminated or truncated)
      self._last_obs[i] = obs
      self._dones[i] = done
      obs_batch.append(jnp.asarray(obs, dtype=jnp.float32))
      reward_batch.append(float(reward))
      done_batch.append(done)

    # Stack to produce batched structures. DiscoRL typically expects
    # observations to be a mapping of arrays with leading batch dimension.
    obs_map = {'observation': jnp.stack(obs_batch)}
    rewards = jnp.asarray(reward_batch, dtype=jnp.float32)
    is_terminal = jnp.asarray(done_batch)
    # Use jnp.where so scalar StepType values broadcast to the array shape.
    step_type = jnp.where(is_terminal, StepType.LAST, StepType.MID)

    timestep = types.EnvironmentTimestep(observation=obs_map, step_type=step_type, reward=rewards)
    return None, timestep

  def reset(self, rng_key=None) -> Tuple[Any, types.EnvironmentTimestep]:
    obs_batch = []
    reward_batch = []
    done_batch = []
    for i, env in enumerate(self._envs):
      # Gym reset returns (obs, info) in Gymnasium; support both
      out = env.reset()
      if isinstance(out, tuple) and len(out) >= 1:
        obs = out[0]
      else:
        obs = out
      self._last_obs[i] = obs
      self._dones[i] = False
      obs_batch.append(jnp.asarray(obs, dtype=jnp.float32))
      reward_batch.append(0.0)
      done_batch.append(False)

    obs_map = {'observation': jnp.stack(obs_batch)}
    rewards = jnp.asarray(reward_batch, dtype=jnp.float32)
    is_terminal = jnp.asarray(done_batch)
    # Use jnp.where so scalar StepType values broadcast to the array shape.
    step_type = jnp.where(is_terminal, StepType.LAST, StepType.MID)
    timestep = types.EnvironmentTimestep(observation=obs_map, step_type=step_type, reward=rewards)
    return None, timestep