## Copyright (C) 2025 Swayam Shah <swayamshah66@gmail.com>
## Copyright (C) 2025 Andreas Bertsatos <abertsatos@biol.uoa.gr>
##
## This file is part of the statistics package for GNU Octave.
##
## This program is free software; you can redistribute it and/or modify it under
## the terms of the GNU General Public License as published by the Free Software
## Foundation; either version 3 of the License, or (at your option) any later
## version.
##
## This program is distributed in the hope that it will be useful, but WITHOUT
## ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
## FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
## details.
##
## You should have received a copy of the GNU General Public License along with
## this program; if not, see <http://www.gnu.org/licenses/>.

classdef hnswSearcher
  ## -*- texinfo -*-
  ## @deftp {statistics} hnswSearcher
  ##
  ## Hierarchical Navigable Small World (HNSW) nearest neighbor searcher class.
  ##
  ## The @code{hnswSearcher} class implements the HNSW algorithm for efficient
  ## nearest neighbor queries.  It stores training data and supports various
  ## distance metrics for performing searches.  The HNSW algorithm builds a
  ## multilayer graph structure that enables fast approximate nearest neighbor
  ## searches by navigating through the graph.  It facilitates nearest neighbor
  ## queries search using @code{knnsearch}.
  ##
  ## You can either use the @code{hnswSearcher} class constructor or the
  ## @code{createns} function to create an @qcode{hnswSearcher} object.
  ##
  ## @seealso{createns, ExhaustiveSearcher, KDTreeSearcher, knnsearch}
  ## @end deftp

  properties(SetAccess = private, Hidden)
    HNSWGraph # HNSW graph structure
  endproperties

  properties(SetAccess = private)
    ## -*- texinfo -*-
    ## @deftp {hnswSearcher} {property} Distance
    ##
    ## Distance metric
    ##
    ## Distance metric used for searches, specified as a character vector (e.g.,
    ## @qcode{'euclidean'}, @qcode{'minkowski'}, @qcode{'cityblock'}). Default
    ## is @qcode{'euclidean'}. Supported metrics align with those in
    ## @code{pdist2}.  This property is private and cannot be modified after
    ## object creation.
    ##
    ## @end deftp
    Distance = 'euclidean'

    ## -*- texinfo -*-
    ## @deftp {hnswSearcher} {property} DistParameter
    ##
    ## Distance parameter
    ##
    ## The type and value of the distance parameter depends on the selected
    ## @qcode{Distance} metric and can be any of the following:
    ##
    ## @itemize
    ## @item For @qcode{'minkowski'}, a positive scalar exponent (default 2).
    ## @item For @qcode{'seuclidean'}, a nonnegative vector of scaling factors
    ## matching the number of columns in @qcode{X} (default is standard
    ## deviation of @qcode{X}).
    ## @item For @qcode{'mahalanobis'}, a positive definite covariance matrix
    ## matching the dimensions of @qcode{X} (default is @code{cov (@var{X})}).
    ## @item Empty for other metrics.
    ## @end itemize
    ##
    ## This property is private and cannot be modified after object creation.
    ##
    ## @end deftp
    DistParameter = []

    ## -*- texinfo -*-
    ## @deftp {hnswSearcher} {property} MaxNumLinksPerNode
    ##
    ## Number of connections created for each node
    ##
    ## Maximum number of neighbors per node in the HNSW graph. Affects graph
    ## connectivity and search accuracy.  Default value is 16.  This property is
    ## private and cannot be modified after object creation.
    ##
    ## @end deftp
    MaxNumLinksPerNode = 16

    ## -*- texinfo -*-
    ## @deftp {hnswSearcher} {property} TrainSetSize
    ##
    ## Number of potential nearest neighbors
    ##
    ## Size of the dynamic candidate list during graph construction.  Higher
    ## values improve accuracy at the cost of construction time.  Default value
    ## is 200.  This property is private and cannot be modified after object
    ## creation.
    ##
    ## @end deftp
    TrainSetSize = 200

    ## -*- texinfo -*-
    ## @deftp {hnswSearcher} {property} X
    ##
    ## Point data
    ##
    ## Point data, specified as an @math{N*P} numeric matrix where each row is
    ## an observation and each column is a feature.  This property is private
    ## and cannot be modified after object creation.
    ##
    ## Data of class @qcode{single} is stored and searched in single
    ## precision, any other numeric class is converted to @qcode{double}.
    ##
    ## @end deftp
    X = []
  endproperties

  methods (Hidden)

    ## Custom display
    function display (this)
      in_name = inputname (1);
      if (! isempty (in_name))
        fprintf ('%s =\n', in_name);
      endif
      disp (this);
    endfunction

    ## Custom display
    function disp (this)
      if (isscalar (this))
        fprintf ("\n  hnswSearcher with properties:\n\n");
        fprintf ("%+25s: %d\n", 'MaxNumLinksPerNode', this.MaxNumLinksPerNode);
        fprintf ("%+25s: %d\n", 'TrainSetSize', this.TrainSetSize);
        fprintf ("%+25s: '%s'\n", 'Distance', this.Distance);
        if (! isempty (this.DistParameter))
          if (isscalar (this.DistParameter))
            fprintf ("%+25s: %g\n", 'DistParameter', this.DistParameter);
          elseif (isvector (this.DistParameter))
            fprintf ("%+25s: %s\n", 'DistParameter', ...
                     mat2str (this.DistParameter));
          else
            fprintf ("%+25s: [%dx%d %s]\n", 'DistParameter', ...
                     size (this.DistParameter), class (this.DistParameter));
          endif
        else
          fprintf ("%+25s: []\n", 'DistParameter');
        endif
        fprintf ("%+25s: [%dx%d %s]\n", 'X', size (this.X), class (this.X));
      endif
    endfunction

    ## Class specific subscripted reference
    function varargout = subsref (this, s)
      chain_s = s(2:end);
      s = s(1);
      switch (s.type)
        case '()'
          error ("hnswSearcher.subsref: () indexing not supported.");
        case '{}'
          error ("hnswSearcher.subsref: {} indexing not supported.");
        case '.'
          if (! ischar (s.subs))
            error (strcat ("hnswSearcher.subsref: property", ...
                           " name must be a character vector."));
          endif
          try
            out = this.(s.subs);
          catch
            error (strcat ("hnswSearcher.subsref: unrecognized", ...
                           " property: '%s'."), s.subs);
          end_try_catch
      endswitch
      ## Chained references
      if (! isempty (chain_s))
        out = subsref (out, chain_s);
      endif
      varargout{1} = out;
    endfunction

    ## Class specific subscripted assignment
    function this = subsasgn (this, s, val)
      switch s.type
        case '()'
          error ("hnswSearcher.subsasgn: () indexing not supported.");
        case '{}'
          error ("hnswSearcher.subsasgn: {} indexing not supported.");
        case '.'
          if (! ischar (s.subs))
            error (strcat ("hnswSearcher.subsasgn: property", ...
                           " name must be a character vector."));
          endif
          switch (s.subs)
            case 'X'
              error (strcat ("hnswSearcher.subsasgn: 'X' is", ...
                             " read-only and cannot be modified."));
            case 'HNSWGraph'
              error (strcat ("hnswSearcher.subsasgn: 'HNSWGraph' is", ...
                             " read-only and cannot be modified."));
            case 'Distance'
              error (strcat ("hnswSearcher.subsasgn: 'Distance' is", ...
                             " read-only and cannot be modified."));
            case 'DistParameter'
              error (strcat ("hnswSearcher.subsasgn: 'DistParameter'", ...
                             " is read-only and cannot be modified."));
            case 'MaxNumLinksPerNode'
              error (strcat ("hnswSearcher.subsasgn: 'MaxNumLinksPerNode'", ...
                             " is read-only and cannot be modified."));
            case 'TrainSetSize'
              error (strcat ("hnswSearcher.subsasgn: 'TrainSetSize'", ...
                             " is read-only and cannot be modified."));
            otherwise
              error (strcat ("hnswSearcher.subsasgn:", ...
                             " unrecognized property: '%s'."), s.subs);
          endswitch
      endswitch
    endfunction

  endmethods

  methods

    ## -*- texinfo -*-
    ## @deftypefn  {hnswSearcher} {@var{obj} =} hnswSearcher (@var{X})
    ## @deftypefnx {hnswSearcher} {@var{obj} =} hnswSearcher (@var{X}, @var{name}, @var{value})
    ##
    ## Create an @qcode{hnswSearcher} object for approximate nearest neighbor
    ## searches.
    ##
    ## @code{@var{obj} = hnswSearcher (@var{X})} constructs an
    ## @qcode{hnswSearcher} object with training data @var{X} using the
    ## default @qcode{'euclidean'} distance metric. @var{X} must be an
    ## @math{N*P} numeric matrix, where rows represent observations and columns
    ## represent features.
    ##
    ## @code{@var{obj} = hnswSearcher (@var{X}, @var{name}, @var{value})}
    ## allows customization through name-value pairs:
    ##
    ## @multitable @columnfractions 0.18 0.8
    ## @headitem @var{Name} @tab @var{Value}
    ##
    ## @item @qcode{'Distance'} @tab Distance metric, specified as a
    ## character vector (e.g., @qcode{'euclidean'}, @qcode{'minkowski'},
    ## @qcode{'cityblock'}). Default is @qcode{'euclidean'}. See @code{pdist2}
    ## for supported metrics.
    ##
    ## @item @qcode{'P'} @tab Minkowski distance exponent, a positive
    ## scalar. Valid only when @qcode{'Distance'} is @qcode{'minkowski'}.
    ## Default is 2.
    ##
    ## @item @qcode{'Scale'} @tab Nonnegative vector of scaling factors
    ## matching the number of columns in @var{X}. Valid only when
    ## @qcode{'Distance'} is @qcode{'seuclidean'}. Default is @code{std (X)}.
    ##
    ## @item @qcode{'Cov'} @tab Positive definite covariance matrix
    ## matching the number of columns in @var{X}. Valid only when
    ## @qcode{'Distance'} is @qcode{'mahalanobis'}. Default is @code{cov (X)}.
    ##
    ## @item @qcode{'MaxNumLinksPerNode'} @tab Maximum number of neighbors
    ## per node in the HNSW graph, a positive integer. Default is 16.
    ##
    ## @item @qcode{'TrainSetSize'} @tab Size of the dynamic candidate
    ## list during graph construction, a positive integer. Default is 200.
    ## @end multitable
    ##
    ## @seealso{hnswSearcher, knnsearch, createns, pdist2}
    ## @end deftypefn
    function obj = hnswSearcher (X, varargin)
      ## Initial input validation
      if (nargin < 1)
        error ("hnswSearcher: too few input arguments.");
      endif
      if (mod (numel (varargin), 2) != 0)
        error ("hnswSearcher: Name-Value arguments must be in pairs.");
      endif

      ## Validate X
      if (isempty (X))
        error ("hnswSearcher: X cannot be empty.");
      endif
      if (! (isnumeric (X) && ismatrix (X) && all (isfinite (X(:)))))
        error ("hnswSearcher: X must be a finite numeric matrix.");
      endif

      ## Single precision is carried through, but every other class is
      ## converted up to double, as MATLAB does.  Integer data would
      ## otherwise round each coordinate difference and corrupt distances.
      if (! isa (X, "single"))
        X = double (X);
      endif

      obj.X = X;
      N = size (X, 1);

      ## Default values
      Distance = 'euclidean';
      P = [];
      S = [];
      C = [];
      MaxNumLinksPerNode = min (16, N);
      TrainSetSize = min (200, N);
      TrainSetSize = max (TrainSetSize, MaxNumLinksPerNode);

      ## Parse optional parameters
      while (numel (varargin) > 0)
        switch (lower (varargin{1}))
          case 'distance'
            Distance = varargin{2};
          case 'p'
            P = varargin{2};
          case 'scale'
            S = varargin{2};
          case 'cov'
            C = varargin{2};
          case 'maxnumlinkspernode'
            MaxNumLinksPerNode = varargin{2};
          case 'trainsetsize'
            TrainSetSize = varargin{2};
          otherwise
            error ("hnswSearcher: invalid parameter name: '%s'.", varargin{1});
        endswitch
        varargin(1:2) = [];
      endwhile

      ## Validate Distance
      valid_metrics = {'euclidean', 'minkowski', 'seuclidean', ...
                       'mahalanobis', 'cityblock', 'manhattan', ...
                       'chebychev', 'cosine', 'correlation', ...
                       'spearman', 'hamming', 'jaccard'};
      if (ischar (Distance))
        if (! any (strcmpi (valid_metrics, Distance)))
          error ("hnswSearcher: unsupported distance metric '%s'.", Distance);
        endif
        obj.Distance = Distance;
      else
        error ("hnswSearcher: 'Distance' must be a string.");
      endif

      ## Set DistParameter
      if (strcmpi (obj.Distance, 'minkowski'))
        if (isempty (P))
          obj.DistParameter = 2;
        else
          if (! (isscalar (P) && isnumeric (P) && P > 0 && isfinite (P)))
            error ("hnswSearcher: 'P' must be a positive finite scalar.");
          endif
          obj.DistParameter = P;
        endif
      elseif (strcmpi (obj.Distance, 'seuclidean'))
        if (isempty (S))
          obj.DistParameter = std (X, [], 1);
        else
          if (! (isvector (S) && isnumeric (S) && all (S >= 0)
                                               && all (isfinite (S))
                                               && length (S) == columns (X)))
            error (strcat ("hnswSearcher: 'Scale' must be a", ...
                           " nonnegative vector matching X columns."));
          endif
          obj.DistParameter = S;
        endif
      elseif (strcmpi (obj.Distance, 'mahalanobis'))
        if (isempty (C))
          obj.DistParameter = cov (X);
        else
          if (! (ismatrix (C) && isnumeric (C) && all (isfinite (C)(:))
                                               && rows (C) == columns (C)
                                               && rows (C) == columns (X)))
            error (strcat ("hnswSearcher: 'Cov' must be a square", ...
                           " matrix matching X columns."));
          endif
          if (! issymmetric (C))
            error (strcat ("hnswSearcher: 'Cov' must be symmetric", ...
                           " for mahalanobis."));
          endif
          [~, p] = chol (C);
          if (p != 0)
            error (strcat ("hnswSearcher: 'Cov' must be positive", ...
                           " definite for mahalanobis."));
          endif
          obj.DistParameter = C;
        endif
      else
        obj.DistParameter = [];
      endif

      ## Validate MaxNumLinksPerNode and TrainSetSize
      if (! (isscalar (MaxNumLinksPerNode) &&
             isnumeric (MaxNumLinksPerNode) &&
             MaxNumLinksPerNode > 0 &&
             MaxNumLinksPerNode == fix (MaxNumLinksPerNode)))
        error (strcat ("hnswSearcher: 'MaxNumLinksPerNode'", ...
                       " must be a positive integer."));
      endif
      if (! (isscalar (TrainSetSize) &&
             isnumeric (TrainSetSize) &&
             TrainSetSize > 0 &&
             TrainSetSize == fix (TrainSetSize)))
        error ("hnswSearcher: 'TrainSetSize' must be a positive integer.");
      endif
      if (TrainSetSize > N)
        error (strcat ("hnswSearcher: 'TrainSetSize' cannot", ...
                       " exceed the number of rows in X."));
      endif
      if (MaxNumLinksPerNode > TrainSetSize)
        error (strcat ("hnswSearcher: 'MaxNumLinksPerNode'", ...
                       " cannot exceed 'TrainSetSize'."));
      endif
      obj.MaxNumLinksPerNode = MaxNumLinksPerNode;
      obj.TrainSetSize = TrainSetSize;

      ## Build HNSW graph
      obj.HNSWGraph = build_hnsw (X, obj.Distance, obj.DistParameter, ...
                                  MaxNumLinksPerNode, TrainSetSize);
    endfunction

    ## -*- texinfo -*-
    ## @deftypefn  {hnswSearcher} {[@var{idx}, @var{D}] =} knnsearch (@var{obj}, @var{Y})
    ## @deftypefnx {hnswSearcher} {[@var{idx}, @var{D}] =} knnsearch (@var{obj}, @var{Y}, @var{name}, @var{value})
    ##
    ## Find the nearest neighbors in the training data to query points.
    ##
    ## @code{[@var{idx}, @var{D}] = knnsearch (@var{obj}, @var{Y})} returns the
    ## indices @var{idx} and distances @var{D} of the nearest neighbor in
    ## @var{obj.X} to each point in @var{Y}, using the distance metric specified
    ## in @var{obj.Distance}.
    ##
    ## @itemize
    ## @item @var{obj} is an @qcode{hnswSearcher} object.
    ## @item @var{Y} is an @math{M*P} numeric matrix of query points, where
    ## @math{P} must match the number of columns in @var{obj.X}.
    ## @item @var{idx} contains the indices of the nearest neighbors in
    ## @var{obj.X}.
    ## @item @var{D} contains the corresponding distances.
    ## @end itemize
    ##
    ## @var{idx} is always of class @qcode{double}.  @var{D} is of class
    ## @qcode{single} when either @var{obj.X} or @var{Y} is @qcode{single},
    ## in which case the distances are computed in single precision, and of
    ## class @qcode{double} otherwise.
    ##
    ## @code{[@var{idx}, @var{D}] = knnsearch (@var{obj}, @var{Y}, @var{name},
    ## @var{value})}
    ## allows additional options via name-value pairs:
    ##
    ## @multitable @columnfractions 0.18 0.8
    ## @headitem @var{Name} @tab @var{Value}
    ##
    ## @item @qcode{'K'} @tab A positive integer specifying the number of
    ## nearest neighbors to find. Default is 1.  A value larger than the
    ## number of observations in the training data is answered with all of
    ## them, since there are no more neighbors to return.
    ##
    ## @item @qcode{'SearchSetSize'} @tab A positive integer specifying the
    ## size of the candidate list of nearest neighbors for a single query point
    ## during the search process.  Default is @qcode{max (10, @var{C})}, where
    ## @var{C} is the number of columns in @var{obj.X}.  @qcode{'SearchSetSize'}
    ## must be at least @var{C} and no more than the number of rows in training
    ## data @var{obj.X}.
    ## @end multitable
    ##
    ## @seealso{hnswSearcher, pdist2}
    ## @end deftypefn
    function [idx, D] = knnsearch (obj, Y, varargin)
      ## Initial input validation
      if (nargin < 2)
        error ("hnswSearcher.knnsearch: too few input arguments.");
      endif
      if (mod (numel (varargin), 2) != 0)
        error (strcat ("hnswSearcher.knnsearch:", ...
                       " Name-Value arguments must be in pairs."));
      endif

      ## Get training data size
      [N, C] = size (obj.X);

      ## Validate Y.  A non-finite query is answered rather than refused: its
      ## distances come back Inf or NaN, which is the honest result and what
      ## MATLAB returns.  An empty query is answered with an empty result.
      if (! (isnumeric (Y) && isreal (Y) && ismatrix (Y)))
        error ("hnswSearcher.knnsearch: Y must be a real numeric matrix.");
      endif
      if (C != size (Y, 2))
        error (strcat ("hnswSearcher.knnsearch: Y must have the same", ...
                       " number of columns as the training data in OBJ.X."));
      endif

      ## Default values
      K = 1;
      SearchSetSize = max (10, C);

      ## Parse options
      while (numel (varargin) > 0)
        switch (lower (varargin{1}))
          case 'k'
            K = varargin{2};
            if (! (isscalar (K) && isnumeric (K) &&
                   K >= 1 && K == fix (K) && isfinite (K)))
              error ("hnswSearcher.knnsearch: 'K' must be a positive integer.");
            endif
          case 'searchsetsize'
            SearchSetSize = varargin{2};
            if (! (isscalar (SearchSetSize) &&
                   isnumeric (SearchSetSize) &&
                   SearchSetSize >= 1 &&
                   SearchSetSize == fix (SearchSetSize) &&
                   isfinite (SearchSetSize)))
              error (strcat ("hnswSearcher.knnsearch: 'SearchSetSize' must", ...
                             " be a positive integer."));
            endif
            if (SearchSetSize < C || SearchSetSize > N)
              error (strcat ("hnswSearcher.knnsearch: 'SearchSetSize' must", ...
                             " be at least the number of features but no", ...
                             " more than the sample size of the training data."));
            endif
          otherwise
            error (strcat ("hnswSearcher.knnsearch: invalid", ...
                           " parameter name: '%s'."), varargin{1});
        endswitch
        varargin(1:2) = [];
      endwhile

      ## More neighbours cannot be returned than there are points to return,
      ## so asking for more is answered with all of them rather than with a
      ## matrix wider than the training data padded out with NaN.
      K = min (K, N);

      ## The search only ever inspects SearchSetSize candidates, so it cannot
      ## produce K neighbours while that is the smaller of the two: asking for
      ## every point used to come back mostly NaN.
      SearchSetSize = min (max (SearchSetSize, K), N);

      ## Distances are computed in single precision when either the training
      ## data or the query is single, and in double otherwise.  The indices
      ## are always double.
      cls = "double";
      if (isa (obj.X, "single") || isa (Y, "single"))
        cls = "single";
      endif
      Y = cast (Y, cls);

      if (isempty (Y))
        idx = zeros (rows (Y), K);
        D = zeros (rows (Y), K, cls);
        return;
      endif

      ## Search HNSW graph
      idx = cell (rows (Y), 1);
      D = cell (rows (Y), 1);
      for i = 1:rows (Y)
        [temp_idx, temp_D] = search_hnsw (obj.HNSWGraph, Y(i,:), obj.X, ...
                                          obj.Distance, obj.DistParameter, ...
                                          K, SearchSetSize);
        [sorted_D, sort_idx] = sort (temp_D);
        idx{i} = temp_idx(sort_idx);
        D{i} = sorted_D;
      endfor

      idx_mat = NaN (rows (Y), K);
      D_mat = NaN (rows (Y), K, cls);
      for i = 1:rows (Y)
        len = min (K, length (idx{i}));
        idx_mat(i, 1:len) = idx{i}(1:len);
        D_mat(i, 1:len) = D{i}(1:len);
      endfor
      idx = idx_mat;
      D = D_mat;
    endfunction

  endmethods

