import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torchvision.transforms import Resize
from skimage.metrics import peak_signal_noise_ratio as psnr

# 生成示例图像（以单通道为例）
def generate_ground_truth(resolution=256):
    x = np.linspace(0, 1, resolution)
    y = np.linspace(0, 1, resolution)
    xx, yy = np.meshgrid(x, y)
    image = np.sin(10 * np.pi * xx) * np.cos(10 * np.pi * yy)  # 高频纹理
    return torch.tensor(image, dtype=torch.float32).unsqueeze(0).unsqueeze(0)  # [1,1,H,W]

# 定义隐式神经表示（INR）模型
class INRModel(nn.Module):
    def __init__(self, hidden_dim=256, freq_bands=10):
        super().__init__()
        self.freq_bands = freq_bands
        
        # 傅里叶坐标编码
        self.embedding = nn.Sequential(
            nn.Linear(2, 2 + 4 * freq_bands),
            nn.LayerNorm(2 + 4 * freq_bands)
        )
        
        # MLP主干网络
        self.mlp = nn.Sequential(
            nn.Linear(2 + 4 * freq_bands, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, coords):
        # coords: [N, 2], 输入坐标范围[-1,1]
        embedded = self.embedding(coords)
        return self.mlp(embedded).squeeze(-1)  # [N]

# 实验1：参数冻结测试
def freeze_test():
    # 生成真实数据
    hr_res = 256  # 高分辨率
    lr_res = 64   # 低分辨率
    hr_image = generate_ground_truth(hr_res)  # [1,1,256,256]

    # 保存原始高分辨率图像
    gt_np = hr_image.squeeze().numpy()
    gt_img = Image.fromarray(((gt_np + 1) * 127.5).astype(np.uint8))  # [-1,1] -> [0,255]
    gt_img.save("results/gt_hr.png")  # Added for saving images
        
    # 创建低分辨率输入
    lr_image = Resize((lr_res, lr_res))(hr_image)  # 下采样

    # 保存低分辨率输入图像
    lr_np = lr_image.squeeze().numpy()
    lr_img = Image.fromarray(((lr_np + 1) * 127.5).astype(np.uint8))  # [-1,1] -> [0,255]
    lr_img = lr_img.resize((256, 256), Image.NEAREST)  # 放大以便查看
    lr_img.save("results/input_lr.png")  # Added for saving images
    
    # 初始化模型和优化器
    model = INRModel()
    optimizer = optim.Adam(model.parameters(), lr=1e-4)
    
    # 训练阶段（在低分辨率上训练）
    print("Training on low-resolution...")
    coord_grid = torch.stack(torch.meshgrid(
        torch.linspace(-1, 1, lr_res),
        torch.linspace(-1, 1, lr_res)
    ), dim=-1).reshape(-1, 2)  # [lr_res^2, 2]
    target = lr_image.reshape(-1)  # [lr_res^2]
    
    for epoch in range(1000):
        pred = model(coord_grid)
        loss = F.mse_loss(pred, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if epoch % 100 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
    
    # 测试阶段（冻结参数，生成高分辨率）
    print("\nTesting on high-resolution...")
    with torch.no_grad():
        hr_coord = torch.stack(torch.meshgrid(
            torch.linspace(-1, 1, hr_res),
            torch.linspace(-1, 1, hr_res)
        , dim=-1).reshape(-1, 2)
        pred_hr = model(hr_coord).reshape(hr_res, hr_res).numpy()

    # 保存预测结果
    pred_img = Image.fromarray(((pred_hr + 1) * 127.5).astype(np.uint8))  # [-1,1] -> [0,255]
    pred_img.save("results/pred_hr.png")  # Added for saving images
    
    # 计算PSNR
    gt_np = hr_image.squeeze().numpy()
    psnr_value = psnr(gt_np, pred_hr, data_range=1.0)
    print(f"PSNR: {psnr_value:.2f} dB (Threshold >30 dB: {'Pass' if psnr_value>30 else 'Fail'})")

# 实验2：梯度敏感性分析
def gradient_analysis():
    model = INRModel()
    
    # 生成不同分辨率输入
    res1, res2 = 64, 128
    coord1 = torch.rand(res1*res1, 2) * 2 - 1  # [-1,1]
    coord2 = torch.rand(res2*res2, 2) * 2 - 1
    
    # 前向传播
    pred1 = model(coord1)
    pred2 = model(coord2)
    
    # 计算梯度差异
    loss1 = pred1.mean()
    loss2 = pred2.mean()
    
    # 第一次反向传播
    model.zero_grad()
    loss1.backward(retain_graph=True)
    grad1 = [p.grad.clone() for p in model.parameters()]
    
    # 第二次反向传播
    model.zero_grad()
    loss2.backward()
    grad2 = [p.grad.clone() for p in model.parameters()]
    
    # 计算梯度差异
    diff = sum([torch.norm(g1 - g2).item() for g1,g2 in zip(grad1, grad2)])
    print(f"\nGradient difference between resolutions {res1}x{res1} and {res2}x{res2}: {diff:.6f}")

if __name__ == "__main__":
    freeze_test()       # 实验1：参数冻结测试
    gradient_analysis() # 实验2：梯度敏感性分析
报错
  File "val_1.py", line 91
    pred_hr = model(hr_coord).reshape(hr_res, hr_res).numpy()
          ^
SyntaxError: invalid syntax