Count-min sketch.

export type Step = {
	/** read: look at a counter. bump: raise it. keep: leave a counter already high enough. */
	kind: 'read' | 'bump' | 'keep';
	row: number;
	column: number;
	before: number;
	after: number;
};

const MAX_COUNTER = 0xffffffff;
const encoder = new TextEncoder();

// FNV-1a over the UTF-8 bytes, then a final mix, so Go computes the same 32-bit hash.
function hash(text: string): number {
	let h = 0x811c9dc5;
	for (const byte of encoder.encode(text)) {
		h ^= byte;
		h = Math.imul(h, 0x01000193) >>> 0;
	}
	return mix(h);
}

function mix(value: number): number {
	let h = value >>> 0;
	h ^= h >>> 16;
	h = Math.imul(h, 0x85ebca6b) >>> 0;
	h ^= h >>> 13;
	h = Math.imul(h, 0xc2b2ae35) >>> 0;
	h ^= h >>> 16;
	return h >>> 0;
}

// A count-min sketch estimates how often each item appeared using a fixed grid of counters:
// `depth` rows, each `width` columns wide. Adding an item hashes it to one column in every row
// and raises those counters. Its estimate is the smallest of them. Other items can share a
// counter and push it up, but nothing ever pushes one down, so an estimate is never too low.
export class CountMinSketch {
	readonly width: number;
	readonly depth: number;
	/** Raise only the counters that are below the new estimate, which overcounts less. */
	readonly conservative: boolean;
	#counters: Uint32Array;
	#total = 0;

	constructor({
		width,
		depth,
		conservative = false
	}: {
		width: number;
		depth: number;
		conservative?: boolean;
	}) {
		if (!Number.isInteger(width) || width < 1 || width > 1_000_000)
			throw new RangeError('width must be a whole number from 1 to 1000000');
		if (!Number.isInteger(depth) || depth < 1 || depth > 16)
			throw new RangeError('depth must be a whole number from 1 to 16');
		this.width = width;
		this.depth = depth;
		this.conservative = conservative;
		this.#counters = new Uint32Array(width * depth);
	}

	/** The column an item uses in each row: two hashes combined, one per row. */
	columns(item: string): number[] {
		const first = hash(item);
		const step = (mix(first ^ 0x9e3779b9) | 1) >>> 0;
		return Array.from(
			{ length: this.depth },
			(_, row) => ((first + Math.imul(row, step)) >>> 0) % this.width
		);
	}