endclassdef

## Private Function to Build HNSW graph
function graph = build_hnsw (X, dist, param, MaxNumLinksPerNode, TrainSetSize)
  N = size (X, 1);
  if (N < 1)
    error ("build_hnsw: X must have at least one point.");
  endif
  max_layers = floor (log2 (max (N, 2))) + 1;
  graph.layers = cell (max_layers, 1);
  graph.entry_point = 1;
  mL = 1 / log (MaxNumLinksPerNode);

  ## Initialize graph with empty adjacency lists
  for L = 1:max_layers
    graph.layers{L} = cell (N, 1);
  endfor

  ## Handle single-point case
  if (N == 1)
    graph.layers{1}{1} = [];
    return;
  endif

  ## Add points to graph
  for i = 1:N
    layer = min (max_layers - 1, floor (-log (rand ()) * mL));
    for L = 0:layer
      if (i == 1)
        graph.layers{L+1}{i} = [];
        continue;
      endif
      ## Find nearest neighbors in current layer
      [neighbors, dists] = search_hnsw_layer (graph, X(i,:), X, dist, param, ...
                                              MaxNumLinksPerNode, ...
                                              TrainSetSize, L, ...
                                              graph.entry_point);
      graph.layers{L+1}{i} = neighbors;
      ## Update neighbors' connections
      for j = neighbors
        if (length (graph.layers{L+1}{j}) < MaxNumLinksPerNode)
          graph.layers{L+1}{j} = [graph.layers{L+1}{j}, i];
        else
          ## Select MaxNumLinksPerNode closest neighbors
          all_neighbors = [graph.layers{L+1}{j}, i];
          dists_j = pdist2 (X(all_neighbors,:), X(j,:), dist, param);
          [~, sort_idx] = sort (dists_j);
          graph.layers{L+1}{j} = all_neighbors(sort_idx(1:MaxNumLinksPerNode));
        endif
      endfor
      if (L == layer && i > 1)
        graph.entry_point = i;
      endif
    endfor
  endfor
