from elasticsearch import Elasticsearch
from PIL import Image
import torch
from transformers import CLIPProcessor, CLIPModel
import numpy as np

es = Elasticsearch("http://localhost:9200")
index_name = "patents"

# Load CLIP model for both text and image embeddings
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")

# CLIP text embedding
def embed_text(text):
    inputs = clip_processor(text=[text], return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        text_features = clip_model.get_text_features(**inputs)
    return text_features[0].cpu().numpy().astype(np.float32).tolist()

# CLIP image embedding
def embed_image(image_path):
    image = Image.open(image_path).convert("RGB")
    inputs = clip_processor(images=image, return_tensors="pt")
    with torch.no_grad():
        image_features = clip_model.get_image_features(**inputs)
    return image_features[0].cpu().numpy().astype(np.float32).tolist()

def search_by_embedding(embedding, dims=512, top_k=5):
    print(f"Embedding length: {len(embedding)} (should be {dims})")
    if len(embedding) != dims:
        print(f"Error: Embedding length {len(embedding)} does not match expected dims {dims}.")
        return
    body = {
        "size": top_k,
        "query": {
            "script_score": {
                "query": {"match_all": {}},
                "script": {
                    "source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0",
                    "params": {"query_vector": embedding}
                }
            }
        }
    }
    try:
        res = es.search(index=index_name, body=body)
        for hit in res['hits']['hits']:
            doc = hit['_source']
            print(f"Patent ID: {doc.get('doc_numbers', [''])[0] if doc.get('doc_numbers') else ''}")
            print(f"Invention Title: {doc.get('invention_title', '')}")
            print(f"Assignors: {doc.get('assignors', '')}")
            print(f"Assignees: {doc.get('assignees', '')}")
            print(f"Conveyance Text: {doc.get('conveyance_text', '')}")
            print(f"Reel No: {doc.get('reel_no', '')}")
            print(f"Frame No: {doc.get('frame_no', '')}")
            print("-"*40)
    except Exception as e:
        print(f"Elasticsearch error: {e}")

if __name__ == "__main__":
    mode = input("Enter 'text' to search by prompt or 'image' to search by image: ").strip().lower()
    if mode == 'text':
        query = input("Enter your search prompt: ")
        emb = embed_text(query)
        search_by_embedding(emb, dims=512)
    elif mode == 'image':
        image_path = input("Enter the path to your image: ")
        emb = embed_image(image_path)
        search_by_embedding(emb, dims=512)
    else:
        print("Invalid mode.")