"""Stage 4 training: RL training loop."""

import torch
import numpy as np
from typing import Dict

from saltts.rl.agent import SACAgent
from saltts.rl.reward import RewardFunction
from saltts.rl.state_encoder import StateEncoder


class Stage4Trainer:
    """Trainer for Stage 4: RL training."""

    def __init__(
        self,
        state_dim: int = 256,
        action_dim: int = 5,
        device: str = "cuda",
    ):
        """Initialize Stage 4 trainer.

        Args:
            state_dim: State dimension
            action_dim: Action dimension
            device: Device
        """
        self.device = device

        # Initialize RL agent
        self.agent = SACAgent(
            state_dim=state_dim,
            action_dim=action_dim,
            device=device,
        )

        # Reward function
        self.reward_fn = RewardFunction()

        # State encoder
        self.state_encoder = StateEncoder().to(device)

    def train_episode(self, env) -> Dict[str, float]:
        """Train for one episode.

        Args:
            env: Environment

        Returns:
            Episode metrics
        """
        state = env.reset()
        episode_reward = 0.0
        episode_length = 0

        done = False
        while not done:
            # Select action
            action = self.agent.select_action(state)

            # Step environment
            next_state, reward, done, info = env.step(action)

            # Store transition
            self.agent.replay_buffer.push(state, action, reward, next_state, done)

            # Update agent
            if len(self.agent.replay_buffer) > 256:
                losses = self.agent.update(batch_size=256)

            state = next_state
            episode_reward += reward
            episode_length += 1

        return {
            "episode_reward": episode_reward,
            "episode_length": episode_length,
        }

    def save(self, path: str) -> None:
        """Save agent."""
        self.agent.save(path)

    def load(self, path: str) -> None:
        """Load agent."""
        self.agent.load(path)