endfunction

## Private Function to Search HNSW graph for k nearest neighbors
function [indices, distances] = search_hnsw (graph, Y, X, dist, param, ...
                                             K, SearchSetSize)
  max_layers = length (graph.layers);
  current_entry = graph.entry_point;

  ## Navigate to the lowest layer (from the top layer down to layer 1)
  for L = max_layers:-1:2
    [new_candidates, new_dists] = search_hnsw_layer (graph, Y, X, ...
                                                     dist, param, ...
                                                     1, SearchSetSize, L-1, ...
                                                     current_entry);
    current_entry = new_candidates(1);
  endfor

  ## Search in the base layer
  [indices, distances] = search_hnsw_layer (graph, Y, X, dist, param, ...
                                            K, SearchSetSize, 0, current_entry);
endfunction

## Private Function Search a single HNSW layer (optimized)
function [indices, distances] = search_hnsw_layer (graph, Y, X, dist, param, ...
                                                   Points, SetSize, L, ...
                                                   entry_points)
  N = size (X, 1);
  visited = false (N, 1);

  ## Pre-allocate fixed-size arrays with generous capacity
  capacity = N;

  ## Candidate list (nodes to explore)
  cand_nodes = zeros (1, capacity);
  cand_dists = zeros (1, capacity);
  cand_count = 0;

  ## Best results list (top SetSize closest nodes, kept sorted)
  best_nodes = zeros (1, capacity);
  best_dists_arr = zeros (1, capacity);
  best_count = 0;

  ## Initialize with entry points
  entry_points = entry_points(:)';
  num_entries = length (entry_points);
  init_dists = pdist2 (X(entry_points,:), Y, dist, param);

  for i = 1:num_entries
    ep = entry_points(i);
    visited(ep) = true;
    cand_count = cand_count + 1;
    cand_nodes(cand_count) = ep;
    cand_dists(cand_count) = init_dists(i);
    best_count = best_count + 1;
    best_nodes(best_count) = ep;
    best_dists_arr(best_count) = init_dists(i);
  endfor

  ## Cache layer adjacency list reference
  layer_adj = graph.layers{L+1};

  ## Main search loop
  while (cand_count > 0)
    ## Find minimum distance candidate
    [~, idx] = min (cand_dists(1:cand_count));
    closest = cand_nodes(idx);

    ## Swap-and-pop removal (O(1) instead of O(n) array shift)
    cand_nodes(idx) = cand_nodes(cand_count);
    cand_dists(idx) = cand_dists(cand_count);
    cand_count = cand_count - 1;

    ## Early termination check
    if (best_count > 0 && cand_count > 0)
      if (max (best_dists_arr(1:best_count)) < min (cand_dists(1:cand_count)))
        break;
      endif
    endif

    ## Get neighbors of closest node
    neighbors = layer_adj{closest};
    if (isempty (neighbors))
      continue;
    endif

    ## Filter to unvisited neighbors only
    unvisited_mask = ! visited(neighbors);
    new_neighbors = neighbors(unvisited_mask);

    if (isempty (new_neighbors))
      continue;
    endif

    ## Mark all new neighbors as visited
    visited(new_neighbors) = true;

    ## Batch distance computation (single pdist2 call)
    new_dists = pdist2 (X(new_neighbors,:), Y, dist, param);

    ## Add each new neighbor to candidates and best lists
    num_new = length (new_neighbors);
    for i = 1:num_new
      n = new_neighbors(i);
      d = new_dists(i);

      ## Add to candidates
      cand_count = cand_count + 1;
      cand_nodes(cand_count) = n;
      cand_dists(cand_count) = d;

      ## Add to best list
      best_count = best_count + 1;
      best_nodes(best_count) = n;
      best_dists_arr(best_count) = d;
    endfor

    ## Trim best list once per batch instead of per element
    if (best_count > SetSize)
      [best_dists_arr(1:best_count), sort_idx] = sort (best_dists_arr(1:best_count));
      best_nodes(1:best_count) = best_nodes(sort_idx);
      best_count = SetSize;
    endif
  endwhile

  ## Return top Points results
  if (best_count > Points)
    [best_dists_arr(1:best_count), sort_idx] = sort (best_dists_arr(1:best_count));
    best_nodes(1:best_count) = best_nodes(sort_idx);
    indices = best_nodes(1:Points);
    distances = best_dists_arr(1:Points);
  else
    indices = best_nodes(1:best_count);
    distances = best_dists_arr(1:best_count);
  endif
