from typing import Tuple, Optional
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal
from torch.optim import Adam
import GANmodels

class AdversaryPolicy(nn.Module):
    def __init__(self, in_channels=3, eps=8/255.0):
        super().__init__()
        self.backbone = GANmodels.Generator(in_channels, in_channels)
        self.eps = eps
        self._log_std_min, self._log_std_max = -5.0, 2.0
        self.log_std = nn.Parameter(torch.tensor(-2.0))

    def forward_raw(self, x):
        delta_mean = self.backbone(x)   
        return delta_mean.view(x.size(0), -1)

    def _dist(self, loc):
        log_std = self.log_std
        log_std = torch.where(torch.isfinite(log_std), log_std, torch.zeros_like(log_std))
        log_std = log_std.clamp(self._log_std_min, self._log_std_max)
        std = log_std.exp()
        std = torch.clamp(std, min=1e-6)
        std = torch.ones_like(loc) * std
        return Normal(loc, std)

    def _raw_to_delta(self, raw, img_shape):
        B, C, H, W = img_shape
        dimg = raw.view(B, C, H, W)
        dimg = torch.tanh(dimg) * self.eps
        return dimg

    def sample(self, x):
        loc = self.forward_raw(x)
        dist = self._dist(loc)
        raw = dist.rsample()
        raw = torch.where(torch.isfinite(raw), raw, torch.zeros_like(raw))
        logp = dist.log_prob(raw).sum(-1)
        ent = dist.entropy().sum(-1)
        delta_img = self._raw_to_delta(raw, x.shape)
        return delta_img, logp, ent, raw

    def log_prob_of(self, x, raw):
        loc = self.forward_raw(x)
        raw = torch.where(torch.isfinite(raw), raw, torch.zeros_like(raw))
        dist = self._dist(loc)
        logp = dist.log_prob(raw).sum(-1)
        ent = dist.entropy().sum(-1)
        return logp, ent

class ConvEncoder(nn.Module):
    def __init__(self, in_ch=3, feat_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_ch, 32, 8, 4), nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, 4, 2),    nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, 1),    nn.ReLU(inplace=True),
            nn.Flatten()
        )
        self.fc = nn.Linear(64 * 7 * 7, feat_dim)
        self.out_dim = feat_dim

    def forward(self, x):
        h = self.net(x)
        if h.shape[-1] != self.fc.in_features:
            self.fc = nn.Linear(h.shape[-1], self.out_dim).to(h.device)
        return torch.relu(self.fc(h))


class AdversaryValue(nn.Module):
    def __init__(self, in_channels=3, latent_dim=256):
        super().__init__()
        self.encoder = ConvEncoder(in_ch=in_channels, feat_dim=latent_dim)
        self.v = nn.Linear(latent_dim, 1)

    def forward(self, x):
        z = self.encoder(x)
        return self.v(z).squeeze(-1)


class NonRLActor:
    def __init__(self, eps: float, K: int = 4, sigma: float = 0.02):
        self.eps = eps
        self.K = K
        self.sigma = sigma

    @torch.no_grad()
    def refine(self, x, raw_center, victim_agent, env):
        device = x.device
        B, C, H, W = x.shape

        def raw_to_delta(r):
            dvec = torch.tanh(r) * self.eps
            return dvec[:, :C].unsqueeze(-1).unsqueeze(-1).expand(B, C, H, W).contiguous()

        def score(delta_img):
            x_adv = (x + delta_img).clamp(0.0, 1.0)
            with torch.no_grad():
                a = victim_agent.act(x_adv)
            _, r, _, _ = env.step(a)
            return -float(r)

        best_raw = raw_center.clone()
        best_delta = raw_to_delta(best_raw)
        try:
            best_score = score(best_delta)
        except Exception:
            best_score = float("inf")

        for _ in range(self.K):
            cand_raw = raw_center + torch.randn_like(raw_center, device=device) * self.sigma
            cand_delta = raw_to_delta(cand_raw)
            try:
                s = score(cand_delta)
            except Exception:
                s = float("inf")
            if s < best_score:
                best_score, best_raw, best_delta = s, cand_raw, cand_delta
        return best_delta

