import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
import math

from GANmodels import Generator

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

    def forward(self, x):
        m = self.backbone(x)
        s = torch.exp(self.log_std.clamp(self._log_std_min, self._log_std_max))
        return m, s


    def sample(self, x):
        m, s = self.forward(x)
        z = torch.randn_like(m)
        raw = m + s * z
        delta = torch.tanh(raw) * self.eps
        logp_elem = -0.5 * ((raw - m) / (s + 1e-8))**2 - torch.log(s + 1e-8) - 0.5 * math.log(2.0 * math.pi)
        logp = logp_elem.view(raw.size(0), -1).sum(dim=-1)
        numel = m[0].numel()
        ent_per_elem = 0.5 * (1.0 + math.log(2.0 * math.pi)) + torch.log(s + 1e-8)
        ent = (ent_per_elem * numel).expand(raw.size(0))
        return delta, logp, ent, raw


    def log_prob_of(self, x, raw):
        m, s = self.forward(x)
        logp_elem = -0.5 * ((raw - m) / (s + 1e-8))**2 - torch.log(s + 1e-8) - 0.5 * math.log(2.0 * math.pi)
        logp = logp_elem.view(raw.size(0), -1).sum(dim=-1)
        numel = m[0].numel()
        ent_per_elem = 0.5 * (1.0 + math.log(2.0 * math.pi)) + torch.log(s + 1e-8)
        ent = (ent_per_elem * numel).expand(raw.size(0))
        return logp, ent



class AdversaryValue(nn.Module):
    def __init__(self, in_channels, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, hidden, 3, 2, 1), nn.ReLU(),
            nn.Conv2d(hidden, hidden, 3, 1, 1), nn.ReLU(),
            nn.ConvTranspose2d(hidden, hidden, 4, 2, 1), nn.ReLU(),
            nn.Conv2d(hidden, 1, 3, 1, 1)
        )

    def forward(self, x):
        v_map = self.net(x)
        return v_map.mean(dim=[1,2,3], keepdim=False)

class OptimalAdversaryTrainer:
    def __init__(self, policy: AdversaryPolicy, victim_agent, device,
                 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):
        self.pi_adv = policy.to(device)
        self.v_net = AdversaryValue(in_channels=policy.out_channels).to(device)
        self.victim = victim_agent
        self.device = device

        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.gae_lambda = gae_lambda
        self.ppo_epochs = ppo_epochs
        self.minibatch_size = minibatch_size

    @torch.no_grad()
    def _proj_box(self, x):
        return x.clamp(0.0, 1.0)

    @torch.no_grad()
    def _collect_batch(self, env, B):
        obs_list, raw_list, logp_list, ent_list = [], [], [], []
        rew_list, val_list = [], []
        next_obs = env.reset()
        for _ in range(B):
            x = torch.as_tensor(next_obs, device=self.device, dtype=torch.float32).unsqueeze(0) / 255.0
            v = self.v_net(x)
            delta, logp, ent, raw = self.pi_adv.sample(x)
            x_adv = self._proj_box(x + delta)
            x_adv = torch.where(torch.isfinite(x_adv), x_adv, x)

            input_victim = x_adv * 255.0
            dist = self.victim.actor(input_victim)
            a = dist.mean[0].cpu().numpy()

            s_next, r, done, _ = env.step(a)

            obs_list.append(x)
            raw_list.append(raw.detach())
            logp_list.append(logp.detach())
            ent_list.append(ent.detach())
            rew_list.append(torch.tensor([-r], device=self.device, dtype=torch.float32))  # maximize -r
            val_list.append(v.detach())

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

        obs = torch.cat(obs_list, dim=0)
        raw = torch.cat(raw_list, dim=0)
        old_logp = torch.stack(logp_list, dim=0)
        old_ent = torch.stack(ent_list, dim=0)
        rewards = torch.stack(rew_list, dim=0).squeeze(-1)
        values = torch.cat(val_list, dim=0)
        return obs, raw, old_logp, old_ent, rewards, values

    def _compute_gae(self, rewards, values):
        with torch.no_grad():
            deltas = rewards + self.gamma * torch.cat([values[1:], values[-1:]], dim=0) - values
            adv = torch.zeros_like(rewards)
            gae = 0.0
            for t in reversed(range(len(rewards))):
                gae = deltas[t] + self.gamma * self.gae_lambda * gae
                adv[t] = gae
            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)
        obs, raw, old_logp, adv, ret = obs[idx], raw[idx], old_logp[idx], adv[idx], ret[idx]
        
        for _ in range(self.ppo_epochs):
            for start in range(0, N, self.minibatch_size):
                end = min(start + self.minibatch_size, N)
                mb_obs  = obs[start:end]
                mb_raw  = raw[start:end]
                mb_oldp = old_logp[start:end]
                mb_adv  = adv[start:end]
                mb_ret  = ret[start:end]

                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, B, logger=None):
        for i in range(n_iter):
            obs, raw, old_logp, old_ent, 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("opt_adv/reward_mean", rewards.mean(), i)
                logger.log("opt_adv/adv_std", adv.std(), i)

def build_optimal_adversary(in_channels, eps, victim_agent, device, lr=1e-4):
    pi = AdversaryPolicy(in_channels=in_channels, eps=eps)
    return OptimalAdversaryTrainer(
        policy=pi,
        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
    )