import torch
import pickle

class ExperienceCollector:
    def __init__(self):
        self.states = []
        self.actions = []
        self.rewards = []
        self.advantages = []
        self._current_episode_states = []
        self._current_episode_actions = []
        self._current_episode_estimated_values = []

    def begin_episode(self):
        self._current_episode_states = []
        self._current_episode_actions = []
        self._current_episode_estimated_values = []

    def record_decision(self, state, action, estimated_value=0):
        self._current_episode_states.append(state)
        self._current_episode_actions.append(action)
        self._current_episode_estimated_values.append(estimated_value)

    def complete_episode(self, reward):
        num_states = len(self._current_episode_states)
        self.states += self._current_episode_states
        self.actions += self._current_episode_actions
        self.rewards += [reward for _ in range(num_states)]

        for i in range(num_states):
            advantage = reward - self._current_episode_estimated_values[i]
            self.advantages.append(advantage)

        self._current_episode_states = []
        self._current_episode_actions = []
        self._current_episode_estimated_values = []


class ExperienceBuffer:
    def __init__(self, states, actions, rewards, advantages):
        self.states = states
        self.actions = actions
        self.rewards = rewards
        self.advantages = advantages

    def save(self, file_path):
        with open(file_path, 'wb') as f:
            pickle.dump((self.states, self.actions, self.rewards, self.advantages), f)

    @classmethod
    def load(cls, file_path):
        with open(file_path, 'rb') as f:
            states, actions, rewards, advantages = pickle.load(f)
        return cls(
            states=states,
            actions=actions,
            rewards=rewards,
            advantages=advantages)

def combine_experience(collectors):
    combined_states = torch.cat([torch.tensor(c.states) for c in collectors])
    combined_actions = torch.cat([torch.tensor(c.actions) for c in collectors])
    combined_rewards = torch.cat([torch.tensor(c.rewards) for c in collectors])
    combined_advantages = torch.cat([torch.tensor(c.advantages) for c in collectors])

    return ExperienceBuffer(
        combined_states,
        combined_actions,
        combined_rewards,
        combined_advantages)