#!/usr/bin/env python3
"""
Simple script to verify a single image using a trained model
"""
import sys
import os

# Add current directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from PIL import Image
from common import load_model, predict

def main():
    if len(sys.argv) < 3:
        print("Usage: python3 verify_image.py <model_path> <image_path>")
        print("Example: python3 verify_image.py models/best_model_ep30_pos99.35_neg94.75_20251125_182545.pt /home/fam/emblem/scans/358626/sbs.jpg")
        sys.exit(1)
    
    model_path = sys.argv[1]
    image_path = sys.argv[2]
    
    if not os.path.exists(model_path):
        print(f"Error: Model file not found: {model_path}")
        sys.exit(1)
    
    if not os.path.exists(image_path):
        print(f"Error: Image file not found: {image_path}")
        sys.exit(1)
    
    print(f"Loading model from {model_path}...")
    model, transforms = load_model(model_path)
    # Ensure model is not compiled (torch.compile can cause issues when loading)
    if hasattr(model, '_orig_mod'):
        model = model._orig_mod
    print("Model loaded successfully")
    
    print(f"Loading image from {image_path}...")
    image = Image.open(image_path).convert('RGB')
    print(f"Image size: {image.size}")
    
    print("Running prediction...")
    predicted_class, probabilities = predict(model, transforms, image)
    
    # probabilities is a list of [neg_prob, pos_prob] for each cell
    # Sum up probabilities across all cells
    neg_sum = sum([p[0] for p in probabilities])
    pos_sum = sum([p[1] for p in probabilities])
    total = neg_sum + pos_sum
    neg_prob = neg_sum / total if total > 0 else 0
    pos_prob = pos_sum / total if total > 0 else 0
    
    print("\n" + "="*50)
    print("Prediction Results:")
    print("="*50)
    print(f"Predicted class: {'POSITIVE' if predicted_class == 1 else 'NEGATIVE'}")
    print(f"Negative probability: {neg_prob:.2%}")
    print(f"Positive probability: {pos_prob:.2%}")
    print(f"Number of cells evaluated: {len(probabilities)}")
    print("="*50)

if __name__ == '__main__':
    main()