import deglib

# Build and optimize a DynamicExplorationGraph directly from a numpy array in one call
deglib.build_from_data(
    data,                                      # 2D numpy array [N, D] of feature vectors
    labels=None,                               # Optional 1D uint32 array [N] of external labels (default: 0..N-1)
    edges_per_vertex=32,                       # Number of outgoing edges per vertex (must be even)
    capacity=-1,                               # Max graph capacity (-1 = data.shape[0])
    metric=Metric.FP32_L2,                     # Distance metric (Metric enum or string e.g. "FP32_L2")
    instruction="Auto",                        # SIMD instruction set ("Auto", "AVX2", "AVX512", "Scalar")
    optimization_target=OptimizationTarget.LowLID, # OptimizationTarget (LowLID, HighLID, StreamingData)
    extend_k=0,                                # Number of neighbors during graph extension (0 = edges_per_vertex)
    extend_eps=0.1,                            # Search expansion epsilon during graph extension
    improve_k=0,                               # Number of neighbors during edge improvement
    improve_eps=0.001,                         # Search expansion epsilon during edge improvement
    max_path_length=5,                         # Maximum edge swaps before canceling an improvement try
    improve_tries=0,                           # Number of improvement attempts per build step
    thread_count=0,                            # Worker thread count (0 = all available CPU cores)
    seed=None,                                 # Random seed (None = default seed)
    callback=None,                             # Optional callback function(BuilderStatus) for progress reporting
    show_progress=False                        # If True and no callback, prints a progress bar to stdout
) -> DynamicExplorationGraph

# Create empty mutable graph with fixed capacity in preallocated memory (SizeBoundedGraph)
deglib.create_empty(capacity, feature_space, edges_per_vertex=32) -> DynamicExplorationGraph

# Create empty mutable graph with chunk-based memory that grows dynamically (DynamicGraph)
deglib.create_dynamic_empty(feature_space, edges_per_vertex=32, chunk_size=1024) -> DynamicExplorationGraph

# Create a randomized initial exploration graph from a numpy feature array
deglib.create_random_graph(features, feature_space, edges_per_vertex=32, seed=7) -> DynamicExplorationGraph

# Load saved graphs from disk
deglib.load_readonly_graph(path) -> DynamicExplorationGraph
deglib.load_dynamic_graph(path, chunk_size=1024) -> DynamicExplorationGraph
deglib.load_mutable_graph(path, capacity=0) -> DynamicExplorationGraph

class DynamicExplorationGraph:
    # --- Factories & Loading ---
    # Create empty mutable graph with fixed capacity in preallocated memory (SizeBoundedGraph)
    create_empty(capacity, feature_space, edges_per_vertex=32) -> DynamicExplorationGraph

    # Create empty mutable graph with chunk-based memory that grows dynamically (DynamicGraph)
    create_dynamic_empty(feature_space, edges_per_vertex=32, chunk_size=1024) -> DynamicExplorationGraph

    # Create a randomized initial exploration graph from a numpy feature array
    create_random_graph(features, feature_space, edges_per_vertex=32, seed=7) -> DynamicExplorationGraph

    # Load a saved read-only graph file optimized for fast search
    load_readonly_graph(path) -> DynamicExplorationGraph

    # Load a saved graph file as a mutable dynamic chunk-allocated graph
    load_dynamic_graph(path, chunk_size=1024) -> DynamicExplorationGraph

    # --- Search & Querying ---
    # Search nearest neighbors for 1D (single) or 2D (batch) query vectors.
    # query: 1D [D] or 2D [Q, D] numpy array
    # eps: Search expansion factor (higher = more accurate, slower)
    # k: Number of nearest neighbors to return per query
    # filter_labels: Filter object or numpy array of valid external labels
    # max_distance_computation_count: Budget cap for distance calculations (0 = unlimited)
    # threads: Worker thread count for batch queries (0 = all CPU cores)
    # return_distances: If True, returns (indices, distances) tuple; if False, returns only indices
    # unsorted: If True, skips sorting results in ascending distance order
    search(query, eps=0.0, k=10, filter_labels=None, max_distance_computation_count=0, threads=0, return_distances=True, unsorted=False)

    # Explore graph starting from entry vertex label(s) (supports single int or batch array).
    # Returns (indices, distances) if return_distances=True, otherwise indices array.
    explore(entry_external_label, k, max_distance_computation_count=0, eps=0.0, include_entry=True, threads=1, filter_labels=None, return_distances=True, unsorted=False)

    # Returns list of neighbor external labels connected to a given vertex
    get_neighbors(external_label) -> list[int]

    # --- State & Inspection ---
    size() -> int                              # Number of active vertices
    get_edges_per_vertex() -> int              # Outgoing edges per vertex
    get_feature_space() -> FloatSpace          # Metric and dimensionality
    has_vertex(external_label) -> bool         # True if label exists in graph
    is_mutable() -> bool                       # True if graph can be modified via GraphBuilder

   # --- Traversal Tuning ---
   # Auto-tune traversal prefetch parameters (po, pl, nl) by timing traversal over sampled vertices.
   # The graph self-samples its own features as queries; call once after construction.
   # sample_count: vertices sampled as queries (capped at graph size); k: expected result count;
   # ef: beam width (clamped to >= k); seed: sampling seed.
   optimize(sample_count=50, k=100, ef=200, seed=7)

    # --- Conversions & Saving ---
    # Convert to compact, immutable graph (ReadOnlyGraph) for maximum query throughput
    to_readonly(feature_space=None, custom_features=None) -> DynamicExplorationGraph

    # Convert to mutable fixed-capacity graph (SizeBoundedGraph) with optional new metric/features
    to_mutable(feature_space=None, custom_features=None, capacity=0) -> DynamicExplorationGraph

    # Convert to mutable chunk-allocated graph (DynamicGraph) that expands dynamically
    to_dynamic(feature_space=None, custom_features=None, chunk_size=1024) -> DynamicExplorationGraph

    # Save graph topology and features to disk (graph must be mutable)
    save_graph(path)

