from pinecone.grpc import PineconeGRPC, GRPCClientConfig
from pinecone import ServerlessSpec

# Initialize a client.
# API key is required, but the value does not matter.
# Host and port of the Pinecone Local instance
# is required when starting without indexes. 
pc = PineconeGRPC(
    api_key="pclocal", 
    host="http://localhost:5080" 
)                                    

# Create two indexes, one dense and one sparse
dense_index_name = "dense-index"
sparse_index_name = "sparse-index"

if not pc.has_index(dense_index_name):  
    dense_index_model = pc.create_index(
        name=dense_index_name,
        vector_type="dense",
        dimension=2,
        metric="cosine",
        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
        deletion_protection="disabled",
        tags={"environment": "development"}
    )

print("Dense index model:\n", dense_index_model)

if not pc.has_index(sparse_index_name):  
    sparse_index_model = pc.create_index(
        name=sparse_index_name,
        vector_type="sparse",
        metric="dotproduct",
        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
        deletion_protection="disabled",
        tags={"environment": "development"}
    )

print("\nSparse index model:\n", sparse_index_model)

# Target each index, disabling tls
dense_index_host = pc.describe_index(name=dense_index_name).host
dense_index = pc.Index(host=dense_index_host, grpc_config=GRPCClientConfig(secure=False))
sparse_index_host = pc.describe_index(name=sparse_index_name).host
sparse_index = pc.Index(host=sparse_index_host, grpc_config=GRPCClientConfig(secure=False))

# Upsert records into the dense index
dense_index.upsert(
    vectors=[
        {
            "id": "vec1", 
            "values": [1.0, -2.5],
            "metadata": {"genre": "drama"}
        },
        {
            "id": "vec2", 
            "values": [3.0, -2.0],
            "metadata": {"genre": "documentary"}
        },
        {
            "id": "vec3", 
            "values": [0.5, -1.5],
            "metadata": {"genre": "documentary"}
        }
    ],
    namespace="example-namespace"
)

# Upsert records into the sparse index
sparse_index.upsert(
    namespace="example-namespace",
    vectors=[
        {
            "id": "vec1",
            "sparse_values": {
                "values": [1.7958984, 0.41577148, 2.828125, 2.8027344, 2.8691406, 1.6533203, 5.3671875, 1.3046875, 0.49780273, 0.5722656, 2.71875, 3.0820312, 2.5019531, 4.4414062, 3.3554688],
                "indices": [822745112, 1009084850, 1221765879, 1408993854, 1504846510, 1596856843, 1640781426, 1656251611, 1807131503, 2543655733, 2902766088, 2909307736, 3246437992, 3517203014, 3590924191]
            },
            "metadata": {
                "chunk_text": "AAPL reported a year-over-year revenue increase, expecting stronger Q3 demand for its flagship phones.",
                "category": "technology",
                "quarter": "Q3"
            }
        },
        {
            "id": "vec2",
            "sparse_values": {
                "values": [0.4362793, 3.3457031, 2.7714844, 3.0273438, 3.3164062, 5.6015625, 2.4863281, 0.38134766, 1.25, 2.9609375, 0.34179688, 1.4306641, 0.34375, 3.3613281, 1.4404297, 2.2558594, 2.2597656, 4.8710938, 0.5605469],
                "indices": [131900689, 592326839, 710158994, 838729363, 1304885087, 1640781426, 1690623792, 1807131503, 2066971792, 2428553208, 2548600401, 2577534050, 3162218338, 3319279674, 3343062801, 3476647774, 3485013322, 3517203014, 4283091697]
            },
            "metadata": {
                "chunk_text": "Analysts suggest that AAPL'\''s upcoming Q4 product launch event might solidify its position in the premium smartphone market.",
                "category": "technology",
                "quarter": "Q4"
            }
        },
        {
            "id": "vec3",
            "sparse_values": {
                "values": [2.6875, 4.2929688, 3.609375, 3.0722656, 2.1152344, 5.78125, 3.7460938, 3.7363281, 1.2695312, 3.4824219, 0.7207031, 0.0826416, 4.671875, 3.7011719, 2.796875, 0.61621094],
                "indices": [8661920, 350356213, 391213188, 554637446, 1024951234, 1640781426, 1780689102, 1799010313, 2194093370, 2632344667, 2641553256, 2779594451, 3517203014, 3543799498, 3837503950, 4283091697]
            },
            "metadata": {
                "chunk_text": "AAPL'\''s strategic Q3 partnerships with semiconductor suppliers could mitigate component risks and stabilize iPhone production",
                "category": "technology",
                "quarter": "Q3"
            }
        }
    ]
)

# Check the number of records in each index
print("\nDense index stats:\n", dense_index.describe_index_stats())
print("\nSparse index stats:\n", sparse_index.describe_index_stats())

# Query the dense index with a metadata filter
dense_response = dense_index.query(
    namespace="example-namespace",
    vector=[3.0, -2.0],
    filter={"genre": {"$eq": "documentary"}},
    top_k=1,
    include_values=False,
    include_metadata=True
)

print("\nDense query response:\n", dense_response)

# Query the sparse index with a metadata filter
sparse_response = sparse_index.query(
    namespace="example-namespace",
    sparse_vector={
      "values": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
      "indices": [767227209, 1640781426, 1690623792, 2021799277, 2152645940, 2295025838, 2443437770, 2779594451, 2956155693, 3476647774, 3818127854, 4283091697]
    }, 
    filter={
        "quarter": {"$eq": "Q4"}
    },
    top_k=1,
    include_values=False,
    include_metadata=True
)

print("/nSparse query response:\n", sparse_response)

# Delete the indexes
pc.delete_index(name=dense_index_name)
pc.delete_index(name=sparse_index_name)

docker compose down