#!/usr/bin/env python3
"""
使用优化后的扩散模型生成IC版图图像
"""

import os
import sys
import torch
import torch.nn.functional as F
from pathlib import Path
import logging
import argparse
import yaml
from PIL import Image
import numpy as np

# 导入优化后的模块
from ic_layout_diffusion_optimized import (
    ManhattanAwareUNet,
    OptimizedNoiseScheduler,
    OptimizedDiffusionTrainer,
    manhattan_post_process,
    manhattan_regularization_loss
)

def setup_logging():
    """设置日志"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(sys.stdout),
            logging.FileHandler('diffusion_generation.log')
        ]
    )
    return logging.getLogger(__name__)

def load_model(checkpoint_path, device):
    """加载训练好的模型"""
    logger = logging.getLogger(__name__)

    # 加载检查点
    checkpoint = torch.load(checkpoint_path, map_location=device)

    # 从检查点中获取配置信息（如果有）
    config = checkpoint.get('config', {})

    # 创建模型
    model = ManhattanAwareUNet(
        in_channels=1,
        out_channels=1,
        use_edge_condition=config.get('edge_condition', False)
    ).to(device)

    # 加载权重
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()

    logger.info(f"模型已加载: {checkpoint_path}")
    logger.info(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")

    return model, config

def ddim_sample(model, scheduler, num_samples, image_size, device, num_steps=50, eta=0.0):
    """DDIM采样，比标准DDPM更快"""
    model.eval()

    # 从纯噪声开始
    x = torch.randn(num_samples, 1, image_size, image_size).to(device)

    # 选择时间步
    timesteps = torch.linspace(scheduler.num_timesteps - 1, 0, num_steps).long().to(device)

    with torch.no_grad():
        for i, t in enumerate(timesteps):
            t_batch = torch.full((num_samples,), t, device=device)

            # 预测噪声
            predicted_noise = model(x, t_batch)

            # 计算原始图像的估计
            alpha_t = scheduler.alphas[t].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
            alpha_cumprod_t = scheduler.alphas_cumprod[t].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
            beta_t = scheduler.betas[t].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
            sqrt_one_minus_alpha_cumprod_t = scheduler.sqrt_one_minus_alphas_cumprod[t].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)

            # 计算x_0的估计
            x_0_pred = (x - sqrt_one_minus_alpha_cumprod_t * predicted_noise) / torch.sqrt(alpha_cumprod_t)

            # 计算前一时间步的方向
            if i < len(timesteps) - 1:
                alpha_t_prev = scheduler.alphas[timesteps[i+1]]
                alpha_cumprod_t_prev = scheduler.alphas_cumprod[timesteps[i+1]]
                sqrt_alpha_cumprod_t_prev = torch.sqrt(alpha_cumprod_t_prev).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
                sqrt_one_minus_alpha_cumprod_t_prev = torch.sqrt(1 - alpha_cumprod_t_prev).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)

                # 计算方差
                variance = eta * torch.sqrt(beta_t).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)

                # 计算前一时间步的x
                x = sqrt_alpha_cumprod_t_prev * x_0_pred + torch.sqrt(1 - alpha_cumprod_t_prev - variance**2) * predicted_noise

                if eta > 0:
                    noise = torch.randn_like(x)
                    x += variance * noise
            else:
                x = x_0_pred

            # 限制范围
            x = torch.clamp(x, -2.0, 2.0)

    return torch.clamp(x, 0.0, 1.0)

def generate_with_guidance(model, scheduler, num_samples, image_size, device,
                          guidance_scale=1.0, num_steps=50, use_ddim=True):
    """带引导的采样（可扩展为classifier-free guidance）"""

    if use_ddim:
        # 使用DDIM采样
        samples = ddim_sample(model, scheduler, num_samples, image_size, device, num_steps)
    else:
        # 使用标准DDPM采样
        trainer = OptimizedDiffusionTrainer(model, scheduler, device)
        samples = trainer.generate(num_samples, image_size, save_dir=None, use_post_process=False)

    return samples

def evaluate_generation_quality(samples, device):
    """评估生成质量"""
    logger = logging.getLogger(__name__)

    quality_metrics = {}

    # 1. 曼哈顿几何合规性
    manhattan_loss = manhattan_regularization_loss(samples, device)
    quality_metrics['manhattan_compliance'] = float(manhattan_loss.item())

    # 2. 边缘锐度
    sobel_x = torch.tensor([[[[-1,0,1],[-2,0,2],[-1,0,1]]]], device=device, dtype=samples.dtype)
    sobel_y = torch.tensor([[[[-1,-2,-1],[0,0,0],[1,2,1]]]], device=device, dtype=samples.dtype)

    edge_x = F.conv2d(samples, sobel_x, padding=1)
    edge_y = F.conv2d(samples, sobel_y, padding=1)
    edge_magnitude = torch.sqrt(edge_x**2 + edge_y**2)

    quality_metrics['edge_sharpness'] = float(torch.mean(edge_magnitude).item())

    # 3. 对比度
    quality_metrics['contrast'] = float(torch.std(samples).item())

    # 4. 稀疏性（IC版图通常是稀疏的）
    quality_metrics['sparsity'] = float((samples < 0.1).float().mean().item())

    logger.info("生成质量评估:")
    for metric, value in quality_metrics.items():
        logger.info(f"  {metric}: {value:.4f}")

    return quality_metrics

def generate_optimized_samples(args):
    """生成优化样本的主函数"""
    logger = setup_logging()

    # 设备检查
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info(f"使用设备: {device}")

    # 设置随机种子
    torch.manual_seed(args.seed)
    if device.type == 'cuda':
        torch.cuda.manual_seed(args.seed)

    # 创建输出目录
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 加载模型
    model, config = load_model(args.checkpoint, device)

    # 创建调度器
    scheduler = OptimizedNoiseScheduler(
        num_timesteps=config.get('timesteps', args.timesteps),
        schedule_type=config.get('schedule_type', args.schedule_type)
    )

    # 生成参数
    generation_config = {
        'num_samples': args.num_samples,
        'image_size': args.image_size,
        'guidance_scale': args.guidance_scale,
        'num_steps': args.num_steps,
        'use_ddim': args.use_ddim,
        'eta': args.eta,
        'seed': args.seed
    }

    # 保存生成配置
    with open(output_dir / 'generation_config.yaml', 'w') as f:
        yaml.dump(generation_config, f, default_flow_style=False)

    logger.info(f"开始生成 {args.num_samples} 个样本...")
    logger.info(f"采样步数: {args.num_steps}, DDIM: {args.use_ddim}, ETA: {args.eta}")

    # 分批生成以避免内存不足
    all_samples = []
    batch_size = min(args.batch_size, args.num_samples)
    num_batches = (args.num_samples + batch_size - 1) // batch_size

    for batch_idx in range(num_batches):
        start_idx = batch_idx * batch_size
        end_idx = min(start_idx + batch_size, args.num_samples)
        current_batch_size = end_idx - start_idx

        logger.info(f"生成批次 {batch_idx + 1}/{num_batches} ({current_batch_size} 个样本)")

        # 生成样本
        with torch.no_grad():
            samples = generate_with_guidance(
                model, scheduler, current_batch_size, args.image_size, device,
                args.guidance_scale, args.num_steps, args.use_ddim
            )

            # 后处理
            if args.use_post_process:
                samples = manhattan_post_process(samples, threshold=args.post_process_threshold)

            all_samples.append(samples)

            # 立即保存当前批次
            batch_dir = output_dir / f"batch_{batch_idx + 1}"
            batch_dir.mkdir(exist_ok=True)

            for i in range(current_batch_size):
                img_tensor = samples[i].cpu()
                img_array = (img_tensor.squeeze().numpy() * 255).astype(np.uint8)
                img = Image.fromarray(img_array, mode='L')
                img.save(batch_dir / f"sample_{start_idx + i:06d}.png")

    # 合并所有样本
    all_samples = torch.cat(all_samples, dim=0)

    # 评估生成质量
    quality_metrics = evaluate_generation_quality(all_samples, device)

    # 保存质量评估结果
    with open(output_dir / 'quality_metrics.yaml', 'w') as f:
        yaml.dump(quality_metrics, f, default_flow_style=False)

    # 创建质量报告
    report_content = f"""
