import torch
import torch.nn as nn
import torch.nn.functional as F


class SelfAttention(nn.Module):
    def __init__(self,seq_length):
        super(SelfAttention,self).__init__()
        self.input_size = seq_length
        # 定义三个权重矩阵
        self.Wq=nn.Linear(seq_length,seq_length)#不改变形状的线性变换
        self.Wk=nn.Linear(seq_length,seq_length)
        self.Wv=nn.Linear(seq_length,seq_length)
        
    def forward(self,input):
        # 计算Q,K,V 三个矩阵
        q = self.Wq(input)
        k = self.Wk(input)
        v = self.Wv(input)
        
        # 计算QK^T，即向量之间的相关度 ; 这里可以理解dk了：torch.tensor(float(self.input_size))，是Wk的维度。
        attention_scores = torch.matmul(q, k.transpose(-1,-2))/torch.sqrt(torch.tensor(float(self.input_size)))
        # 计算向量权重，softmax归一化
        attention_weight = F.softmax(attention_scores, dim=-1)
        # 计算输出
        output = torch.matmul(attention_weight, v)
        return output

import torch
import torch.nn as nn
import torch.nn.functional as F

class Unet(nn.Module):
    def __init__(self, in_channels, out_channels, ch=128):
        super().__init__()
        self.down = nn.ModuleList([
            Block(in_channels, ch, ch * 2),
            Block(ch * 2, ch * 2, ch * 4),
            Block(ch * 4, ch * 4, ch * 8),
            Block(ch * 8, ch * 8, ch * 16),
        ])
        self.up = nn.ModuleList([
            Block(ch * 16, ch * 8, ch * 8),
            Block(ch * 8, ch * 4, ch * 4),
            Block(ch * 4, ch * 2, ch * 2),
            Block(ch * 2, ch, ch),
        ])
        self.final_conv = nn.Conv2d(ch, out_channels, kernel_size=1)

    def forward(self, x):
        skip_connections = []
        for down_block in self.down:
            x = down_block(x)
            skip_connections.append(x)
        skip_connections = skip_connections[::-1]
        for i, up_block in enumerate(self.up):
            x = up_block(x, skip_connections[i])
        return self.final_conv(x)

class Block(nn.Module):
    def __init__(self, in_channels, out_channels, hidden_channels=None):
        super().__init__()
        hidden_channels = hidden_channels or out_channels
        self.conv1 = nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(hidden_channels, out_channels, kernel_size=3, padding=1)
        self.time_emb = nn.Linear(1, out_channels)

    def forward(self, x, skip=None):
        t = torch.randint(0, 1000, (x.shape[0], 1, 1, 1), device=x.device).float() / 1000.
        t_emb = self.time_emb(t)
        h = F.relu(self.conv1(x))
        h = self.conv2(h)
        if skip is not None:
            h = torch.cat([h, skip], dim=1)
        return F.relu(h + t_emb)

class DDPM(nn.Module):
    def __init__(self, image_size, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        super().__init__()
        self.image_size = image_size
        self.timesteps = timesteps
        self.beta = torch.linspace(beta_start, beta_end, timesteps)
        self.alpha = 1. - self.beta
        self.alpha_bar = torch.cumprod(self.alpha, dim=0)
        self.unet = Unet(3, 3)

    def forward(self, x_0, t):
        # Sample noise from a Gaussian distribution
        epsilon = torch.randn_like(x_0)
        # Apply the forward diffusion process
        x_t = torch.sqrt(self.alpha_bar[t]) * x_0 + torch.sqrt(1. - self.alpha_bar[t]) * epsilon
        # Predict the noise using the U-Net
        epsilon_theta = self.unet(x_t, t)
        return epsilon_theta

    def sample(self, batch_size=16):
        # Sample noise from a Gaussian distribution
        x_t = torch.randn(batch_size, 3, self.image_size, self.image_size)
        # Iterate over the reverse diffusion process
        for t in reversed(range(self.timesteps)):
            # Predict the noise using the U-Net
            epsilon_theta = self.unet(x_t, t)
            # Apply the reverse diffusion process
            x_t = (x_t - torch.sqrt(1. - self.alpha[t]) * epsilon_theta) / torch.sqrt(self.alpha[t])
        return x_t

# Create a DDPM model
model = DDPM(image_size=32)

# Train the model
# ...

# Sample images from the model
samples = model.sample()

# Display the sampled images
# ...

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

#定义Encoder and Decoder（只用简单的线性层）
class Encoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(Encoder, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        mu = self.fc_mu(x)
        logvar = self.fc_logvar(x)
        return mu, logvar

class Decoder(nn.Module):
    def __init__(self, latent_dim, hidden_dim, output_dim):
        super(Decoder, self).__init__()
        self.fc1 = nn.Linear(latent_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, z):
       

 z = torch.relu(self.fc1(z))
        x_hat = torch.sigmoid(self.fc2(z))
        return x_hat

#定义模型
class VAE(nn.Module):
    def __init__(self, encoder, decoder):
        super(VAE, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
    
    def forward(self, x):
        mu, logvar = self.encoder(x)
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mu + eps * std
        x_hat = self.decoder(z)
        return x_hat, mu, logvar

#训练
def train_vae(model, dataloader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    for batch in dataloader:
        batch = batch.to(device)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(batch.view(-1, 784))
        loss = criterion(recon_batch, batch.view(-1, 784), mu, logvar)
        loss.backward()
        running_loss += loss.item()
        optimizer.step()
    return running_loss / len(dataloader.dataset)

# 定义超参数
input_dim = 784
hidden_dim = 256
latent_dim = 20
output_dim = 784
batch_size = 128
epochs = 20

# 创建数据加载器
transform = transforms.Compose([transforms.ToTensor()])
# 这一步加载的是MNIST数据集，可以尝试更换其他想要的数据集。
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 创建VAE模型
encoder = Encoder(input_dim, hidden_dim, latent_dim)
decoder = Decoder(latent_dim, hidden_dim, output_dim)
vae = VAE(encoder, decoder)

# 定义优化器和损失函数
optimizer = optim.Adam(vae.parameters(), lr=0.001)
criterion = VAE_loss()

# 训练VAE模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vae.to(device)
for epoch in range(epochs):
    train_loss = train_vae(vae, train_loader, optimizer, criterion, device)
    print(f"Epoch {epoch+1}/{epochs}, Loss: {train_loss:.4f}")

# 生成手写数字图像
with torch.no_grad():
    z = torch.randn(64, latent_dim).to(device)
    generated_images = vae.decoder(z).view(-1, 1, 28, 28).cpu()

# 显示生成的图像
fig, axes = plt.subplots(8, 8, figsize=(10, 10))
for i, ax in enumerate(axes.flatten()):
    ax.imshow(generated_images[i][0], cmap='gray')
    ax.axis('off')
plt.show()