MinHash and locality-sensitive hashing.

export type MinHashErrorCode =
	| 'bad-options'
	| 'empty-document-list'
	| 'too-many-documents'
	| 'bad-document-id'
	| 'duplicate-document-id'
	| 'bad-document-text'
	| 'too-many-shingles'
	| 'no-shingles';

export class MinHashError extends Error {
	readonly code: MinHashErrorCode;

	constructor(code: MinHashErrorCode, message: string) {
		super(message);
		this.name = 'MinHashError';
		this.code = code;
	}
}

export interface Document {
	id: string;
	text: string;
}

export interface MinHashOptions {
	shingleSize?: number;
	permutations?: number;
	bands?: number;
}

export interface MinHashConfig {
	readonly shingleSize: number;
	readonly permutations: number;
	readonly bands: number;
	readonly rowsPerBand: number;
}

export interface Signature {
	id: string;
	shingles: string[];
	values: number[];
}

export interface PairComparison {
	left: string;
	right: string;
	sharedShingles: number;
	unionShingles: number;
	matchingRows: number;
	exactJaccard: number;
	estimatedJaccard: number;
	candidate: boolean;
}

export interface BandBucket {
	band: number;
	key: string;
	ids: string[];
}

export interface SearchResult {
	signatures: Signature[];
	buckets: BandBucket[];
	candidates: string[];
}

function validateOptions(options: MinHashOptions): MinHashConfig {
	const shingleSize = options.shingleSize ?? 3;
	const permutations = options.permutations ?? 24;
	const bands = options.bands ?? 6;
	if (
		!Number.isInteger(shingleSize) ||
		shingleSize < 2 ||
		shingleSize > MAX_SHINGLE_SIZE ||
		!Number.isInteger(permutations) ||
		permutations < 8 ||
		permutations > MAX_PERMUTATIONS ||
		!Number.isInteger(bands) ||
		bands < 1 ||
		bands > MAX_BANDS ||
		permutations % bands !== 0 ||
		permutations / bands > MAX_ROWS_PER_BAND
	)
		throw new MinHashError(
			'bad-options',
			`Use shingles 2-${MAX_SHINGLE_SIZE}, ${8}-${MAX_PERMUTATIONS} permutations, and bands that divide permutations with at most ${MAX_ROWS_PER_BAND} rows.`
		);
	return { shingleSize, permutations, bands, rowsPerBand: permutations / bands };
}

function validateText(text: string): void {
	if (
		typeof text !== 'string' ||
		text.length === 0 ||
		text.length > MAX_TEXT_LENGTH ||
		!ASCII_TEXT.test(text)
	)
		throw new MinHashError(
			'bad-document-text',
			`Document text must be 1-${MAX_TEXT_LENGTH} printable ASCII characters.`
		);
}

function validateDocument(document: Document): void {
	if (typeof document.id !== 'string' || !DOCUMENT_ID.test(document.id))
		throw new MinHashError('bad-document-id', 'Document ids must be lowercase slugs.');
	validateText(document.text);
}

function validateDocuments(documents: Document[]): void {
	if (documents.length === 0)
		throw new MinHashError('empty-document-list', 'At least one document is required.');
	if (documents.length > MAX_DOCUMENTS)
		throw new MinHashError(
			'too-many-documents',
			`At most ${MAX_DOCUMENTS} documents are accepted.`
		);
	const ids = new Set<string>();
	for (const document of documents) {
		validateDocument(document);
		if (ids.has(document.id))
			throw new MinHashError('duplicate-document-id', 'Document ids must be unique.');
		ids.add(document.id);
	}
}

// FNV-1a plus a final avalanche. ASCII shingles make this 32-bit hash match Go byte for byte.
function hash32(value: string, seed: number): number {
	let hash = (0x811c9dc5 ^ Math.imul(seed + 1, 0x9e3779b9)) >>> 0;
	for (let index = 0; index < value.length; index += 1) {
		hash ^= value.charCodeAt(index);
		hash = Math.imul(hash, 0x01000193) >>> 0;
	}
	hash ^= hash >>> 16;
	hash = Math.imul(hash, 0x85ebca6b) >>> 0;
	hash ^= hash >>> 13;
	hash = Math.imul(hash, 0xc2b2ae35) >>> 0;
	hash ^= hash >>> 16;
	return hash >>> 0;
}

export function tokenize(text: string): string[] {
	validateText(text);
	return text.toLowerCase().match(WORD) ?? [];
}

