"""Soft Actor-Critic (SAC) agent for adaptive mixing."""

import torch
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from typing import Tuple, Optional

from saltts.rl.networks import Actor, Critic
from saltts.rl.replay_buffer import ReplayBuffer
from saltts.rl.hnsw_memory import HNSWMemory
from saltts.rl.state_encoder import StateEncoder


class SACAgent:
    """Soft Actor-Critic agent for adaptive mixing.

    Learns to mix Generator A and B outputs based on context.
    """

    def __init__(
        self,
        state_dim: int = 256,
        action_dim: int = 5,
        lr: float = 3e-4,
        gamma: float = 0.99,
        tau: float = 0.005,
        alpha: float = 0.2,
        buffer_capacity: int = 1000000,
        device: str = "cuda",
    ):
        """Initialize SAC agent.

        Args:
            state_dim: State dimension
            action_dim: Action dimension
            lr: Learning rate
            gamma: Discount factor
            tau: Soft update coefficient
            alpha: Entropy coefficient
            buffer_capacity: Replay buffer capacity
            device: Device (cuda or cpu)
        """
        self.device = device
        self.gamma = gamma
        self.tau = tau
        self.alpha = alpha

        # Networks
        self.actor = Actor(state_dim, action_dim).to(device)
        self.critic = Critic(state_dim, action_dim).to(device)
        self.critic_target = Critic(state_dim, action_dim).to(device)
        self.critic_target.load_state_dict(self.critic.state_dict())

        # Optimizers
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr)

        # Replay buffer
        self.replay_buffer = ReplayBuffer(buffer_capacity)

        # HNSW memory
        self.hnsw_memory = HNSWMemory(dim=state_dim)

    def select_action(self, state: np.ndarray, evaluate: bool = False) -> np.ndarray:
        """Select action.

        Args:
            state: State [state_dim]
            evaluate: If True, use deterministic policy

        Returns:
            Action [action_dim]
        """
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)

        with torch.no_grad():
            if evaluate:
                mean, _ = self.actor(state_tensor)
                action = torch.tanh(mean)
            else:
                action, _ = self.actor.sample(state_tensor)

        return action.cpu().numpy()[0]

    def update(self, batch_size: int = 256) -> dict:
        """Update networks.

        Args:
            batch_size: Batch size

        Returns:
            Dictionary of losses
        """
        if len(self.replay_buffer) < batch_size:
            return {}

        # Sample batch
        states, actions, rewards, next_states, dones = self.replay_buffer.sample(batch_size)

        states = torch.FloatTensor(states).to(self.device)
        actions = torch.FloatTensor(actions).to(self.device)
        rewards = torch.FloatTensor(rewards).unsqueeze(1).to(self.device)
        next_states = torch.FloatTensor(next_states).to(self.device)
        dones = torch.FloatTensor(dones).unsqueeze(1).to(self.device)

        # Update critic
        with torch.no_grad():
            next_actions, next_log_probs = self.actor.sample(next_states)
            q1_next, q2_next = self.critic_target(next_states, next_actions)
            q_next = torch.min(q1_next, q2_next) - self.alpha * next_log_probs
            q_target = rewards + (1 - dones) * self.gamma * q_next

        q1, q2 = self.critic(states, actions)
        critic_loss = F.mse_loss(q1, q_target) + F.mse_loss(q2, q_target)

        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # Update actor
        new_actions, log_probs = self.actor.sample(states)
        q1_new, q2_new = self.critic(states, new_actions)
        q_new = torch.min(q1_new, q2_new)
        actor_loss = (self.alpha * log_probs - q_new).mean()

        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # Soft update target network
        self._soft_update()

        return {
            "critic_loss": critic_loss.item(),
            "actor_loss": actor_loss.item(),
        }

    def _soft_update(self) -> None:
        """Soft update target network."""
        for target_param, param in zip(
            self.critic_target.parameters(), self.critic.parameters()
        ):
            target_param.data.copy_(
                target_param.data * (1.0 - self.tau) + param.data * self.tau
            )

    def save(self, path: str) -> None:
        """Save agent."""
        torch.save({
            "actor": self.actor.state_dict(),
            "critic": self.critic.state_dict(),
            "critic_target": self.critic_target.state_dict(),
        }, path)

    def load(self, path: str) -> None:
        """Load agent."""
        checkpoint = torch.load(path)
        self.actor.load_state_dict(checkpoint["actor"])
        self.critic.load_state_dict(checkpoint["critic"])
        self.critic_target.load_state_dict(checkpoint["critic_target"])
