import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from deepfake_detector import DeepfakeDetector
from image_processor import ImageProcessor
from labeling import ThreatLabeler
from heatmap_generator import HeatmapGenerator
from combined_visualization import CombinedVisualizer

def ensure_dir(directory):
    """Ensure directory exists"""
    if not os.path.exists(directory):
        os.makedirs(directory)

def demo_combined_visualization(image1_path, image2_path, output_dir, threshold=30, min_area=100):
    """
    Demonstrate the combined visualization that overlays difference image, 
    low and medium threat heatmaps, and bounding boxes.
    
    Args:
        image1_path: Path to original image
        image2_path: Path to modified image
        output_dir: Directory to save outputs
        threshold: Threshold for difference detection
        min_area: Minimum area for region detection
    """
    # Initialize components
    detector = DeepfakeDetector()
    labeler = ThreatLabeler()
    heatmap_gen = HeatmapGenerator()
    img_processor = ImageProcessor()
    visualizer = CombinedVisualizer()
    
    # Create output directory
    ensure_dir(output_dir)
    
    # Get base filename for outputs
    base_name = os.path.splitext(os.path.basename(image1_path))[0]
    
    print(f"Processing images: {image1_path} and {image2_path}")
    
    # Step 1: Verification Module - Process the image pair
    results = detector.process_image_pair(image1_path, image2_path, threshold, min_area)
    
    # Step 2: Labeling System - Label detected regions by threat level
    original_image = img_processor.load_image(image1_path)
    modified_image = img_processor.load_image(image2_path)
    labeled_image, labeled_regions = labeler.label_regions(
        original_image, results['difference_image'], results['bounding_boxes'])
    
    # Get threat summary
    threat_summary = labeler.get_threat_summary(labeled_regions)
    
    # Step 3: Generate multi-level heatmaps
    multi_heatmaps = heatmap_gen.generate_multi_level_heatmap(original_image, labeled_regions)
    
    # Prepare results for combined visualization
    image_pair_results = {
        'original_image': original_image,
        'modified_image': modified_image,
        'difference_image': results['difference_image'],
        'threshold_image': results['threshold_image'],
        'annotated_image': results['annotated_image'],
        'labeled_image': labeled_image,
        'multi_heatmaps': multi_heatmaps,
        'bounding_boxes': results['bounding_boxes'],
        'labeled_regions': labeled_regions,
        'threat_summary': threat_summary,
        'smi_score': results['smi_score']
    }
    
    # Create combined visualization
    output_path = os.path.join(output_dir, f"{base_name}_combined_overlay.png")
    visualizer.create_combined_visualization(
        image_pair_results, 
        output_path, 
        alpha_diff=0.4,  # Transparency for difference image
        alpha_low=0.3,   # Transparency for low threat heatmap
        alpha_medium=0.4 # Transparency for medium threat heatmap
    )
    
    # Print summary information
    print(f"\nAnalysis Results:")
    print(f"SMI Score: {results['smi_score']:.4f} (1.0 = identical, 0.0 = completely different)")
    print(f"Total regions detected: {threat_summary['total_regions']}")
    print(f"Threat counts: Low={threat_summary['threat_counts']['low']}, "
          f"Medium={threat_summary['threat_counts']['medium']}, "
          f"High={threat_summary['threat_counts']['high']}")
    if threat_summary['max_threat']:
        print(f"Maximum threat: {threat_summary['max_threat']['level'].upper()} "
              f"({threat_summary['max_threat']['percentage']:.1f}%)")
    print(f"Average difference: {threat_summary['average_difference']:.1f}%")
    
    print(f"\nCombined visualization saved to: {output_path}")
    return output_path

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 3:
        print("Usage: python demo_combined_visualization.py <original_image> <modified_image> [output_dir]")
        sys.exit(1)
    
    image1_path = sys.argv[1]
    image2_path = sys.argv[2]
    output_dir = sys.argv[3] if len(sys.argv) > 3 else "./output"
    
    demo_combined_visualization(image1_path, image2_path, output_dir)