IC版图扩散模型生成报告
======================

生成配置:
- 模型检查点: {args.checkpoint}
- 样本数量: {args.num_samples}
- 图像尺寸: {args.image_size}x{args.image_size}
- 采样步数: {args.num_steps}
- DDIM采样: {args.use_ddim}
- 后处理: {args.use_post_process}

质量指标:
- 曼哈顿几何合规性: {quality_metrics['manhattan_compliance']:.4f} (越低越好)
- 边缘锐度: {quality_metrics['edge_sharpness']:.4f}
- 对比度: {quality_metrics['contrast']:.4f}
- 稀疏性: {quality_metrics['sparsity']:.4f}

输出目录: {output_dir}
"""

    with open(output_dir / 'generation_report.txt', 'w') as f:
        f.write(report_content)

    logger.info("生成完成!")
    logger.info(f"样本保存目录: {output_dir}")
    logger.info(f"质量报告: {output_dir / 'generation_report.txt'}")


def main():
    parser = argparse.ArgumentParser(description="使用优化的扩散模型生成IC版图")

    # 必需参数
    parser.add_argument('--checkpoint', type=str, required=True, help='模型检查点路径')
    parser.add_argument('--output_dir', type=str, required=True, help='输出目录')

    # 生成参数
    parser.add_argument('--num_samples', type=int, default=200, help='生成样本数量')
    parser.add_argument('--image_size', type=int, default=256, help='图像尺寸')
    parser.add_argument('--seed', type=int, default=42, help='随机种子')
    parser.add_argument('--batch_size', type=int, default=8, help='批次大小')

    # 采样参数
    parser.add_argument('--num_steps', type=int, default=50, help='采样步数')
    parser.add_argument('--use_ddim', action='store_true', default=True, help='使用DDIM采样')
    parser.add_argument('--guidance_scale', type=float, default=1.0, help='引导尺度')
    parser.add_argument('--eta', type=float, default=0.0, help='DDIM eta参数 (0=确定性, 1=随机)')

    # 后处理参数
    parser.add_argument('--use_post_process', action='store_true', default=True, help='启用后处理')
    parser.add_argument('--post_process_threshold', type=float, default=0.5, help='后处理阈值')

    # 模型配置（用于覆盖检查点中的配置）
    parser.add_argument('--timesteps', type=int, default=1000, help='扩散时间步数')
    parser.add_argument('--schedule_type', type=str, default='cosine',
                       choices=['linear', 'cosine'], help='噪声调度类型')

    args = parser.parse_args()

    # 检查检查点文件
    if not Path(args.checkpoint).exists():
        print(f"错误: 检查点文件不存在: {args.checkpoint}")
        sys.exit(1)

    # 开始生成
    generate_optimized_samples(args)


if __name__ == "__main__":
    main()