import numpy as np
import torch
import torch.nn.functional as F
from torch_geometric.data import Batch
from torch_scatter import scatter_add


class RolloutBuffer:

    def __init__(self, buffer_size: int,
        gae_lambda: float,
        discount_factor: float,
        device=None,
    ):
        self.buffer_size = buffer_size
        self.gae_lambda = gae_lambda
        self.discount_factor = discount_factor
        self.device = device
        self.reset()

    def reset(self):
        self.observations = []
        self.actions = []
        self.log_probs = []
        self.rewards = []  # per-agent rewards (list of tensors, varying shapes)
        self.values = []  # per-agent values (list of tensors, varying shapes)
        self.dones = []
        self.agent_mappings = []  # mapping from new → old agent indices per step
        self.pos = 0

    def add(
        self, observation, actions, reward, done, value, log_probs,
        agent_mapping=None,
    ):
        dev = self.device
        self.observations.append(observation.to(dev))
        self.actions.append(actions.to(dev))
        self.log_probs.append(log_probs.to(dev))
        self.rewards.append(torch.as_tensor(reward, dtype=torch.float32, device=dev).flatten())
        self.values.append(value.flatten().to(dev))
        self.dones.append(float(done))
        self.agent_mappings.append(
            torch.as_tensor(agent_mapping, dtype=torch.long, device=dev).flatten()
        )
        self.pos += 1

    def compute_returns_and_advantage(self, last_value):
        """Single-path GAE: potential-shaped per-agent reward with scatter_add for mesh refinement."""
        last_value = last_value.to(self.device).flatten()
        n = self.buffer_size

        dones = torch.as_tensor(self.dones, device=self.device)

        # ---- 0. Normalize rewards to unit scale ----
        all_rews = torch.cat([r.flatten() for r in self.rewards])
        rew_mean = all_rews.mean()
        rew_std = all_rews.std()
        if rew_std > 1e-8:
            self.rewards = [(r - rew_mean) / rew_std for r in self.rewards]

        # ---- 1. Per-agent GAE (scatter_add for mesh refinement) ----
        advantages = [None] * n
        deltas = []
        next_values = self.values[1:] + [last_value]

        for step in range(n):
            if dones[step]:
                next_val = self.values[step]
            else:
                next_val = scatter_add(next_values[step], self.agent_mappings[step], dim=0)
            delta = self.rewards[step] + (0 if dones[step] else self.discount_factor * next_val) - self.values[step]
            deltas.append(delta)

        last_gae = torch.zeros_like(self.agent_mappings[-1], dtype=torch.float32, device=self.device)
        for step in reversed(range(n)):
            if dones[step]:
                last_gae = deltas[step]
            else:
                last_gae = deltas[step] + self.discount_factor * self.gae_lambda * scatter_add(last_gae, self.agent_mappings[step], dim=0)
            advantages[step] = last_gae

        self.returns = [adv + val for adv, val in zip(advantages, self.values)]

        # ---- 2. Normalize advantages (per-batch, zero-mean unit-std) ----
        all_advs = torch.cat([a.flatten() for a in advantages])
        adv_mean = all_advs.mean()
        adv_std = all_advs.std()
        if adv_std > 1e-8:
            advantages = [(a - adv_mean) / adv_std for a in advantages]
        # NOTE: returns and values keep their original scale — no unit-scale normalization,
        # so the value network sees a stable regression target across iterations.

        self.advantages = [ret - val for ret, val in zip(self.returns, self.values)]

    def get(self, batch_size: int):
        """Yield random minibatches from the buffer."""
        indices = np.random.permutation(self.buffer_size)
        start = 0
        while start < self.buffer_size:
            batch_idx = indices[start : start + batch_size]
            start += batch_size

            obs_batch = Batch.from_data_list([self.observations[i] for i in batch_idx])
            acts = torch.cat([self.actions[i] for i in batch_idx], dim=0)
            lps = torch.cat([self.log_probs[i].flatten() for i in batch_idx], dim=0)
            vals = torch.cat([self.values[i].flatten() for i in batch_idx], dim=0)
            advs = torch.cat([self.advantages[i].flatten() for i in batch_idx], dim=0)
            rets = torch.cat([self.returns[i].flatten() for i in batch_idx], dim=0)

            obs_batch, acts, lps, vals, advs, rets = (
                x.to(self.device) for x in (obs_batch, acts, lps, vals, advs, rets)
            )
            yield obs_batch, acts, lps, vals, advs, rets

    @property
    def full(self):
        return self.pos >= self.buffer_size

    @property
    def explained_variance(self):
        all_vals = torch.cat([v.flatten() for v in self.values])
        all_rets = torch.cat([r.flatten() for r in self.returns])
        var_ret = torch.var(all_rets)
        if var_ret < 1e-12:
            return 0.0
        return float(1.0 - torch.var(all_rets - all_vals) / var_ret)


# ── PPO losses ────────────────────────────────────────────
def policy_loss(advantages: torch.Tensor, ratio: torch.Tensor, clip_range: float) -> torch.Tensor:
    """Clipped PPO policy loss."""
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
    loss1 = advantages * ratio
    loss2 = advantages * torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range)
    return -torch.min(loss1, loss2).mean()


