/**
 * Scoring Pipeline: Stage 1 (Semantic + Frequency), Grammar Filter, Stage 2 (LM Re-ranking).
 *
 * Operates synchronously on pre-loaded data. Each stage gracefully degrades
 * when its required data isn't available (cold → warm → full warm).
 */

import posTagsData from './data/combined_l2_l3_pos_tags.json';
import ghostPosTagsData from './data/ghost_pos_tags.json';
import grammarTransitionsData from './data/grammar_transitions_10k.json';

// ─── Types ──────────────────────────────────────────────────

/** Metadata returned by the grammar filter for debug logging in the caller. */

// ─── Scoring Constants ──────────────────────────────────────

const ALPHA = 0.5;
const BETA = 0.5;
const NEUTRAL_SCORE = 0.5;
export const STAGE1_WEIGHT = 0.35;
export const STAGE2_WEIGHT = 0.65;
export const MIN_STAGE1_SCORE = 0.35;
const L1_SESSION_CAP = 1.2;
// Minimum prefix-payload max LM probability before Stage 2 activates.
// Below this threshold the LM signal is too weak to suppress Stage 1 — finalScore
// falls back to stage1Score directly. Prevents weak prefixes (e.g. "ins" → "instances"
// at 0.00024) from triggering re-ranking.
const LM_GATE_THRESHOLD = 0.0005;

// ─── Grammar Data (loaded once on import) ───────────────────

const posTags = new Map(Object.entries(posTagsData));
const grammarTransitions = grammarTransitionsData;

/**
 * Precomputed map from each POS tag to the set of allowed next POS tags.
 * Built once at module load from grammarTransitions so applyGrammarFilter
 * never re-iterates the transition rules per call.
 */
const precomputedAllowedByPos = new Map(Object.entries(grammarTransitions.transitions).map(([pos, rule]) => [pos, new Set(rule.allowed)]));

// ─── Math ───────────────────────────────────────────────────

function cosineSimilarity(a, b) {
  let dot = 0;
  let normA = 0;
  let normB = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    normA += a[i] * a[i];
    normB += b[i] * b[i];
  }
  const dNormA = Math.sqrt(normA);
  const dNormB = Math.sqrt(normB);
  if (dNormA === 0 || dNormB === 0) {
    return NEUTRAL_SCORE;
  }
  return (1 + dot / (dNormA * dNormB)) / 2;
}

// ─── Stage 1: Semantic + Frequency ──────────────────────────

function scoreStage1(candidate, contextVector, getWordVector, maxTenantFreq) {
  // 1. Calculate Base Global Score (Normalized Log)
  const maxPossibleLog = Math.log10(maxTenantFreq + 1);

  // Diversity Adjustment
  const diversityRaw = (Math.log10(candidate.tenantFreq + 1) * 0.5 + Math.log10(candidate.docFreq + 1) * 0.25 + Math.log10(candidate.authorFreq + 1) * 0.25) / maxPossibleLog;
  const sessionMultiplier = candidate.sessionFreq > 0 ? 1 + Math.log10(candidate.sessionFreq + 1) * 2.5 : 1;

  // Apply multiplier; capped at L1_SESSION_CAP (default 1.2) to prevent excessive over-indexing
  const freqScore = Math.min(diversityRaw * sessionMultiplier, L1_SESSION_CAP);

  // 3. Semantic Scoring
  let semanticScore = NEUTRAL_SCORE;
  if (contextVector) {
    const wordVec = getWordVector(candidate.word);
    semanticScore = wordVec ? cosineSimilarity(contextVector, wordVec) : NEUTRAL_SCORE;
  }
  return {
    semanticScore,
    freqScore,
    stage1Score: ALPHA * semanticScore + BETA * freqScore
  };
}

// ─── Grammar Filter ─────────────────────────────────────────

/**
 * GHOST POS DICTIONARY
 * A hardcoded mapping of common structural English words that were stripped
 * from the main domain vocabulary. This allows the grammar filter to understand
 * context without suggesting these words to the user.
 */