from deglib.builder import GraphBuilder, OptimizationTarget, build_from_data

# Optimization strategy based on dataset characteristics
class OptimizationTarget:
    StreamingData  # Optimized for streaming or shifting distributions (single-threaded insertion)
    HighLID        # Optimized for high intrinsic dimensionality (> 15, parallel extension)
    LowLID         # Optimized for low intrinsic dimensionality (<= 15, parallel extension)


# Status object passed to build progress callbacks
class BuilderStatus:
    step: int                 # Number of graph manipulation steps completed
    added: int                # Total number of added vertices
    deleted: int              # Total number of deleted vertices
    improved: int             # Total number of successful edge improvements
    tries: int                # Total number of improvement attempts
    step_added_ids: list[int] # External labels added in the current step
    step_deleted_ids: list[int]# External labels deleted in the current step
    total_added_ids: list[int]# All added external labels across the entire build
    total_deleted_ids: list[int]# All deleted external labels across the entire build


class GraphBuilder:
    def __init__(
        graph,                                 # Mutable DynamicExplorationGraph
        seed=None,                             # Random seed
        optimization_target=OptimizationTarget.LowLID, # Optimization target strategy
        extend_k=0,                            # Neighbors during extension (0 = graph's edges_per_vertex)
        extend_eps=0.1,                        # Epsilon for neighbor search during extension
        improve_k=0,                           # Neighbors during edge improvement
        improve_eps=0.001,                     # Epsilon for neighbor search during improvement
        max_path_length=5,                     # Max edge swaps per improvement attempt
        improve_tries=0                        # Improvement attempts per build step
    )

    # Queue a single vector [D] or batch [N, D] for addition
    add_entry(external_label, feature)

    # Queue an external label for deletion
    remove_entry(external_label)

    # Queue status inspection
    get_num_new_entries() -> int               # Count of pending additions
    get_num_remove_entries() -> int            # Count of pending deletions

    # Multithreading configuration
    set_thread_count(thread_count)             # Set worker thread count (0 = all CPU cores)
    set_batch_size(tasks_per_batch, task_size) # Configure batch granularity
    get_batch_size() -> int                    # Current total batch size (threads * tasks * task_size)

    # Execute build/optimization loop.
    # callback: function(BuilderStatus) for custom progress reporting
    # show_progress: If True and no callback, prints a console progress bar
    # infinite: If True, runs indefinitely until stop() is called from another thread
    build(callback=None, show_progress=False, infinite=False) -> BuilderStatus

    # Request graceful termination of the build loop between steps
    stop()

from deglib.distances import FloatSpace, Metric, floats_to_fp16, fp16_to_floats

