Extract dates from free text

"""
Extract dates from text: one regular expression per format, then a real
calendar check.

Rung N0. Deterministic, standard library only, and the whole of it fits on a
screen.

The regular expression is the easy half. It finds three digits groups and
knows nothing else: 31/02/2024 matches it perfectly, and so does 29/02/2023.
The second half is what makes the difference, and it is one line long, because
`datetime.date` already owns the calendar — month lengths, leap years, and the
century rule that makes 1900 a common year.

What the digits cannot say is whether 03/04/2024 is 3 April or 4 March. No
amount of pattern matching settles that, so the caller settles it once.
"""

import re
import unicodedata
from datetime import date

# Month names in French and English, the two a French-language document mixes.
MONTHS = {name: number for number, names in enumerate(
    ("janvier january", "fevrier february", "mars march", "avril april", "mai may",
     "juin june", "juillet july", "aout august", "septembre september",
     "octobre october", "novembre november", "decembre december"), 1) for name in names.split()}

# "3 avril 2024", "1er mars 2024".
TEXTUAL = re.compile(r"(?<!\d)(\d{1,2})(?:er)?\s+([^\W\d_]+)\s+(\d{4})(?!\d)")

# "12/03/2024", "12.03.24", "12-03-2024", and the ISO "2024-03-12". Years are
# two or four digits, never three: that is what keeps "1.2.3" out.
NUMERIC = re.compile(r"(?<!\d)(\d{1,2}|\d{4})[/.-](\d{1,2})[/.-](\d{2}|\d{4})(?!\d)")


def _fold(word: str) -> str:
    """Drop accents, so that « février » and « fevrier » reach the same entry."""
    decomposed = unicodedata.normalize("NFKD", word.lower())
    return "".join(c for c in decomposed if not unicodedata.combining(c))


def _to_date(year: int, month: int, day: int) -> date | None:
    """
    Real calendar validation, and the point of this rung.

    A regular expression accepts 31 February; `date` does not. Leap years come
    with it, century rule included.
    """
    try:
        return date(year, month, day)
    except ValueError:
        return None


def _full_year(year: int) -> int:
    """Two-digit years on the usual pivot: 69 reads as 2069, 70 as 1970."""
    return year if year >= 100 else year + (2000 if year < 70 else 1900)


def _read_textual(match: re.Match) -> date | None:
    month = MONTHS.get(_fold(match.group(2)))
    return _to_date(int(match.group(3)), month, int(match.group(1))) if month else None


def _read_numeric(match: re.Match, day_first: bool) -> date | None:
    first, second, third = (int(group) for group in match.groups())
    if len(match.group(1)) == 4:  # ISO order, whatever the local habit is
        return _to_date(first, second, third)
    day, month = (first, second) if day_first else (second, first)
    return _to_date(_full_year(third), month, day)


def extract_dates(text: str, day_first: bool = True) -> list[tuple[str, date]]:
    """
    Return every real date in `text`, as (what was written, what it means).

    `day_first` says how to read 03/04/2024. The digits cannot say, so the
    caller decides once, for a whole document, and lives with it.
    """
    found = []
    for pattern in (TEXTUAL, NUMERIC):
        for match in pattern.finditer(text):
            # A match that overlaps an accepted one is a second reading of the
            # same characters, not a second date.
            if any(start < match.end() and match.start() < end for start, end, _, _ in found):
                continue
            value = _read_textual(match) if pattern is TEXTUAL else _read_numeric(match, day_first)
            if value is not None:
                found.append((match.start(), match.end(), match.group(0), value))
    found.sort(key=lambda item: item[0])
    return [(written, value) for _, _, written, value in found]

