#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试数据生成器
用于生成模拟的测试数据，方便测试模型评估功能

作者: AI Assistant
创建时间: 2024
"""

import os
import cv2
import numpy as np
import json
import random
from datetime import datetime
from typing import List, Dict, Tuple

class TestDataGenerator:
    """
    测试数据生成器
    生成模拟的UAV和鸟类图像用于测试
    """
    
    def __init__(self, output_dir: str = None):
        """
        初始化生成器
        
        Args:
            output_dir: 输出目录
        """
        if output_dir is None:
            output_dir = os.path.join(os.path.dirname(__file__), 'Data')
        
        self.output_dir = output_dir
        self.image_size = (640, 480)  # 默认图像尺寸
        
        # 创建目录结构
        self.birds_dir = os.path.join(output_dir, 'BIRDS')
        self.uav_dir = os.path.join(output_dir, 'UAV')
        self.others_dir = os.path.join(output_dir, 'OTHERS')
        
        os.makedirs(self.birds_dir, exist_ok=True)
        os.makedirs(self.uav_dir, exist_ok=True)
        os.makedirs(self.others_dir, exist_ok=True)
    
    def generate_bird_image(self, image_id: int) -> str:
        """
        生成鸟类图像（模拟）
        
        Args:
            image_id: 图像ID
            
        Returns:
            生成的图像路径
        """
        # 创建蓝天背景
        img = np.ones((self.image_size[1], self.image_size[0], 3), dtype=np.uint8)
        img[:, :] = [135, 206, 235]  # 天蓝色背景
        
        # 添加一些云朵
        for _ in range(random.randint(2, 5)):
            center = (random.randint(50, self.image_size[0]-50), 
                     random.randint(50, self.image_size[1]-50))
            radius = random.randint(20, 60)
            cv2.circle(img, center, radius, (255, 255, 255), -1)
        
        # 绘制鸟类形状（简化的V形）
        bird_x = random.randint(100, self.image_size[0]-100)
        bird_y = random.randint(100, self.image_size[1]-100)
        bird_size = random.randint(15, 40)
        
        # 绘制鸟的翅膀（V形）
        pts1 = np.array([[bird_x-bird_size, bird_y], 
                        [bird_x, bird_y-bird_size//2], 
                        [bird_x, bird_y+bird_size//2]], np.int32)
        pts2 = np.array([[bird_x+bird_size, bird_y], 
                        [bird_x, bird_y-bird_size//2], 
                        [bird_x, bird_y+bird_size//2]], np.int32)
        
        cv2.fillPoly(img, [pts1], (50, 50, 50))  # 深灰色鸟
        cv2.fillPoly(img, [pts2], (50, 50, 50))
        
        # 添加噪声
        noise = np.random.normal(0, 10, img.shape).astype(np.uint8)
        img = cv2.add(img, noise)
        
        # 保存图像
        filename = f"bird_{image_id:04d}.jpg"
        filepath = os.path.join(self.birds_dir, filename)
        cv2.imwrite(filepath, img)
        
        return filepath
    
    def generate_uav_image(self, image_id: int) -> str:
        """
        生成无人机图像（模拟）
        
        Args:
            image_id: 图像ID
            
        Returns:
            生成的图像路径
        """
        # 创建蓝天背景
        img = np.ones((self.image_size[1], self.image_size[0], 3), dtype=np.uint8)
        img[:, :] = [135, 206, 235]  # 天蓝色背景
        
        # 添加一些云朵
        for _ in range(random.randint(1, 3)):
            center = (random.randint(50, self.image_size[0]-50), 
                     random.randint(50, self.image_size[1]-50))
            radius = random.randint(30, 80)
            cv2.circle(img, center, radius, (255, 255, 255), -1)
        
        # 绘制无人机形状（简化的十字形）
        uav_x = random.randint(100, self.image_size[0]-100)
        uav_y = random.randint(100, self.image_size[1]-100)
        uav_size = random.randint(20, 50)
        
        # 绘制无人机机身（矩形）
        cv2.rectangle(img, 
                     (uav_x-uav_size//4, uav_y-uav_size//8), 
                     (uav_x+uav_size//4, uav_y+uav_size//8), 
                     (80, 80, 80), -1)
        
        # 绘制螺旋桨臂
        cv2.line(img, (uav_x-uav_size, uav_y), (uav_x+uav_size, uav_y), (60, 60, 60), 3)
        cv2.line(img, (uav_x, uav_y-uav_size), (uav_x, uav_y+uav_size), (60, 60, 60), 3)
        
        # 绘制螺旋桨（圆形）
        for dx, dy in [(-uav_size, 0), (uav_size, 0), (0, -uav_size), (0, uav_size)]:
            cv2.circle(img, (uav_x+dx, uav_y+dy), uav_size//6, (40, 40, 40), -1)
        
        # 添加噪声
        noise = np.random.normal(0, 8, img.shape).astype(np.uint8)
        img = cv2.add(img, noise)
        
        # 保存图像
        filename = f"uav_{image_id:04d}.jpg"
        filepath = os.path.join(self.uav_dir, filename)
        cv2.imwrite(filepath, img)
        
        return filepath
    
    def generate_other_image(self, image_id: int) -> str:
        """
        生成其他类型图像（负样本）
        
        Args:
            image_id: 图像ID
            
        Returns:
            生成的图像路径
        """
        # 随机选择背景类型
        bg_type = random.choice(['sky', 'landscape', 'urban'])
        
        if bg_type == 'sky':
            # 纯天空背景
            img = np.ones((self.image_size[1], self.image_size[0], 3), dtype=np.uint8)
            img[:, :] = [135, 206, 235]  # 天蓝色
            
            # 添加云朵
            for _ in range(random.randint(3, 8)):
                center = (random.randint(0, self.image_size[0]), 
                         random.randint(0, self.image_size[1]))
                radius = random.randint(40, 100)
                cv2.circle(img, center, radius, (255, 255, 255), -1)
        
        elif bg_type == 'landscape':
            # 风景背景
            img = np.ones((self.image_size[1], self.image_size[0], 3), dtype=np.uint8)
            
            # 天空部分
            img[:self.image_size[1]//2, :] = [135, 206, 235]
            
            # 地面部分
            img[self.image_size[1]//2:, :] = [34, 139, 34]  # 森林绿
            
            # 添加一些树木形状
            for _ in range(random.randint(5, 10)):
                tree_x = random.randint(0, self.image_size[0])
                tree_y = random.randint(self.image_size[1]//2, self.image_size[1])
                tree_height = random.randint(30, 80)
                cv2.rectangle(img, 
                             (tree_x-5, tree_y), 
                             (tree_x+5, tree_y-tree_height), 
                             (139, 69, 19), -1)  # 棕色树干
                cv2.circle(img, (tree_x, tree_y-tree_height), 15, (0, 100, 0), -1)  # 绿色树冠
        
        else:  # urban
            # 城市背景
            img = np.ones((self.image_size[1], self.image_size[0], 3), dtype=np.uint8)
            img[:, :] = [169, 169, 169]  # 灰色背景
            
            # 添加建筑物
            for _ in range(random.randint(3, 6)):
                building_x = random.randint(0, self.image_size[0]-50)
                building_y = random.randint(self.image_size[1]//3, self.image_size[1])
                building_w = random.randint(30, 80)
                building_h = random.randint(50, 150)
                
                cv2.rectangle(img, 
                             (building_x, building_y), 
                             (building_x+building_w, building_y-building_h), 
                             (105, 105, 105), -1)  # 深灰色建筑
        
        # 添加噪声
        noise = np.random.normal(0, 12, img.shape).astype(np.uint8)
        img = cv2.add(img, noise)
        
        # 保存图像
        filename = f"other_{image_id:04d}.jpg"
        filepath = os.path.join(self.others_dir, filename)
        cv2.imwrite(filepath, img)
        
        return filepath
    
    def generate_dataset(self, num_birds: int = 50, num_uavs: int = 50, num_others: int = 100) -> Dict:
        """
        生成完整的测试数据集
        
        Args:
            num_birds: 鸟类图像数量
            num_uavs: 无人机图像数量
            num_others: 其他图像数量
            
        Returns:
            数据集信息字典
        """
        print(f"开始生成测试数据集...")
        print(f"鸟类图像: {num_birds} 张")
        print(f"无人机图像: {num_uavs} 张")
        print(f"其他图像: {num_others} 张")
        print(f"输出目录: {self.output_dir}")
        
        dataset_info = {
            'creation_time': datetime.now().isoformat(),
            'total_images': num_birds + num_uavs + num_others,
            'categories': {
                'birds': {'count': num_birds, 'files': []},
                'uavs': {'count': num_uavs, 'files': []},
                'others': {'count': num_others, 'files': []}
            }
        }
        
        # 生成鸟类图像
        print("\n生成鸟类图像...")
        for i in range(num_birds):
            filepath = self.generate_bird_image(i)
            dataset_info['categories']['birds']['files'].append(filepath)
            if (i + 1) % 10 == 0:
                print(f"已生成 {i + 1}/{num_birds} 张鸟类图像")
        
        # 生成无人机图像
        print("\n生成无人机图像...")
        for i in range(num_uavs):
            filepath = self.generate_uav_image(i)
            dataset_info['categories']['uavs']['files'].append(filepath)
            if (i + 1) % 10 == 0:
                print(f"已生成 {i + 1}/{num_uavs} 张无人机图像")
        
        # 生成其他图像
        print("\n生成其他图像...")
        for i in range(num_others):
            filepath = self.generate_other_image(i)
            dataset_info['categories']['others']['files'].append(filepath)
            if (i + 1) % 20 == 0:
                print(f"已生成 {i + 1}/{num_others} 张其他图像")
        
        # 保存数据集信息
        info_file = os.path.join(self.output_dir, 'dataset_info.json')
        with open(info_file, 'w', encoding='utf-8') as f:
            json.dump(dataset_info, f, ensure_ascii=False, indent=2)
        
        print(f"\n数据集生成完成！")
        print(f"总计生成 {dataset_info['total_images']} 张图像")
        print(f"数据集信息已保存到: {info_file}")
        
        return dataset_info
    
    def create_annotation_file(self, dataset_info: Dict) -> str:
        """
        创建标注文件
        
        Args:
            dataset_info: 数据集信息
            
        Returns:
            标注文件路径
        """
        annotations = {
            'info': {
                'description': 'Generated test dataset for UAV and bird detection',
                'version': '1.0',
                'creation_date': dataset_info['creation_time']
            },
            'images': []
        }
        
        # 添加鸟类图像标注
        for filepath in dataset_info['categories']['birds']['files']:
            annotations['images'].append({
                'path': filepath,
                'label': 'bird',
                'has_target': True,
                'bboxes': []  # 简化版本，不包含具体边界框
            })
        
        # 添加无人机图像标注
        for filepath in dataset_info['categories']['uavs']['files']:
            annotations['images'].append({
                'path': filepath,
                'label': 'uav',
                'has_target': True,
                'bboxes': []
            })
        
        # 添加其他图像标注
        for filepath in dataset_info['categories']['others']['files']:
            annotations['images'].append({
                'path': filepath,
                'label': 'others',
                'has_target': False,
                'bboxes': []
            })
        
        # 保存标注文件
        annotation_file = os.path.join(self.output_dir, 'annotations.json')
        with open(annotation_file, 'w', encoding='utf-8') as f:
            json.dump(annotations, f, ensure_ascii=False, indent=2)
        
        print(f"标注文件已保存到: {annotation_file}")
        return annotation_file

def main():
    """
    主函数 - 生成测试数据
    """
    print("测试数据生成器")
    print("=" * 50)
    
    # 配置参数
    output_dir = r"d:\PythonProject\R360-UAVmodelTool\Data"
    num_birds = 30
    num_uavs = 30
    num_others = 60
    
    # 初始化生成器
    generator = TestDataGenerator(output_dir)
    
    # 生成数据集
    dataset_info = generator.generate_dataset(num_birds, num_uavs, num_others)
    
    # 创建标注文件
    annotation_file = generator.create_annotation_file(dataset_info)
    
    print("\n数据集结构:")
    print(f"├── BIRDS/     ({num_birds} 张图像)")
    print(f"├── UAV/       ({num_uavs} 张图像)")
    print(f"├── OTHERS/    ({num_others} 张图像)")
    print(f"├── dataset_info.json")
    print(f"└── annotations.json")
    
    print("\n现在可以运行 model_evaluator.py 进行模型评估！")

if __name__ == "__main__":
    main()