import requests
import base64
from io import BytesIO
from PIL import Image
import json
import matplotlib.pyplot as plt
import numpy as np

# This is an example client that demonstrates how to use the Gradio API
# You would replace this URL with the actual URL of your deployed Huggingface Space
GRADIO_API_URL = "https://your-username-nomic-vision-embedding.hf.space/api/predict"

def embed_image_from_url(image_url):
    """
    Generate embeddings for an image using the Gradio API
    
    Args:
        image_url: URL of the image to embed
        
    Returns:
        The embedding vector and its dimension
    """
    try:
        # Download the image
        response = requests.get(image_url)
        image = Image.open(BytesIO(response.content))
        
        # Convert image to bytes
        img_byte_arr = BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()
        
        # Prepare the request
        files = {
            'data': ('image.png', img_byte_arr, 'image/png')
        }
        
        # Send the request to the Gradio API
        response = requests.post(GRADIO_API_URL, files=files)
        
        # Parse the response
        if response.status_code == 200:
            result = response.json()
            embedding_data = result['data'][0]
            embedding_dim = result['data'][1]
            
            return embedding_data, embedding_dim
        else:
            print(f"Error: HTTP {response.status_code}")
            print(response.text)
            return None, None
            
    except Exception as e:
        print(f"Error: {str(e)}")
        return None, None

def embed_image_from_file(image_path):
    """
    Generate embeddings for an image using the Gradio API
    
    Args:
        image_path: Path to the image file
        
    Returns:
        The embedding vector and its dimension
    """
    try:
        # Load the image
        image = Image.open(image_path)
        
        # Convert image to bytes
        img_byte_arr = BytesIO()
        image.save(img_byte_arr, format=image.format if image.format else 'PNG')
        img_byte_arr = img_byte_arr.getvalue()
        
        # Prepare the request
        files = {
            'data': ('image.png', img_byte_arr, 'image/png')
        }
        
        # Send the request to the Gradio API
        response = requests.post(GRADIO_API_URL, files=files)
        
        # Parse the response
        if response.status_code == 200:
            result = response.json()
            embedding_data = result['data'][0]
            embedding_dim = result['data'][1]
            
            return embedding_data, embedding_dim
        else:
            print(f"Error: HTTP {response.status_code}")
            print(response.text)
            return None, None
            
    except Exception as e:
        print(f"Error: {str(e)}")
        return None, None

def visualize_embedding(embedding):
    """
    Visualize the embedding vector
    
    Args:
        embedding: The embedding vector
    """
    # Convert the embedding to a numpy array
    embedding_array = np.array(embedding)
    
    # Plot the embedding
    plt.figure(figsize=(10, 5))
    plt.plot(embedding_array)
    plt.title("Embedding Vector")
    plt.xlabel("Dimension")
    plt.ylabel("Value")
    plt.grid(True)
    plt.show()
    
    # Plot the histogram of the embedding
    plt.figure(figsize=(10, 5))
    plt.hist(embedding_array, bins=50)
    plt.title("Embedding Histogram")
    plt.xlabel("Value")
    plt.ylabel("Frequency")
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    # Example usage with an image URL
    image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/bert-architects.png"
    print(f"Generating embedding for image: {image_url}")
    
    embedding_data, embedding_dim = embed_image_from_url(image_url)
    
    if embedding_data:
        print(f"Embedding dimension: {embedding_dim}")
        print(f"First 10 values of embedding: {embedding_data['embedding'][:10]}...")
        
        # Visualize the embedding
        visualize_embedding(embedding_data['embedding'])
    
    # Example usage with a local image file
    # Uncomment the following lines to use a local image file
    # image_path = "path/to/your/image.jpg"
    # print(f"Generating embedding for image: {image_path}")
    # embedding_data, embedding_dim = embed_image_from_file(image_path)
    # 
    # if embedding_data:
    #     print(f"Embedding dimension: {embedding_dim}")
    #     print(f"First 10 values of embedding: {embedding_data['embedding'][:10]}...")
    #     
    #     # Visualize the embedding
    #     visualize_embedding(embedding_data['embedding'])