import flatnav
from flatnav.index import index_factory
from flatnav.index import L2Index, IPIndex
from typing import Union 

def create_index(
    distance_type: str, dim: int, dataset_size: int, max_edges_per_node: int
) -> Union[L2Index, IPIndex]:
    index = index_factory(
        distance_type=distance_type,
        dim=dim,
        dataset_size=dataset_size,
        max_edges_per_node=max_edges_per_node,
        verbose=True,
    )
    if not (isinstance(index, L2Index) or isinstance(index, IPIndex)):
        raise RuntimeError("Invalid index.")

    return index

flatnav_index = create_index(distance_type="l2", dataset_size=60000, max_edges_per_node=32)