class PAADTrainer:
    def __init__(self, policy: AdversaryPolicy, value: AdversaryValue, victim_agent,
                 device="cpu", lr_pi=1e-4, lr_v=1e-4, clip_ratio=0.2, ent_coef=0.0,
                 gamma=0.99, gae_lambda=0.95, ppo_epochs=4, minibatch_size=64,
                 nonrl_K=4, nonrl_sigma=0.02):
        self.device = torch.device(device)
        self.pi_adv = policy.to(self.device)
        self.v_net  = value.to(self.device)
        self.victim_agent = victim_agent

        self.opt_pi = Adam(self.pi_adv.parameters(), lr=lr_pi)
        self.opt_v  = Adam(self.v_net.parameters(), lr=lr_v)

        self.clip_ratio = clip_ratio
        self.ent_coef   = ent_coef
        self.gamma      = gamma
        self.lam        = gae_lambda
        self.ppo_epochs = ppo_epochs
        self.minibatch  = minibatch_size

        self.actor = NonRLActor(eps=self.pi_adv.eps, K=nonrl_K, sigma=nonrl_sigma) if nonrl_K > 0 else None

    def _to_tensor(self, obs_np):
        if isinstance(obs_np, dict):
            obs_np = next(iter(obs_np.values()))
        if obs_np.dtype == np.uint8:
            x = obs_np.astype(np.float32) / 255.0
            if x.ndim == 3 and x.shape[-1] in (1, 3):
                x = np.transpose(x, (2, 0, 1))
        else:
            x = obs_np.astype(np.float32)
            if x.ndim == 3 and x.shape[-1] in (1, 3):
                x = np.transpose(x, (2, 0, 1))
        return torch.from_numpy(x).unsqueeze(0).to(self.device)

    def _collect_batch(self, env, B: int):
        obs_buf, raw_buf, old_logp_buf, rew_buf, val_buf = [], [], [], [], []
        obs = env.reset()
        for _ in range(B):
            x = self._to_tensor(obs)
            delta_img, logp, ent, raw = self.pi_adv.sample(x)

            if self.actor is not None:
                delta_img = self.actor.refine(x, raw, self.victim_agent, env)

            x_adv = (x + delta_img).clamp(0.0, 1.0)
            with torch.no_grad():
                a = self.victim_agent.act(x_adv)
            next_obs, r, done, _ = env.step(a)

            obs_buf.append(x)
            raw_buf.append(raw.detach())
            old_logp_buf.append(logp.detach())
            rew_buf.append(torch.tensor([r], dtype=torch.float32, device=self.device))
            val_buf.append(self.v_net(x).detach())

            obs = next_obs if not done else env.reset()

        obs_t = torch.cat(obs_buf, dim=0)
        raw_t = torch.cat(raw_buf, dim=0)
        logp_t = torch.cat(old_logp_buf, dim=0).squeeze(-1)
        rews_t = torch.cat(rew_buf, dim=0).squeeze(-1)
        values_t = torch.cat(val_buf, dim=0).squeeze(-1)
        return obs_t, raw_t, logp_t, rews_t, values_t

    def _compute_gae(self, rewards, values):
        T = rewards.size(0)
        adv = torch.zeros_like(rewards)
        lastgaelam = 0.0
        for t in reversed(range(T)):
            next_v = values[t+1] if t+1 < T else 0.0
            delta = rewards[t] + self.gamma * next_v - values[t]
            lastgaelam = delta + self.gamma * self.lam * lastgaelam
            adv[t] = lastgaelam
        ret = adv + values
        adv = (adv - adv.mean()) / (adv.std() + 1e-8)
        return adv, ret

    def _ppo_update(self, obs, raw, old_logp, adv, ret):
        N = obs.size(0)
        idx = torch.randperm(N, device=self.device)
        for _ in range(self.ppo_epochs):
            for s in range(0, N, self.minibatch):
                e = min(s + self.minibatch, N)
                mb = idx[s:e]
                mb_obs, mb_raw, mb_oldp = obs[mb], raw[mb], old_logp[mb]
                mb_adv, mb_ret = adv[mb], ret[mb]

                new_logp, ent = self.pi_adv.log_prob_of(mb_obs, mb_raw)
                ratio = (new_logp - mb_oldp).exp()
                surr1 = ratio * mb_adv
                surr2 = torch.clamp(ratio, 1.0 - self.clip_ratio, 1.0 + self.clip_ratio) * mb_adv
                pi_loss = -(torch.min(surr1, surr2)).mean() - self.ent_coef * ent.mean()

                v_pred = self.v_net(mb_obs)
                v_loss = F.mse_loss(v_pred, mb_ret)

                self.opt_pi.zero_grad()
                pi_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.pi_adv.parameters(), 1.0)
                self.opt_pi.step()
                with torch.no_grad():
                    self.pi_adv.log_std.clamp_(self.pi_adv._log_std_min, self.pi_adv._log_std_max)

                self.opt_v.zero_grad()
                v_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.v_net.parameters(), 1.0)
                self.opt_v.step()

    def train(self, env, n_iter: int, B: int, logger=None):
        for i in range(n_iter):
            obs, raw, old_logp, rewards, values = self._collect_batch(env, B)
            adv, ret = self._compute_gae(rewards, values)
            self._ppo_update(obs, raw, old_logp, adv, ret)
            if logger is not None:
                logger.log("paad/reward_mean", rewards.mean(), i)
                logger.log("paad/adv_std", adv.std(), i)


def build_paad_adversary(in_channels, eps, victim_agent, device, lr=1e-4):
    pi = AdversaryPolicy(in_channels=in_channels, eps=eps)
    v  = AdversaryValue(in_channels=in_channels)
    return PAADTrainer(
        policy=pi, value=v, victim_agent=victim_agent, device=device,
        lr_pi=lr, lr_v=lr, clip_ratio=0.2, ent_coef=0.0,
        gamma=0.99, gae_lambda=0.95, ppo_epochs=4, minibatch_size=64,
        nonrl_K=4, nonrl_sigma=0.02
    )