# Supported distance metrics
class Metric:
    FP32_L2            # Single-precision Euclidean L2 distance (dtype: np.float32)
    FP32_InnerProduct  # Single-precision Inner Product (1 - <a,b>, dtype: np.float32)
    Uint8_L2           # 8-bit unsigned integer L2 distance (dtype: np.uint8)
    Uint8_InnerProduct # 8-bit unsigned integer Inner Product (dtype: np.uint8)
    FP16_L2            # Half-precision FP16 Euclidean L2 distance (dtype: np.uint16)
    FP16_InnerProduct  # Half-precision FP16 Inner Product (dtype: np.uint16)
    EVP_InnerProduct   # Extreme Vector Quantization binary/ternary IP (dtype: np.uint8)
    Int8_InnerProduct  # 8-bit signed integer Inner Product (dtype: np.int8)
    Int8_L2            # 8-bit signed integer L2 distance (dtype: np.int8)

    get_dtype() -> np.dtype # Expected numpy data type for feature vectors


class FloatSpace:
    # Create feature space with dimension, metric, and optional SIMD instruction set
    create(dim, metric, instruction=InstructionSet.Auto) -> FloatSpace

    dim() -> int                               # Dimensionality
    metric() -> Metric                         # Configured metric enum
    get_data_size() -> int                     # Byte size of a single vector
    get_instruction() -> InstructionSet        # Active SIMD instruction set ("AVX512", "AVX2", "Scalar")

    # Compute distance between two individual 1D feature vectors
    compute_distance(vec1, vec2) -> float

    # Compute 1D array of distances between query [D] and a batch of targets [N, D]
    compute_distances(query, targets) -> np.ndarray


# --- Type Conversion ---

# Convert float32 numpy array to uint16-packed FP16 representation
floats_to_fp16(floats) -> np.ndarray

# Convert uint16-packed FP16 array back to float32
fp16_to_floats(fp16_vals) -> np.ndarray

from deglib.search import Filter, rerank, Reranker

class Filter:
    def __init__(valid_labels, max_value=-1, max_label_count=-1)
    """
    Creates a search label filter restricting results to allowed IDs.

    valid_labels: 1D int32 numpy array containing all allowed external labels.
    max_value: Maximum label value (computed automatically if -1).
    max_label_count: Total dataset size (defaults to graph size).
    """

    # Static factory supporting None, numpy array, or existing Filter
    create_filter(filter_labels, graph_size) -> cpp_filter


# Multi-threaded exact distance candidate reranking in C++.
# space: FloatSpace used to compute distances (metric + SIMD instruction set)
# queries: 2D array [Q, D]
# candidate_indices: 2D uint32 array [Q, K_cand]
# base_vectors: 2D array [N, D] (defaults to queries if None)
# k_top: Number of nearest candidates to return per query (0 = all)
# num_threads: Number of threads (0 = all CPU cores)
# return_distances: If True, returns (indices, distances) tuple
# unsorted: If True, skips sorting the resulting candidates
# Returns 2D uint32 array [Q, k_top] of candidate IDs, or (indices, distances) if return_distances is True.
rerank(space, queries, candidate_indices, base_vectors=None, k_top=0, num_threads=0, return_distances=False, unsorted=False)


# Stateful reranker holding base vectors for repeated exact-distance reranking.
class Reranker:
    # Constructor: Reranker(space, base_vectors)
    # space: FloatSpace (metric + SIMD instruction set)
    # base_vectors: 2D array [N, D] the candidates are scored against

    # Auto-tune the reranker's prefetch parameters (po, pl) by timing reranking over sampled base
    # vectors. Self-samples its own base vectors; call once after construction to tune for the
    # expected operating point.
    # sample_count: base vectors used as sample queries (capped at the base count)
    # num_candidates: candidates scored per query (clamped to the base count)
    # k: candidates kept per query after reranking (clamped to num_candidates)
    optimize(sample_count=50, num_candidates=100, k=10)

    # Read the active prefetch knobs: po is the lookahead offset in candidates (0 disables
    # prefetching), pl the cache lines per candidate (0 prefetches the whole vector).
    get_po() -> int
    get_pl() -> int

    # Override the prefetch knobs without auto-tuning. Prefetching is on by default (po=8, pl=0),
    # which is the right starting point whenever the base vectors do not fit in cache.
    # po/pl values <= 0 disable the respective knob.
    set_prefetch(po, pl)

    # Re-evaluate and rank candidate indices per query against the base vectors.
    # Returns 2D uint32 array [Q, k_top], or (indices, distances) if return_distances is True.
    rerank(queries, candidate_indices, k_top=0, num_threads=0, return_distances=False, unsorted=False) -> np.ndarray | tuple[np.ndarray, np.ndarray]

    # Number of base vectors the reranker scores candidates against.
    get_num_base_vectors() -> int