	// Returns the item's estimate after adding. Counters stop at 4,294,967,295.
	add(item: string, count = 1, steps?: Step[]): number {
		if (!Number.isInteger(count) || count < 1 || count > 1_000_000)
			throw new RangeError('count must be a whole number from 1 to 1000000');
		const columns = this.columns(item);
		let estimate = MAX_COUNTER;
		if (this.conservative) {
			// Every counter for this item is at least its true count, so raising the ones below
			// the new estimate is enough.
			estimate = Math.min(MAX_COUNTER, this.#read(columns, steps) + count);
			columns.forEach((column, row) => {
				const before = this.#counters[row * this.width + column];
				const after = Math.max(before, estimate);
				this.#counters[row * this.width + column] = after;
				steps?.push({ kind: after === before ? 'keep' : 'bump', row, column, before, after });
			});
		} else {
			columns.forEach((column, row) => {
				const before = this.#counters[row * this.width + column];
				const after = Math.min(MAX_COUNTER, before + count);
				this.#counters[row * this.width + column] = after;
				steps?.push({ kind: 'bump', row, column, before, after });
				estimate = Math.min(estimate, after);
			});
		}
		this.#total += count;
		return estimate;
	}

	/** The smallest counter the item hashes to: never below its true count. */
	estimate(item: string, steps?: Step[]): number {
		return this.#read(this.columns(item), steps);
	}

	/** Every counter, row by row, for drawing. */
	rows(): number[][] {
		return Array.from({ length: this.depth }, (_, row) =>
			Array.from(this.#counters.subarray(row * this.width, (row + 1) * this.width))
		);
	}

	/** The sum of every count added. */
	get total(): number {
		return this.#total;
	}

	#read(columns: number[], steps?: Step[]): number {
		let smallest = MAX_COUNTER;
		columns.forEach((column, row) => {
			const value = this.#counters[row * this.width + column];
			steps?.push({ kind: 'read', row, column, before: value, after: value });
			smallest = Math.min(smallest, value);
		});
		return smallest;
	}
}

// Step is one counter touched. Kind is "read" (look at a counter), "bump" (raise it), or
// "keep" (leave a counter already high enough). Before and After are its values.
type Step struct {
	Kind   string `json:"kind"`
	Row    int    `json:"row"`
	Column int    `json:"column"`
	Before uint32 `json:"before"`
	After  uint32 `json:"after"`
	Tag    string `json:"tag,omitempty"`
}

const maxCounter = ^uint32(0)

// hash is FNV-1a over the UTF-8 bytes, then a final mix, so TypeScript computes the same
// 32-bit hash.
func hash(text string) uint32 {
	h := uint32(0x811c9dc5)
	for i := 0; i < len(text); i++ {
		h ^= uint32(text[i])
		h *= 0x01000193
	}
	return mix(h)
}

func mix(h uint32) uint32 {
	h ^= h >> 16
	h *= 0x85ebca6b
	h ^= h >> 13
	h *= 0xc2b2ae35
	h ^= h >> 16
	return h
}

func record(steps *[]Step, step Step) {
	if steps != nil {
		*steps = append(*steps, step)
	}
}

// CountMinSketch estimates how often each item appeared using a fixed grid of counters:
// depth rows, each width columns wide. Adding an item hashes it to one column in every row
// and raises those counters. Its estimate is the smallest of them. Other items can share a
// counter and push it up, but nothing ever pushes one down, so an estimate is never too low.
type CountMinSketch struct {
	width, depth int
	// conservative raises only the counters below the new estimate, which overcounts less.
	conservative bool
	counters     []uint32
	total        uint64
}

func NewCountMinSketch(width, depth int, conservative bool) (*CountMinSketch, error) {
	if width < 1 || width > 1_000_000 {
		return nil, errors.New("width must be a whole number from 1 to 1000000")
	}
	if depth < 1 || depth > 16 {
		return nil, errors.New("depth must be a whole number from 1 to 16")
	}
	return &CountMinSketch{width: width, depth: depth, conservative: conservative, counters: make([]uint32, width*depth)}, nil
}

// Columns returns the column an item uses in each row: two hashes combined, one per row.
func (s *CountMinSketch) Columns(item string) []int {
	first := hash(item)
	step := mix(first^0x9e3779b9) | 1
	columns := make([]int, s.depth)
	for row := range columns {
		columns[row] = int((first + uint32(row)*step) % uint32(s.width))
	}
	return columns
}

// Add returns the item's estimate after adding. Counters stop at 4,294,967,295.
func (s *CountMinSketch) Add(item string, count int, steps *[]Step) (uint32, error) {
	if count < 1 || count > 1_000_000 {
		return 0, errors.New("count must be a whole number from 1 to 1000000")
	}
	columns := s.Columns(item)
	estimate := maxCounter
	if s.conservative {
		// Every counter for this item is at least its true count, so raising the ones below
		// the new estimate is enough.
		estimate = addCapped(s.read(columns, steps), count)
		for row, column := range columns {
			before := s.counters[row*s.width+column]
			after := max(before, estimate)
			s.counters[row*s.width+column] = after
			kind := "bump"
			if after == before {
				kind = "keep"
			}
			record(steps, Step{Kind: kind, Row: row, Column: column, Before: before, After: after})
		}
	} else {
		for row, column := range columns {
			before := s.counters[row*s.width+column]
			after := addCapped(before, count)
			s.counters[row*s.width+column] = after
			record(steps, Step{Kind: "bump", Row: row, Column: column, Before: before, After: after})
			estimate = min(estimate, after)
		}
	}
	s.total += uint64(count)
	return estimate, nil
}

// Estimate returns the smallest counter the item hashes to: never below its true count.
func (s *CountMinSketch) Estimate(item string, steps *[]Step) uint32 {
	return s.read(s.Columns(item), steps)
}

// Rows returns every counter, row by row, for drawing.
func (s *CountMinSketch) Rows() [][]uint32 {
	rows := make([][]uint32, s.depth)
	for row := range rows {
		rows[row] = slices.Clone(s.counters[row*s.width : (row+1)*s.width])
	}
	return rows
}

// Total is the sum of every count added.
func (s *CountMinSketch) Total() uint64 { return s.total }

func (s *CountMinSketch) read(columns []int, steps *[]Step) uint32 {
	smallest := maxCounter
	for row, column := range columns {
		value := s.counters[row*s.width+column]
		record(steps, Step{Kind: "read", Row: row, Column: column, Before: value, After: value})
		smallest = min(smallest, value)
	}
	return smallest
}

func addCapped(value uint32, count int) uint32 {
	if uint64(value)+uint64(count) > uint64(maxCounter) {
		return maxCounter
	}
	return value + uint32(count)
}