export type DistanceMetric = "euclidean" | "cosine";

export interface KDPoint<T = unknown> {
	vector: number[];
	payload: T;
}

export interface KNNResult<T = unknown> {
	point: KDPoint<T>;
	distance: number;
}

interface KDNode<T> {
	point: KDPoint<T>;
	axis: number;
	left: KDNode<T> | null;
	right: KDNode<T> | null;
}

// ─── Distance helpers ─────────────────────────────────────────────────────────

function euclidean(a: number[], b: number[]): number {
	let sum = 0;
	for (let i = 0; i < a.length; i++) {
		const d = a[i] - b[i];
		sum += d * d;
	}
	return Math.sqrt(sum);
}

function cosine(a: number[], b: number[]): number {
	let dot = 0, normA = 0, normB = 0;
	for (let i = 0; i < a.length; i++) {
		dot   += a[i] * b[i];
		normA += a[i] * a[i];
		normB += b[i] * b[i];
	}
	const denom = Math.sqrt(normA) * Math.sqrt(normB);
	return denom === 0 ? 1 : 1 - dot / denom; // distance = 1 - similarity
}

/**
 * Keeps the k closest candidates in memory, evicts the furthest when full
 */
class BoundedMaxHeap<T> {
	private heap: KNNResult<T>[] = [];

	constructor(private readonly k: number) {}

	get size(): number { return this.heap.length; }

	get worstDistance(): number {
		return this.heap.length < this.k ? Infinity : this.heap[0].distance;
	}

	push(item: KNNResult<T>): void {
		if (this.heap.length < this.k) {
			this.heap.push(item);
			this.bubbleUp(this.heap.length - 1);
		} else if (item.distance < this.heap[0].distance) {
			this.heap[0] = item;
			this.sinkDown(0);
		}
	}

	toSortedArray(): KNNResult<T>[] {
		return [...this.heap].sort((a, b) => a.distance - b.distance);
	}

	private bubbleUp(i: number): void {
		while (i > 0) {
			const parent = (i - 1) >> 1;
			if (this.heap[parent].distance >= this.heap[i].distance) break;
			[this.heap[parent], this.heap[i]] = [this.heap[i], this.heap[parent]];
			i = parent;
		}
	}

	private sinkDown(i: number): void {
		const n = this.heap.length;
		while (true) {
			let largest = i;
			const l = 2 * i + 1, r = 2 * i + 2;
			if (l < n && this.heap[l].distance > this.heap[largest].distance) largest = l;
			if (r < n && this.heap[r].distance > this.heap[largest].distance) largest = r;
			if (largest === i) break;
			[this.heap[largest], this.heap[i]] = [this.heap[i], this.heap[largest]];
			i = largest;
		}
	}
}

/**
 * K-D Tree for efficient nearest-neighbor search over high-dimensional vectors / embeddings.
 *
 * Supports:
 *  - Insertion of labeled points
 *  - k-nearest-neighbor (KNN) search
 *  - Radius search (all points within a given distance)
 *  - Euclidean and cosine distance metrics
 *  - Bulk construction (balanced tree) for best query performance
 */
export class KDTree<T = unknown> {
	private root: KDNode<T> | null = null;
	private _size = 0;
	private readonly distanceFn: (a: number[], b: number[]) => number;

	readonly dims: number;

	/**
	 * @param dims           Dimensionality of all vectors (must be consistent).
	 * @param metric         Distance metric to use. Default: "euclidean".
	 * @param points         Optional initial set of points. Builds a balanced tree
	 *                       in O(n log² n) — prefer this over inserting one-by-one
	 *                       when you have a large corpus.
	 */
	constructor(
		dims: number,
		metric: DistanceMetric = "euclidean",
		points?: KDPoint<T>[]
	) {
		this.dims = dims;
		this.distanceFn = metric === "cosine" ? cosine : euclidean;

		if (points && points.length > 0) {
			this.validateAll(points);
			this.root = this.buildBalanced([...points], 0);
			this._size = points.length;
		}
	}

	/** Total number of points stored in the tree. */
	get size(): number { return this._size; }

	// ── Insertion ──────────────────────────────────────────────────────────────

	/**
	 * Insert a single point. O(log n) average, O(n) worst case on skewed data.
	 * For bulk loading prefer passing points to the constructor.
	 */
	insert(point: KDPoint<T>): void {
		this.validate(point);
		this.root = this.insertNode(this.root, point, 0);
		this._size++;
	}

	// ── KNN search ─────────────────────────────────────────────────────────────

	/**
	 * Find the k nearest neighbors to `query`.
	 * Returns results sorted by distance ascending.
	 */
	knn(query: number[], k: number): KNNResult<T>[] {
		if (k <= 0) throw new RangeError("k must be a positive integer");
		this.validateVector(query);

		const heap = new BoundedMaxHeap<T>(k);
		this.searchKNN(this.root, query, k, heap, 0);
		return heap.toSortedArray();
	}