endfunction

%!demo
%! ## Create an hnswSearcher with Euclidean distance
%! X = [1, 2; 3, 4; 5, 6];
%! obj = hnswSearcher (X);
%! ## Find the nearest neighbor to [2, 3]
%! Y = [2, 3];
%! [idx, D] = knnsearch (obj, Y, 'K', 1);
%! disp ('Nearest neighbor index:');
%! disp (idx);
%! disp ('Distance:');
%! disp (D);

%!demo
%! ## Create an hnswSearcher with Minkowski distance (P=3)
%! X = [0, 0; 1, 0; 2, 0];
%! obj = hnswSearcher (X, 'Distance', 'minkowski', 'P', 3);
%! ## Find the nearest neighbor to [1, 0]
%! Y = [1, 0];
%! [idx, D] = knnsearch (obj, Y, 'K', 1);
%! disp ('Nearest neighbor index:');
%! disp (idx);
%! disp ('Distance:');
%! disp (D);

## Test Cases

%!test
%! ## the graph is drawn from rand, so seed it: the search is approximate
%! rand ("seed", 42);
%! load fisheriris
%! X = meas;
%! obj = hnswSearcher (X, 'Distance', 'chebychev');
%! Y = X(30:35,:);
%! [idx, D] = knnsearch (obj, Y, 'K', 4);
%! ## Under chebychev these queries have several equidistant neighbours, so
%! ## which of them is listed first is not defined.  Assert the distances,
%! ## which are, and that each returned index really sits at the distance
%! ## reported for it.
%! assert_equal (D, [[0 0.1000 0.1000 0.2000]; [0 0.1000 0.1000 0.1000]; [0 0.2000 ...
%!              0.2000 0.2000]; [0 0.3000 0.3000 0.3000]; [0 0.2000 0.3000 ...
%!              0.3000]; [0 0.1000 0.1000 0.1000]], 5e-15)
%! for i = 1:rows (Y)
%!   assert_equal (numel (unique (idx(i,:))), 4);
%!   for j = 1:4
%!     assert_equal (max (abs (X(idx(i,j),:) - Y(i,:))), D(i,j), 5e-15);
%!   endfor
%! endfor

