/*
 * lm.c - a small word-level trigram language model with back-off (C89).
 *
 * Storage:
 *   - a vocabulary of interned word strings (open-addressed string table)
 *   - unigram counts (cumulative for sampling)
 *   - bigram counts: prefixes (w1) -> bucket of (w2, count)
 *   - trigram counts: prefixes (w1,w2) -> bucket of (w3, count)
 *
 * Generation samples the trigram distribution, backing off to the bigram and
 * then the unigram distribution when a context was never seen.
 *
 * Author: DeepSeek (deepseek-v4.1-flash), October 2026.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "lm.h"

typedef struct {
    int word;       /* next-word id */
    int cnt;        /* count */
    int next;       /* next entry in bucket, -1 = end */
} NEntry;

struct LM {
    /* vocabulary */
    char **words;
    int vocab;
    int vcap;
    int *vhash;          /* string table -> word id, -1 empty */
    int vhcap;

    /* unigram counts and cumulative */
    unsigned *uni;
    unsigned long long *unicum;   /* size vocab+1 */
    unsigned long long utotal;

    /* bigram prefixes */
    int *bk;             /* key = w1, -1 empty */
    int *bh;             /* head entry */
    int bcap;
    NEntry *bie;
    long nbie, cbie;

    /* trigram prefixes */
    unsigned long long *tk;
    int *th;
    int tcap;
    NEntry *tie;
    long ntie, ctie;

    unsigned long long tokens;
    unsigned long long corpus_bytes;
};

static unsigned long rng_state = 0;

static unsigned long rnd(void)
{
    if (rng_state == 0)
        rng_state = (unsigned long)time((time_t *)0) ^ 0x9e3779b9UL;
    rng_state = rng_state * 1103515245UL + 12345UL;
    return (rng_state >> 16) & 0xffffffffUL;
}

static unsigned hash_str(const char *s)
{
    unsigned h = 5381u;
    while (*s)
        h = h * 33u + (unsigned char)*s++;
    return h;
}

static unsigned hash_int(unsigned x)
{
    x ^= x >> 16;
    x *= 0x7feb352dUL;
    x ^= x >> 15;
    x *= 0x846ca68bUL;
    x ^= x >> 16;
    return x;
}

static unsigned long long hash_ll(unsigned long long x)
{
    x ^= x >> 33;
    x *= (((unsigned long long)0xff51afd7u << 32) | (unsigned long long)0xed558ccdu);
    x ^= x >> 33;
    return x;
}

/* ---- vocabulary ---- */

static void vocab_rehash(LM *m, int newcap)
{
    int i;
    int *t = (int *)malloc((size_t)newcap * sizeof(int));
    for (i = 0; i < newcap; i++) t[i] = -1;
    for (i = 0; i < m->vocab; i++) {
        unsigned h = hash_str(m->words[i]) & (unsigned)(newcap - 1);
        while (t[h] >= 0) h = (h + 1) & (unsigned)(newcap - 1);
        t[h] = i;
    }
    free(m->vhash);
    m->vhash = t;
    m->vhcap = newcap;
}

static int vocab_lookup(LM *m, const char *s)
{
    unsigned h = hash_str(s) & (unsigned)(m->vhcap - 1);
    while (m->vhash[h] >= 0) {
        if (strcmp(m->words[m->vhash[h]], s) == 0)
            return m->vhash[h];
        h = (h + 1) & (unsigned)(m->vhcap - 1);
    }
    return -1;
}

static int vocab_intern(LM *m, const char *s, int slen)
{
    unsigned h = hash_str(s) & (unsigned)(m->vhcap - 1);
    int id;
    while (m->vhash[h] >= 0) {
        if (strcmp(m->words[m->vhash[h]], s) == 0)
            return m->vhash[h];
        h = (h + 1) & (unsigned)(m->vhcap - 1);
    }
    if (m->vocab >= m->vcap) {
        m->vcap = m->vcap ? m->vcap * 2 : 4096;
        m->words = (char **)realloc(m->words, (size_t)m->vcap * sizeof(char *));
        m->uni = (unsigned *)realloc(m->uni, (size_t)m->vcap * sizeof(unsigned));
    }
    id = m->vocab++;
    m->words[id] = (char *)malloc((size_t)slen + 1);
    memcpy(m->words[id], s, (size_t)slen);
    m->words[id][slen] = '\0';
    m->uni[id] = 0;
    m->vhash[h] = id;
    if ((long)m->vocab * 10 > (long)m->vhcap * 7)
        vocab_rehash(m, m->vhcap * 2);
    return id;
}

