import {
  PAPER_MEMORY_LEXICAL_VERSION,
  type PaperChunk,
  type PaperEmbeddingIndex,
  type PaperIndexGeneration,
  type PaperLexicalIndex,
} from './types';

const WORD_RE = /[\p{L}\p{N}]+\^[-+]?\d+|[\p{L}\p{N}]+(?:[_./:+-][\p{L}\p{N}]+)*/gu;
const COMPARISON_RE = /\b(?:compare|comparison|contrast|differ|difference|disagree|contradict|methodolog|datasets?|conclusions?|across papers|between papers)\b/iu;

function normalizeToken(value: string): string {
  return value.normalize('NFKC').toLocaleLowerCase('en-US');
}

export function tokenizeScientificText(value: string): string[] {
  const normalized = value.normalize('NFKC');
  const tokens = normalized.match(WORD_RE) ?? [];
  const result: string[] = [];
  for (const token of tokens) {
    const compact = normalizeToken(token);
    if (!compact) continue;
    result.push(compact);
    // Preserve useful compound identifiers while also making their components
    // discoverable. This keeps exact scientific tokens such as ResNet-50,
    // beta_2 and 10^-4 searchable without sacrificing ordinary word recall.
    for (const part of compact.split(/[_./:+^-]+/u)) {
      if (part.length >= 2 && part !== compact) result.push(part);
    }
  }
  return result;
}

export function buildPaperLexicalIndex(chunks: readonly PaperChunk[]): PaperLexicalIndex {
  const postings = new Map<string, number[]>();
  const documentLengths = chunks.map((chunk, chunkIndex) => {
    const tokens = tokenizeScientificText(`${chunk.heading ?? ''} ${chunk.section ?? ''} ${chunk.text}`);
    for (const term of new Set(tokens)) {
      const values = postings.get(term) ?? [];
      values.push(chunkIndex);
      postings.set(term, values);
    }
    return Math.max(1, tokens.length);
  });
  const averageDocumentLength = documentLengths.length
    ? documentLengths.reduce((sum, value) => sum + value, 0) / documentLengths.length
    : 0;
  return {
    version: PAPER_MEMORY_LEXICAL_VERSION,
    documentLengths,
    averageDocumentLength,
    postings: [...postings.entries()]
      .sort(([a], [b]) => a.localeCompare(b))
      .map(([term, chunkIndexes]) => ({ term, chunkIndexes })),
  };
}

function postingMap(index: PaperLexicalIndex): Map<string, readonly number[]> {
  return new Map(index.postings.map((posting) => [posting.term, posting.chunkIndexes] as const));
}

function termFrequency(text: string, term: string): number {
  let count = 0;
  for (const token of tokenizeScientificText(text)) if (token === term) count += 1;
  return count;
}

export function rankPaperLexical(
  generation: PaperIndexGeneration,
  query: string,
): Array<{ chunkIndex: number; score: number }> {
  const lexical = generation.lexical ?? buildPaperLexicalIndex(generation.chunks);
  if (!generation.chunks.length) return [];
  const queryTerms = [...new Set(tokenizeScientificText(query))];
  if (!queryTerms.length) return [];
  const postings = postingMap(lexical);
  const scores = new Map<number, number>();
  const documentCount = generation.chunks.length;
  const avg = lexical.averageDocumentLength || 1;
  const k1 = 1.2;
  const b = 0.75;

  for (const term of queryTerms) {
    const matches = postings.get(term) ?? [];
    if (!matches.length) continue;
    const idf = Math.log(1 + (documentCount - matches.length + 0.5) / (matches.length + 0.5));
    for (const chunkIndex of matches) {
      const chunk = generation.chunks[chunkIndex];
      if (!chunk) continue;
      const haystack = `${chunk.heading ?? ''} ${chunk.section ?? ''} ${chunk.text}`;
      const tf = termFrequency(haystack, term);
      const length = lexical.documentLengths[chunkIndex] ?? 1;
      const denominator = tf + k1 * (1 - b + b * length / avg);
      const score = idf * (tf * (k1 + 1)) / Math.max(0.0001, denominator);
      scores.set(chunkIndex, (scores.get(chunkIndex) ?? 0) + score);
    }
  }

  // Exact phrase/token bonus matters disproportionately in scientific search.
  const exactNeedle = normalizeToken(query.trim());
  if (exactNeedle) {
    generation.chunks.forEach((chunk, chunkIndex) => {
      const haystack = normalizeToken(`${chunk.heading ?? ''} ${chunk.section ?? ''} ${chunk.text}`);
      if (haystack.includes(exactNeedle)) scores.set(chunkIndex, (scores.get(chunkIndex) ?? 0) + 4);
    });
  }
  return [...scores.entries()]
    .map(([chunkIndex, score]) => ({ chunkIndex, score }))
    .sort((a, b) => b.score - a.score);
}