%!test
%! ## the graph is drawn from rand, so seed it: the search is approximate
%! rand ("seed", 42);
%! load fisheriris
%! X = meas;
%! C = cov (X);
%! obj = hnswSearcher (X, 'Distance', 'mahalanobis', 'Cov', C);
%! Y = X(120:125,:);
%! [idx, D] = knnsearch (obj, Y, 'K', 2);
%! assert_equal (idx(1, :), [120 82])
%! assert_equal (idx(4, :), [123 106])
%! assert_equal (idx(5, :), [124 127])
%! assert_equal (idx(6, :), [125 57])
%! assert_equal (D(1, :), [0 0.7734], 1e-4)
%! assert_equal (D(4, :), [0 0.8452], 1e-4)
%! assert_equal (D(5, :), [0 0.4152], 1e-4)
%! assert_equal (D(6, :), [0 0.7322], 1e-4)

%!test
%! ## Basic constructor with default Euclidean
%! X = [1, 2; 3, 4; 5, 6];
%! obj = hnswSearcher (X);
%! assert_equal (obj.X, X);
%! assert_equal (obj.Distance, "euclidean");
%! assert_equal (isempty (obj.DistParameter), true);

%!test
%! ## Minkowski distance with custom P
%! X = [0, 0; 1, 1; 2, 2];
%! obj = hnswSearcher (X, 'Distance', 'minkowski', 'P', 3);
%! assert_equal (obj.Distance, "minkowski");
%! assert_equal (obj.DistParameter, 3);

