# coding=utf-8
# Copyright 2026 The Perch Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Brute force search and reranking utilities."""

import concurrent
import threading
from typing import Any, cast

import numpy as np
from perch_hoplite.db import interface
from perch_hoplite.db import score_functions
from perch_hoplite.db import search_results


def worker_initializer(state):
  name = threading.current_thread().name
  state[name + 'db'] = state['db'].thread_split()


def brute_search_worker_fn(emb_ids: np.ndarray, state: dict[str, Any]):
  name = threading.current_thread().name
  embeddings = state[name + 'db'].get_embeddings_batch(emb_ids)
  scores = state['score_fn'](embeddings, state['query_embedding'])
  top_locs = np.argpartition(scores, state['search_list_size'], axis=-1)
  return emb_ids[top_locs], scores[top_locs]


def threaded_brute_search(
    db: interface.HopliteDBInterface,
    query_embedding: np.ndarray,
    search_list_size: int,
    score_fn: score_functions.ScoreFn,
    batch_size: int = 1024,
    max_workers: int = 8,
    sample_size: int | float | None = None,
    rng_seed: int | None = None,
) -> search_results.TopKSearchResults:
  """Performs a brute-force search for neighbors of the query embedding.

  Args:
    db: Graph DB instance.
    query_embedding: Query embedding vector.
    search_list_size: Number of results to return.
    score_fn: Scoring function to use for ranking results.
    batch_size: Number of embeddings to score in each thread.
    max_workers: Maximum number of threads to use for the search.
    sample_size: Number of embeddings to compare with. If None, all embeddings
      are used. For floats between 0 and 1, sample a proportion of the database.
      For ints, sample the specified number of embeddings.
    rng_seed: Random number generator seed to use for sampled search.

  Returns:
    A TopKSearchResults object containing the search results.
  """
  state = {}
  state['search_list_size'] = search_list_size
  state['db'] = db
  state['query_embedding'] = query_embedding
  state['score_fn'] = score_fn

  results = search_results.TopKSearchResults(search_list_size)
  # Commit the DB, since we are about to create views in multiple threads.
  db.commit()
  with concurrent.futures.ThreadPoolExecutor(
      max_workers=max_workers,
      initializer=worker_initializer,
      initargs=(state,),
  ) as executor:
    ids = get_brute_search_ids(db, sample_size, rng_seed)
    futures = []
    for q in range(0, ids.shape[0], batch_size):
      futures.append(
          executor.submit(
              brute_search_worker_fn, ids[q : q + batch_size], state
          )
      )
    for f in futures:
      idxes, scores = f.result()
      for idx, score in zip(idxes, scores):
        if not results.will_filter(idx, score):
          results.update(
              search_results.SearchResult(idx, score), force_insert=True
          )
  return results


def brute_search(
    db: interface.HopliteDBInterface,
    query_embedding: np.ndarray,
    search_list_size: int,
    score_fn: score_functions.ScoreFn,
    sample_size: int | float | None = None,
    rng_seed: int | None = None,
) -> search_results.TopKSearchResults:
  """Performs a brute-force search for neighbors of the query embedding.

  Args:
    db: Graph DB instance.
    query_embedding: Query embedding vector.
    search_list_size: Number of results to return.
    score_fn: Scoring function to use for ranking results.
    sample_size: Number of embeddings to compare with. If None, all embeddings
      are used. For floats between 0 and 1, sample a proportion of the database.
      For ints, sample the specified number of embeddings.
    rng_seed: Random number generator seed to use for sampled search.

  Returns:
    A TopKSearchResults object containing the search results.
  """
  results = search_results.TopKSearchResults(search_list_size)
  ids = get_brute_search_ids(db, sample_size, rng_seed)
  for idx in ids:
    target_embedding = db.get_embedding(idx)
    score = score_fn(query_embedding, target_embedding)
    # Check filtering and then force insert to avoid creating a SearchResult
    # object for discarded objects. This saves a small amount of time in the
    # inner loop.
    if not results.will_filter(idx, score):
      results.update(search_results.SearchResult(idx, score), force_insert=True)
  return results


def get_brute_search_ids(
    db: interface.HopliteDBInterface,
    sample_size: int | float | None = None,
    rng_seed: int | None = None,
):
  """Get IDs for brute force search, subsampling as needed."""
  ids = np.array(db.match_window_ids())
  if sample_size is None:
    return ids
  if isinstance(sample_size, float):
    if sample_size < 0 or sample_size > 1:
      raise ValueError('Float sample size must be between 0 and 1.')
    sample_size = int(sample_size * db.count_embeddings())
  np.random.default_rng(rng_seed).shuffle(ids)
  return ids[:sample_size]


def get_random_embedding_scores(
    db: interface.HopliteDBInterface,
    query_embedding: np.ndarray,
    score_fn: score_functions.ScoreFn,
    sample_size: int | float | None = 1024,
    rng_seed: int | None = None,
) -> np.ndarray:
  """Get scores for a random sample of embeddings."""
  ids = get_brute_search_ids(db, sample_size, rng_seed)
  embeddings = db.get_embeddings_batch(ids)
  scores = score_fn(embeddings, query_embedding)
  return cast(np.ndarray, scores)


def rerank(
    query_embedding: np.ndarray,
    results: search_results.TopKSearchResults,
    db: interface.HopliteDBInterface,
    score_fn: score_functions.ScoreFn,
) -> search_results.TopKSearchResults:
  """Rescore the search results using a different score function."""
  new_results = search_results.TopKSearchResults(results.top_k)
  for r in results:
    new_results.update(
        search_results.SearchResult(
            r.window_id,
            score_fn(query_embedding, db.get_embedding(r.window_id)),
        )
    )
  return new_results
