import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np
from typing import Dict, Tuple, Optional
import logging
from collections import deque

logger = logging.getLogger(__name__)


class ActorCriticNetwork(nn.Module):
    def __init__(self, state_dim: int = 200, num_actions: int = 25, hidden_dim: int = 256):
        super().__init__()
        self.feature_net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
        )
        self.actor = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, num_actions),
        )
        self.critic = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 1),
        )

    def forward(self, state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        features = self.feature_net(state)
        logits = self.actor(features)
        policy = torch.softmax(logits, dim=-1)
        value = self.critic(features)
        return policy, value

    def get_action(self, state: torch.Tensor, deterministic: bool = False) -> Tuple[int, float, float]:
        policy, value = self.forward(state)
        dist = Categorical(policy)
        if deterministic:
            action = policy.argmax(dim=-1)
        else:
            action = dist.sample()
        log_prob = dist.log_prob(action)
        return action.item(), log_prob.item(), value.item()


class PPOAgent:
    _sd_cache: dict = {}

    @classmethod
    def _get_cached_state_dict(cls, path: str) -> Optional[dict]:
        import os as _os
        path = _os.path.abspath(_os.path.expanduser(path))
        if path not in cls._sd_cache:
            if not _os.path.exists(path):
                return None
            try:
                cls._sd_cache[path] = torch.load(path, map_location='cpu', weights_only=False)
            except Exception as e:
                logger.warning(f"[PPOAgent] Failed to load {path}: {e}")
                return None
        return cls._sd_cache.get(path)

    @classmethod
    def _clear_cache(cls):
        cls._sd_cache.clear()

    def __init__(self, state_dim: int = 200, num_actions: int = 25, config: Dict = None):
        self.config = config or {}
        self.state_dim = state_dim
        self.num_actions = num_actions

        hidden_dim = self.config.get("hidden_dim", 256)
        self.network = ActorCriticNetwork(state_dim, num_actions, hidden_dim)

        lr = self.config.get("learning_rate", 3e-4)
        self.optimizer = optim.Adam(self.network.parameters(), lr=lr)

        self.gamma = self.config.get("gamma", 0.99)
        self.gae_lambda = self.config.get("gae_lambda", 0.95)
        self.clip_epsilon = self.config.get("clip_epsilon", 0.2)
        self.entropy_coef = self.config.get("entropy_coef", 0.01)
        self.value_coef = self.config.get("value_coef", 0.5)
        self.max_grad_norm = self.config.get("max_grad_norm", 0.5)
        self.update_epochs = self.config.get("update_epochs", 10)
        self.batch_size = self.config.get("batch_size", 64)

        from core.rl.replay_buffer import EpisodeBuffer
        self.buffer = EpisodeBuffer()
        self.replay_buffer = self.buffer

        self.training_stats = {
            "policy_loss": deque(maxlen=100),
            "value_loss": deque(maxlen=100),
            "entropy": deque(maxlen=100),
            "total_loss": deque(maxlen=100),
        }

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.network.to(self.device)
        logger.info(f"PPOAgent initialized: state_dim={state_dim}, actions={num_actions}, device={self.device}")

    def select_action(self, state: np.ndarray, deterministic: bool = False) -> Tuple[int, float, float]:
        tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        with torch.no_grad():
            was_training = self.network.training
            if deterministic:
                self.network.eval()
            action, log_prob, value = self.network.get_action(tensor, deterministic=deterministic)
            if deterministic:
                self.network.train(was_training)
        return action, log_prob, value

    def store_transition(self, state: np.ndarray, action: int, reward: float,
                         done: bool, log_prob: float, value: float):
        self.buffer.push(state, action, reward, done, log_prob, value)

    def train(self) -> Optional[Dict]:
        if len(self.buffer) < self.batch_size:
            return {"total_loss": 0.0, "policy_loss": 0.0, "value_loss": 0.0, "entropy": 0.0}

        batch = self.buffer.get()
        states = torch.FloatTensor(batch["states"]).to(self.device)
        actions = torch.LongTensor(batch["actions"]).to(self.device)
        old_log = torch.FloatTensor(batch["log_probs"]).to(self.device)
        rewards = torch.FloatTensor(batch["rewards"]).to(self.device)
        dones = torch.FloatTensor(batch["dones"]).to(self.device)
        old_vals = torch.FloatTensor(batch["values"]).to(self.device)

        advantages, returns = self._compute_gae(rewards, old_vals, dones)
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        for _ in range(self.update_epochs):
            policy, values = self.network(states)
            values = values.squeeze()
            dist = Categorical(policy)
            log_probs = dist.log_prob(actions)
            entropy = dist.entropy().mean()

            ratio = torch.exp(log_probs - old_log)
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1.0 - self.clip_epsilon, 1.0 + self.clip_epsilon) * advantages
            policy_loss = -torch.min(surr1, surr2).mean()

            value_loss = nn.MSELoss()(values, returns)

            loss = policy_loss + self.value_coef * value_loss - self.entropy_coef * entropy

            self.optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.network.parameters(), self.max_grad_norm)
            self.optimizer.step()

        self.training_stats["policy_loss"].append(policy_loss.item())
        self.training_stats["value_loss"].append(value_loss.item())
        self.training_stats["entropy"].append(entropy.item())
        self.training_stats["total_loss"].append(loss.item())
        self.buffer.clear()

        return {
            "policy_loss": policy_loss.item(),
            "value_loss": value_loss.item(),
            "entropy": entropy.item(),
            "total_loss": loss.item(),
        }

    def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor, dones: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        advantages = []
        gae = 0.0
        for t in reversed(range(len(rewards))):
            next_val = 0 if t == len(rewards) - 1 else values[t + 1]
            delta = rewards[t] + self.gamma * next_val * (1 - dones[t]) - values[t]
            gae = delta + self.gamma * self.gae_lambda * (1 - dones[t]) * gae
            advantages.insert(0, gae)
        advantages = torch.FloatTensor(advantages).to(self.device)
        return advantages, advantages + values

    def save(self, path: str):
        import os as _os
        sd = {
            "network_state": self.network.state_dict(),
            "optimizer_state": self.optimizer.state_dict(),
            "config": self.config,
            "training_stats": dict(self.training_stats),
        }
        _os.makedirs(_os.path.dirname(path), exist_ok=True)
        torch.save(sd, path)
        self.__class__._sd_cache[_os.path.abspath(path)] = sd
        logger.info(f"Model saved to {path}")

    def load(self, path: str, weights_only: bool = False) -> bool:
        sd = self._get_cached_state_dict(path)
        if sd is None:
            logger.warning(f"Checkpoint not found: {path}")
            return False
        try:
            self.network.load_state_dict(sd["network_state"])
            if "optimizer_state" in sd:
                try:
                    self.optimizer.load_state_dict(sd["optimizer_state"])
                except Exception:
                    pass
            logger.info(f"Model loaded from {path}")
            return True
        except Exception as e:
            logger.warning(f"Failed to load checkpoint {path}: {e}")
            return False

    def get_stats(self) -> Dict:
        return {
            k: {"mean": float(np.mean(v)) if v else 0.0, "std": float(np.std(v)) if v else 0.0}
            for k, v in self.training_stats.items()
        }
