from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Dict, Iterable, Optional

import torch
from torch import Tensor


@dataclass
class Episode:
    """Container for episodic transitions.

    Each episode is stored as parallel tensors for states, actions, and rewards.
    Optional fields include `dones` and `infos` for compatibility with common
    RL interfaces. All tensors share the leading time dimension.
    """

    states: Tensor
    actions: Tensor
    rewards: Tensor
    dones: Optional[Tensor] = None
    infos: Optional[Iterable[Dict[str, Any]]] = None
    # Optional post-processed targets
    returns: Optional[Tensor] = None
    advantages: Optional[Tensor] = None

    def __post_init__(self) -> None:
        time_dim = self.states.shape[0]
        for name, tensor in [("actions", self.actions), ("rewards", self.rewards)]:
            if tensor.shape[0] != time_dim:
                raise ValueError(f"Episode {name} length {tensor.shape[0]} != {time_dim}")
        if self.dones is not None and self.dones.shape[0] != time_dim:
            raise ValueError("Episode dones length mismatch")
        if self.infos is not None:
            if not isinstance(self.infos, list):
                self.infos = list(self.infos)
            if len(self.infos) != time_dim:
                raise ValueError("Episode infos length mismatch")

    def __len__(self) -> int:  # pragma: no cover - trivial
        return self.states.shape[0]

    @property
    def device(self) -> torch.device:
        return self.states.device

    @property
    def total_reward(self) -> float:
        return float(self.rewards.sum().item())

    def slice(self, start: int, end: int) -> "Episode":
        """Return a shallow slice preserving tensor types."""
        return Episode(
            states=self.states[start:end],
            actions=self.actions[start:end],
            rewards=self.rewards[start:end],
            dones=self.dones[start:end] if self.dones is not None else None,
            infos=list(self.infos)[start:end] if self.infos is not None else None,
        )

    def to(self, device: torch.device) -> "Episode":
        """Move all tensor fields to the target device."""
        return Episode(
            states=self.states.to(device),
            actions=self.actions.to(device),
            rewards=self.rewards.to(device),
            dones=self.dones.to(device) if self.dones is not None else None,
            infos=self.infos,
        )

    def to_dict(self) -> Dict[str, Any]:
        """Serialize episode tensors for checkpointing."""
        return {
            "states": self.states,
            "actions": self.actions,
            "rewards": self.rewards,
            "dones": self.dones,
            "infos": self.infos,
            "returns": self.returns,
            "advantages": self.advantages,
        }

    @classmethod
    def from_dict(cls, payload: Dict[str, Any]) -> "Episode":
        return cls(
            states=payload["states"],
            actions=payload["actions"],
            rewards=payload["rewards"],
            dones=payload.get("dones"),
            infos=payload.get("infos"),
            returns=payload.get("returns"),
            advantages=payload.get("advantages"),
        )

    def compute_returns(
        self,
        gamma: float,
        n_steps: Optional[int] = None,
        *,
        value_fn: Optional[Callable[[Tensor], Tensor]] = None,
        gae_lambda: Optional[float] = None,
        last_value: float | Tensor = 0.0,
        store: bool = True,
    ) -> Tensor:
        """Compute discounted returns and optionally GAE advantages.

        Args:
            gamma: Discount factor in [0, 1].
            n_steps: Horizon for n-step bootstrapping. If ``None`` or >= ``T``,
                compute full-episode Monte Carlo (no bootstrapping).
            value_fn: Optional callback mapping states ``[T, ...]`` to values
                ``[T]`` (or ``[T, 1]``). Used for n-step bootstrapping and GAE.
            gae_lambda: When provided along with ``value_fn``, compute Generalized
                Advantage Estimation. In this mode, ``returns = advantages + values``.
            last_value: Bootstrap value for the state following the final transition.
                Used by GAE (as ``V_{T}``) and for n-step bootstrapping past the end.
            store: If True, cache results in ``self.returns`` (and
                ``self.advantages`` when GAE is used).

        Returns:
            Tensor of shape ``[T]`` containing the computed returns.
        """
        T = self.rewards.shape[0]
        device = self.rewards.device
        dtype = self.rewards.dtype

        dones = (
            self.dones.to(device=device)
            if self.dones is not None
            else torch.zeros(T, dtype=torch.bool, device=device)
        )

        # When a value function is provided, pre-compute per-step values once.
        values: Optional[Tensor] = None
        if value_fn is not None:
            with torch.no_grad():
                v = value_fn(self.states)
            values = torch.as_tensor(v, dtype=dtype, device=device).squeeze(-1)
            if values.dim() != 1 or values.shape[0] != T:
                raise ValueError("value_fn must return shape [T] or [T,1] values")

        # GAE path takes precedence when lambda provided alongside a value function.
        if gae_lambda is not None:
            if values is None:
                raise ValueError("GAE requires value_fn to be provided")
            lv = torch.as_tensor(last_value, dtype=dtype, device=device)
            next_values = torch.empty(T, dtype=dtype, device=device)
            if T > 1:
                next_values[:-1] = values[1:]
            next_values[-1] = lv

            deltas = self.rewards + gamma * next_values * (~dones).to(dtype) - values
            advantages = torch.zeros(T, dtype=dtype, device=device)
            gae = torch.tensor(0.0, dtype=dtype, device=device)
            for t in reversed(range(T)):
                mask = 1.0 - dones[t].to(dtype)
                gae = deltas[t] + gamma * gae_lambda * mask * gae
                advantages[t] = gae
            returns = advantages + values
            if store:
                self.advantages = advantages
                self.returns = returns
            return returns

        # n-step (including MC when n_steps is None or large)
        if n_steps is None or n_steps <= 0 or n_steps >= T:
            # Full Monte Carlo without bootstrapping.
            returns = torch.zeros(T, dtype=dtype, device=device)
            running = torch.tensor(0.0, dtype=dtype, device=device)
            for t in reversed(range(T)):
                running = self.rewards[t] + gamma * running * (1.0 - dones[t].to(dtype))
                returns[t] = running
            if store:
                self.returns = returns
            return returns

        # General n-step with optional value bootstrapping.
        horizon = int(n_steps)
        returns = torch.zeros(T, dtype=dtype, device=device)
        # Optional V for bootstrapping beyond the observed rewards
        if values is None:
            # Treat missing value_fn as zero bootstrap
            values = torch.zeros(T, dtype=dtype, device=device)
        lv = torch.as_tensor(last_value, dtype=dtype, device=device)

        for t in range(T):
            ret = torch.tensor(0.0, dtype=dtype, device=device)
            discount = torch.tensor(1.0, dtype=dtype, device=device)
            done_reached = False
            for k in range(horizon):
                idx = t + k
                if idx >= T:
                    break
                ret = ret + discount * self.rewards[idx]
                if dones[idx]:
                    done_reached = True
                    break
                discount = discount * gamma
            # Bootstrapping: if horizon not exhausted due to terminal or end,
            # and we have a value estimate for s_{t+horizon}
            target_idx = t + horizon
            if not done_reached:
                if target_idx < T:
                    ret = ret + discount * values[target_idx]
                elif target_idx == T:
                    ret = ret + discount * lv
            returns[t] = ret

        if store:
            self.returns = returns
        return returns


__all__ = ["Episode"]
