import os
import sys
import torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification

def verify():
    custom_dir = "custom_model"
    if not os.path.exists(custom_dir):
        print(f"Error: custom_model directory '{custom_dir}' does not exist.")
        sys.exit(1)

    print("-> Loading custom fine-tuned model...")
    try:
        processor = AutoImageProcessor.from_pretrained(custom_dir)
        model = AutoModelForImageClassification.from_pretrained(custom_dir)
        print("[OK] Custom model loaded successfully!")
    except Exception as e:
        print(f"[ERROR] Failed to load custom model: {e}")
        sys.exit(1)

    # Print labels
    labels = list(model.config.id2label.values())
    print(f"Model Labels: {labels}")
    assert "HEN" in labels, "HEN label missing"
    assert "PEACOCK" in labels, "PEACOCK label missing"
    assert "OTHER" in labels, "OTHER label missing"

    # Define test images
    test_cases = [
        ("dataset/HEN/hen1.jpg", "HEN"),
        ("dataset/PEACOCK/peacock1.jpg", "PEACOCK"),
    ]

    # Find a file in OTHER to test
    other_dir = "dataset/OTHER"
    if os.path.exists(other_dir):
        other_files = [f for f in os.listdir(other_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
        if other_files:
            test_cases.append((os.path.join(other_dir, other_files[0]), "OTHER"))

    print("\n-> Running prediction tests...")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    all_passed = True
    for img_path, expected_class in test_cases:
        if not os.path.exists(img_path):
            print(f"[Warning] Test image '{img_path}' not found. Skipping test.")
            continue

        try:
            image = Image.open(img_path).convert("RGB")
            inputs = processor(images=image, return_tensors="pt").to(device)
            
            with torch.no_grad():
                outputs = model(**inputs)
            
            logits = outputs.logits
            probs = torch.softmax(logits, dim=-1)
            pred_idx = torch.argmax(probs, dim=-1).item()
            pred_label = model.config.id2label[pred_idx]
            confidence = probs[0][pred_idx].item()

            print(f"Image: {img_path}")
            print(f"  Expected: {expected_class}")
            print(f"  Predicted: {pred_label} (confidence: {confidence:.4f})")
            
            if pred_label == expected_class:
                print("  [PASS]")
            else:
                print("  [FAIL] (Mismatch)")
                all_passed = False
        except Exception as e:
            print(f"  [FAIL] (Error: {e})")
            all_passed = False

    if all_passed:
        print("\nAll verification tests passed successfully!")
    else:
        print("\nSome verification tests failed. Please inspect the outputs.")

if __name__ == "__main__":
    verify()