function cosine(left: Float32Array, right: Float32Array): number {
  if (!left.length || left.length !== right.length) return Number.NEGATIVE_INFINITY;
  let dot = 0;
  let leftNorm = 0;
  let rightNorm = 0;
  for (let index = 0; index < left.length; index += 1) {
    const a = left[index]!;
    const b = right[index]!;
    dot += a * b;
    leftNorm += a * a;
    rightNorm += b * b;
  }
  const denominator = Math.sqrt(leftNorm) * Math.sqrt(rightNorm);
  return denominator > 0 ? dot / denominator : Number.NEGATIVE_INFINITY;
}

export function embeddingVectorAt(embedding: PaperEmbeddingIndex, chunkIndex: number): Float32Array {
  const bytesPerVector = embedding.dimensions * Float32Array.BYTES_PER_ELEMENT;
  const start = chunkIndex * bytesPerVector;
  if (start < 0 || start + bytesPerVector > embedding.vectors.byteLength) return new Float32Array();
  return new Float32Array(embedding.vectors, start, embedding.dimensions);
}

export function rankPaperSemantic(
  generation: PaperIndexGeneration,
  queryVector: Float32Array,
): Array<{ chunkIndex: number; score: number }> {
  const embedding = generation.embedding;
  if (!embedding || queryVector.length !== embedding.dimensions) return [];
  return generation.chunks
    .map((_, chunkIndex) => ({ chunkIndex, score: cosine(queryVector, embeddingVectorAt(embedding, chunkIndex)) }))
    .filter((item) => Number.isFinite(item.score))
    .sort((a, b) => b.score - a.score);
}

export function isPaperComparisonQuery(query: string): boolean {
  return COMPARISON_RE.test(query);
}

export type PaperRankCandidate = {
  paperId: string;
  chunkIndex: number;
  lexicalScore: number;
  semanticScore?: number;
  score: number;
};

export function combinePaperRanks(
  inputs: readonly {
    paperId: string;
    lexical: readonly { chunkIndex: number; score: number }[];
    semantic?: readonly { chunkIndex: number; score: number }[];
  }[],
): PaperRankCandidate[] {
  const combined: PaperRankCandidate[] = [];
  for (const input of inputs) {
    const lexicalMax = Math.max(0, ...input.lexical.map((item) => item.score));
    const semanticByIndex = new Map(input.semantic?.map((item) => [item.chunkIndex, item.score] as const) ?? []);
    const lexicalByIndex = new Map(input.lexical.map((item) => [item.chunkIndex, item.score] as const));
    const indexes = new Set<number>([...lexicalByIndex.keys(), ...semanticByIndex.keys()]);
    for (const chunkIndex of indexes) {
      const lexicalRaw = lexicalByIndex.get(chunkIndex) ?? 0;
      const lexicalScore = lexicalMax > 0 ? lexicalRaw / lexicalMax : 0;
      const semanticScore = semanticByIndex.get(chunkIndex);
      const semanticNormalized = semanticScore === undefined ? undefined : Math.max(0, Math.min(1, (semanticScore + 1) / 2));
      const score = semanticNormalized === undefined
        ? lexicalScore
        : semanticNormalized * 0.68 + lexicalScore * 0.32;
      combined.push({ paperId: input.paperId, chunkIndex, lexicalScore, ...(semanticScore !== undefined ? { semanticScore } : {}), score });
    }
  }
  return combined.sort((a, b) => b.score - a.score);
}

export function diversifyPaperRanks(
  ranked: readonly PaperRankCandidate[],
  limit: number,
  perPaperCap = 3,
): PaperRankCandidate[] {
  const output: PaperRankCandidate[] = [];
  const counts = new Map<string, number>();
  const globalHead = Math.min(2, limit, ranked.length);
  for (let index = 0; index < globalHead; index += 1) {
    const candidate = ranked[index]!;
    output.push(candidate);
    counts.set(candidate.paperId, (counts.get(candidate.paperId) ?? 0) + 1);
  }
  for (const candidate of ranked.slice(globalHead)) {
    if (output.length >= limit) break;
    if ((counts.get(candidate.paperId) ?? 0) >= perPaperCap) continue;
    output.push(candidate);
    counts.set(candidate.paperId, (counts.get(candidate.paperId) ?? 0) + 1);
  }
  return output;
}