from octanedb import OctaneDB

# Initialize with text embedding support
db = OctaneDB(
    dimension=384,  # Will be auto-set by embedding model
    embedding_model="all-MiniLM-L6-v2"
)

# Create a collection
collection = db.create_collection("documents")
db.use_collection("documents")

# Add text documents (ChromaDB-compatible!)
result = db.add(
    ids=["doc1", "doc2"],
    documents=[
        "This is a document about pineapple",
        "This is a document about oranges"
    ],
    metadatas=[
        {"category": "tropical", "color": "yellow"},
        {"category": "citrus", "color": "orange"}
    ]
)

# Search by text query
results = db.search_text(
    query_text="fruit",
    k=2,
    filter="category == 'tropical'",
    include_metadata=True
)

for doc_id, distance, metadata in results:
    print(f"Document: {db.get_document(doc_id)}")
    print(f"Distance: {distance:.4f}")
    print(f"Metadata: {metadata}")

from octanedb import OctaneDB

# Initialize database with text embeddings
db = OctaneDB(
    dimension=384,  # sentence-transformers default dimension
    storage_mode="in-memory",
    enable_text_embeddings=True,
    embedding_model="all-MiniLM-L6-v2"  # Lightweight model
)

# Create a collection
db.create_collection("fruits")
db.use_collection("fruits")

# Add some fruit documents
fruits_data = [
    {"id": "apple", "text": "Apple is a sweet and crunchy fruit that grows on trees.", "category": "temperate"},
    {"id": "banana", "text": "Banana is a yellow tropical fruit rich in potassium.", "category": "tropical"},
    {"id": "mango", "text": "Mango is a sweet tropical fruit with a large seed.", "category": "tropical"},
    {"id": "orange", "text": "Orange is a citrus fruit with a bright orange peel.", "category": "citrus"}
]

for fruit in fruits_data:
    db.add(
        ids=[fruit["id"]],
        documents=[fruit["text"]],
        metadatas=[{"category": fruit["category"], "type": "fruit"}]
    )

# Simple text search
results = db.search_text(query_text="sweet", k=2, include_metadata=True)
print("Sweet fruits:")
for doc_id, distance, metadata in results:
    print(f"  • {doc_id}: {metadata.get('document', 'N/A')[:50]}...")

# Text search with filter
results = db.search_text(
    query_text="fruit", 
    k=2, 
    filter="category == 'tropical'",
    include_metadata=True
)
print("\nTropical fruits:")
for doc_id, distance, metadata in results:
    print(f"  • {doc_id}: {metadata.get('document', 'N/A')[:50]}...")

# Old ChromaDB code
# collection.add(
#     ids=["id1", "id2"],
#     documents=["doc1", "doc2"]
# )

# New OctaneDB code (identical API!)
db.add(
    ids=["id1", "id2"],
    documents=["doc1", "doc2"]
)

# Batch text search
query_texts = ["machine learning", "artificial intelligence", "data science"]
batch_results = db.search_text_batch(
    query_texts=query_texts,
    k=5,
    include_metadata=True
)

# Change embedding models
db.change_embedding_model("all-mpnet-base-v2")  # Higher quality, 768 dimensions

# Get available models
models = db.get_available_models()
print(f"Available models: {models}")

# Use pre-computed embeddings
custom_embeddings = np.random.randn(100, 384).astype(np.float32)
result = db.add(
    ids=[f"vec_{i}" for i in range(100)],
    embeddings=custom_embeddings,
    metadatas=[{"source": "custom"} for _ in range(100)]
)

# Optimize for speed vs. accuracy
db = OctaneDB(
    dimension=384,
    m=8,              # Fewer connections = faster, less accurate
    ef_construction=100,  # Lower = faster build
    ef_search=50      # Lower = faster search
)

# Persistent storage
db = OctaneDB(
    dimension=384,
    storage_path="./data",
    embedding_model="all-MiniLM-L6-v2"
)

# Save and load
db.save("./my_database.h5")
loaded_db = OctaneDB.load("./my_database.h5")

# Complex filters
results = db.search_text(
    query_text="technology",
    k=10,
    filter={
        "$and": [
            {"category": "tech"},
            {"$or": [
                {"year": {"$gte": 2020}},
                {"priority": "high"}
            ]}
        ]
    }
)

# This will work correctly:
results = db.search_text(
    query_text="fruit", 
    k=2, 
    filter="category == 'tropical'",
    include_metadata=True  # Important!
)

# Process results correctly:
for doc_id, distance, metadata in results:
    print(f"ID: {doc_id}, Distance: {distance:.4f}")
    if metadata:
        print(f"  Document: {metadata.get('document', 'N/A')}")
        print(f"  Category: {metadata.get('category', 'N/A')}")