import torch
from net import create_food_cnn, create_mobile_food_cnn
import os


def main():
    print("开始转换模型为移动端格式...")
    
    # 1. 加载训练好的基础模型权重
    base_model = create_food_cnn()
    model_path = "../model/10/best_food_model.pth"
    
    if not os.path.exists(model_path):
        print(f"错误：模型文件不存在 {model_path}")
        return
    
    base_model.load_state_dict(torch.load(model_path, map_location='cpu'))
    print("✓ 基础模型权重加载成功")
    
    # 2. 创建移动端模型并复制权重
    mobile_model = create_mobile_food_cnn()
    mobile_model.load_state_dict(base_model.state_dict())
    mobile_model.eval()
    print("✓ 移动端模型创建成功")
    
    # 3. 测试模型
    # 创建示例输入 (模拟Android端输入：1x3x224x224，值范围0-255)
    example_input = torch.randint(0, 256, (1, 3, 224, 224), dtype=torch.float32)
    
    with torch.no_grad():
        output = mobile_model(example_input)
        print(f"✓ 模型测试成功，输出形状: {output.shape}")
        print(f"  输出值范围: [{output.min():.3f}, {output.max():.3f}]")
    
    # 4. 转换为 TorchScript
    try:
        traced_model = torch.jit.trace(mobile_model, example_input)
        
        # 保存模型
        output_path = "../model/10/best_food_model_mobile.pt"
        traced_model.save(output_path)
        print(f"✓ TorchScript模型保存成功: {output_path}")
        
        # 验证保存的模型
        loaded_model = torch.jit.load(output_path)
        with torch.no_grad():
            loaded_output = loaded_model(example_input)
            print(f"✓ 保存的模型验证成功，输出形状: {loaded_output.shape}")
            
            # 检查输出是否一致
            if torch.allclose(output, loaded_output, atol=1e-6):
                print("✓ 模型输出一致性验证通过")
            else:
                print("⚠ 警告：模型输出存在差异")
        
    except Exception as e:
        print(f"✗ TorchScript转换失败: {e}")
        return
    
    print("\n" + "="*50)
    print("移动端模型转换完成！")
    print("="*50)
    print("使用说明：")
    print("1. Android端输入图片应为 [batch, 3, height, width] 格式")
    print("2. 像素值范围：0-255 (uint8 或 float32)")
    print("3. 模型会自动处理：")
    print("   - 归一化到 [0,1]")
    print("   - 缩放到 32x32")
    print("   - ImageNet标准化")
    print("4. 输出：3个类别的logits")
    print("="*50)


if __name__ == "__main__":
    main()