/* ---- n-gram prefix tables ---- */

static void bi_rehash(LM *m, int newcap)
{
    int i;
    int *k = (int *)malloc((size_t)newcap * sizeof(int));
    int *h = (int *)malloc((size_t)newcap * sizeof(int));
    for (i = 0; i < newcap; i++) { k[i] = -1; h[i] = -1; }
    for (i = 0; i < m->bcap; i++) {
        if (m->bk[i] >= 0) {
            unsigned p = hash_int((unsigned)m->bk[i]) & (unsigned)(newcap - 1);
            while (k[p] >= 0) p = (p + 1) & (unsigned)(newcap - 1);
            k[p] = m->bk[i];
            h[p] = m->bh[i];
        }
    }
    free(m->bk); free(m->bh);
    m->bk = k; m->bh = h; m->bcap = newcap;
}

static void tri_rehash(LM *m, int newcap)
{
    int i;
    unsigned long long *k = (unsigned long long *)malloc((size_t)newcap * sizeof(unsigned long long));
    int *h = (int *)malloc((size_t)newcap * sizeof(int));
    for (i = 0; i < newcap; i++) { k[i] = (unsigned long long)0; h[i] = -1; }
    for (i = 0; i < m->tcap; i++) {
        if (m->tk[i] != (unsigned long long)0) {
            unsigned p = (unsigned)(hash_ll(m->tk[i]) & (unsigned long long)(newcap - 1));
            while (k[p] != (unsigned long long)0) p = (p + 1) & (unsigned)(newcap - 1);
            k[p] = m->tk[i];
            h[p] = m->th[i];
        }
    }
    free(m->tk); free(m->th);
    m->tk = k; m->th = h; m->tcap = newcap;
}

static int is_word_char(int c);

/* Extract the next token from text[*pos..len). Returns token length, 0 at end.
 * Words are lowercased; punctuation becomes a single-character token. */
static int next_token(const char *text, long len, long *pos, char *tok, int tokmax)
{
    long i = *pos;
    int tl = 0;
    while (i < len) {
        int c = (unsigned char)text[i];
        if (c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '\f' || c == '\v') {
            i++;
            continue;
        }
        break;
    }
    if (i >= len) { *pos = i; return 0; }
    if (is_word_char((unsigned char)text[i])) {
        while (i < len && is_word_char((unsigned char)text[i])) {
            int ch = (unsigned char)text[i];
            if (ch >= 'A' && ch <= 'Z') ch = ch - 'A' + 'a';
            if (tl < tokmax - 1) tok[tl++] = (char)ch;
            i++;
        }
    } else {
        tok[tl++] = (char)text[i];
        i++;
    }
    tok[tl] = '\0';
    *pos = i;
    return tl;
}

static void add_bigram(LM *m, int w1, int w2)
{
    unsigned p = hash_int((unsigned)w1) & (unsigned)(m->bcap - 1);
    int e;
    while (m->bk[p] >= 0 && m->bk[p] != w1)
        p = (p + 1) & (unsigned)(m->bcap - 1);
    if (m->bk[p] < 0) {
        m->bk[p] = w1;
        m->bh[p] = -1;
        if ((long)(m->nbie) >= 0 && (long)m->bcap * 7 < (long)(m->nbie + 1) * 10)
            bi_rehash(m, m->bcap * 2);
        p = hash_int((unsigned)w1) & (unsigned)(m->bcap - 1);
        while (m->bk[p] >= 0 && m->bk[p] != w1)
            p = (p + 1) & (unsigned)(m->bcap - 1);
    }
    for (e = m->bh[p]; e >= 0; e = m->bie[e].next)
        if (m->bie[e].word == w2) { m->bie[e].cnt++; return; }
    if (m->nbie >= m->cbie) {
        m->cbie = m->cbie ? m->cbie * 2 : 4096;
        m->bie = (NEntry *)realloc(m->bie, (size_t)m->cbie * sizeof(NEntry));
    }
    m->bie[m->nbie].word = w2;
    m->bie[m->nbie].cnt = 1;
    m->bie[m->nbie].next = m->bh[p];
    m->bh[p] = (int)m->nbie;
    m->nbie++;
}

