import torch
import torch.nn.functional as F
 
def model_inversion_basic(model, target_class, image_size=(3, 224, 224),
                          steps=2000, lr=0.01):
    """重建目標類別的代表性輸入。"""
    x = torch.randn(1, *image_size, requires_grad=True)
    optimizer = torch.optim.Adam([x], lr=lr)
 
    for step in range(steps):
        optimizer.zero_grad()
        output = model(x)
 
        # 最大化目標類別的機率
        loss = -F.log_softmax(output, dim=1)[0, target_class]
 
        # 正規化:平滑度的總變異
        tv_loss = torch.sum(torch.abs(x[:, :, :, :-1] - x[:, :, :, 1:])) + \
                  torch.sum(torch.abs(x[:, :, :-1, :] - x[:, :, 1:, :]))
 
        total_loss = loss + 0.001 * tv_loss
        total_loss.backward()
        optimizer.step()
 
        with torch.no_grad():
            x.clamp_(0, 1)
 
    return x.detach()

def gradient_inversion(model, target_gradient, batch_size=1,
                       image_size=(3, 32, 32), steps=5000):
    """從共享梯度重建訓練資料。"""
    dummy_x = torch.randn(batch_size, *image_size, requires_grad=True)
    dummy_y = torch.randint(0, 10, (batch_size,))
 
    optimizer = torch.optim.LBFGS([dummy_x], lr=0.1)
 
    for step in range(steps):
        def closure():
            optimizer.zero_grad()
            output = model(dummy_x)
            loss = F.cross_entropy(output, dummy_y)
            dummy_gradient = torch.autograd.grad(loss, model.parameters(),
                                                  create_graph=True)
 
            # 最小化假梯度與目標梯度之間的距離
            grad_loss = sum(
                ((dg - tg) ** 2).sum()
                for dg, tg in zip(dummy_gradient, target_gradient)
            )
            grad_loss.backward()
            return grad_loss
 
        optimizer.step(closure)
 
    return dummy_x.detach()

def gan_inversion(target_model, generator, target_class,
                  latent_dim=512, steps=3000, n_starts=5):
    """使用 GAN 先驗進行模型反演以獲得逼真重建。"""
    best_image = None
    best_conf = -float('inf')
 
    for _ in range(n_starts):
        z = torch.randn(1, latent_dim, requires_grad=True)
        optimizer = torch.optim.Adam([z], lr=0.02)
 
        for step in range(steps):
            optimizer.zero_grad()
            fake_image = generator(z)
            output = target_model(fake_image)
 
            conf = F.softmax(output, dim=1)[0, target_class]
            loss = -torch.log(conf)
            loss.backward()
            optimizer.step()
 
        final_image = generator(z).detach()
        final_conf = F.softmax(target_model(final_image), dim=1)[0, target_class].item()
 
        if final_conf > best_conf:
            best_conf = final_conf
            best_image = final_image
 
    return best_image, best_conf