const ghostPosTags = ghostPosTagsData;
function applyGrammarFilter(candidates, previousWord) {
  if (!previousWord) return {
    filtered: candidates,
    grammarMeta: null
  };
  const lowerPrev = previousWord.toLowerCase();
  const prevTags = ghostPosTags[lowerPrev] || posTags.get(lowerPrev);
  if (!prevTags || prevTags.length === 0) {
    return {
      filtered: candidates,
      grammarMeta: null
    };
  }
  let allowedNextTags;
  if (prevTags.length === 1) {
    var _precomputedAllowedBy;
    // Common case: single POS tag — reuse the precomputed Set directly (no allocation)
    allowedNextTags = (_precomputedAllowedBy = precomputedAllowedByPos.get(prevTags[0])) !== null && _precomputedAllowedBy !== void 0 ? _precomputedAllowedBy : new Set();
  } else {
    allowedNextTags = new Set();
    for (const pt of prevTags) {
      const allowed = precomputedAllowedByPos.get(pt);
      if (allowed) allowed.forEach(tag => allowedNextTags.add(tag));
    }
  }
  const filtered = [];
  const dropped = [];
  for (const entry of candidates) {
    const candidateTags = posTags.get(entry.candidate.word.toLowerCase());

    // If candidate has no tags (unknown word), let it pass to be safe
    if (!candidateTags || candidateTags.length === 0) {
      filtered.push(entry);
      continue;
    }
    if (candidateTags.some(ct => allowedNextTags.has(ct))) {
      filtered.push(entry);
    } else {
      dropped.push(entry.candidate.word);
    }
  }

  // Grammar is authoritative.
  return {
    filtered: filtered,
    grammarMeta: {
      prevWord: lowerPrev,
      prevTags,
      before: candidates.length,
      after: filtered.length,
      dropped: dropped
    }
  };
}

// ─── Stage 2: LM Re-ranking ────────────────────────────────
function getLmScore(word, lmLogits) {
  if (!lmLogits) return 0;

  // Look up the word directly! No more tokens.
  const val = lmLogits[word.toLowerCase()];
  if (typeof val === 'number') {
    return val;
  }
  return 0;
}

// ─── Public API ─────────────────────────────────────────────

export function rankCandidates(candidates, contextVector, getWordVector, lmLogits, maxTenantFreq, previousWord) {
  var _grammarMeta$dropped;
  // Stage 1
  const stage1Results = candidates.map(candidate => {
    const {
      semanticScore,
      freqScore,
      stage1Score
    } = scoreStage1(candidate, contextVector, getWordVector, maxTenantFreq);
    return {
      candidate,
      semanticScore,
      freqScore,
      stage1Score
    };
  });
  const stage1Survivors = [];
  const stage1Rejected = [];
  for (const entry of stage1Results) {
    if (entry.stage1Score >= MIN_STAGE1_SCORE) {
      stage1Survivors.push(entry);
    } else {
      stage1Rejected.push(entry.candidate.word);
    }
  }

  // Grammar Filter
  const {
    filtered,
    grammarMeta
  } = applyGrammarFilter(stage1Survivors, previousWord);

  // Stage 2 + final assembly
  let lmMax = 0;
  if (lmLogits) {
    const values = Object.values(lmLogits);
    if (values.length > 0) lmMax = Math.max(...values);
  }
  const scored = filtered.map(entry => {
    let lmScore = 0;
    let finalScore = entry.stage1Score;
    if (lmLogits && lmMax >= LM_GATE_THRESHOLD) {
      const rawLm = getLmScore(entry.candidate.word, lmLogits);
      if (rawLm !== 0) {
        const logitDiff = Math.log(rawLm) - Math.log(lmMax);
        lmScore = Math.exp(logitDiff);
      }
      // Words absent from the prefix-filtered payload get lmScore = 0,
      // not 0.05, so they don't outrank genuine LM predictions.

      finalScore = STAGE1_WEIGHT * entry.stage1Score + STAGE2_WEIGHT * lmScore;
    }
    return {
      word: entry.candidate.word,
      freqScore: entry.freqScore,
      semanticScore: entry.semanticScore,
      lmScore,
      finalScore
    };
  });
  scored.sort((a, b) => {
    if (b.finalScore !== a.finalScore) return b.finalScore - a.finalScore;
    return a.word.length - b.word.length;
  });
  return {
    candidates: scored,
    grammarMeta,
    pipelineDebug: {
      initial: candidates.length,
      stage1Rejected,
      grammarRejected: (_grammarMeta$dropped = grammarMeta === null || grammarMeta === void 0 ? void 0 : grammarMeta.dropped) !== null && _grammarMeta$dropped !== void 0 ? _grammarMeta$dropped : [],
      final: scored.length
    }
  };
}