class VectorStore
  def initialize(distance_metric = :cosine)
    @vectors = {}
    @distance_metric = distance_metric
  end

  def add(id, vector)
    validate_vector(vector)
    @vectors[id] = vector
  end

  def get(id)
    @vectors[id]
  end

  def search(query_vector, k = 5)
    validate_vector(query_vector)
    @vectors.map do |id, vector|
      distance = calculate_distance(query_vector, vector)
      [id, distance]
    end.sort_by { |_, distance| distance }
      .take(k)
  end

  private

  def validate_vector(vector)
    raise ArgumentError, "Vector must be an Array" unless vector.is_a?(Array)
    raise ArgumentError, "Vector must contain only numbers" unless vector.all? { |v| v.is_a?(Numeric) }
  end

  def calculate_distance(vec1, vec2)
    case @distance_metric
    when :cosine then cosine_distance(vec1, vec2)
    when :euclidean then euclidean_distance(vec1, vec2)
    when :manhattan then manhattan_distance(vec1, vec2)
    else
      raise ArgumentError, "Unsupported distance metric: #{@distance_metric}"
    end
  end

  def cosine_distance(vec1, vec2)
    dot_product = vec1.zip(vec2).sum { |a, b| a * b }
    magnitude1 = Math.sqrt(vec1.map { |v| v**2 }.sum)
    magnitude2 = Math.sqrt(vec2.map { |v| v**2 }.sum)
    return 1.0 if magnitude1.zero? || magnitude2.zero?
    1 - (dot_product / (magnitude1 * magnitude2))
  end

  def euclidean_distance(vec1, vec2)
    Math.sqrt(vec1.zip(vec2).sum { |a, b| (a - b)**2 })
  end

  def manhattan_distance(vec1, vec2)
    vec1.zip(vec2).sum { |a, b| (a - b).abs }
  end
end

store = VectorStore.new(:cosine)
store.add("doc1", [0.2, 0.3, 0.4, 0.1])
store.add("doc2", [0.3, 0.2, 0.1, 0.4])
store.add("doc3", [0.1, 0.2, 0.3, 0.4])
store.add("doc4", [0.5, 0.5, 0.2, 0.1])

query = [0.2, 0.2, 0.3, 0.3]
results = store.search(query, 2)

puts "Kết quả tìm kiếm cho #{query}:"
results.each { |id, distance| puts " - #{id}: #{distance}" }