#include "grader.h"
#include <bits/stdc++.h>
using namespace std;
static constexpr int K = 56;
static constexpr int S = 65536; // symbols: 1..65535
static constexpr int B = 1 << 18; // bigram hash bins (262144)
// Tunables (these matter)
static constexpr double ALPHA1 = 0.25; // unigram smoothing
static constexpr double ALPHA2 = 0.25; // bigram smoothing
static constexpr double BETA = 1.0; // prior smoothing (avoid log(0))
static constexpr double W2 = 0.9; // weight for bigram score vs unigram
// Unigram model
static uint32_t cnt1[K][S];
static uint32_t tot1[K];
static uint32_t types1[K];
static uint32_t ex[K];
static bitset<S> seen1[K];
// Bigram model (hashed)
static uint32_t cnt2[K][B];
static uint32_t tot2[K];
static uint32_t types2[K];
static bitset<B> seen2[K];
static inline uint32_t h2(uint32_t a, uint32_t b) {
// decent mix, fast
uint32_t x = a * 0x9e3779b1u ^ (b + 0x85ebca6bu);
x ^= x >> 16;
x *= 0x7feb352du;
x ^= x >> 15;
return x & (B - 1);
}
static inline void learn(int L, const int E[100]) {
ex[L]++;
// unigrams
for (int i = 0; i < 100; i++) {
int x = E[i];
tot1[L]++;
if (!seen1[L].test(x)) {
seen1[L].set(x);
types1[L]++;
}
cnt1[L][x]++;
}
// bigrams
for (int i = 0; i + 1 < 100; i++) {
uint32_t id = h2((uint32_t)E[i], (uint32_t)E[i + 1]);
tot2[L]++;
if (!seen2[L].test(id)) {
seen2[L].set(id);
types2[L]++;
}
cnt2[L][id]++;
}
}
void excerpt(int E[100]) {
// Pre-hash bigrams of this excerpt once
uint32_t big[99];
for (int i = 0; i < 99; i++) big[i] = h2((uint32_t)E[i], (uint32_t)E[i + 1]);
int bestL = 0;
double bestScore = -1e300;
// total examples is constant for argmax if you normalize, so we skip it
for (int L = 0; L < K; L++) {
// Prior (smoothed)
double score = log((double)ex[L] + BETA);
// --- unigram log-likelihood with "UNK bucket" ---
double denom1 = (double)tot1[L] + ALPHA1 * ((double)types1[L] + 1.0);
score -= 100.0 * log(denom1);
for (int i = 0; i < 100; i++) {
int x = E[i];
if (seen1[L].test(x)) score += log((double)cnt1[L][x] + ALPHA1);
else score += log(ALPHA1);
}
// --- bigram log-likelihood (hashed) ---
double score2 = 0.0;
double denom2 = (double)tot2[L] + ALPHA2 * ((double)types2[L] + 1.0);
score2 -= 99.0 * log(denom2);
for (int i = 0; i < 99; i++) {
uint32_t id = big[i];
if (seen2[L].test(id)) score2 += log((double)cnt2[L][id] + ALPHA2);
else score2 += log(ALPHA2);
}
score += W2 * score2;
if (score > bestScore) {
bestScore = score;
bestL = L;
}
}
int correct = language(bestL);
learn(correct, E);
}