def value_loss(
    returns: torch.Tensor, values: torch.Tensor,
    old_values: torch.Tensor, clip_range: float,
) -> torch.Tensor:
    """Clipped value function loss."""
    vf_loss = F.mse_loss(returns, values)
    if clip_range > 0:
        v_clipped = old_values + (values - old_values).clamp(-clip_range, clip_range)
        vf_loss = torch.max(vf_loss, F.mse_loss(returns, v_clipped))
    return vf_loss


def entropy_loss(entropy) -> torch.Tensor:
    """Entropy bonus for exploration."""
    return -torch.mean(entropy)


class PPOTrainer:

    def __init__(self, actor_critic, environment, config: dict, device=None):
        self.policy = actor_critic
        self.env = environment
        self.device = device

        ppo_cfg = config.get("ppo", {})
        self.num_rollout_steps = ppo_cfg.get("num_rollout_steps", 256)
        self.epochs_per_iteration = ppo_cfg.get("epochs_per_iteration", 5)
        self.batch_size = config.get("batch_size", 32)
        self.clip_range = ppo_cfg.get("clip_range", 0.2)
        self.max_grad_norm = ppo_cfg.get("max_grad_norm", 0.5)
        self.entropy_coef = ppo_cfg.get("entropy_coefficient", 0.0)
        self.vf_coef = ppo_cfg.get("value_function_coefficient", 0.5)
        self.vf_clip_range = ppo_cfg.get("value_function_clip_range", 0.2)
        self.gae_lambda = ppo_cfg.get("gae_lambda", 0.95)
        self.discount_factor = config.get("discount_factor", 1.0)

        self.buffer = RolloutBuffer(
            buffer_size=self.num_rollout_steps,
            gae_lambda=self.gae_lambda,
            discount_factor=self.discount_factor,
            device=device,
        )

    def collect_rollouts(self):
        self.policy.eval()
        self.buffer.reset()
        obs = self.env.reset()
        step_rewards, step_num_agents = [], []
        _rho_keys = ("rho_int_mean", "rho_jump_mean", "rho_sbc_mean",
                      "w_rho_int", "w_rho_jump", "w_rho_sbc")
        rho_accum = {k: 0.0 for k in _rho_keys}
        diag_keys = ("neg_action_ratio", "eligible_ratio", "selected_count",
                      "dorfler_tail_ratio", "dorfler_floor_active")
        diag_accum = {k: 0.0 for k in diag_keys}
        diag_steps = 0

        for _ in range(self.num_rollout_steps):
            with torch.no_grad():
                actions, values, log_probs = self.policy(
                    Batch.from_data_list([obs]), deterministic=False
                )
            values = values.flatten()
            next_obs, reward, done, info = self.env.step(actions.cpu().numpy())
            step_rewards.append(float(np.sum(reward)))
            step_num_agents.append(int(len(reward)))
            for k in _rho_keys:
                if k in info:
                    rho_accum[k] += float(info[k])
            for k in diag_keys:
                if k in info:
                    diag_accum[k] += float(info[k])
            diag_steps += 1

            self.buffer.add(
                observation=obs, actions=actions, reward=reward,
                done=float(done), value=values, log_probs=log_probs,
                agent_mapping=self.env.agent_mapping,
            )
            obs = self.env.reset() if done else next_obs

        with torch.no_grad():
            _, last_value, _ = self.policy(Batch.from_data_list([obs]), deterministic=True)
        last_value = last_value.squeeze(-1).flatten()
        self.buffer.compute_returns_and_advantage(last_value)

        n = max(1, self.num_rollout_steps)
        metrics = {
            "num_agents": step_num_agents[-1], "reward": step_rewards[-1],
            "avg_agents": np.mean(step_num_agents),
            "avg_reward": np.mean(step_rewards),
            "min_reward": np.min(step_rewards),
            "max_reward": np.max(step_rewards),
            "sum_reward": np.sum(step_rewards),
        }
        # rho diagnostics for weight calibration (averaged over rollout)
        for k in _rho_keys:
            metrics[k] = rho_accum[k] / n
        # score-based refinement diagnostics
        n_diag = max(1, diag_steps)
        for k in diag_keys:
            metrics[k] = diag_accum[k] / n_diag
        return metrics

    def train_step(self):
        self.policy.train()
        total_losses = []
        for _ in range(self.epochs_per_iteration):
            for obs_batch, acts, old_lp, old_vals, advs, rets in self.buffer.get(self.batch_size):
                values, log_probs, entropy = self.policy.evaluate_actions(obs_batch, acts)
                values = values.squeeze(-1)

                ratio = torch.exp(log_probs - old_lp)

                pl = policy_loss(advs, ratio, self.clip_range)
                vl = self.vf_coef * value_loss(rets, values, old_vals, self.vf_clip_range)
                el = self.entropy_coef * entropy_loss(entropy)
                loss = pl + vl + el

                self.policy.optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
                self.policy.optimizer.step()
                if self.policy.log_std is not None:
                    self.policy.log_std.data.clamp_(-3.0, -1.0)  # σ ∈ [0.05, 0.37]
                total_losses.append(loss.item())

        if self.policy.lr_scheduler is not None:
            self.policy.lr_scheduler.step()

        return {
            "loss": np.mean(total_losses) if total_losses else 0.0,
            "explained_variance": self.buffer.explained_variance,
        }

    def fit_iteration(self):
        metrics = self.collect_rollouts()
        metrics.update(self.train_step())
        return metrics