	/**
	 * Nearest single neighbor. Convenience wrapper around knn(query, 1).
	 * Returns null if the tree is empty.
	 */
	nearest(query: number[]): KNNResult<T> | null {
		const results = this.knn(query, 1);
		return results[0] ?? null;
	}

	// ── Radius search ──────────────────────────────────────────────────────────

	/**
	 * Return all points whose distance to `query` is ≤ `radius`,
	 * sorted by distance ascending.
	 */
	radiusSearch(query: number[], radius: number): KNNResult<T>[] {
		if (radius < 0) throw new RangeError("radius must be non-negative");
		this.validateVector(query);

		const results: KNNResult<T>[] = [];
		this.searchRadius(this.root, query, radius, results, 0);
		results.sort((a, b) => a.distance - b.distance);
		return results;
	}

	// ── Conversion ─────────────────────────────────────────────────────────────

	/** Collect all points in the tree (order not guaranteed). */
	toArray(): KDPoint<T>[] {
		const out: KDPoint<T>[] = [];
		this.collect(this.root, out);
		return out;
	}

	/**
	 * Rebuild the tree from its current points as a balanced tree.
	 * Useful after many individual insertions to restore O(log n) query time.
	 */
	rebalance(): void {
		const points = this.toArray();
		this.root = points.length ? this.buildBalanced(points, 0) : null;
	}

	// ── Private: build ─────────────────────────────────────────────────────────

	private buildBalanced(points: KDPoint<T>[], depth: number): KDNode<T> {
		const axis = depth % this.dims;
		points.sort((a, b) => a.vector[axis] - b.vector[axis]);

		const mid = Math.floor(points.length / 2);
		return {
			point: points[mid],
			axis,
			left:  points.slice(0, mid).length
				? this.buildBalanced(points.slice(0, mid), depth + 1)
				: null,
			right: points.slice(mid + 1).length
				? this.buildBalanced(points.slice(mid + 1), depth + 1)
				: null,
		};
	}

	// ── Private: insert ────────────────────────────────────────────────────────

	private insertNode(
		node: KDNode<T> | null,
		point: KDPoint<T>,
		depth: number
	): KDNode<T> {
		if (node === null) {
			return { point, axis: depth % this.dims, left: null, right: null };
		}
		const axis = depth % this.dims;
		if (point.vector[axis] < node.point.vector[axis]) {
			node.left = this.insertNode(node.left, point, depth + 1);
		} else {
			node.right = this.insertNode(node.right, point, depth + 1);
		}
		return node;
	}

	// ── Private: KNN traversal ─────────────────────────────────────────────────

	private searchKNN(
		node: KDNode<T> | null,
		query: number[],
		k: number,
		heap: BoundedMaxHeap<T>,
		depth: number
	): void {
		if (node === null) return;

		const dist = this.distanceFn(query, node.point.vector);
		heap.push({ point: node.point, distance: dist });

		const axis = node.axis;
		const diff = query[axis] - node.point.vector[axis];
		const [near, far] = diff <= 0
			? [node.left, node.right]
			: [node.right, node.left];

		this.searchKNN(near, query, k, heap, depth + 1);

		// Only explore the far side if it could contain a closer point.
		// For cosine distance we can't prune by axis gap alone, so always explore.
		const shouldExplore =
			this.distanceFn === cosine
				? true
				: Math.abs(diff) < heap.worstDistance;

		if (shouldExplore) {
			this.searchKNN(far, query, k, heap, depth + 1);
		}
	}

	// ── Private: radius traversal ──────────────────────────────────────────────

	private searchRadius(
		node: KDNode<T> | null,
		query: number[],
		radius: number,
		results: KNNResult<T>[],
		depth: number
	): void {
		if (node === null) return;

		const dist = this.distanceFn(query, node.point.vector);
		if (dist <= radius) {
			results.push({ point: node.point, distance: dist });
		}

		const axis = node.axis;
		const diff = query[axis] - node.point.vector[axis];
		const [near, far] = diff <= 0
			? [node.left, node.right]
			: [node.right, node.left];

		this.searchRadius(near, query, radius, results, depth + 1);

		const shouldExplore =
			this.distanceFn === cosine ? true : Math.abs(diff) <= radius;

		if (shouldExplore) {
			this.searchRadius(far, query, radius, results, depth + 1);
		}
	}

	// ── Private: collect ───────────────────────────────────────────────────────

	private collect(node: KDNode<T> | null, out: KDPoint<T>[]): void {
		if (node === null) return;
		out.push(node.point);
		this.collect(node.left, out);
		this.collect(node.right, out);
	}

	// ── Private: validation ────────────────────────────────────────────────────

	private validateVector(v: number[]): void {
		if (v.length !== this.dims) {
			throw new TypeError(
				`Vector length ${v.length} does not match tree dimensionality ${this.dims}`
			);
		}
	}

	private validate(point: KDPoint<T>): void {
		this.validateVector(point.vector);
	}

	private validateAll(points: KDPoint<T>[]): void {
		for (const p of points) this.validate(p);
	}
}