# High-performance searcher with optional on-the-fly quantization and exact reranking.
class Searcher:
    # Constructor via create_searcher(graph, quantizer=None, refine_space=None, refine_data=None)

    # Entry vertex optimization via k-means medoids
    optimize(n_clusters=256, n_iter=20, sample_size=0, k=100, seed=7, num_threads=1)

    # Search for nearest neighbors of a single query (1D) or a batch of queries (2D).
    # query: 1D vector [D] or 2D batch [N, D]
    # k: Number of nearest neighbors to return
    # eps_or_ef: Exploration parameter. Values >= 1.0 are treated as fixed pool size (ef),
    #            values < 1.0 are treated as relative distance margin (eps). Default: 0.1.
    # rerank_factor: Candidate expansion factor for exact reranking (retrieves ceil(k * rerank_factor) candidates)
    # threads: Number of worker threads for batch queries
    # return_distances: If True, returns (indices, distances) tuple
    # unsorted: If True, returns candidates in heap order instead of ascending distance order
    search(query, k, eps_or_ef=0.1, rerank_factor=1.0, threads=1, return_distances=False, unsorted=False) -> np.ndarray | tuple[np.ndarray, np.ndarray]

create_searcher(graph, quantizer=None, refine_space=None, refine_data=None) -> Searcher

from deglib.optimization import (
    prune_non_rng_edges,
    prune_worst_edges,
    presort,
    mips_l2_transform,
    mips_l2_transform_query,
    EvpQuantizer,
    ScalarQuantizerInt8,
    ScalarQuantizerInt8PerDim,
    ScalarQuantizerUint8,
    ScalarQuantizerUint8PerDim,
    make_evp_quantizer,
    make_scalar_quantizer_int8,
    make_scalar_quantizer_int8_perdim,
    make_scalar_quantizer_uint8,
    make_scalar_quantizer_uint8_perdim,
)

# 1. State-aware Scalar Quantization (All classes support float32 & float16)
# INT8 Quantizer [-127, 127] (Cosine / InnerProduct)
quantizer = make_scalar_quantizer_int8(base_vectors, drop_ratio=0.0)
quant_base = quantizer.quantize(base_vectors)
quant_query = quantizer.quantize(query_vectors)

# INT8 Per-Dimension Quantizer [-127, 127]
int8_pdim = make_scalar_quantizer_int8_perdim(base_vectors)
quant_int8_pdim = int8_pdim.quantize(base_vectors)

# UINT8 Quantizer [0, 255] (Euclidean L2)
l2_quantizer = make_scalar_quantizer_uint8(base_vectors)
quant_l2_base = l2_quantizer.quantize(base_vectors)

# UINT8 Per-Dimension Quantizer [0, 255]
perdim_quantizer = make_scalar_quantizer_uint8_perdim(base_vectors)
quant_perdim_base = perdim_quantizer.quantize(base_vectors)

# 2. EVP Quantization (float32 or float16 vectors to byte-packed EVP format)
# One shared quantizer for database and query vectors guarantees matching non_zeros.
quantizer = EvpQuantizer(non_zeros)  # or make_evp_quantizer(non_zeros)
quant_base = quantizer.quantize(base_vectors, num_threads=0)
quant_query = quantizer.quantize(query_vectors, num_threads=0)

# Remove all edges violating the Relative Neighborhood Graph (RNG) rule. Returns removed edge count.
prune_non_rng_edges(graph, num_threads=0) -> int
# Prune the worst (longest/highest-weight) neighbors per vertex by replacing them with self-loops
prune_worst_edges(graph, prune_worst, num_threads=0)

