"""
Image Quality Assessment Analyzer
Main module for analyzing image quality using comprehensive metrics.
"""

import os
import json
import pandas as pd
from pathlib import Path
from typing import Dict, List, Tuple
from datetime import datetime

from .metrics import MetricsManager
from .logger_config import get_logger
from .batch_loader import BatchProcessor

logger = get_logger(__name__)


class IQAAnalyzer:
    """Main IQA analyzer using comprehensive metrics."""
    
    def __init__(self, image_dir: str, output_dir: str = "results", 
                 enable_deqa: bool = True, enable_traditional: bool = True,
                 enable_pyiqa: bool = True, pyiqa_selected_metrics: List[str] | None = None,
                 batch_size: int = 8, max_image_size: tuple = (1024, 1024),
                 enable_memory_monitoring: bool = True):
        """
        Initialize the IQA analyzer.
        
        Args:
            image_dir: Directory containing images to analyze
            output_dir: Directory to save results
            enable_deqa: Whether to enable DeQA metric
            enable_traditional: Whether to enable traditional metrics
            enable_pyiqa: Whether to enable PyIQA metrics
            pyiqa_selected_metrics: List of specific PyIQA metrics to use
            batch_size: Number of images to process per batch
            max_image_size: Maximum image dimensions for preprocessing
            enable_memory_monitoring: Whether to monitor memory usage
        """
        self.image_dir = Path(image_dir)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # Store batch processing settings
        self.batch_size = batch_size
        self.max_image_size = max_image_size
        self.enable_memory_monitoring = enable_memory_monitoring
        
        # Initialize metrics manager
        self.metrics_manager = MetricsManager(
            enable_deqa=enable_deqa,
            enable_traditional=enable_traditional,
            enable_pyiqa=enable_pyiqa,
            pyiqa_selected_metrics=pyiqa_selected_metrics
        )
        
        # Initialize batch processor (only if DeQA is enabled for memory efficiency)
        self.batch_processor = None
        if enable_deqa:
            self.batch_processor = BatchProcessor(
                batch_size=batch_size,
                max_image_size=max_image_size,
                enable_memory_monitoring=enable_memory_monitoring
            )
        
        # Results storage
        self.results = {}
        self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        logger.info(f"IQA Analyzer initialized for: {image_dir}")
        logger.info(f"Available metrics: {self.metrics_manager.get_available_metrics()}")
        logger.info(f"Batch processing: {'enabled' if self.batch_processor else 'disabled'} "
                   f"(batch_size={batch_size}, max_size={max_image_size})")
    
    def _get_image_files(self) -> List[Path]:
        """Get all image files from the directory."""
        image_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.PNG', '.JPG', '.JPEG', '.BMP', '.TIFF'}
        
        image_files = []
        for ext in image_extensions:
            image_files.extend(self.image_dir.glob(f"*{ext}"))
        
        return sorted(image_files)
    
    def analyze_images(self) -> Dict[str, Dict]:
        """
        Analyze all images using comprehensive metrics.
        
        Returns:
            Dictionary containing analysis results
        """
        logger.info("Starting image quality analysis...")
        
        # Get all image files
        image_files = self._get_image_files()
        
        if not image_files:
            logger.error("No image files found!")
            return {}
        
        logger.info(f"Found {len(image_files)} images to analyze")
        
        # Prepare image paths for analysis
        image_paths = [str(img_path) for img_path in image_files]
        
        # Use batch processing if available and only DeQA is enabled
        available_metrics = self.metrics_manager.get_available_metrics()
        use_batch_processing = (self.batch_processor is not None and 
                               available_metrics == ['deqa'])
        
        if use_batch_processing:
            logger.info(f"Using memory-efficient batch processing (batch_size={self.batch_size})...")
            
            # Get DeQA scorer
            deqa_scorer = self.metrics_manager.get_deqa_scorer()
            if deqa_scorer is None:
                logger.error("DeQA scorer not available for batch processing")
                return {}
            
            # Process with batch loader
            def progress_callback(processed, total):
                logger.info(f"Batch progress: {processed}/{total} images processed")
            
            self.results = self.batch_processor.process_images_in_batches(
                image_paths, deqa_scorer, progress_callback
            )
        else:
            # Fall back to original processing
            logger.info("Using standard processing...")
            self.results = self.metrics_manager.calculate_metrics_batch(image_paths)
        
        logger.info(f"Analysis completed for {len(self.results)} images")
        return self.results
    
    def generate_report(self) -> str:
        """Generate comprehensive analysis report."""
        if not self.results:
            return "No results to report. Run analyze_images() first."
        
        # Create summary statistics
        valid_results = {k: v for k, v in self.results.items() if 'error' not in v}
        
        if not valid_results:
            return "No valid results to report."
        
        # Extract scores
        deqa_scores = [v.get('deqa_score', 0) for v in valid_results.values() if v.get('deqa_score') is not None]
        musiq_scores = [v.get('musiq_score', 0) for v in valid_results.values() if v.get('musiq_score') is not None]
        
        # Get available metrics info
        metrics_info = self.metrics_manager.get_metric_info()
        
        # Calculate statistics
        stats = {
            'total_images': len(self.results),
            'valid_images': len(valid_results),
            'deqa_score_stats': {
                'mean': round(sum(deqa_scores) / len(deqa_scores), 3) if deqa_scores else 0,
                'min': round(min(deqa_scores), 3) if deqa_scores else 0,
                'max': round(max(deqa_scores), 3) if deqa_scores else 0,
                'std': round(pd.Series(deqa_scores).std(), 3) if len(deqa_scores) > 1 else 0
            }
        }
        
        # Generate report
        report = f"""
Image Quality Assessment Report
Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
Image Directory: {self.image_dir}

Summary Statistics:
==================
Total Images: {stats['total_images']}
Valid Images: {stats['valid_images']}
Available Metrics: {', '.join(metrics_info.keys())}

DeQA Score Statistics (0-5 scale):
- Mean: {stats['deqa_score_stats']['mean']}
- Min: {stats['deqa_score_stats']['min']}
- Max: {stats['deqa_score_stats']['max']}
- Std Dev: {stats['deqa_score_stats']['std']}

MUSIQ Score Statistics (0-100 scale):
- Mean: {round(sum(musiq_scores) / len(musiq_scores), 3) if musiq_scores else 0}
- Min: {round(min(musiq_scores), 3) if musiq_scores else 0}
- Max: {round(max(musiq_scores), 3) if musiq_scores else 0}

Quality Distribution (DeQA):
- Excellent (4.5-5.0): {len([s for s in deqa_scores if s >= 4.5])}
- Good (3.5-4.4): {len([s for s in deqa_scores if 3.5 <= s < 4.5])}
- Fair (2.5-3.4): {len([s for s in deqa_scores if 2.5 <= s < 3.5])}
- Poor (1.5-2.4): {len([s for s in deqa_scores if 1.5 <= s < 2.5])}
- Very Poor (0.0-1.4): {len([s for s in deqa_scores if s < 1.5])}

Top 10 Highest Quality Images (DeQA):
"""
        
        # Sort by DeQA score
        sorted_results = sorted(
            valid_results.items(), 
            key=lambda x: x[1].get('deqa_score', 0), 
            reverse=True
        )
        
        for i, (img_name, result) in enumerate(sorted_results[:10]):
            deqa = result.get('deqa_score', 'N/A')
            musiq = result.get('musiq_score', 'N/A')
            size_mb = result.get('file_size_mb', 'N/A')
            report += f"{i+1}. {img_name}: DeQA: {deqa}, MUSIQ: {musiq}, Size: {size_mb} MB\n"
        
        report += "\nBottom 10 Lowest Quality Images (DeQA):\n"
        for i, (img_name, result) in enumerate(sorted_results[-10:]):
            deqa = result.get('deqa_score', 'N/A')
            musiq = result.get('musiq_score', 'N/A')
            size_mb = result.get('file_size_mb', 'N/A')
            report += f"{i+1}. {img_name}: DeQA: {deqa}, MUSIQ: {musiq}, Size: {size_mb} MB\n"
        
        # Add traditional metrics summary if available
        if 'traditional' in metrics_info:
            report += f"\nTraditional Metrics Available:\n"
            for metric_name in metrics_info['traditional']['metrics'][:10]:  # Show first 10
                report += f"- {metric_name}\n"
            if len(metrics_info['traditional']['metrics']) > 10:
                report += f"... and {len(metrics_info['traditional']['metrics']) - 10} more\n"
        
        return report
    
    def save_results(self, format: str = 'json') -> str:
        """Save results to file."""
        if not self.results:
            return "No results to save. Run analyze_images() first."
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        if format.lower() == 'json':
            output_file = self.output_dir / f"iqa_results_{timestamp}.json"
            with open(output_file, 'w') as f:
                json.dump(self.results, f, indent=2, default=str)
        
        elif format.lower() == 'csv':
            # Flatten results for CSV with all available metrics
            flat_results = []
            for img_name, result in self.results.items():
                if 'error' not in result:
                    row = {'image_name': img_name}
                    # Add all available metrics
                    for key, value in result.items():
                        if key not in ['file_path']:  # Skip file path for CSV
                            row[key] = value
                    flat_results.append(row)
            
            df = pd.DataFrame(flat_results)
            output_file = self.output_dir / f"iqa_results_{timestamp}.csv"
            df.to_csv(output_file, index=False)
        
        elif format.lower() == 'txt':
            # Always save DeQA
            output_file = self.output_dir / f"deqa-score-{timestamp}.txt"
            with open(output_file, 'w') as f:
                sorted_results = sorted(
                    self.results.items(),
                    key=lambda x: x[1].get('deqa_score', 0) if 'error' not in x[1] else 0,
                    reverse=True
                )
                for img_name, result in sorted_results:
                    if 'error' not in result and result.get('deqa_score') is not None:
                        f.write(f"{result['deqa_score']:.1f}  - {img_name}\n")

            # For PyIQA selected metrics, emit one TXT per metric ONLY if any value exists
            selected = self.metrics_manager.get_selected_pyiqa_metrics() or []
            # Precompute which metrics have at least one value
            metrics_with_values = set()
            for metric_name in selected:
                key = f"pyiqa_{metric_name}"
                for result in self.results.values():
                    if 'error' not in result and result.get(key) is not None:
                        metrics_with_values.add(metric_name)
                        break
            for metric_name in selected:
                if metric_name not in metrics_with_values:
                    # Skip creating a file when no values are available
                    continue
                key = f"pyiqa_{metric_name}"
                path = self.output_dir / f"{metric_name}-score-{timestamp}.txt"
                with open(path, 'w') as f:
                    # Sort by this metric when available
                    sorted_items = sorted(
                        [item for item in self.results.items() if 'error' not in item[1] and item[1].get(key) is not None],
                        key=lambda x: x[1].get(key, float('-inf')),
                        reverse=True
                    )
                    for img_name, result in sorted_items:
                        val = result.get(key)
                        f.write(f"{float(val):.3f}  - {img_name}\n")
        
        else:
            return f"Unsupported format: {format}"
        
        return str(output_file)
    
    def run_analysis(self) -> Tuple[Dict, str]:
        """
        Run complete analysis and generate report.
        
        Returns:
            Tuple of (results_dict, report_string)
        """
        # Analyze images
        results = self.analyze_images()
        
        # Generate report
        report = self.generate_report()
        
        # Do not auto-save JSON anymore; TXT will be saved via main's choice
        return results, report
