"""
Main application file for the Image Evaluator tool.
This module integrates all components and provides a Gradio interface.
"""

import os
import gradio as gr
import numpy as np
import pandas as pd
import torch
import glob
from PIL import Image
import json
import tempfile
import shutil
from datetime import datetime

# Import custom modules
from modules.metadata_extractor import MetadataExtractor
from modules.technical_metrics import TechnicalMetrics
from modules.aesthetic_metrics import AestheticMetrics
from modules.aggregator import ResultsAggregator
from modules.visualizer import Visualizer


class ImageEvaluator:
    """Main class for the Image Evaluator application."""
    
    def __init__(self):
        """Initialize the Image Evaluator."""
        self.results_dir = os.path.join(os.getcwd(), "results")
        os.makedirs(self.results_dir, exist_ok=True)
        
        # Initialize components
        self.metadata_extractor = MetadataExtractor()
        self.technical_metrics = TechnicalMetrics()
        self.aesthetic_metrics = AestheticMetrics()
        self.aggregator = ResultsAggregator()
        self.visualizer = Visualizer(self.results_dir)
        
        # Storage for results
        self.evaluation_results = {}
        self.metadata_cache = {}
        self.current_comparison = None
        
    def process_images(self, image_files, progress=None):
        """
        Process a list of image files and extract metadata.
        
        Args:
            image_files: list of image file paths
            progress: optional gradio Progress object
            
        Returns:
            tuple: (metadata_by_model, metadata_by_prompt)
        """
        metadata_list = []
        
        total_files = len(image_files)
        for i, img_path in enumerate(image_files):
            # Safe progress update without accessing internal attributes
            if progress is not None:
                try:
                    progress((i + 1) / total_files, f"Processing image {i+1}/{total_files}")
                except Exception as e:
                    print(f"Progress update error (non-critical): {e}")
            
            # Extract metadata
            metadata = self.metadata_extractor.extract_metadata(img_path)
            metadata_list.append((img_path, metadata))
            
            # Cache metadata
            self.metadata_cache[img_path] = metadata
        
        # Group by model and prompt
        metadata_by_model = self.metadata_extractor.group_images_by_model(metadata_list)
        metadata_by_prompt = self.metadata_extractor.group_images_by_prompt(metadata_list)
        
        return metadata_by_model, metadata_by_prompt
    
    def evaluate_images(self, image_files, progress=None):
        """
        Evaluate a list of image files using all metrics.
        
        Args:
            image_files: list of image file paths
            progress: optional gradio Progress object
            
        Returns:
            dict: evaluation results by image path
        """
        results = {}
        
        total_files = len(image_files)
        for i, img_path in enumerate(image_files):
            # Safe progress update without accessing internal attributes
            if progress is not None:
                try:
                    progress((i + 1) / total_files, f"Evaluating image {i+1}/{total_files}")
                except Exception as e:
                    print(f"Progress update error (non-critical): {e}")
            
            # Get metadata if available
            metadata = self.metadata_cache.get(img_path, {})
            prompt = metadata.get('prompt', '')
            
            # Calculate technical metrics
            tech_metrics = self.technical_metrics.calculate_all_metrics(img_path)
            
            # Calculate aesthetic metrics
            aesthetic_metrics = self.aesthetic_metrics.calculate_all_metrics(img_path, prompt)
            
            # Combine results
            combined_metrics = {**tech_metrics, **aesthetic_metrics}
            
            # Store results
            results[img_path] = combined_metrics
        
        return results
    
    def compare_models(self, evaluation_results, metadata_by_model):
        """
        Compare different models based on evaluation results.
        
        Args:
            evaluation_results: dictionary with image paths as keys and metrics as values
            metadata_by_model: dictionary with model names as keys and lists of image paths as values
            
        Returns:
            tuple: (comparison_df, visualizations)
        """
        # Group results by model
        results_by_model = {}
        for model, image_paths in metadata_by_model.items():
            model_results = [evaluation_results[img] for img in image_paths if img in evaluation_results]
            results_by_model[model] = model_results
        
        # Compare models
        comparison = self.aggregator.compare_models(results_by_model)
        
        # Create comparison dataframe
        comparison_df = self.aggregator.create_comparison_dataframe(comparison)
        
        # Store current comparison
        self.current_comparison = comparison_df
        
        # Create visualizations
        visualizations = {}
        
        # Create heatmap
        heatmap_path = self.visualizer.plot_heatmap(comparison_df)
        visualizations['Model Comparison Heatmap'] = heatmap_path
        
        # Create radar chart for key metrics
        key_metrics = ['aesthetic_score', 'sharpness', 'noise', 'contrast', 'color_harmony', 'prompt_similarity']
        available_metrics = [m for m in key_metrics if m in comparison_df.columns]
        if available_metrics:
            radar_path = self.visualizer.plot_radar_chart(comparison_df, available_metrics)
            visualizations['Model Comparison Radar Chart'] = radar_path
        
        # Create bar charts for important metrics
        for metric in ['overall_score', 'aesthetic_score', 'prompt_similarity']:
            if metric in comparison_df.columns:
                bar_path = self.visualizer.plot_metric_comparison(comparison_df, metric)
                visualizations[f'{metric} Comparison'] = bar_path
        
        return comparison_df, visualizations
    
    def export_results(self, format='csv'):
        """
        Export current comparison results.
        
        Args:
            format: export format ('csv', 'excel', or 'html')
            
        Returns:
            str: path to exported file
        """
        if self.current_comparison is not None:
            return self.visualizer.export_comparison_table(self.current_comparison, format)
        return None
    
    def generate_report(self, comparison_df, visualizations):
        """
        Generate a comprehensive HTML report.
        
        Args:
            comparison_df: pandas DataFrame with comparison data
            visualizations: dictionary of visualization paths
            
        Returns:
            str: path to HTML report
        """
        metrics_list = comparison_df.columns.tolist()
        return self.visualizer.generate_html_report(comparison_df, visualizations, metrics_list)


