from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Tuple

import torch
from torch import Tensor, nn
from torch.distributions import Normal


@dataclass
class SACConfig:
    state_dim: int
    action_dim: int
    hidden_dim: int = 256
    gamma: float = 0.99
    tau: float = 0.005
    lr: float = 3e-4
    alpha: float = 0.2  # fixed entropy coefficient for simplicity
    action_limit: float = 1.0
    device: str = "cpu"


class _MLP(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, hidden: int) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, out_dim),
        )

    def forward(self, x: Tensor) -> Tensor:  # noqa: D401 - trivial
        return self.net(x)


class _Actor(nn.Module):
    def __init__(self, state_dim: int, action_dim: int, hidden: int) -> None:
        super().__init__()
        self.backbone = _MLP(state_dim, 2 * action_dim, hidden)
        self.action_dim = action_dim

    def forward(self, state: Tensor) -> Tuple[Tensor, Tensor]:
        out = self.backbone(state)
        mu, log_std = out[..., : self.action_dim], out[..., self.action_dim :]
        log_std = torch.clamp(log_std, -20.0, 2.0)
        return mu, log_std

    def sample(self, state: Tensor, *, deterministic: bool = False) -> Tuple[Tensor, Tensor]:
        mu, log_std = self(state)
        if deterministic:
            z = mu
        else:
            std = log_std.exp()
            dist = Normal(mu, std)
            z = dist.rsample()
        # Tanh squashing
        action = torch.tanh(z)
        # Compute log prob with tanh correction when not deterministic
        if deterministic:
            log_prob = torch.zeros(action.shape[0], 1, device=action.device)
        else:
            std = log_std.exp()
            dist = Normal(mu, std)
            log_prob = dist.log_prob(z)
            # sum over action dims
            log_prob = log_prob.sum(dim=-1, keepdim=True)
            # Change of variables for tanh
            log_prob -= torch.log(1 - action.pow(2) + 1e-6).sum(dim=-1, keepdim=True)
        return action, log_prob


class _QNet(nn.Module):
    def __init__(self, state_dim: int, action_dim: int, hidden: int) -> None:
        super().__init__()
        self.net = _MLP(state_dim + action_dim, 1, hidden)

    def forward(self, state: Tensor, action: Tensor) -> Tensor:  # noqa: D401 - trivial
        x = torch.cat([state, action], dim=-1)
        return self.net(x)


class SACAgent:
    """Minimal SAC agent for continuous actions.

    Expects replay batches with keys: states, actions, rewards, next_states, dones.
    Actions are assumed to be in [-action_limit, action_limit].
    """

    def __init__(self, cfg: SACConfig) -> None:
        self.cfg = cfg
        self.device = torch.device(cfg.device)
        self.actor = _Actor(cfg.state_dim, cfg.action_dim, cfg.hidden_dim).to(self.device)
        self.q1 = _QNet(cfg.state_dim, cfg.action_dim, cfg.hidden_dim).to(self.device)
        self.q2 = _QNet(cfg.state_dim, cfg.action_dim, cfg.hidden_dim).to(self.device)
        self.q1_target = _QNet(cfg.state_dim, cfg.action_dim, cfg.hidden_dim).to(self.device)
        self.q2_target = _QNet(cfg.state_dim, cfg.action_dim, cfg.hidden_dim).to(self.device)
        self.q1_target.load_state_dict(self.q1.state_dict())
        self.q2_target.load_state_dict(self.q2.state_dict())

        self.pi_optim = torch.optim.Adam(self.actor.parameters(), lr=cfg.lr)
        self.q1_optim = torch.optim.Adam(self.q1.parameters(), lr=cfg.lr)
        self.q2_optim = torch.optim.Adam(self.q2.parameters(), lr=cfg.lr)

        self._action_limit = float(cfg.action_limit)

    @torch.no_grad()
    def act(self, state, explore: bool = True):  # -> np.ndarray | Tensor
        state_t = (
            state if isinstance(state, Tensor) else torch.as_tensor(state, dtype=torch.float32)
        )
        if state_t.dim() == 1:
            state_t = state_t.unsqueeze(0)
        state_t = state_t.to(self.device).float()
        a, _ = self.actor.sample(state_t, deterministic=not explore)
        a = a[0] if a.dim() > 1 else a
        return (a * self._action_limit).cpu().numpy()

    def update(self, batch: Dict[str, Tensor]) -> Dict[str, float]:
        states = batch["states"].to(self.device).float()
        actions = batch["actions"].to(self.device).float()
        rewards = batch["rewards"].to(self.device).float().view(-1, 1)
        next_states = batch["next_states"].to(self.device).float()
        dones = batch["dones"].to(self.device).float().view(-1, 1)

        # Critic update
        with torch.no_grad():
            a_next, logp_next = self.actor.sample(next_states, deterministic=False)
            a_next = a_next.clamp(-1.0, 1.0) * self._action_limit
            q1_t = self.q1_target(next_states, a_next)
            q2_t = self.q2_target(next_states, a_next)
            q_t_min = torch.min(q1_t, q2_t)
            target = rewards + self.cfg.gamma * (1.0 - dones) * (
                q_t_min - self.cfg.alpha * logp_next
            )

        q1 = self.q1(states, actions)
        q2 = self.q2(states, actions)
        q1_loss = nn.functional.mse_loss(q1, target)
        q2_loss = nn.functional.mse_loss(q2, target)

        self.q1_optim.zero_grad()
        q1_loss.backward()
        self.q1_optim.step()

        self.q2_optim.zero_grad()
        q2_loss.backward()
        self.q2_optim.step()

        # Actor update
        a_pi, logp = self.actor.sample(states, deterministic=False)
        a_pi = a_pi.clamp(-1.0, 1.0) * self._action_limit
        q1_pi = self.q1(states, a_pi)
        q2_pi = self.q2(states, a_pi)
        q_pi = torch.min(q1_pi, q2_pi)
        pi_loss = (self.cfg.alpha * logp - q_pi).mean()

        self.pi_optim.zero_grad()
        pi_loss.backward()
        self.pi_optim.step()

        # Soft target updates
        self._soft_update(self.q1_target, self.q1)
        self._soft_update(self.q2_target, self.q2)

        return {
            "q1_loss": float(q1_loss.item()),
            "q2_loss": float(q2_loss.item()),
            "pi_loss": float(pi_loss.item()),
        }

    def _soft_update(self, target: nn.Module, source: nn.Module) -> None:
        with torch.no_grad():
            for p_t, p in zip(target.parameters(), source.parameters()):
                p_t.data.mul_(1.0 - self.cfg.tau).add_(p.data, alpha=self.cfg.tau)


__all__ = ["SACAgent", "SACConfig"]