export function makeShingles(text: string, shingleSize = 3): string[] {
	if (!Number.isInteger(shingleSize) || shingleSize < 2 || shingleSize > MAX_SHINGLE_SIZE)
		throw new MinHashError(
			'bad-options',
			`Shingle size must be a whole number from 2 to ${MAX_SHINGLE_SIZE}.`
		);
	const words = tokenize(text);
	const shingles = new Set<string>();
	for (let index = 0; index + shingleSize <= words.length; index += 1)
		shingles.add(words.slice(index, index + shingleSize).join(' '));
	if (shingles.size > MAX_TEXT_LENGTH)
		throw new MinHashError(
			'too-many-shingles',
			`A document may contain at most ${MAX_TEXT_LENGTH} shingles.`
		);
	return [...shingles].sort();
}

function setJaccard(
	left: string[],
	right: string[]
): { shared: number; union: number; score: number } {
	const rightSet = new Set(right);
	const leftSet = new Set(left);
	let shared = 0;
	for (const shingle of leftSet) if (rightSet.has(shingle)) shared += 1;
	const union = new Set([...leftSet, ...rightSet]).size;
	// signature() refuses empty shingle sets, so union is never 0 here; if it were, report 0
	// rather than calling two empty documents identical.
	return { shared, union, score: union === 0 ? 0 : shared / union };
}

// The MinHash step: for each of the deterministic hash rows, keep the smallest hash over the
// document's shingles. Two documents agree on a row with probability equal to their Jaccard
// similarity. A document with no shingles has no minimum, so it is refused.
function minHashValues(shingles: string[], permutations: number): number[] {
	if (shingles.length === 0)
		throw new MinHashError(
			'no-shingles',
			'A document needs at least as many words as the shingle size to form one shingle.'
		);
	return Array.from({ length: permutations }, (_, seed) =>
		shingles.reduce((minimum, shingle) => Math.min(minimum, hash32(shingle, seed)), 0xffffffff)
	);
}

function bandKey(values: number[], start: number, length: number): string {
	return values.slice(start, start + length).join(':');
}

type MinHashErrorCode string

const (
	BadOptions          MinHashErrorCode = "bad-options"
	EmptyDocumentList   MinHashErrorCode = "empty-document-list"
	TooManyDocuments    MinHashErrorCode = "too-many-documents"
	BadDocumentID       MinHashErrorCode = "bad-document-id"
	DuplicateDocumentID MinHashErrorCode = "duplicate-document-id"
	BadDocumentText     MinHashErrorCode = "bad-document-text"
	NoShingles          MinHashErrorCode = "no-shingles"
)

type MinHashError struct {
	Code    MinHashErrorCode
	Message string
}

func (e *MinHashError) Error() string { return e.Message }

type Document struct {
	ID   string
	Text string
}

type MinHashOptions struct {
	ShingleSize  int
	Permutations int
	Bands        int
}

type MinHashConfig struct {
	ShingleSize  int
	Permutations int
	Bands        int
	RowsPerBand  int
}

type Signature struct {
	ID       string
	Shingles []string
	Values   []uint32
}

type PairComparison struct {
	Left             string
	Right            string
	SharedShingles   int
	UnionShingles    int
	MatchingRows     int
	ExactJaccard     float64
	EstimatedJaccard float64
	Candidate        bool
}

type BandBucket struct {
	Band int
	Key  string
	IDs  []string
}

type SearchResult struct {
	Signatures []Signature
	Buckets    []BandBucket
	Candidates []string
}

func validateOptions(options MinHashOptions) (MinHashConfig, error) {
	shingleSize := options.ShingleSize
	if shingleSize == 0 {
		shingleSize = defaultShingle
	}
	permutations := options.Permutations
	if permutations == 0 {
		permutations = defaultPermutation
	}
	bands := options.Bands
	if bands == 0 {
		bands = defaultBand
	}
	if shingleSize < 2 || shingleSize > maxShingleSize ||
		permutations < 8 || permutations > maxPermutations ||
		bands < 1 || bands > maxBands || permutations%bands != 0 ||
		permutations/bands > maxRowsPerBand {
		return MinHashConfig{}, &MinHashError{
			Code:    BadOptions,
			Message: fmt.Sprintf("use shingles 2-%d, 8-%d permutations, and bands that divide permutations with at most %d rows", maxShingleSize, maxPermutations, maxRowsPerBand),
		}
	}
	return MinHashConfig{ShingleSize: shingleSize, Permutations: permutations, Bands: bands, RowsPerBand: permutations / bands}, nil
}

func validSlug(id string) bool {
	if len(id) < 1 || len(id) > 32 || id[0] < 'a' || id[0] > 'z' {
		return false
	}
	for _, char := range []byte(id[1:]) {
		if !(char >= 'a' && char <= 'z') && !(char >= '0' && char <= '9') && char != '-' {
			return false
		}
	}
	return true
}