# 1D dataset pre-sorting using Fast Linear Assignment Sorter (FLAS).
# vectors: 2D float32 array [N, D]
# space_or_metric: FloatSpace or Metric enum (default: Metric.FP32_L2)
# radius_decay: Neighborhood decay factor (default: 0.9)
# threads: Worker threads (0 = all CPU cores)
# callback: function(progress_float) for custom progress reporting
# show_progress: If True and no callback, prints progress to stdout
# Returns 1D uint32 permutation array [N] of sorted vector indices.
presort(vectors, space_or_metric=None, radius_decay=0.9, threads=0, callback=None, *, show_progress=False) -> np.ndarray

# Transform database vectors from d dimensions to (d+1) dimensions for MIPS via L2 distance.
# database: 2D float32 array [N, d]
# Returns tuple (transformed_database [N, d+1], max_norm float M).
mips_l2_transform(database) -> tuple[np.ndarray, float]

# Pad query vectors with 0 at dimension (d+1) for MIPS queries against an L2-transformed database.
# queries: 1D [d] or 2D [Q, d] float32 array
# Returns transformed array of shape [d+1] or [Q, d+1].
mips_l2_transform_query(queries) -> np.ndarray

from deglib.analysis import (
    analyze_graph,
    check_graph_regularity,
    check_graph_weights,
    check_graph_connectivity,
    calc_avg_edge_weight,
    calc_edge_weight_histogram,
    calc_non_rng_edges,
    calc_search_reachability,
    calc_exploration_reach,
)

# Comprehensive graph analysis returning a dictionary of topology and health metrics
analyze_graph(graph) -> dict
# Returns dictionary:
# {
#     "vertex_count": int,              # Total active vertices
#     "edge_count": int,                # Total valid edges
#     "feature_dims": int,              # Vector dimensionality
#     "edges_per_vertex": int,          # Outgoing edge capacity per vertex
#     "avg_out_degree": float,          # Average out-degree
#     "min_out_degree": int,            # Minimum out-degree
#     "max_out_degree": int,            # Maximum out-degree
#     "avg_in_degree": float,           # Average in-degree
#     "min_in_degree": int,             # Minimum in-degree
#     "max_in_degree": int,             # Maximum in-degree
#     "source_vertices": int,           # Vertices with in-degree = 0 (unreachable as targets)
#     "search_reachability": float,     # Ratio of vertices reachable from entry points (0.0 to 1.0)
#     "exploration_reachability": float,# Average exploration reachability across all vertices (0.0 to 1.0)
#     "memory_bytes": int               # Estimated memory footprint in bytes
# }

# Verify that vertex count, ascending edge ordering, and neighbor uniqueness are valid
check_graph_regularity(graph, expected_vertices, check_back_link=False) -> bool

# Verify that stored edge weights match exact feature distance calculations
check_graph_weights(graph) -> bool

# Verify that the graph consists of a single connected component
check_graph_connectivity(graph) -> bool

# Calculate average edge weight across all graph edges
calc_avg_edge_weight(graph, scale=1) -> float

# Calculate 10-bin histogram of graph edge weights
calc_edge_weight_histogram(graph, sort, scale=1) -> list[float]

# Count total non-RNG conform edges
calc_non_rng_edges(graph) -> int

# Fraction of vertices reachable from graph entry points via BFS (0.0 to 1.0)
calc_search_reachability(graph) -> float

# Average exploration reachability across all starting vertices (0.0 to 1.0)
calc_exploration_reach(graph) -> float

# Average graph quality against a precomputed base_top groundtruth (0.0 to 1.0)
calc_graph_quality(graph, base_top, sample_size=0, num_threads=0) -> float

from deglib.cpu import InstructionSet, has_avx2, has_avx_vnni, has_avx512, has_avx512_vnni

class InstructionSet:
    Auto         # Automatically select highest supported instruction set
    Scalar       # Standard scalar operations
    AVX2         # 256-bit AVX2 + FMA SIMD
    AVX2_VNNI    # 256-bit AVX2 with VNNI acceleration
    AVX512       # 512-bit AVX-512 SIMD
    AVX512_VNNI  # 512-bit AVX-512 with VNNI acceleration

# Runtime hardware support checks
has_avx2() -> bool         # Returns True if host CPU supports AVX2
has_avx_vnni() -> bool     # Returns True if host CPU supports AVX-VNNI (256-bit)
has_avx512() -> bool       # Returns True if host CPU supports AVX-512
has_avx512_vnni() -> bool  # Returns True if host CPU supports AVX-512 VNNI