# Create Gradio interface
def create_interface():
    """Create and configure the Gradio interface."""
    
    # Initialize evaluator
    evaluator = ImageEvaluator()
    
    # Track state
    state = {
        'uploaded_images': [],
        'metadata_by_model': {},
        'metadata_by_prompt': {},
        'evaluation_results': {},
        'comparison_df': None,
        'visualizations': {},
        'report_path': None
    }
    
    def upload_images(files):
        """Handle image upload and processing."""
        # Reset state
        state['uploaded_images'] = []
        state['metadata_by_model'] = {}
        state['metadata_by_prompt'] = {}
        state['evaluation_results'] = {}
        state['comparison_df'] = None
        state['visualizations'] = {}
        state['report_path'] = None
        
        # Process uploaded files
        image_paths = [f.name for f in files]
        state['uploaded_images'] = image_paths
        
        # Extract metadata and group images
        # Use a simple progress message instead of Gradio Progress object
        print("Extracting metadata...")
        metadata_by_model, metadata_by_prompt = evaluator.process_images(image_paths)
        state['metadata_by_model'] = metadata_by_model
        state['metadata_by_prompt'] = metadata_by_prompt
        
        # Create model summary
        model_summary = []
        for model, images in metadata_by_model.items():
            model_summary.append(f"- {model}: {len(images)} images")
        
        # Create prompt summary
        prompt_summary = []
        for prompt, images in metadata_by_prompt.items():
            prompt_summary.append(f"- {prompt}: {len(images)} images")
        
        return (
            f"Processed {len(image_paths)} images.\n\n"
            f"Found {len(metadata_by_model)} models:\n" + "\n".join(model_summary) + "\n\n"
            f"Found {len(metadata_by_prompt)} unique prompts."
        )
    
    def evaluate_images():
        """Evaluate all uploaded images."""
        if not state['uploaded_images']:
            return "No images uploaded. Please upload images first."
        
        # Evaluate images
        # Use a simple progress message instead of Gradio Progress object
        print("Evaluating images...")
        evaluation_results = evaluator.evaluate_images(state['uploaded_images'])
        state['evaluation_results'] = evaluation_results
        
        return f"Evaluated {len(evaluation_results)} images with all metrics."
    
    def compare_models():
        """Compare models based on evaluation results."""
        if not state['evaluation_results'] or not state['metadata_by_model']:
            return "No evaluation results available. Please evaluate images first.", None, None
        
        # Compare models
        comparison_df, visualizations = evaluator.compare_models(
            state['evaluation_results'], state['metadata_by_model']
        )
        state['comparison_df'] = comparison_df
        state['visualizations'] = visualizations
        
        # Generate report
        report_path = evaluator.generate_report(comparison_df, visualizations)
        state['report_path'] = report_path
        
        # Get visualization paths
        heatmap_path = visualizations.get('Model Comparison Heatmap')
        radar_path = visualizations.get('Model Comparison Radar Chart')
        overall_score_path = visualizations.get('overall_score Comparison')
        
        # Convert DataFrame to markdown for display
        df_markdown = comparison_df.to_markdown()
        
        return df_markdown, heatmap_path, radar_path
    
    def export_results(format):
        """Export results in the specified format."""
        if state['comparison_df'] is None:
            return "No comparison results available. Please compare models first."
        
        export_path = evaluator.export_results(format)
        if export_path:
            return f"Results exported to {export_path}"
        else:
            return "Failed to export results."
    
    def view_report():
        """View the generated HTML report."""
        if state['report_path'] and os.path.exists(state['report_path']):
            return state['report_path']
        else:
            return "No report available. Please compare models first."
    
    # Create interface
    with gr.Blocks(title="Image Model Evaluator") as interface:
        gr.Markdown("# Image Model Evaluator")
        gr.Markdown("Upload images generated by different AI models to compare their quality and performance.")
        
        with gr.Tab("Upload & Process"):
            with gr.Row():
                with gr.Column():
                    upload_input = gr.File(
                        label="Upload Images (PNG format)",
                        file_count="multiple",
                        type="filepath"  # Changed from 'file' to 'filepath'
                    )
                    upload_button = gr.Button("Process Uploaded Images")
                
                with gr.Column():
                    upload_output = gr.Textbox(
                        label="Processing Results",
                        lines=10,
                        interactive=False
                    )
            
            evaluate_button = gr.Button("Evaluate Images")
            evaluate_output = gr.Textbox(
                label="Evaluation Status",
                lines=2,
                interactive=False
            )
        
        with gr.Tab("Compare Models"):
            compare_button = gr.Button("Compare Models")
            
            with gr.Row():
                comparison_output = gr.Markdown(
                    label="Comparison Results"
                )
            
            with gr.Row():
                with gr.Column():
                    heatmap_output = gr.Image(
                        label="Model Comparison Heatmap",
                        interactive=False
                    )
                
                with gr.Column():
                    radar_output = gr.Image(
                        label="Model Comparison Radar Chart",
                        interactive=False
                    )
        
        with gr.Tab("Export & Report"):
            with gr.Row():
                with gr.Column():
                    export_format = gr.Radio(
                        label="Export Format",
                        choices=["csv", "excel", "html"],
                        value="csv"
                    )
                    export_button = gr.Button("Export Results")
                    export_output = gr.Textbox(
                        label="Export Status",
                        lines=2,
                        interactive=False
                    )
                
                with gr.Column():
                    report_button = gr.Button("View Full Report")
                    report_output = gr.HTML(
                        label="Full Report"
                    )
        
        # Set up event handlers
        upload_button.click(
            upload_images,
            inputs=[upload_input],
            outputs=[upload_output]
        )
        
        evaluate_button.click(
            evaluate_images,
            inputs=[],
            outputs=[evaluate_output]
        )
        
        compare_button.click(
            compare_models,
            inputs=[],
            outputs=[comparison_output, heatmap_output, radar_output]
        )
        
        export_button.click(
            export_results,
            inputs=[export_format],
            outputs=[export_output]
        )
        
        report_button.click(
            view_report,
            inputs=[],
            outputs=[report_output]
        )
    
    return interface


# Launch the application
if __name__ == "__main__":
    interface = create_interface()
    # Remove share=True for HuggingFace Spaces
    interface.launch()