from pymilvus import MilvusClient, DataType

# 1. Set up a Milvus client
client = MilvusClient(
    uri="http://localhost:19530"
)
client.get_collection_stats('nuggets_so400m')
# %%

schema = MilvusClient.create_schema(
    auto_id=False,
    enable_dynamic_field=False,
)
schema.add_field(field_name="primary_id",datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="filepath", datatype=DataType.VARCHAR, max_length=128)
schema.add_field(field_name="frame_number", datatype=DataType.INT32)
schema.add_field(field_name="so400m", datatype=DataType.FLOAT16_VECTOR, dim=1152)


index_params = client.prepare_index_params()


index_params.add_index(
    field_name="primary_id",
    index_type="STL_SORT")

index_params.add_index(
    field_name="filepath",
    index_type="Trie")

index_params.add_index(
    field_name="so400m", 
    index_type="IVF_FLAT",
    metric_type="COSINE",
    params={ "nlist": 128 })



client.create_collection(
    collection_name="nuggets_so400m",
    schema=schema,
    index_params=index_params
)

# %%
res = client.get_load_state(
    collection_name="nuggets_so400m"
)


res = client.load_collection(collection_name="nuggets_so400m")