Match misspelled company names

"""
Match two company names: normalise, drop the legal form, then Jaro-Winkler.

Rung N0. Deterministic, standard library only, and short enough to read in one
sitting. Two decisions carry the whole result.

First, the legal form is removed rather than compared. « Boulangerie Martin
SARL » and « Boulangerie Martin SAS » are one trading name under two statuses;
leaving SARL and SAS inside the strings would push them apart for a reason
nobody cares about.

Second, Jaro-Winkler rather than a plain edit distance. It rewards a shared
opening, which is how company names actually vary: the head is the brand, the
tail is a form, a city or a scrap of punctuation.
"""

import re
import unicodedata

# Legal forms, French and foreign. This is business knowledge, not example
# data: every real matching job carries a list like this one, and grows it.
LEGAL_FORMS = frozenset(("sarl sas sasu sa eurl sci snc ltd limited plc gmbh "
                         "ag inc llc corp bv nv spa srl").split())

# Winkler looks at the first four characters only, and never gives back more
# than a tenth of the score Jaro withheld.
PREFIX_LENGTH, PREFIX_SCALING = 4, 0.1


def normalise(name: str) -> str:
    """Lowercase, strip accents and punctuation, drop the legal form."""
    folded = unicodedata.normalize("NFKD", name.lower())
    plain = "".join(c for c in folded if not unicodedata.combining(c))
    words = re.findall(r"[a-z0-9]+", plain)
    kept = [w for w in words if w not in LEGAL_FORMS]
    # A name made of nothing but a legal form keeps it. Emptying it would make
    # it match every other emptied name perfectly, which is worse than useless.
    return " ".join(kept or words)