static void add_trigram(LM *m, int w1, int w2, int w3)
{
    unsigned long long key = (((unsigned long long)w1 << 32) | (unsigned)w2) + (unsigned long long)1;
    unsigned p = (unsigned)(hash_ll(key) & (unsigned long long)(m->tcap - 1));
    int e;
    while (m->tk[p] != (unsigned long long)0 && m->tk[p] != key)
        p = (p + 1) & (unsigned)(m->tcap - 1);
    if (m->tk[p] == (unsigned long long)0) {
        m->tk[p] = key;
        m->th[p] = -1;
        if ((long)m->tcap * 7 < (long)(m->ntie + 1) * 10) {
            tri_rehash(m, m->tcap * 2);
            p = (unsigned)(hash_ll(key) & (unsigned long long)(m->tcap - 1));
            while (m->tk[p] != (unsigned long long)0 && m->tk[p] != key)
                p = (p + 1) & (unsigned)(m->tcap - 1);
        }
    }
    for (e = m->th[p]; e >= 0; e = m->tie[e].next)
        if (m->tie[e].word == w3) { m->tie[e].cnt++; return; }
    if (m->ntie >= m->ctie) {
        m->ctie = m->ctie ? m->ctie * 2 : 4096;
        m->tie = (NEntry *)realloc(m->tie, (size_t)m->ctie * sizeof(NEntry));
    }
    m->tie[m->ntie].word = w3;
    m->tie[m->ntie].cnt = 1;
    m->tie[m->ntie].next = m->th[p];
    m->th[p] = (int)m->ntie;
    m->ntie++;
}

/* ---- construction ---- */

LM *lm_create(void)
{
    LM *m = (LM *)calloc(1, sizeof(LM));
    m->vhcap = 1 << 16;
    m->vhash = (int *)malloc((size_t)m->vhcap * sizeof(int));
    { int i; for (i = 0; i < m->vhcap; i++) m->vhash[i] = -1; }
    m->bcap = 1 << 16;
    m->bk = (int *)malloc((size_t)m->bcap * sizeof(int));
    m->bh = (int *)malloc((size_t)m->bcap * sizeof(int));
    { int i; for (i = 0; i < m->bcap; i++) { m->bk[i] = -1; m->bh[i] = -1; } }
    m->tcap = 1 << 16;
    m->tk = (unsigned long long *)calloc((size_t)m->tcap, sizeof(unsigned long long));
    m->th = (int *)malloc((size_t)m->tcap * sizeof(int));
    { int i; for (i = 0; i < m->tcap; i++) m->th[i] = -1; }
    return m;
}

static int is_word_char(int c)
{
    return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
           (c >= '0' && c <= '9') || c == '\'' || c == '_';
}

void lm_add_text(LM *m, const char *text, long len)
{
    long pos = 0;
    int prev1 = -1, prev2 = -1;
    char tok[64];
    int tl;

    while ((tl = next_token(text, len, &pos, tok, (int)sizeof tok)) > 0) {
        int id = vocab_intern(m, tok, tl);
        m->uni[id]++;
        m->tokens++;
        if (prev1 >= 0) add_bigram(m, prev1, id);
        if (prev2 >= 0) add_trigram(m, prev2, prev1, id);
        prev2 = prev1;
        prev1 = id;
    }
    m->corpus_bytes += (unsigned long long)len;
}

void lm_finish(LM *m)
{
    int i;
    m->unicum = (unsigned long long *)malloc((size_t)(m->vocab + 1) * sizeof(unsigned long long));
    m->unicum[0] = 0;
    for (i = 0; i < m->vocab; i++)
        m->unicum[i + 1] = m->unicum[i] + (unsigned long long)m->uni[i];
    m->utotal = m->unicum[m->vocab];
}

void lm_stats(const LM *m, LMStats *s)
{
    s->vocab = m->vocab;
    s->tokens = m->tokens;
    s->distinct_bigrams = m->nbie;
    s->distinct_trigrams = m->ntie;
    s->corpus_bytes = m->corpus_bytes;
}

int lm_vocab(const LM *m) { return m->vocab; }
const char *lm_word(const LM *m, int id) { return m->words[id]; }

/* ---- generation ---- */

