#!/usr/bin/env python3
"""
Data Preparation Module for YOLOv8 Training
"""
import os
import sys
import yaml
import argparse
from pathlib import Path
import logging
import cv2
import numpy as np
from PIL import Image
import random

# Import config
sys.path.append(str(Path(__file__).parent.parent))
from config import DATA_YAML_PATH

logger = logging.getLogger(__name__)

class DataPreparator:
    """
    Data Preparation for YOLOv8 Training
    """
    
    def __init__(self, data_yaml_path: str = None):
        """
        Initialize Data Preparator
        
        Args:
            data_yaml_path: Path to data.yaml file (optional, uses default if None)
        """
        self.data_yaml_path = Path(data_yaml_path) if data_yaml_path else DATA_YAML_PATH
        self.data_config = self._load_data_config()
    
    def _load_data_config(self):
        """Load data configuration from YAML file"""
        if not self.data_yaml_path.exists():
            raise FileNotFoundError(f"data.yaml not found at {self.data_yaml_path}")
        
        with open(self.data_yaml_path, 'r') as f:
            config = yaml.safe_load(f)
        
        return config
    
    def check_data_structure(self):
        """Check data structure and validate paths"""
        logger.info("Checking data structure...")
        
        # Check training data
        train_path = Path(self.data_config['train'])
        if train_path.exists():
            train_images = list(train_path.glob('*.jpg')) + list(train_path.glob('*.jpeg')) + list(train_path.glob('*.png'))
            train_labels = list(train_path.glob('*.txt'))
            logger.info(f"Training data: {len(train_images)} images, {len(train_labels)} labels")
        else:
            logger.warning(f"Training path does not exist: {train_path}")
        
        # Check validation data
        val_path = Path(self.data_config['val'])
        if val_path.exists():
            val_images = list(val_path.glob('*.jpg')) + list(val_path.glob('*.jpeg')) + list(val_path.glob('*.png'))
            val_labels = list(val_path.glob('*.txt'))
            logger.info(f"Validation data: {len(val_images)} images, {len(val_labels)} labels")
        else:
            logger.warning(f"Validation path does not exist: {val_path}")
        
        # Check test data
        if 'test' in self.data_config:
            test_path = Path(self.data_config['test'])
            if test_path.exists():
                test_images = list(test_path.glob('*.jpg')) + list(test_path.glob('*.jpeg')) + list(test_path.glob('*.png'))
                test_labels = list(test_path.glob('*.txt'))
                logger.info(f"Test data: {len(test_images)} images, {len(test_labels)} labels")
            else:
                logger.warning(f"Test path does not exist: {test_path}")
        
        # Check class information
        logger.info(f"Number of classes: {self.data_config['nc']}")
        logger.info(f"Class names: {self.data_config['names']}")
    
    def validate_labels(self, split='train'):
        """Validate YOLO format labels"""
        logger.info(f"Validating {split} labels...")
        
        if split == 'train':
            images_path = Path(self.data_config['train'])
        elif split == 'val':
            images_path = Path(self.data_config['val'])
        elif split == 'test' and 'test' in self.data_config:
            images_path = Path(self.data_config['test'])
        else:
            logger.error(f"Invalid split: {split}")
            return
        
        if not images_path.exists():
            logger.error(f"Path does not exist: {images_path}")
            return
        
        # Get all image files
        image_files = list(images_path.glob('*.jpg')) + list(images_path.glob('*.jpeg')) + list(images_path.glob('*.png'))
        
        valid_images = 0
        invalid_images = 0
        total_annotations = 0
        
        for img_file in image_files:
            # Check if corresponding label file exists
            label_file = img_file.with_suffix('.txt')
            
            if not label_file.exists():
                logger.warning(f"No label file for {img_file.name}")
                invalid_images += 1
                continue
            
            # Validate label format
            try:
                with open(label_file, 'r') as f:
                    lines = f.readlines()
                
                # Check each annotation
                for line_num, line in enumerate(lines, 1):
                    parts = line.strip().split()
                    if len(parts) != 5:
                        logger.warning(f"Invalid annotation format in {label_file.name}, line {line_num}")
                        continue
                    
                    # Check class index
                    class_idx = int(parts[0])
                    if class_idx >= self.data_config['nc']:
                        logger.warning(f"Invalid class index {class_idx} in {label_file.name}, line {line_num}")
                        continue
                    
                    # Check coordinates (should be normalized between 0 and 1)
                    coords = [float(x) for x in parts[1:]]
                    if any(coord < 0 or coord > 1 for coord in coords):
                        logger.warning(f"Invalid coordinates in {label_file.name}, line {line_num}")
                        continue
                    
                    total_annotations += 1
                
                valid_images += 1
                
            except Exception as e:
                logger.error(f"Error reading {label_file}: {e}")
                invalid_images += 1
        
        logger.info(f"{split} validation results:")
        logger.info(f"  - Valid images: {valid_images}")
        logger.info(f"  - Invalid images: {invalid_images}")
        logger.info(f"  - Total annotations: {total_annotations}")
    
    def check_image_quality(self, split='train', sample_size=50):
        """Check image quality and statistics"""
        logger.info(f"Checking {split} image quality...")
        
        if split == 'train':
            images_path = Path(self.data_config['train'])
        elif split == 'val':
            images_path = Path(self.data_config['val'])
        elif split == 'test' and 'test' in self.data_config:
            images_path = Path(self.data_config['test'])
        else:
            logger.error(f"Invalid split: {split}")
            return
        
        if not images_path.exists():
            logger.error(f"Path does not exist: {images_path}")
            return
        
        # Get all image files
        image_files = list(images_path.glob('*.jpg')) + list(images_path.glob('*.jpeg')) + list(images_path.glob('*.png'))
        
        if len(image_files) == 0:
            logger.warning(f"No images found in {images_path}")
            return
        
        # Sample images for analysis
        sample_files = random.sample(image_files, min(sample_size, len(image_files)))
        
        widths = []
        heights = []
        channels = []
        
        for img_file in sample_files:
            try:
                # Read image
                img = cv2.imread(str(img_file))
                if img is None:
                    logger.warning(f"Could not read image: {img_file}")
                    continue
                
                height, width = img.shape[:2]
                channel_count = img.shape[2] if len(img.shape) == 3 else 1
                
                widths.append(width)
                heights.append(height)
                channels.append(channel_count)
                
            except Exception as e:
                logger.error(f"Error reading {img_file}: {e}")
        
        if widths:
            logger.info(f"Image statistics (sample of {len(widths)} images):")
            logger.info(f"  - Width: min={min(widths)}, max={max(widths)}, avg={sum(widths)/len(widths):.1f}")
            logger.info(f"  - Height: min={min(heights)}, max={max(heights)}, avg={sum(heights)/len(heights):.1f}")
            logger.info(f"  - Channels: {set(channels)}")
    
    def run_full_validation(self):
        """Run complete data validation"""
        logger.info("Running complete data validation...")
        
        # Check data structure
        self.check_data_structure()
        
        # Validate labels for each split
        for split in ['train', 'val']:
            self.validate_labels(split)
        
        # Check image quality
        for split in ['train', 'val']:
            self.check_image_quality(split)
        
        logger.info("Data validation completed!")
        return True 