%!test
%! ## Seuclidean distance with custom Scale
%! X = [1, 2; 3, 4; 5, 6];
%! S = [1, 2];
%! obj = hnswSearcher (X, 'Distance', 'seuclidean', 'Scale', S);
%! assert_equal (obj.Distance, "seuclidean");
%! assert_equal (obj.DistParameter, S);

%!test
%! ## Mahalanobis distance with custom Cov
%! X = [1, 2; 3, 4; 5, 6];
%! C = [1, 0; 0, 1];
%! obj = hnswSearcher (X, 'Distance', 'mahalanobis', 'Cov', C);
%! assert_equal (obj.Distance, "mahalanobis");
%! assert_equal (obj.DistParameter, C);

%!test
%! ## knnsearch with Euclidean distance
%! X = [1, 2; 3, 4; 5, 6];
%! obj = hnswSearcher (X);
%! Y = [2, 3];
%! [idx, D] = knnsearch (obj, Y, 'K', 1);
%! assert_equal (ismember (idx, [2]), true);
%! assert_equal (abs (D - sqrt (2)) < 1e-2, true);

%!test
%! ## knnsearch with Cityblock distance
%! X = [0, 0; 1, 1; 2, 2];
%! obj = hnswSearcher (X, 'Distance', 'cityblock');
%! Y = [1, 0];
%! [idx, D] = knnsearch (obj, Y, 'K', 1);
%! assert_equal (ismember (idx, [1, 2]), true);
%! assert_equal (abs (D - 1) < 1e-2, true);

%!test
%! ## knnsearch with Chebychev distance
%! X = [1, 1; 2, 3; 4, 2];
%! obj = hnswSearcher (X, 'Distance', 'chebychev');
%! Y = [2, 2];
%! [idx, D] = knnsearch (obj, Y, 'K', 1);
%! assert_equal (ismember (idx, [1, 2]), true);
%! assert_equal (abs (D - 1) < 1e-2, true);

%!test
%! ## knnsearch with Minkowski P=3
%! X = [0, 0; 1, 0; 2, 0];
%! obj = hnswSearcher (X, 'Distance', 'minkowski', 'P', 3);
%! Y = [1, 0];
%! [idx, D] = knnsearch (obj, Y, 'K', 1);
%! assert_equal (ismember (idx, [2]), true);
%! assert_equal (abs (D - 0) < 1e-2, true);

%!test
%! ## Diverse dataset with Euclidean
%! X = [0, 10; 5, 5; 10, 0];
%! obj = hnswSearcher (X);
%! Y = [5, 5];
%! [idx, D] = knnsearch (obj, Y, 'K', 1);
%! assert_equal (ismember (idx, [2]), true);
%! assert_equal (abs (D - 0) < 1e-2, true);

%!test
%! ## High-dimensional data with Cityblock
%! X = [1, 2, 3; 4, 5, 6; 7, 8, 9];
%! obj = hnswSearcher (X, 'Distance', 'cityblock');
%! Y = [4, 5, 6];
%! [idx, D] = knnsearch (obj, Y, 'K', 1);
%! assert_equal (ismember (idx, [2]), true);
%! assert_equal (abs (D - 0) < 1e-2, true);