def _jaro(a: str, b: str) -> float:
    """
    Share of characters found on both sides, discounted by their disorder.

    A character counts as found only if its twin sits within half the length
    of the longer name. That window is what separates Jaro from a plain count
    of common letters.
    """
    if a == b:
        return 1.0
    if not a or not b:
        return 0.0
    window = max(max(len(a), len(b)) // 2 - 1, 0)
    hit_a, hit_b = [False] * len(a), [False] * len(b)
    for i, char in enumerate(a):
        for j in range(max(0, i - window), min(len(b), i + window + 1)):
            if not hit_b[j] and b[j] == char:
                hit_a[i] = hit_b[j] = True
                break
    matched_a = [c for c, hit in zip(a, hit_a) if hit]
    matched_b = [c for c, hit in zip(b, hit_b) if hit]
    if not matched_a:
        return 0.0
    m = len(matched_a)
    swaps = sum(x != y for x, y in zip(matched_a, matched_b)) // 2
    return (m / len(a) + m / len(b) + (m - swaps) / m) / 3


def jaro_winkler(a: str, b: str) -> float:
    """Jaro, raised towards 1 in proportion to the shared opening."""
    score = _jaro(a, b)
    prefix = 0
    for x, y in zip(a[:PREFIX_LENGTH], b[:PREFIX_LENGTH]):
        if x != y:
            break
        prefix += 1
    return score + prefix * PREFIX_SCALING * (1 - score)


def similarity(left: str, right: str) -> float:
    """Similarity of two company names, from 0 to 1."""
    return jaro_winkler(normalise(left), normalise(right))

/**
 * Match two company names: normalise, drop the legal form, then Jaro-Winkler.
 *
 * Rung N0. Deterministic, no dependency, and short enough to read in one
 * sitting. Two decisions carry the whole result.
 *
 * First, the legal form is removed rather than compared. « Boulangerie Martin
 * SARL » and « Boulangerie Martin SAS » are one trading name under two
 * statuses; leaving SARL and SAS inside the strings would push them apart for
 * a reason nobody cares about.
 *
 * Second, Jaro-Winkler rather than a plain edit distance. It rewards a shared
 * opening, which is how company names actually vary: the head is the brand,
 * the tail is a form, a city or a scrap of punctuation.
 *
 * Written out in full rather than pulled from a package: the algorithm is
 * thirty lines, and that is the argument of this entry.
 */

// Legal forms, French and foreign. This is business knowledge, not example
// data: every real matching job carries a list like this one, and grows it.
const LEGAL_FORMS = new Set(
  'sarl sas sasu sa eurl sci snc ltd limited plc gmbh ag inc llc corp bv nv spa srl'.split(' '),
);

// Winkler looks at the first four characters only, and never gives back more
// than a tenth of the score Jaro withheld.
const PREFIX_LENGTH = 4;
const PREFIX_SCALING = 0.1;

/** Lowercase, strip accents and punctuation, drop the legal form. */
export function normalise(name) {
  const plain = name.toLowerCase().normalize('NFKD').replace(/\p{M}/gu, '');
  const words = plain.match(/[a-z0-9]+/g) ?? [];
  const kept = words.filter((w) => !LEGAL_FORMS.has(w));
  // A name made of nothing but a legal form keeps it. Emptying it would make
  // it match every other emptied name perfectly, which is worse than useless.
  return (kept.length ? kept : words).join(' ');
}

/**
 * Share of characters found on both sides, discounted by their disorder.
 *
 * A character counts as found only if its twin sits within half the length of
 * the longer name. That window is what separates Jaro from a plain count of
 * common letters.
 */
function jaro(a, b) {
  if (a === b) return 1;
  if (!a.length || !b.length) return 0;
  const window = Math.max(Math.floor(Math.max(a.length, b.length) / 2) - 1, 0);
  const hitA = new Array(a.length).fill(false);
  const hitB = new Array(b.length).fill(false);
  for (let i = 0; i < a.length; i += 1) {
    for (let j = Math.max(0, i - window); j < Math.min(b.length, i + window + 1); j += 1) {
      if (!hitB[j] && b[j] === a[i]) {
        hitA[i] = true;
        hitB[j] = true;
        break;
      }
    }
  }
  const matchedA = [...a].filter((_, i) => hitA[i]);
  const matchedB = [...b].filter((_, j) => hitB[j]);
  const m = matchedA.length;
  if (!m) return 0;
  const swaps = Math.floor(matchedA.filter((c, i) => c !== matchedB[i]).length / 2);
  return (m / a.length + m / b.length + (m - swaps) / m) / 3;
}

/** Jaro, raised towards 1 in proportion to the shared opening. */
export function jaroWinkler(a, b) {
  const score = jaro(a, b);
  let prefix = 0;
  while (prefix < PREFIX_LENGTH && prefix < a.length && prefix < b.length && a[prefix] === b[prefix]) {
    prefix += 1;
  }
  return score + prefix * PREFIX_SCALING * (1 - score);
}

/** Similarity of two company names, from 0 to 1. */
export function similarity(left, right) {
  return jaroWinkler(normalise(left), normalise(right));
}

"""
Match company names on character n-grams weighted by TF-IDF.

Rung N1. N0 compares two strings. This compares one name against a whole
register, and it does so by looking at the fragments a name is made of.

Two changes matter.

A rare fragment now weighs more than a common one. « boulangerie » appears in
half the register and tells you almost nothing; « quiquengrogne » appears once
and settles the question. TF-IDF is exactly that arithmetic, and N0 has no
equivalent: Jaro-Winkler treats every character alike.

And because a name becomes a vector, the search is a matrix product rather
than a loop over every possible pair, which is what makes a whole register
searchable at all.

`char_wb` keeps n-grams inside word boundaries, so a fragment never straddles
two words. « Martin Dubois » and « Dubois Martin » still meet, because word
order costs nothing here — unlike N0, where it costs almost everything.
"""

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer

# Two to four characters: long enough to be a syllable, short enough to
# survive a typo somewhere else in the word.
NGRAM_RANGE = (2, 4)


def build_index(names: list[str]) -> dict:
    """Fit the vocabulary on the register, once, and vectorise it."""
    vectoriser = TfidfVectorizer(analyzer="char_wb", ngram_range=NGRAM_RANGE)
    return {"names": list(names), "vectoriser": vectoriser,
            "matrix": vectoriser.fit_transform(names)}


def match(index: dict, query: str, top_k: int = 3) -> list[tuple[str, float]]:
    """
    The nearest names in the register, best first, with their cosine score.

    TfidfVectorizer returns rows of length one, so the cosine similarity is
    just the dot product. No normalisation to write, and none to get wrong.
    """
    vector = index["vectoriser"].transform([query])
    scores = (index["matrix"] @ vector.T).toarray().ravel()
    # A stable sort, so two names with the same score always come back in
    # register order. A matching run has to be replayable.
    order = np.argsort(-scores, kind="stable")[:top_k]
    return [(index["names"][i], float(scores[i])) for i in order]

/**
 * Match company names on character n-grams weighted by TF-IDF.
 *
 * Rung N1. N0 compares two strings. This compares one name against a whole
 * register, and it does so by looking at the fragments a name is made of.
 *
 * A rare fragment now weighs more than a common one. « boulangerie » appears
 * in half the register and tells you almost nothing; « quiquengrogne »
 * appears once and settles the question. TF-IDF is exactly that arithmetic,
 * and N0 has no equivalent: Jaro-Winkler treats every character alike.
 *
 * The n-grams are taken inside word boundaries, so a fragment never straddles
 * two words. « Martin Dubois » and « Dubois Martin » still meet, because word
 * order costs nothing here — unlike N0, where it costs almost everything.
 *
 * This is the same arithmetic scikit-learn performs, written out: raw counts,
 * an inverse document frequency smoothed on both sides, then rows brought to
 * length one so a cosine is a dot product.
 */

// Two to four characters: long enough to be a syllable, short enough to
// survive a typo somewhere else in the word.
const MIN_N = 2;
const MAX_N = 4;

/** Every n-gram of every word, each word padded with a space at both ends. */
function ngrams(text) {
  const out = [];
  for (const word of text.toLowerCase().split(/\s+/).filter(Boolean)) {
    const padded = ` ${word} `;
    for (let n = MIN_N; n <= MAX_N; n += 1) {
      // A word shorter than the window is counted once, whole, and no wider
      // window can tell you anything more about it.
      if (padded.length <= n) {
        out.push(padded);
        break;
      }
      for (let i = 0; i + n <= padded.length; i += 1) out.push(padded.slice(i, i + n));
    }
  }
  return out;
}

/** Counts times idf, brought to length one. Unknown n-grams are dropped. */
function weigh(grams, idf) {
  const counts = new Map();
  for (const g of grams) if (idf.has(g)) counts.set(g, (counts.get(g) ?? 0) + 1);
  const vector = new Map();
  let sum = 0;
  // Alphabetical order, so a name and the same name with its words swapped
  // add their weights in the same order and land on the very same float.
  for (const g of [...counts.keys()].sort()) {
    const tf = counts.get(g);
    const w = tf * idf.get(g);
    vector.set(g, w);
    sum += w * w;
  }
  const norm = Math.sqrt(sum);
  if (norm) for (const [g, w] of vector) vector.set(g, w / norm);
  return vector;
}

/** Fit the vocabulary on the register, once, and vectorise it. */
export function buildIndex(names) {
  const grams = names.map(ngrams);
  const df = new Map();
  for (const g of grams) for (const gram of new Set(g)) df.set(gram, (df.get(gram) ?? 0) + 1);
  const idf = new Map();
  for (const [gram, d] of df) idf.set(gram, Math.log((1 + names.length) / (1 + d)) + 1);
  return { names: [...names], idf, vectors: grams.map((g) => weigh(g, idf)) };
}

/** The nearest names in the register, best first, with their cosine score. */
export function match(index, query, topK = 3) {
  const vector = weigh(ngrams(query), index.idf);
  const scored = index.names.map((name, i) => {
    let dot = 0;
    for (const [gram, w] of vector) dot += w * (index.vectors[i].get(gram) ?? 0);
    return [name, dot];
  });
  // A stable sort, so two names with the same score always come back in
  // register order. A matching run has to be replayable.
  return scored.sort((a, b) => b[1] - a[1]).slice(0, topK);
}

"""
Match company names by meaning, with a self-hosted encoder.

Rung N2. N0 and N1 both compare characters, so both miss the pair this entry
keeps coming back to: an acronym and the name it stands for share almost no
letters. An encoder maps each name to a vector by meaning rather than by
spelling, which is the only way that pair can ever meet.

What it costs: a model file to ship and keep in sync, a warm process to hold
it, and a score you cannot explain to the colleague who asks why two names
were merged. The register also has to be re-encoded whenever the model is
upgraded, and the scores of the old version are not comparable to the new.

Note what is not here: no legal form is stripped and nothing is lowercased.
The encoder is supposed to handle that itself. Whether it does is exactly the
part a local double cannot prove — see the test.
"""

from __future__ import annotations

MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"


def load_encoder(name: str = MODEL_NAME):
    """The real encoder: fetched once, then held in memory and run locally."""
    from sentence_transformers import SentenceTransformer  # pragma: no cover

    return SentenceTransformer(name)


def build_index(names: list[str], encoder=None) -> dict:
    """
    Encode the whole register once, and keep the encoder for the queries.

    `encoder` is injected so this can be tested without loading a model. Left
    alone, it is the real one above.
    """
    encoder = load_encoder() if encoder is None else encoder
    names = list(names)
    return {"names": names, "encoder": encoder,
            "vectors": [_unit(v) for v in encoder.encode(names)]}


def match(index: dict, query: str, top_k: int = 3) -> list[tuple[str, float]]:
    """The nearest names by meaning, best first, with their cosine score."""
    vector = _unit(index["encoder"].encode([query])[0])
    scored = [(name, _dot(known, vector))
              for name, known in zip(index["names"], index["vectors"])]
    # A stable sort, so two names with the same score always come back in
    # register order. A matching run has to be replayable.
    scored.sort(key=lambda pair: -pair[1])
    return scored[:top_k]


def _unit(vector) -> list[float]:
    """Cosine similarity is a dot product once both sides have length one."""
    values = [float(v) for v in vector]
    norm = sum(v * v for v in values) ** 0.5
    return [v / norm for v in values] if norm else values


def _dot(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b))