static int sample_bucket(LM *m, NEntry *pool, int head)
{
    int e;
    long sum = 0;
    unsigned long r;
    for (e = head; e >= 0; e = pool[e].next) sum += pool[e].cnt;
    if (sum <= 0) return -1;
    r = rnd() % (unsigned long)sum;
    for (e = head; e >= 0; e = pool[e].next) {
        if (r < (unsigned long)pool[e].cnt) return pool[e].word;
        r -= (unsigned long)pool[e].cnt;
    }
    return -1;
}

static int sample_trigram(LM *m, int w1, int w2)
{
    unsigned long long key = (((unsigned long long)w1 << 32) | (unsigned)w2) + (unsigned long long)1;
    unsigned p = (unsigned)(hash_ll(key) & (unsigned long long)(m->tcap - 1));
    while (m->tk[p] != (unsigned long long)0 && m->tk[p] != key)
        p = (p + 1) & (unsigned)(m->tcap - 1);
    if (m->tk[p] == key && m->th[p] >= 0)
        return sample_bucket(m, m->tie, m->th[p]);
    return -1;
}

static int sample_bigram(LM *m, int w1)
{
    unsigned p = hash_int((unsigned)w1) & (unsigned)(m->bcap - 1);
    while (m->bk[p] >= 0 && m->bk[p] != w1)
        p = (p + 1) & (unsigned)(m->bcap - 1);
    if (m->bk[p] == w1 && m->bh[p] >= 0)
        return sample_bucket(m, m->bie, m->bh[p]);
    return -1;
}

static int sample_unigram(LM *m)
{
    unsigned long long r, *base, *lo, *hi;
    if (m->utotal == 0) return -1;
    r = (rnd() % m->utotal);
    base = m->unicum + 1;          /* find first i with unicum[i+1] > r */
    lo = base; hi = base + m->vocab;
    while (lo < hi) {
        unsigned long long *mid = lo + (hi - lo) / 2;
        if (*mid <= r) lo = mid + 1; else hi = mid;
    }
    return (int)(lo - m->unicum) - 1;
}

static int next_word(LM *m, int w1, int w2)
{
    int w = -1;
    if (w1 >= 0 && w2 >= 0) w = sample_trigram(m, w1, w2);
    if (w < 0 && w2 >= 0) w = sample_bigram(m, w2);
    if (w < 0) w = sample_unigram(m);
    return w;
}

/* render a token list with light punctuation spacing */
static void render(LM *m, int *ids, int n, char *out, int outsz)
{
    int i, o = 0;
    for (i = 0; i < n && o < outsz - 1; i++) {
        const char *w = m->words[ids[i]];
        int punct = (w[1] == '\0') && !is_word_char((unsigned char)w[0]);
        int nospace = 0;
        if (punct && (w[0] == '.' || w[0] == ',' || w[0] == ';' || w[0] == ':' ||
                      w[0] == ')' || w[0] == '!' || w[0] == '?' || w[0] == '\''))
            nospace = 1;
        if (i > 0 && !nospace) {
            if (o < outsz - 1) out[o++] = ' ';
        }
        { const char *p = w; while (*p && o < outsz - 1) out[o++] = *p++; }
    }
    out[o] = '\0';
}

int lm_generate(LM *m, const char *seed, int max_words, char *out, int outsz)
{
    int ids[4096];
    int n = 0;
    int w1 = -1, w2 = -1;
    char tok[64];
    int tl;
    long pos = 0;
    long slen = (long)strlen(seed);

    /* tokenize the seed; unknown words become id -1 (back-off handles it) */
    while ((tl = next_token(seed, slen, &pos, tok, (int)sizeof tok)) > 0) {
        int id = vocab_lookup(m, tok);
        w2 = w1;
        w1 = id;
    }

    if (w1 < 0 && max_words > 0) {
        /* empty or all-unknown seed: start from a random word */
        w1 = sample_unigram(m);
    }

    while (n < max_words && n < 4095) {
        int w = next_word(m, w2, w1);
        if (w < 0) break;
        ids[n++] = w;
        w2 = w1;
        w1 = w;
    }

    render(m, ids, n, out, outsz);
    return n;
}

void lm_free(LM *m)
{
    int i;
    if (!m) return;
    for (i = 0; i < m->vocab; i++) free(m->words[i]);
    free(m->words); free(m->uni); free(m->unicum); free(m->vhash);
    free(m->bk); free(m->bh); free(m->bie);
    free(m->tk); free(m->th); free(m->tie);
    free(m);
}