import numpy as np
import torch
import random
from typing import Dict, List, Tuple

class ReplayBuffer:
    """
    A simple sequence replay buffer for Dreamer training.
    Expects episodes of shape (seq_len, dim) to be added.
    """
    def __init__(self, capacity: int = 10000, seq_len: int = 50):
        self.capacity = capacity
        self.seq_len = seq_len
        self.episodes = []
        
    def add_episode(self, obs: np.ndarray, actions: np.ndarray, rewards: np.ndarray):
        """
        obs: (time, obs_dim)
        actions: (time, action_dim)
        rewards: (time,)
        """
        if len(self.episodes) >= self.capacity:
            self.episodes.pop(0)
            
        self.episodes.append({
            'obs': obs,
            'actions': actions,
            'rewards': rewards
        })
        
    def sample_batch(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Samples a batch of sequence chunks.
        Returns tensors of shape (batch, seq_len, dim)
        """
        if len(self.episodes) == 0:
            raise ValueError("Replay buffer is empty.")
            
        obs_batch, action_batch, reward_batch = [], [], []
        
        for _ in range(batch_size):
            ep_idx = random.randint(0, len(self.episodes) - 1)
            ep = self.episodes[ep_idx]
            ep_len = len(ep['obs'])
            
            if ep_len <= self.seq_len:
                start = 0
                pad_len = self.seq_len - ep_len
                obs = np.pad(ep['obs'], ((0, pad_len), (0, 0)), mode='edge')
                acts = np.pad(ep['actions'], ((0, pad_len), (0, 0)), mode='edge')
                rews = np.pad(ep['rewards'], (0, pad_len), mode='constant', constant_values=0)
            else:
                start = random.randint(0, ep_len - self.seq_len)
                obs = ep['obs'][start:start + self.seq_len]
                acts = ep['actions'][start:start + self.seq_len]
                rews = ep['rewards'][start:start + self.seq_len]
                
            obs_batch.append(obs)
            action_batch.append(acts)
            reward_batch.append(rews)
            
        return (
            torch.FloatTensor(np.stack(obs_batch)),
            torch.FloatTensor(np.stack(action_batch)),
            torch.FloatTensor(np.stack(reward_batch))
        )