/**
 * Extract dates from text: one regular expression per format, then a real
 * calendar check.
 *
 * Rung N0. Deterministic, no dependency, and the whole of it fits on a screen.
 *
 * The regular expression is the easy half. It finds three groups of digits and
 * knows nothing else: 31/02/2024 matches it perfectly, and so does 29/02/2023.
 *
 * The second half is what makes the difference, and in JavaScript it needs
 * care: `new Date(2024, 1, 31)` does not fail, it quietly rolls over to
 * 2 March. The only honest check is to build the date and read its parts back.
 *
 * What the digits cannot say is whether 03/04/2024 is 3 April or 4 March. No
 * amount of pattern matching settles that, so the caller settles it once.
 */

// Month names in French and English, the two a French-language document mixes.
const MONTHS = new Map(['janvier january', 'fevrier february', 'mars march', 'avril april',
  'mai may', 'juin june', 'juillet july', 'aout august', 'septembre september',
  'octobre october', 'novembre november', 'decembre december']
  .flatMap((names, index) => names.split(' ').map((name) => [name, index + 1])));

// "3 avril 2024", "1er mars 2024".
const TEXTUAL = /(?<!\d)(\d{1,2})(?:er)?\s+(\p{L}+)\s+(\d{4})(?!\d)/gu;

// "12/03/2024", "12.03.24", "12-03-2024", and the ISO "2024-03-12". Years are
// two or four digits, never three: that is what keeps "1.2.3" out.
const NUMERIC = /(?<!\d)(\d{1,2}|\d{4})[/.-](\d{1,2})[/.-](\d{2}|\d{4})(?!\d)/g;

/** Drop accents, so that "février" and "fevrier" reach the same entry. */
function fold(word) {
  return word.toLowerCase().normalize('NFKD').replace(/\p{M}/gu, '');
}

/**
 * Real calendar validation, and the point of this rung.
 *
 * Building the date is not enough, because JavaScript rolls an impossible one
 * over instead of refusing it. Reading the parts back is the check.
 */
function toDate(year, month, day) {
  const value = new Date(Date.UTC(year, month - 1, day));
  const real = value.getUTCFullYear() === year && value.getUTCMonth() === month - 1 && value.getUTCDate() === day;
  return real ? value : null;
}

/** Two-digit years on the usual pivot: 69 reads as 2069, 70 as 1970. */
function fullYear(year) {
  return year >= 100 ? year : year + (year < 70 ? 2000 : 1900);
}

function readTextual(match) {
  const month = MONTHS.get(fold(match[2]));
  return month ? toDate(Number(match[3]), month, Number(match[1])) : null;
}

function readNumeric(match, dayFirst) {
  const [first, second, third] = match.slice(1).map(Number);
  if (match[1].length === 4) return toDate(first, second, third); // ISO order
  const [day, month] = dayFirst ? [first, second] : [second, first];
  return toDate(fullYear(third), month, day);
}

/**
 * Return every real date in `text`, as { text: what was written, date }.
 *
 * `dayFirst` says how to read 03/04/2024. The digits cannot say, so the caller
 * decides once, for a whole document, and lives with it.
 */
export function extractDates(text, dayFirst = true) {
  const found = [];
  for (const pattern of [TEXTUAL, NUMERIC]) {
    for (const match of text.matchAll(pattern)) {
      const [start, end] = [match.index, match.index + match[0].length];
      // A match that overlaps an accepted one is a second reading of the same
      // characters, not a second date.
      if (found.some((item) => item.start < end && start < item.end)) continue;
      const date = pattern === TEXTUAL ? readTextual(match) : readNumeric(match, dayFirst);
      if (date) found.push({ start, end, text: match[0], date });
    }
  }
  found.sort((a, b) => a.start - b.start);
  return found.map(({ text: written, date }) => ({ text: written, date }));
}

"""
Decide whether 03/04/2024 is 3 April or 4 March, with a light classifier.

Rung N1. Rules still find the candidates, exactly as N0 does: a date is a
shape, and a shape is what regular expressions are for. What a rule cannot do
is read 03/04/2024, because nothing in those digits says which field is the
day. N0 answers by asking the caller to pick one convention for a whole
document, which is wrong the moment a document quotes a supplier from abroad.

The convention is not in the digits, it is in the prose around them. That is a
classification problem, and a few hundred labelled sentences are enough for it.
The model is small enough to keep beside the code, and the rules still settle
every case they can settle on their own.
"""

