import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
from tqdm.auto import tqdm

# --- 1. Simulation Dataset ---
class ToyDataset(Dataset):
    def __init__(self, num_samples, data_dim):
        # 创建 num_samples 个独立的、高斯分布的“原型”数据
        self.prototypes = torch.randn(num_samples, data_dim)

    def __len__(self):
        return len(self.prototypes)

    def __getitem__(self, idx):
        # 对同一个原型添加两次不同的噪声，生成两个视图
        prototype = self.prototypes[idx]
        view1 = prototype + torch.randn_like(prototype) * 0.1 # view1
        view2 = prototype + torch.randn_like(prototype) * 0.1 # view2
        return view1, view2, idx

# --- 2. Simple MLP Model ---
class MLPEncoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.net(x)
        
# --- 3. Contrastive Loss with Memory Bank ---
class NCECriterion(nn.Module):
    def __init__(self, num_samples):
        super().__init__()
        self.num_samples = num_samples

    def forward(self, x, targets, memory_bank, temperature):
        sim = torch.einsum('bd,nd->bn', x, memory_bank) / temperature
        loss = F.cross_entropy(sim, targets)
        return loss

# --- Hyperparameters & Setup ---
torch.manual_seed(42)
np.random.seed(42)

# 数据参数
NUM_SAMPLES = 2048   # 数据集大小，也是 Memory Bank 大小
DATA_DIM = 512       # 原始数据维度
PROJECTION_DIM = 128 # 对比学习的特征维度

# 训练参数
BATCH_SIZE = 128
EPOCHS = 10
LEARNING_RATE = 1e-3
TEMPERATURE = 0.07   # 对比学习温度

def run_experiment():
    # 初始化数据
    dataset = ToyDataset(NUM_SAMPLES, DATA_DIM)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
    
    # 初始化模型、Memory Bank 和损失函数
    model = MLPEncoder(DATA_DIM, 512, PROJECTION_DIM)
    memory_bank = F.normalize(torch.randn(NUM_SAMPLES, PROJECTION_DIM), dim=1)
    criterion = NCECriterion(NUM_SAMPLES)
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
    
    # 用于记录损失
    step_losses = []
    
    print("Starting toy model training...")
    global_step = 0
    for epoch in range(EPOCHS):
        progress_bar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{EPOCHS}")
        for view1, view2, indices in progress_bar:
            model.train()
            
            # --- 前向传播 ---
            q = model(view1) # Query
            k = model(view2) # Key
            
            # 归一化
            q = F.normalize(q, dim=1)
            k = F.normalize(k, dim=1)
            
            # --- 计算损失 ---
            loss1 = criterion(q, indices, memory_bank, TEMPERATURE)
            loss2 = criterion(k, indices, memory_bank, TEMPERATURE)
            loss = (loss1 + loss2) / 2
            
            # --- 反向传播与优化 ---
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # --- 更新 Memory Bank ---
            with torch.no_grad():
                memory_bank[indices] = k.detach()
            
            # --- 记录与显示 ---
            step_losses.append(loss.item())
            progress_bar.set_postfix(loss=f"{loss.item():.4f}")
            global_step += 1

    print("Training finished.")
    return step_losses

training_losses = run_experiment()

# 绘制损失曲线
plt.style.use('seaborn-v0_8-whitegrid')
fig, ax = plt.subplots(figsize=(12, 7))

# 原始损失
ax.plot(training_losses, alpha=1, label='Raw Step Loss')

# 理论初始损失
initial_loss_theory = np.log(NUM_SAMPLES)
ax.axhline(y=initial_loss_theory, color='red', linestyle='--', label=f'Theoretical Initial Loss: log({NUM_SAMPLES}) ≈ {initial_loss_theory:.2f}')

ax.set_title('Contrastive Loss Curve', fontsize=16, pad=20)
ax.set_xlabel('Training Steps', fontsize=12)
ax.set_ylabel('InfoNCE Loss', fontsize=12)
ax.legend(fontsize=11)
ax.set_ylim(bottom=0)

plt.tight_layout()
plt.show()