## Raising 'SearchSetSize' must not make the search worse.  The early
## termination check read the last element of the best list on the assumption
## that the list was sorted, but the list is only sorted when it is trimmed,
## and it is only trimmed once it grows past 'SearchSetSize'.  For a large
## enough 'SearchSetSize' the trim never ran, the check compared whatever
## distance happened to have been appended last, and the walk stopped at the
## entry point -- returning the same neighbour for every query.
%!test
%! ## both the data and the graph are drawn, so seed both: the search is
%! ## approximate and an unlucky graph misses a true neighbour
%! rand ("seed", 42);
%! randn ("seed", 4);
%! X = [randn(20,2); randn(20,2) + 3; randn(20,2) + [0, 6]];
%! Y = randn (8, 2);
%! hn = hnswSearcher (X);
%! es = ExhaustiveSearcher (X);
%! truth = knnsearch (es, Y, 'K', 3);
%! for ess = [5, 16, 31, 32, 33, 60]
%!   assert_equal (knnsearch (hn, Y, 'K', 3, 'SearchSetSize', ess), truth);
%! endfor
%!test
%! ## the boundary that used to break was 2 * MaxNumLinksPerNode
%! rand ("seed", 42);
%! randn ("seed", 4);
%! X = [randn(20,2); randn(20,2) + 3; randn(20,2) + [0, 6]];
%! Y = randn (8, 2);
%! hn = hnswSearcher (X);
%! assert_equal (knnsearch (hn, Y, 'K', 3, 'SearchSetSize', 31), ...
%!               knnsearch (hn, Y, 'K', 3, 'SearchSetSize', 32));
%!test
%! ## distinct queries must not collapse onto one index
%! rand ("seed", 42);
%! randn ("seed", 4);
%! X = [randn(20,2); randn(20,2) + 3; randn(20,2) + [0, 6]];
%! Y = randn (8, 2);
%! idx = knnsearch (hnswSearcher (X), Y, 'K', 3, 'SearchSetSize', 60);
%! assert_equal (numel (unique (idx(:,1))) > 1, true);

## Test Input Validation

%!error<hnswSearcher: too few input arguments.> ...
%! hnswSearcher ()
%!error<hnswSearcher: Name-Value arguments must be in pairs.> ...
%! hnswSearcher (ones (3,2), 'Distance')
%!error<hnswSearcher: X cannot be empty.> ...
%! hnswSearcher ([])
%!error<hnswSearcher: X must be a finite numeric matrix.> ...
%! hnswSearcher ('abc')
%!error<hnswSearcher: X must be a finite numeric matrix.> ...
%! hnswSearcher ([1; Inf; 3])
%!error<hnswSearcher: invalid parameter name: 'foo'.> ...
%! hnswSearcher (ones (3,2), 'foo', 'bar')
%!error<hnswSearcher: unsupported distance metric 'invalid'.> ...
%! hnswSearcher (ones (3,2), 'Distance', 'invalid')
%!error<hnswSearcher: 'Distance' must be a string.> ...
%! hnswSearcher (ones (3,2), 'Distance', 1)
%!error<hnswSearcher: 'P' must be a positive finite scalar.> ...
%! hnswSearcher (ones (3,2), 'Distance', 'minkowski', 'P', -1)
%!error<hnswSearcher: 'Scale' must be a nonnegative vector matching X columns.> ...
%! hnswSearcher (ones (3,2), 'Distance', 'seuclidean', 'Scale', [-1, 1])
%!error<hnswSearcher: 'Cov' must be a square matrix matching X columns.> ...
%! hnswSearcher (ones (3,2), 'Distance', 'mahalanobis', 'Cov', ones (3,3))
%!error<hnswSearcher: 'Cov' must be symmetric for mahalanobis.> ...
%! hnswSearcher (ones (3,2), 'Distance', 'mahalanobis', 'Cov', [1, 2; 3, 4])
%!error<hnswSearcher: 'Cov' must be positive definite for mahalanobis.> ...
%! hnswSearcher (ones (3,2), 'Distance', 'mahalanobis', 'Cov', -eye (2))
%!error<hnswSearcher: 'MaxNumLinksPerNode' must be a positive integer.> ...
%! hnswSearcher (ones (3,2), 'MaxNumLinksPerNode', 0)
%!error<hnswSearcher: 'TrainSetSize' must be a positive integer.> ...
%! hnswSearcher (ones (3,2), 'TrainSetSize', -1)
%!error<hnswSearcher: 'TrainSetSize' cannot exceed the number of rows in X.> ...
%! hnswSearcher (ones (3,2), 'TrainSetSize', 4)
%!error<hnswSearcher: 'TrainSetSize' cannot exceed the number of rows in X.> ...
%! hnswSearcher (ones (3,2), 'MaxNumLinksPerNode', 200, 'TrainSetSize', 100)

%!error<hnswSearcher.knnsearch: too few input arguments.> ...
%! knnsearch (hnswSearcher (ones (3,2)))
%!error<hnswSearcher.knnsearch: Name-Value arguments must be in pairs.> ...
%! knnsearch (hnswSearcher (ones (3,2)), ones (3,2), 'K')
%!error<hnswSearcher.knnsearch: Y must be a real numeric matrix.> ...
%! knnsearch (hnswSearcher (ones (3,2)), {1, 2})
%!error<hnswSearcher.knnsearch: 'K' must be a positive integer.> ...
%! knnsearch (hnswSearcher (ones (3,2)), ones (3,2), 'K', 0)
%!error<hnswSearcher.knnsearch: invalid parameter name: 'foo'.> ...
%! knnsearch (hnswSearcher (ones (3,2)), ones (3,2), 'foo', 'bar')
%!error<hnswSearcher.subsref: \(\) indexing not supported.> ...
%! obj = hnswSearcher (ones (3,2)); obj(1)
%!error<hnswSearcher.subsref: {} indexing not supported.> ...
%! obj = hnswSearcher (ones (3,2)); obj{1}
%!error<hnswSearcher.subsref: unrecognized property: 'invalid'.> ...
%! obj = hnswSearcher (ones (3,2)); obj.invalid