import re
from datetime import date

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

CANDIDATE = re.compile(r"(?<!\d)(\d{1,2})[/.-](\d{1,2})[/.-](\d{4})(?!\d)")
WINDOW = 40  # characters of context kept on each side of a candidate


def context(text: str, span: tuple[int, int]) -> str:
    """
    The words around a date, with every digit removed.

    Removing the digits is what stops the classifier memorising the dates of
    the training set instead of learning the habits of the prose around them.
    """
    start, end = span
    around = text[max(0, start - WINDOW):start] + " " + text[end:end + WINDOW]
    return re.sub(r"\d+", " ", around).lower()


def train(texts: list[str], labels: list[int]):
    """`labels` is 1 when the text writes the day first, 0 when the month comes first."""
    model = make_pipeline(
        TfidfVectorizer(ngram_range=(1, 2), min_df=1),
        LogisticRegression(class_weight="balanced", max_iter=1000),
    )
    model.fit([context(t, CANDIDATE.search(t).span()) for t in texts], labels)
    return model


def _to_date(year: int, month: int, day: int) -> date | None:
    """Real calendar validation, kept from N0: a regular expression accepts 31 February."""
    try:
        return date(year, month, day)
    except ValueError:
        return None


def extract_dates(model, text: str) -> list[tuple[str, date]]:
    """Return every real date in `text`, reading each one the way its context suggests."""
    found = []
    for match in CANDIDATE.finditer(text):
        first, second, year = (int(group) for group in match.groups())
        if second > 12:  # the second field cannot be a month, so it is a day
            day_first = False
        elif first > 12:  # symmetrically, the first field can only be a day
            day_first = True
        else:  # nothing in the digits decides it, so ask the prose
            day_first = model.predict_proba([context(text, match.span())])[0][1] >= 0.5
        day, month = (first, second) if day_first else (second, first)
        value = _to_date(year, month, day)
        if value is not None:
            found.append((match.group(0), value))
    return found

/**
 * Decide whether 03/04/2024 is 3 April or 4 March, with a light classifier.
 *
 * Rung N1. Rules still find the candidates, exactly as N0 does: a date is a
 * shape, and a shape is what regular expressions are for. What a rule cannot
 * do is read 03/04/2024, because nothing in those digits says which field is
 * the day. N0 answers by asking the caller to pick one convention for a whole
 * document, which is wrong the moment a document quotes a supplier from abroad.
 *
 * The convention is not in the digits, it is in the prose around them. That is
 * a classification problem, and a few hundred labelled sentences are enough.
 *
 * Logistic regression on word counts is written out here rather than pulled
 * from a library, because it is short enough to read. That is the whole
 * argument of this rung.
 */

const CANDIDATE = /(?<!\d)(\d{1,2})[/.-](\d{1,2})[/.-](\d{4})(?!\d)/g;
const WINDOW = 40; // characters of context kept on each side of a candidate

/**
 * The words around a date, with every digit removed.
 *
 * Removing the digits is what stops the classifier memorising the dates of the
 * training set instead of learning the habits of the prose around them.
 */
export function context(text, start, end) {
  const around = `${text.slice(Math.max(0, start - WINDOW), start)} ${text.slice(end, end + WINDOW)}`;
  return around.toLowerCase().replace(/\d+/g, ' ');
}

/** Word counts of one context, normalised so that long sentences do not shout. */
function features(text) {
  const counts = new Map();
  for (const word of text.match(/\p{L}+/gu) ?? []) counts.set(word, (counts.get(word) ?? 0) + 1);
  const norm = Math.hypot(...counts.values());
  return new Map([...counts].map(([word, count]) => [word, count / norm]));
}

