ulfbot/src/markov.ts

/**
 * Markov chain for generating sentences from learned text.
 * Uses trigrams (word triplets) for better grammatical coherence.
 */

interface ChainData {
  // Map of "word1 word2" -> possible next words with frequency
  transitions: Map<string, Map<string, number>>;
  // All sentence start patterns
  starts: Map<string, number>;
}

export class MarkovChain {
  private order: number;
  private chain: ChainData;

  constructor(order: number = 1) {
    this.order = order;
    this.chain = {
      transitions: new Map(),
      starts: new Map(),
    };
  }

  /**
   * Tokenize text into words.
   * Keeps special characters and emojis as requested.
   */
  private tokenize(text: string): string[][] {
    // We treat the whole message as one sequence to keep it simple and creative
    const words = text
      .trim()
      .split(/\s+/)
      .filter(w => w.length > 0);

    return words.length > 0 ? [words] : [];
  }

  /**
   * Learn from a text, adding it to the chain.
   * Returns the new transitions and starts for efficient DB updates.
   */
  learn(text: string): { transitions: Array<{ key: string; next: string }>; starts: string[] } {
    const sentences = this.tokenize(text);
    const learned = {
      transitions: [] as Array<{ key: string; next: string }>,
      starts: [] as string[],
    };

    for (const words of sentences) {
      if (words.length < this.order + 1) continue;

      // Mark sentence start
      const startKey = words.slice(0, this.order).join(' ');
      this.chain.starts.set(startKey, (this.chain.starts.get(startKey) || 0) + 1);
      learned.starts.push(startKey);

      // Build transitions
      for (let i = 0; i < words.length - this.order; i++) {
        const key = words.slice(i, i + this.order).join(' ');
        const next = words[i + this.order];

        if (!this.chain.transitions.has(key)) {
          this.chain.transitions.set(key, new Map());
        }

        const transitions = this.chain.transitions.get(key)!;
        transitions.set(next, (transitions.get(next) || 0) + 1);
        learned.transitions.push({ key, next });
      }
    }

    return learned;
  }

  /**
   * Pick a weighted random item from a map of {item: weight}
   */
  private weightedRandom<T>(items: Map<T, number>): T | null {
    const entries = Array.from(items.entries());
    if (entries.length === 0) return null;

    const total = entries.reduce((sum, [, w]) => sum + w, 0);
    let random = Math.random() * total;

    for (const [item, weight] of entries) {
      random -= weight;
      if (random <= 0) return item;
    }

    return entries[0]![0];
  }

  /**
   * Generate a sentence from the learned chain
   */
  generate(maxWords: number = 20): string {
    if (this.chain.starts.size === 0) {
      return 'Ich brauche erst mehr Nachrichten zum Lernen.';
    }

    // Pick random start
    const startKey = this.weightedRandom(this.chain.starts);
    if (!startKey) return '...';

    const words: string[] = startKey.split(' ');
    let currentKey = startKey;

    while (words.length < maxWords) {
      const transitions = this.chain.transitions.get(currentKey);

      if (!transitions || transitions.size === 0) break;

      const next = this.weightedRandom(transitions);
      if (!next) break;

      words.push(next);

      // Update key (slide window)
      const keyWords = words.slice(-this.order);
      currentKey = keyWords.join(' ');
    }

    // Capitalize first letter
    const result = words.join(' ');
    return result.charAt(0).toUpperCase() + result.slice(1);
  }

  /**
   * Check if chain has learned anything
   */
  hasLearned(): boolean {
    return this.chain.starts.size > 0;
  }

  /**
   * Add a single transition (for loading from DB)
   */
  addTransition(key: string, nextWord: string, count: number): void {
    if (!this.chain.transitions.has(key)) {
      this.chain.transitions.set(key, new Map());
    }
    this.chain.transitions.get(key)!.set(nextWord, count);
  }

  /**
   * Add a single start (for loading from DB)
   */
  addStart(key: string, count: number): void {
    this.chain.starts.set(key, count);
  }

  /**
   * Export chain data for persistence
   */
  export(): { transitions: Record<string, Record<string, number>>; starts: Record<string, number> } {
    const transitions: Record<string, Record<string, number>> = {};

    for (const [key, words] of this.chain.transitions) {
      transitions[key] = Object.fromEntries(words);
    }

    return {
      transitions,
      starts: Object.fromEntries(this.chain.starts),
    };
  }
}