"""
Example usage script for the Skin Type Classification model on Hugging Face.
"""

from transformers import AutoModelForImageClassification, AutoImageProcessor
from PIL import Image
import torch
import requests
from io import BytesIO

def load_model(model_name="your-username/skin-type-classifier"):
    """Load the model and processor from Hugging Face."""
    model = AutoModelForImageClassification.from_pretrained(model_name)
    processor = AutoImageProcessor.from_pretrained(model_name)
    return model, processor

def predict_skin_type(image_path_or_url, model, processor):
    """
    Predict skin type from an image.
    
    Args:
        image_path_or_url: Path to local image or URL
        model: The loaded model
        processor: The loaded processor
    
    Returns:
        dict: Prediction results with class and confidence
    """
    # Load image
    if image_path_or_url.startswith(('http://', 'https://')):
        response = requests.get(image_path_or_url)
        image = Image.open(BytesIO(response.content))
    else:
        image = Image.open(image_path_or_url)
    
    # Convert to RGB if needed
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Process image
    inputs = processor(images=image, return_tensors="pt")
    
    # Make prediction
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_class_idx = predictions.argmax().item()
        confidence = predictions[0][predicted_class_idx].item()
    
    # Map to class names
    class_names = {0: "dry", 1: "oily"}
    predicted_class = class_names[predicted_class_idx]
    
    return {
        "predicted_class": predicted_class,
        "confidence": confidence,
        "all_scores": {
            "dry": predictions[0][0].item(),
            "oily": predictions[0][1].item()
        }
    }

def main():
    """Example usage of the skin type classification model."""
    print("🔬 Loading Skin Type Classification Model...")
    
    # Load model and processor
    model, processor = load_model()
    
    print("✅ Model loaded successfully!")
    
    # Example with local image (replace with your image path)
    try:
        image_path = "example_skin_image.jpg"  # Replace with actual image path
        result = predict_skin_type(image_path, model, processor)
        
        print(f"\n📊 Prediction Results:")
        print(f"Predicted Skin Type: {result['predicted_class']}")
        print(f"Confidence: {result['confidence']:.2%}")
        print(f"All Scores: {result['all_scores']}")
        
    except FileNotFoundError:
        print("ℹ️  Please provide a valid image path to test the model")
    
    # Example usage patterns
    print("\n💡 Usage Examples:")
    print("1. Local image: predict_skin_type('path/to/image.jpg', model, processor)")
    print("2. URL image: predict_skin_type('https://example.com/image.jpg', model, processor)")

if __name__ == "__main__":
    main()