import numpy as np
import cv2
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.preprocessing import image
from skimage.metrics import structural_similarity as ssim
import os
import argparse

class ImageCharacterClassifier:
    def __init__(self, similarity_threshold=0.7):
        # Initialize ResNet50 model without top classification layer
        self.model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
        self.similarity_threshold = similarity_threshold

    def load_and_preprocess_image(self, image_path, target_size=(224, 224)):
        # Load and preprocess image for ResNet50
        img = image.load_img(image_path, target_size=target_size)
        img_array = image.img_to_array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = preprocess_input(img_array)
        return img_array

    def extract_features(self, image_path):
        # Extract deep features using ResNet50
        preprocessed_img = self.load_and_preprocess_image(image_path)
        features = self.model.predict(preprocessed_img)
        return features

    def calculate_ssim(self, img1_path, img2_path):
        # Calculate SSIM between two images
        img1 = cv2.imread(img1_path)
        img2 = cv2.imread(img2_path)
        
        # Convert to grayscale if images are in color
        if len(img1.shape) == 3:
            img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
        if len(img2.shape) == 3:
            img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
        
        # Resize images to same dimensions
        img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))
        
        score = ssim(img1, img2)
        return score

    def classify_images(self, reference_image_path, image_folder_path):
        # Extract features from reference image
        reference_features = self.extract_features(reference_image_path)
        
        results = []
        
        # Process each image in the folder
        for image_name in os.listdir(image_folder_path):
            if image_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                image_path = os.path.join(image_folder_path, image_name)
                
                try:
                    # Calculate SSIM
                    ssim_score = self.calculate_ssim(reference_image_path, image_path)
                    
                    # Extract features and calculate similarity
                    image_features = self.extract_features(image_path)
                    
                    # Calculate cosine similarity
                    feature_similarity = np.dot(reference_features.flatten(), 
                                             image_features.flatten()) / (
                        np.linalg.norm(reference_features) * np.linalg.norm(image_features))
                    
                    # Give more weight to feature similarity
                    combined_similarity = (0.3 * ssim_score + 0.7 * feature_similarity)
                    
                    # Classify based on similarity threshold
                    is_similar = combined_similarity >= self.similarity_threshold
                    
                    results.append({
                        'image_name': image_name,
                        'ssim_score': ssim_score,
                        'feature_similarity': feature_similarity,
                        'combined_similarity': combined_similarity,
                        'is_similar': is_similar
                    })
                    
                except Exception as e:
                    print(f"Error processing {image_name}: {str(e)}")
                    continue
        
        return results

def main():
    # Create argument parser
    parser = argparse.ArgumentParser(description='Image Character Classification')
    parser.add_argument('--reference', '-r', 
                       type=str, 
                       required=True,
                       help='Path to reference image')
    parser.add_argument('--folder', '-f', 
                       type=str, 
                       required=True,
                       help='Path to folder containing images to compare')
    parser.add_argument('--threshold', '-t', 
                       type=float, 
                       default=0.5,  # Lowered the default threshold
                       help='Similarity threshold (default: 0.5)')
    
    # Parse arguments
    args = parser.parse_args()
    
    # Initialize classifier
    classifier = ImageCharacterClassifier(similarity_threshold=args.threshold)
    
    # Check if paths exist
    if not os.path.exists(args.reference):
        print(f"Error: Reference image not found at {args.reference}")
        return
    
    if not os.path.exists(args.folder):
        print(f"Error: Image folder not found at {args.folder}")
        return
    
    # Perform classification
    results = classifier.classify_images(args.reference, args.folder)
    
    # Sort results by similarity score
    results.sort(key=lambda x: x['combined_similarity'], reverse=True)
    
    # Print results
    print("\nResults sorted by similarity (highest to lowest):")
    print("-" * 50)
    for result in results:
        print(f"\nImage: {result['image_name']}")
        print(f"SSIM Score: {result['ssim_score']:.3f}")
        print(f"Feature Similarity: {result['feature_similarity']:.3f}")
        print(f"Combined Similarity: {result['combined_similarity']:.3f}")
        print(f"Is Similar: {result['is_similar']}")
        print("-" * 30)

if __name__ == "__main__":
    main()