#!/usr/bin/env python3
import cv2
import numpy as np
import matplotlib.pyplot as plt
import argparse
from skimage.metrics import structural_similarity as ssim
from skimage.feature import hog
from scipy.spatial.distance import cosine

def preprocess_signature(image_path):
    """Preprocess the signature image to prepare for comparison."""
    # Load image and convert to grayscale
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # Threshold to create binary image
    _, binary = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
    
    # Find bounding rectangle of signature content
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None, None
    
    # Combine all contours to get full signature area
    all_points = np.concatenate(contours)
    x, y, w, h = cv2.boundingRect(all_points)
    
    # Crop to the content area
    cropped = binary[y:y+h, x:x+w]
    
    # Resize to standard dimensions for comparison (preserving aspect ratio)
    target_size = (300, 150)
    h, w = cropped.shape
    aspect = w / h
    
    if aspect > target_size[0] / target_size[1]:
        new_w = target_size[0]
        new_h = int(new_w / aspect)
    else:
        new_h = target_size[1]
        new_w = int(new_h * aspect)
    
    resized = cv2.resize(cropped, (new_w, new_h))
    
    # Create a blank white image with target size
    normalized = np.ones((target_size[1], target_size[0]), dtype=np.uint8) * 255
    
    # Calculate position to paste (centered)
    start_x = (target_size[0] - new_w) // 2
    start_y = (target_size[1] - new_h) // 2
    
    # Place the signature in the center
    normalized[start_y:start_y+new_h, start_x:start_x+new_w] = resized
    
    # Invert back to white background, black signature for display
    normalized_display = cv2.bitwise_not(normalized)
    
    return normalized, normalized_display

def compare_signatures(img1, img2):
    """Compare two preprocessed signatures using multiple methods."""
    results = {}
    
    # Method 1: Structural Similarity Index (SSIM)
    ssim_score, _ = ssim(img1, img2, full=True)
    results['ssim'] = ssim_score * 100  # Convert to percentage
    
    # Method 2: Histogram of Oriented Gradients (HOG) with cosine similarity
    hog_features1 = hog(img1, orientations=8, pixels_per_cell=(16, 16),
                       cells_per_block=(1, 1), visualize=False)
    hog_features2 = hog(img2, orientations=8, pixels_per_cell=(16, 16),
                       cells_per_block=(1, 1), visualize=False)
    
    cosine_sim = 1 - cosine(hog_features1, hog_features2)
    results['hog_cosine'] = cosine_sim * 100  # Convert to percentage
    
    # Method 3: Template matching
    result = cv2.matchTemplate(img1, img2, cv2.TM_CCOEFF_NORMED)
    _, max_val, _, _ = cv2.minMaxLoc(result)
    results['template'] = max_val * 100  # Convert to percentage
    
    # Calculate combined score (weighted average)
    combined_score = (results['ssim'] * 0.4 + 
                      results['hog_cosine'] * 0.4 + 
                      results['template'] * 0.2)
    
    results['combined'] = combined_score
    
    return results

def visualize_comparison(img1, img2, result):
    """Generate a visual comparison of the signatures."""
    # Create a diff image
    diff = cv2.absdiff(img1, img2)
    
    # Create figure
    plt.figure(figsize=(12, 4))
    
    # Plot first signature
    plt.subplot(1, 3, 1)
    plt.imshow(img1, cmap='gray')
    plt.title('Signature 1')
    plt.axis('off')
    
    # Plot second signature
    plt.subplot(1, 3, 2)
    plt.imshow(img2, cmap='gray')
    plt.title('Signature 2')
    plt.axis('off')
    
    # Plot difference
    plt.subplot(1, 3, 3)
    plt.imshow(diff, cmap='hot')
    plt.title(f'Differences (Similarity: {result["combined"]:.2f}%)')
    plt.axis('off')
    
    plt.tight_layout()
    
    # Save the comparison image
    comparison_path = 'signature_comparison.png'
    plt.savefig(comparison_path)
    plt.close()
    
    return comparison_path

def main():
    parser = argparse.ArgumentParser(description='Compare two signature images')
    parser.add_argument('--sig1', required=True, help='Path to first signature image')
    parser.add_argument('--sig2', required=True, help='Path to second signature image')
    parser.add_argument('--threshold', type=float, default=75.0, 
                        help='Similarity threshold percentage to consider signatures matching')
    
    args = parser.parse_args()
    
    # Preprocess signatures
    print("Preprocessing signatures...")
    norm_img1, display_img1 = preprocess_signature(args.sig1)
    norm_img2, display_img2 = preprocess_signature(args.sig2)
    
    if norm_img1 is None or norm_img2 is None:
        print("Error: Could not process one or both signature images")
        return
    
    # Compare preprocessed signatures
    print("Comparing signatures...")
    results = compare_signatures(norm_img1, norm_img2)
    
    # Generate visual comparison
    comparison_path = visualize_comparison(display_img1, display_img2, results)
    
    # Print results
    print("\nSignature Comparison Results:")
    print(f"SSIM Similarity: {results['ssim']:.2f}%")
    print(f"HOG Feature Similarity: {results['hog_cosine']:.2f}%")
    print(f"Template Matching: {results['template']:.2f}%")
    print(f"\nCombined Similarity: {results['combined']:.2f}%")
    
    # Determine match status
    if results['combined'] >= args.threshold:
        print("\n✅ Signatures MATCH")
    else:
        print("\n❌ Signatures DON'T MATCH")
    
    print(f"\nComparison visualization saved to: {comparison_path}")

if __name__ == "__main__":
    main()