func validateText(text string) error {
	if len(text) < 1 || len(text) > maxTextLength {
		return &MinHashError{Code: BadDocumentText, Message: fmt.Sprintf("document text must be 1-%d printable ASCII characters", maxTextLength)}
	}
	for _, char := range []byte(text) {
		if char != '\t' && char != '\n' && char != '\r' && (char < 0x20 || char > 0x7e) {
			return &MinHashError{Code: BadDocumentText, Message: fmt.Sprintf("document text must be 1-%d printable ASCII characters", maxTextLength)}
		}
	}
	return nil
}

func validateDocument(document Document) error {
	if !validSlug(document.ID) {
		return &MinHashError{Code: BadDocumentID, Message: "document ids must be lowercase slugs"}
	}
	return validateText(document.Text)
}

func validateDocuments(documents []Document) error {
	if len(documents) == 0 {
		return &MinHashError{Code: EmptyDocumentList, Message: "at least one document is required"}
	}
	if len(documents) > maxDocuments {
		return &MinHashError{Code: TooManyDocuments, Message: fmt.Sprintf("at most %d documents are accepted", maxDocuments)}
	}
	seen := map[string]bool{}
	for _, document := range documents {
		if err := validateDocument(document); err != nil {
			return err
		}
		if seen[document.ID] {
			return &MinHashError{Code: DuplicateDocumentID, Message: "document ids must be unique"}
		}
		seen[document.ID] = true
	}
	return nil
}

func hash32(value string, seed int) uint32 {
	hash := uint32(0x811c9dc5) ^ (uint32(seed+1) * uint32(0x9e3779b9))
	for _, char := range []byte(value) {
		hash ^= uint32(char)
		hash *= uint32(0x01000193)
	}
	hash ^= hash >> 16
	hash *= uint32(0x85ebca6b)
	hash ^= hash >> 13
	hash *= uint32(0xc2b2ae35)
	hash ^= hash >> 16
	return hash
}

func bandKey(values []uint32, start, length int) string {
	parts := make([]string, length)
	for index := range parts {
		parts[index] = fmt.Sprintf("%d", values[start+index])
	}
	return strings.Join(parts, ":")
}

func Tokenize(text string) ([]string, error) {
	if err := validateText(text); err != nil {
		return nil, err
	}
	lower := strings.ToLower(text)
	words := []string{}
	start := -1
	for index := 0; index <= len(lower); index++ {
		isWord := index < len(lower) && ((lower[index] >= 'a' && lower[index] <= 'z') || (lower[index] >= '0' && lower[index] <= '9'))
		if isWord && start < 0 {
			start = index
		} else if !isWord && start >= 0 {
			words = append(words, lower[start:index])
			start = -1
		}
	}
	return words, nil
}

func MakeShingles(text string, size int) ([]string, error) {
	if size < 2 || size > maxShingleSize {
		return nil, &MinHashError{Code: BadOptions, Message: fmt.Sprintf("shingle size must be a whole number from 2 to %d", maxShingleSize)}
	}
	words, err := Tokenize(text)
	if err != nil {
		return nil, err
	}
	unique := map[string]bool{}
	for index := 0; index+size <= len(words); index++ {
		unique[strings.Join(words[index:index+size], " ")] = true
	}
	shingles := make([]string, 0, len(unique))
	for shingle := range unique {
		shingles = append(shingles, shingle)
	}
	sort.Strings(shingles)
	return shingles, nil
}

func setJaccard(left, right []string) (int, int, float64) {
	rightSet := map[string]bool{}
	leftSet := map[string]bool{}
	for _, shingle := range right {
		rightSet[shingle] = true
	}
	for _, shingle := range left {
		leftSet[shingle] = true
	}
	shared := 0
	for shingle := range leftSet {
		if rightSet[shingle] {
			shared++
		}
	}
	unionSet := map[string]bool{}
	for shingle := range leftSet {
		unionSet[shingle] = true
	}
	for shingle := range rightSet {
		unionSet[shingle] = true
	}
	union := len(unionSet)
	// Signature refuses empty shingle sets, so union is never 0 here; if it were, report 0
	// rather than calling two empty documents identical.
	if union == 0 {
		return shared, union, 0
	}
	return shared, union, float64(shared) / float64(union)
}

// The MinHash step: for each of the deterministic hash rows, keep the smallest hash over the
// document's shingles. Two documents agree on a row with probability equal to their Jaccard
// similarity. A document with no shingles has no minimum, so it is refused.
func minHashValues(shingles []string, permutations int) ([]uint32, error) {
	if len(shingles) == 0 {
		return nil, &MinHashError{Code: NoShingles, Message: "a document needs at least as many words as the shingle size to form one shingle"}
	}
	values := make([]uint32, permutations)
	for seed := range values {
		values[seed] = ^uint32(0)
		for _, shingle := range shingles {
			hash := hash32(shingle, seed)
			if hash < values[seed] {
				values[seed] = hash
			}
		}
	}
	return values, nil
}