/**
 * Match company names by meaning, with a self-hosted encoder.
 *
 * Rung N2. N0 and N1 both compare characters, so both miss the pair this
 * entry keeps coming back to: an acronym and the name it stands for share
 * almost no letters. An encoder maps each name to a vector by meaning rather
 * than by spelling, which is the only way that pair can ever meet.
 *
 * What it costs: a model file to ship and keep in sync, a warm process to
 * hold it, and a score you cannot explain to the colleague who asks why two
 * names were merged. The register also has to be re-encoded whenever the
 * model is upgraded, and the old scores are not comparable to the new.
 *
 * Note what is not here: no legal form is stripped and nothing is lowercased.
 * The encoder is supposed to handle that itself. Whether it does is exactly
 * the part a local double cannot prove — see the test.
 */

export const MODEL_NAME = 'Xenova/paraphrase-multilingual-MiniLM-L12-v2';

/** The real encoder: fetched once, then held in memory and run locally. */
export async function loadEncoder(name = MODEL_NAME) {
  const { pipeline } = await import('@xenova/transformers');
  const extract = await pipeline('feature-extraction', name);
  return { encode: async (texts) => (await extract(texts, { pooling: 'mean' })).tolist() };
}

/**
 * Encode the whole register once, and keep the encoder for the queries.
 *
 * `encoder` is injected so this can be tested without loading a model. Left
 * alone, it is the real one above.
 */
export async function buildIndex(names, encoder) {
  const model = encoder ?? (await loadEncoder());
  const kept = [...names];
  return { names: kept, encoder: model, vectors: (await model.encode(kept)).map(unit) };
}

/** The nearest names by meaning, best first, with their cosine score. */
export async function match(index, query, topK = 3) {
  const vector = unit((await index.encoder.encode([query]))[0]);
  const scored = index.names.map((name, i) => [name, dot(index.vectors[i], vector)]);
  // A stable sort, so two names with the same score always come back in
  // register order. A matching run has to be replayable.
  return scored.sort((a, b) => b[1] - a[1]).slice(0, topK);
}

/** Cosine similarity is a dot product once both sides have length one. */
function unit(vector) {
  let sum = 0;
  for (const v of vector) sum += v * v;
  const norm = Math.sqrt(sum);
  return norm ? vector.map((v) => v / norm) : [...vector];
}

function dot(a, b) {
  let total = 0;
  for (let i = 0; i < a.length; i += 1) total += a[i] * b[i];
  return total;
}