%!error<hnswSearcher.subsasgn: \(\) indexing not supported.> ...
%! obj = hnswSearcher (ones (3,2)); obj(1) = 1
%!error<hnswSearcher.subsasgn: {} indexing not supported.> ...
%! obj = hnswSearcher (ones (3,2)); obj{1} = 1
%!error<hnswSearcher.subsasgn: 'X' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2)); obj.X = 1


%!error<hnswSearcher.subsasgn: 'HNSWGraph' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2)); obj.HNSWGraph = 1
%!error<hnswSearcher.subsasgn: 'Distance' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2)); obj.Distance = 'invalid'
%!error<hnswSearcher.subsasgn: 'Distance' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2)); obj.Distance = 1
%!error<hnswSearcher.subsasgn: 'DistParameter' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2), 'Distance', 'minkowski'); obj.DistParameter = -1
%!error<hnswSearcher.subsasgn: 'DistParameter' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2), 'Distance', 'seuclidean'); obj.DistParameter = [-1, 1]
%!error<pdist2: covariance matrix for mahalanobis distance must be symmetric and positive definite.> ...
%! obj = hnswSearcher (ones (3,2), 'Distance', 'mahalanobis'); obj.DistParameter = ones (3,3)
%!error<pdist2: covariance matrix for mahalanobis distance must be symmetric and positive definite.> ...
%! obj = hnswSearcher (ones (3,2), 'Distance', 'mahalanobis'); obj.DistParameter = -eye (2)
%!error<hnswSearcher.subsasgn: 'DistParameter' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2)); obj.DistParameter = 1
%!error<hnswSearcher.subsasgn: 'MaxNumLinksPerNode' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2)); obj.MaxNumLinksPerNode = 0
%!error<hnswSearcher.subsasgn: 'TrainSetSize' is read-only and cannot be modified.> ...
%! obj = hnswSearcher (ones (3,2)); obj.TrainSetSize = -1
%!error<hnswSearcher.subsasgn: unrecognized property: 'efSearch'.> ...
%! obj = hnswSearcher (ones (3,2)); obj.efSearch = 1.5
%!error<hnswSearcher.subsasgn: unrecognized property: 'invalid'.> ...
%! obj = hnswSearcher (ones (3,2)); obj.invalid = 1

## More neighbours than there are points, and queries that are empty or not
## finite, are answered rather than refused.  Verified against MATLAB R2024a.
%!shared Xs, Ys, hs
%! ## the graph is built from rand, so seed it: the search is approximate and
%! ## an unlucky graph misses a true neighbour
%! rand ("seed", 42);
%! randn ("seed", 4);
%! Xs = [randn(20,2); randn(20,2) + 3; randn(20,2) + [0, 6]];
%! Ys = randn (8, 2);
%! hs = hnswSearcher (Xs);

%!test  # asking for every point returns every point, not a padding of NaN
%! [idx, D] = knnsearch (hs, Ys, "K", 60);
%! assert_equal (size (idx), [8, 60]);
%! assert_equal (any (isnan (idx(:))), false);
%! for r = 1:8
%!   assert_equal (sort (idx(r,:)), 1:60);
%! endfor

%!test  # and the neighbours are the right ones, not merely present
%! for K = [1, 5, 30, 60]
%!   [~, Dh] = knnsearch (hs, Ys, "K", K);
%!   [~, De] = knnsearch (ExhaustiveSearcher (Xs), Ys, "K", K);
%!   assert_equal (Dh, De, 1e-12);
%! endfor

%!test  # more neighbours than points gives all of them, not a wider matrix
%! for K = [61, 100, 1000]
%!   [idx, D] = knnsearch (hs, Ys, "K", K);
%!   assert_equal (size (idx), [8, 60]);
%!   assert_equal (size (D), [8, 60]);
%!   assert_equal (any (isnan (idx(:))), false);
%! endfor

%!test  # an empty query is answered with an empty result
%! [idx, D] = knnsearch (hs, zeros (0, 2), "K", 3);
%! assert_equal (size (idx), [0, 3]);
%! assert_equal (size (D), [0, 3]);

%!test  # a query that is not finite is answered, its distances saying so
%! [idx, D] = knnsearch (hs, [Inf, 0], "K", 3);
%! assert_equal (size (idx), [1, 3]);
%! assert_equal (all (isinf (D)), true);
%! [idx, D] = knnsearch (hs, [NaN, NaN], "K", 3);
%! assert_equal (size (idx), [1, 3]);
%! assert_equal (all (isnan (D)), true);

## Single precision is carried through to the distances, every other class is
## converted up to double, and the indices are always double.
%!test
%! rand ("seed", 42);
%! X = [1, 2; 3, 4; 5, 6; 7, 8];  Y = [2, 3; 6, 7];
%! obj = hnswSearcher (single (X));
%! assert_equal (class (obj.X), 'single');
%! [idx, D] = knnsearch (obj, single (Y), "K", 2);
%! assert_equal (class (idx), 'double');
%! assert_equal (class (D), 'single');
%! ## a double query against single data is still computed in single
%! [~, Dm] = knnsearch (obj, Y, "K", 2);
%! assert_equal (class (Dm), 'single');
%! assert_equal (Dm, D);
%! ## and so is a single query against double data
%! [~, Ds] = knnsearch (hnswSearcher (X), single (Y), "K", 2);
%! assert_equal (class (Ds), 'single');
%! ## double throughout stays double
%! [~, Dd] = knnsearch (hnswSearcher (X), Y, "K", 2);
%! assert_equal (class (Dd), 'double');

## Integer data is converted to double on the way in.  Held as int32 the
## coordinate differences were rounded and the distances came out wrong.
%!test
%! rand ("seed", 42);
%! X = int32 ([10, 20; 33, 41; 55, 62]);  Y = [21, 33];
%! obj = hnswSearcher (X);
%! assert_equal (class (obj.X), 'double');
%! [~, D] = knnsearch (obj, Y, "K", 1);
%! assert_equal (D, min (sqrt (sum ((double (X) - Y) .^ 2, 2))), 1e-12);