/** Probability that this context comes from a document writing the day first. */
function probability(model, row) {
  let z = model.bias;
  for (const [word, value] of row) z += (model.weights.get(word) ?? 0) * value;
  return 1 / (1 + Math.exp(-z));
}

function firstContext(text) {
  const [match] = text.matchAll(CANDIDATE);
  return match ? context(text, match.index, match.index + match[0].length) : '';
}

/** `labels` is 1 when the text writes the day first, 0 when the month comes first. */
export function train(texts, labels, { epochs = 300, rate = 0.5 } = {}) {
  const rows = texts.map((text) => features(firstContext(text)));
  const model = { weights: new Map(), bias: 0 };
  for (let epoch = 0; epoch < epochs; epoch += 1) {
    rows.forEach((row, i) => {
      const error = probability(model, row) - labels[i];
      for (const [word, value] of row) {
        model.weights.set(word, (model.weights.get(word) ?? 0) - rate * error * value);
      }
      model.bias -= rate * error;
    });
  }
  return model;
}

/** Real calendar validation, kept from N0: a regular expression accepts 31 February. */
function toDate(year, month, day) {
  const value = new Date(Date.UTC(year, month - 1, day));
  const real = value.getUTCFullYear() === year && value.getUTCMonth() === month - 1 && value.getUTCDate() === day;
  return real ? value : null;
}

/** Return every real date in `text`, reading each one the way its context suggests. */
export function extractDates(model, text) {
  const found = [];
  for (const match of text.matchAll(CANDIDATE)) {
    const [first, second, year] = match.slice(1).map(Number);
    const [start, end] = [match.index, match.index + match[0].length];
    // The rules settle what they can; the classifier only sees what is left.
    const dayFirst = second > 12 ? false : first > 12 ? true
      : probability(model, features(context(text, start, end))) >= 0.5;
    const date = toDate(year, dayFirst ? second : first, dayFirst ? first : second);
    if (date) found.push({ text: match[0], date });
  }
  return found;
}

"""
Extract dates by asking a general-purpose model.

Rung N3. This is the option people reach for first, and it is the only one on
this entry that reads « jeudi prochain ». That is a real capability, and it is
why the rung is here.

Note what the code has to do that N0 did not: pass a reference date, because
the model has no idea what day it is; cap the input size; retry on failure;
parse an answer that is only probably valid JSON; and check the calendar
itself, because a model will answer 2024-02-31 in flawless JSON without
blinking. That plumbing is the real cost of the rung, and it is the part the
tests have to cover, because the model itself is not testable.
"""

import json
from datetime import date

PROMPT = (
    "Find every date mentioned in the text below. Answer with JSON only: a list\n"
    "of objects with keys `text` and `date`, where `text` is the words as written\n"
    "and `date` is the day in ISO format, YYYY-MM-DD. Resolve relative dates such\n"
    "as « next Thursday » against today, which is {today}. If there is no date,\n"
    "answer with an empty list.\n\nText:\n{text}"
)

MAX_CHARACTERS = 8000


class ExtractionUnavailable(Exception):
    """The provider could not be reached, or answered something unusable."""


def extract_dates(text: str, client=None, *, today: date | None = None, attempts: int = 3):
    """
    Return every date the model reports, as (what was written, what it means).

    `client` is injected so this function can be tested without a network call.
    In production it defaults to a real provider client.
    """
    if client is None:  # pragma: no cover - needs a key and a network
        from openai import OpenAI

        client = OpenAI()

    # A model charges by the token. Refusing oversized input is not an
    # optimisation, it is a cost control.
    if len(text) > MAX_CHARACTERS:
        raise ValueError(f"text longer than {MAX_CHARACTERS} characters")

    found = []
    for item in _ask(client, text, today or date.today(), attempts):
        try:
            # The same calendar check as N0, on the model's answer this time.
            value = date.fromisoformat(item["date"])
        except (TypeError, KeyError, ValueError):
            continue  # a day that does not exist is not a date, however fluent
        found.append((item.get("text", ""), value))
    return found


def _ask(client, text: str, today: date, attempts: int) -> list[dict]:
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            answer = client.complete(
                prompt=PROMPT.format(text=text, today=today.isoformat()),
                # Temperature zero: a date that changes between two identical
                # calls cannot be reviewed.
                temperature=0,
            )
            parsed = json.loads(answer)
            if isinstance(parsed, list):
                return parsed
            last_error = ValueError("the model answered something that is not a list")
        except Exception as error:  # noqa: BLE001 - any provider failure is retried
            last_error = error
    raise ExtractionUnavailable(str(last_error))

/**
 * Extract dates by asking a general-purpose model.
 *
 * Rung N3. This is the option people reach for first, and it is the only one
 * on this entry that reads "jeudi prochain". That is a real capability, and it
 * is why the rung is here.
 *
 * Note what the code has to do that N0 did not: pass a reference date, because
 * the model has no idea what day it is; cap the input size; retry on failure;
 * parse an answer that is only probably valid JSON; and check the calendar
 * itself, because a model will answer 2024-02-31 in flawless JSON without
 * blinking. That plumbing is the real cost of the rung, and it is the part the
 * tests have to cover, because the model itself is not testable.
 */

const PROMPT = [
  'Find every date mentioned in the text below. Answer with JSON only: a list',
  'of objects with keys `text` and `date`, where `text` is the words as written',
  'and `date` is the day in ISO format, YYYY-MM-DD. Resolve relative dates such',
  'as "next Thursday" against today, which is {today}. If there is no date,',
  'answer with an empty list.', '', 'Text:',
].join('\n');

export const MAX_CHARACTERS = 8000;

export class ExtractionUnavailable extends Error {}

/** The calendar check of N0, applied to the model's answer this time. */
function parseIsoDay(value) {
  const match = /^(\d{4})-(\d{2})-(\d{2})$/.exec(String(value));
  if (!match) return null;
  const [year, month, day] = match.slice(1).map(Number);
  const date = new Date(Date.UTC(year, month - 1, day));
  return date.getUTCMonth() === month - 1 && date.getUTCDate() === day ? date : null;
}

/**
 * Return every date the model reports, as { text: what was written, date }.
 *
 * @param {string} text
 * @param {object} options
 * @param {{complete: Function}} [options.client] injected so this can be
 *   tested without a network call; defaults to a real provider client
 * @param {Date} [options.today] the day relative dates are resolved against
 * @param {number} [options.attempts]
 */
export async function extractDates(text, { client, today = new Date(), attempts = 3 } = {}) {
  if (!client) {
    // Needs a key and a network, so it is never reached in the tests.
    const { OpenAI } = await import('openai');
    client = new OpenAI();
  }

  // A model charges by the token. Refusing oversized input is not an
  // optimisation, it is a cost control.
  if (text.length > MAX_CHARACTERS) throw new RangeError(`text longer than ${MAX_CHARACTERS} characters`);

  const items = await ask(client, text, today, attempts);
  // A day that does not exist is not a date, however fluent the answer.
  return items.map((item) => ({ text: item?.text ?? '', date: parseIsoDay(item?.date) }))
    .filter((item) => item.date);
}

async function ask(client, text, today, attempts) {
  let lastError;
  for (let i = 0; i < attempts; i += 1) {
    try {
      // Temperature zero: a date that changes between two identical calls
      // cannot be reviewed.
      const prompt = `${PROMPT.replace('{today}', today.toISOString().slice(0, 10))}\n${text}`;
      const answer = await client.complete({ prompt, temperature: 0 });
      const parsed = JSON.parse(answer);
      if (Array.isArray(parsed)) return parsed;
      lastError = new Error('the model answered something that is not a list');
    } catch (error) {
      lastError = error;
    }
  }
  throw new ExtractionUnavailable(String(lastError));
}