src/utils/vectorStorage.ts

// Enhanced Vector Storage with TF-IDF and Semantic Similarity
// Provides better document retrieval and similarity matching

export interface DocumentChunk {
  id: string;
  content: string;
  metadata: {
    source: string;
    pageNumber?: number;
    category?: string;
    chunkIndex: number;
    keywords: string[];
  };
  embedding: number[];
  tfidf: Map<string, number>;
}

export interface SearchResult {
  chunk: DocumentChunk;
  score: number;
  matchedKeywords: string[];
}

class VectorStorage {
  private documents: DocumentChunk[] = [];
  private idfScores: Map<string, number> = new Map();
  private vocabularySize: number = 0;

  /**
   * Add document to the vector storage with chunking
   */
  async addDocument(
    content: string,
    metadata: {
      source: string;
      pageNumber?: number;
      category?: string;
    }
  ): Promise<void> {
    // Split document into chunks for better retrieval
    const chunks = this.chunkDocument(content);
    
    chunks.forEach((chunkText, index) => {
      const keywords = this.extractKeywords(chunkText);
      const tfidf = this.calculateTFIDF(chunkText);
      const embedding = this.createEmbedding(chunkText, tfidf);
      
      const chunk: DocumentChunk = {
        id: `${metadata.source}_chunk_${index}`,
        content: chunkText,
        metadata: {
          ...metadata,
          chunkIndex: index,
          keywords
        },
        embedding,
        tfidf
      };
      
      this.documents.push(chunk);
    });
    
    // Recalculate IDF scores after adding documents
    this.calculateIDF();
  }

  /**
   * Search for relevant documents using enhanced similarity
   */
  search(query: string, options?: {
    topK?: number;
    categoryFilter?: string;
    minScore?: number;
  }): SearchResult[] {
    const topK = options?.topK || 10;
    const minScore = options?.minScore || 0.05; // Lowered from 0.1 to catch more results
    
    // Extract query keywords and create embedding
    const queryKeywords = this.extractKeywords(query);
    const queryTFIDF = this.calculateTFIDF(query);
    const queryEmbedding = this.createEmbedding(query, queryTFIDF);
    
    // Calculate similarity scores for all documents
    const results: SearchResult[] = [];
    
    for (const doc of this.documents) {
      // Apply category filter if specified
      if (options?.categoryFilter && doc.metadata.category !== options.categoryFilter) {
        continue;
      }
      
      // Calculate multiple similarity scores
      const cosineSim = this.cosineSimilarity(queryEmbedding, doc.embedding);
      const keywordSim = this.keywordSimilarity(queryKeywords, doc.metadata.keywords);
      const tfidfSim = this.tfidfSimilarity(queryTFIDF, doc.tfidf);
      
      // Weighted combination - boosted keyword matching from 0.3 to 0.4
      const finalScore = (
        cosineSim * 0.35 +
        keywordSim * 0.40 +  // Increased weight for keyword matches
        tfidfSim * 0.25
      );
      
      if (finalScore >= minScore) {
        const matchedKeywords = queryKeywords.filter(kw => 
          doc.metadata.keywords.some(dk => 
            dk.includes(kw) || kw.includes(dk)
          )
        );
        
        results.push({
          chunk: doc,
          score: finalScore,
          matchedKeywords
        });
      }
    }
    
    // Sort by score and return top K
    return results
      .sort((a, b) => b.score - a.score)
      .slice(0, topK);
  }

  /**
   * Semantic search with query expansion
   */
  semanticSearch(query: string, options?: {
    topK?: number;
    categoryFilter?: string;
    expandQuery?: boolean;
  }): SearchResult[] {
    const topK = options?.topK || 10;
    let searchQuery = query;
    
    // Expand query with synonyms and related terms
    if (options?.expandQuery !== false) {
      searchQuery = this.expandQuery(query);
    }
    
    // Perform regular search with expanded query
    const results = this.search(searchQuery, {
      topK: topK * 2, // Get more results initially
      categoryFilter: options?.categoryFilter,
      minScore: 0.05 // Lower threshold for semantic search
    });
    
    // Re-rank results based on semantic relevance
    const rerankedResults = this.rerankBySemanticRelevance(query, results);
    
    return rerankedResults.slice(0, topK);
  }

  /**
   * Get documents by category
   */
  getDocumentsByCategory(category: string): DocumentChunk[] {
    return this.documents.filter(doc => doc.metadata.category === category);
  }

  /**
   * Get all unique categories
   */
  getCategories(): string[] {
    const categories = new Set<string>();
    this.documents.forEach(doc => {
      if (doc.metadata.category) {
        categories.add(doc.metadata.category);
      }
    });
    return Array.from(categories);
  }

  /**
   * Clear all documents
   */
  clear(): void {
    this.documents = [];
    this.idfScores.clear();
    this.vocabularySize = 0;
  }

  /**
   * Get storage statistics
   */
  getStats(): {
    totalDocuments: number;
    totalChunks: number;
    vocabularySize: number;
    categories: string[];
  } {
    return {
      totalDocuments: new Set(this.documents.map(d => d.metadata.source)).size,
      totalChunks: this.documents.length,
      vocabularySize: this.vocabularySize,
      categories: this.getCategories()
    };
  }

  // Private helper methods

  private chunkDocument(content: string, chunkSize: number = 1000): string[] {
    const chunks: string[] = [];
    
    // Try to split by paragraphs first
    const paragraphs = content.split(/\n\s*\n/);
    let currentChunk = '';
    
    for (const para of paragraphs) {
      if ((currentChunk + para).length < chunkSize) {
        currentChunk += para + '\n\n';
      } else {
        if (currentChunk) {
          chunks.push(currentChunk.trim());
        }
        currentChunk = para + '\n\n';
      }
    }
    
    if (currentChunk) {
      chunks.push(currentChunk.trim());
    }
    
    return chunks.length > 0 ? chunks : [content];
  }

  private extractKeywords(text: string): string[] {
    // Tokenize and clean
    const words = text.toLowerCase()
      .replace(/[^\w\s]/g, ' ')
      .split(/\s+/)
      .filter(w => w.length > 3); // Remove short words
    
    // Remove common stop words
    const stopWords = new Set([
      'the', 'this', 'that', 'these', 'those', 'with', 'from', 'have', 'been',
      'were', 'was', 'will', 'would', 'could', 'should', 'about', 'what', 'when',
      'where', 'which', 'their', 'there', 'them', 'they', 'than', 'then', 'your'
    ]);
    
    const keywords = words.filter(w => !stopWords.has(w));
    
    // Return unique keywords
    return Array.from(new Set(keywords));
  }

  private calculateTFIDF(text: string): Map<string, number> {
    const words = this.extractKeywords(text);
    const tfidf = new Map<string, number>();
    const wordCount = words.length;
    
    // Calculate term frequency
    const termFreq = new Map<string, number>();
    words.forEach(word => {
      termFreq.set(word, (termFreq.get(word) || 0) + 1);
    });
    
    // Calculate TF-IDF
    termFreq.forEach((freq, term) => {
      const tf = freq / wordCount;
      const idf = this.idfScores.get(term) || 0;
      tfidf.set(term, tf * idf);
    });
    
    return tfidf;
  }

  private calculateIDF(): void {
    // Count document frequency for each term
    const docFreq = new Map<string, number>();
    const totalDocs = this.documents.length;
    
    this.documents.forEach(doc => {
      const uniqueWords = new Set(doc.metadata.keywords);
      uniqueWords.forEach(word => {
        docFreq.set(word, (docFreq.get(word) || 0) + 1);
      });
    });
    
    // Calculate IDF scores
    this.idfScores.clear();
    docFreq.forEach((freq, term) => {
      const idf = Math.log((totalDocs + 1) / (freq + 1)) + 1;
      this.idfScores.set(term, idf);
    });
    
    this.vocabularySize = docFreq.size;
  }

  private createEmbedding(text: string, tfidf: Map<string, number>): number[] {
    // Create a simple but effective embedding using TF-IDF scores
    // In a production system, you'd use a pre-trained model like Sentence-BERT
    
    const embedding: number[] = new Array(100).fill(0);
    const words = this.extractKeywords(text);
    
    words.forEach((word, index) => {
      const score = tfidf.get(word) || 0;
      const hashIndex = this.hashToIndex(word, 100);
      embedding[hashIndex] += score;
    });
    
    // Normalize the embedding
    const magnitude = Math.sqrt(embedding.reduce((sum, val) => sum + val * val, 0));
    if (magnitude > 0) {
      for (let i = 0; i < embedding.length; i++) {
        embedding[i] /= magnitude;
      }
    }
    
    return embedding;
  }

  private hashToIndex(str: string, size: number): number {
    let hash = 0;
    for (let i = 0; i < str.length; i++) {
      hash = ((hash << 5) - hash) + str.charCodeAt(i);
      hash = hash & hash; // Convert to 32-bit integer
    }
    return Math.abs(hash) % size;
  }

  private cosineSimilarity(vec1: number[], vec2: number[]): number {
    if (vec1.length !== vec2.length) return 0;
    
    let dotProduct = 0;
    let mag1 = 0;
    let mag2 = 0;
    
    for (let i = 0; i < vec1.length; i++) {
      dotProduct += vec1[i] * vec2[i];
      mag1 += vec1[i] * vec1[i];
      mag2 += vec2[i] * vec2[i];
    }
    
    mag1 = Math.sqrt(mag1);
    mag2 = Math.sqrt(mag2);
    
    if (mag1 === 0 || mag2 === 0) return 0;
    
    return dotProduct / (mag1 * mag2);
  }

  private keywordSimilarity(keywords1: string[], keywords2: string[]): number {
    if (keywords1.length === 0 || keywords2.length === 0) return 0;
    
    const set1 = new Set(keywords1);
    const set2 = new Set(keywords2);
    
    let matches = 0;
    set1.forEach(kw1 => {
      set2.forEach(kw2 => {
        // Partial matching
        if (kw1.includes(kw2) || kw2.includes(kw1)) {
          matches++;
        }
      });
    });
    
    // Jaccard similarity with partial matching bonus
    const union = new Set([...keywords1, ...keywords2]).size;
    return matches / Math.max(keywords1.length, keywords2.length);
  }

  private tfidfSimilarity(tfidf1: Map<string, number>, tfidf2: Map<string, number>): number {
    if (tfidf1.size === 0 || tfidf2.size === 0) return 0;
    
    let dotProduct = 0;
    let mag1 = 0;
    let mag2 = 0;
    
    // Calculate dot product and magnitudes
    const allTerms = new Set([...tfidf1.keys(), ...tfidf2.keys()]);
    
    allTerms.forEach(term => {
      const val1 = tfidf1.get(term) || 0;
      const val2 = tfidf2.get(term) || 0;
      
      dotProduct += val1 * val2;
      mag1 += val1 * val1;
      mag2 += val2 * val2;
    });
    
    mag1 = Math.sqrt(mag1);
    mag2 = Math.sqrt(mag2);
    
    if (mag1 === 0 || mag2 === 0) return 0;
    
    return dotProduct / (mag1 * mag2);
  }

  private expandQuery(query: string): string {
    const expansions = new Map<string, string[]>([
      ['experience', ['background', 'history', 'track record', 'portfolio', 'projects', 'expertise']],
      ['approach', ['methodology', 'process', 'framework', 'strategy', 'method', 'technique']],
      ['team', ['staff', 'personnel', 'resources', 'people', 'workforce', 'employees']],
      ['cost', ['price', 'budget', 'pricing', 'rate', 'fee', 'expense']],
      ['timeline', ['schedule', 'duration', 'timeframe', 'deadline', 'delivery']],
      ['quality', ['excellence', 'standards', 'assurance', 'control', 'reliability']],
      ['services', ['offerings', 'solutions', 'capabilities', 'products', 'deliverables']],
      ['financial', ['fiscal', 'monetary', 'economic', 'revenue', 'budget']],
      ['implementation', ['deployment', 'installation', 'rollout', 'execution', 'delivery']],
      ['support', ['maintenance', 'assistance', 'help', 'service', 'backup']]
    ]);
    
    let expandedQuery = query;
    const queryWords = query.toLowerCase().split(/\s+/);
    
    queryWords.forEach(word => {
      if (expansions.has(word)) {
        const synonyms = expansions.get(word) || [];
        // Add a couple of most relevant synonyms
        expandedQuery += ' ' + synonyms.slice(0, 2).join(' ');
      }
    });
    
    return expandedQuery;
  }

  private rerankBySemanticRelevance(originalQuery: string, results: SearchResult[]): SearchResult[] {
    // Additional semantic analysis for reranking
    const queryLower = originalQuery.toLowerCase();
    
    return results.map(result => {
      let semanticBoost = 0;
      
      // Boost if the chunk contains question-answer patterns
      if (result.chunk.content.toLowerCase().includes('answer:') ||
          result.chunk.content.toLowerCase().includes('response:')) {
        semanticBoost += 0.1;
      }
      
      // Boost if chunk is near the beginning (often contains important info)
      if (result.chunk.metadata.chunkIndex < 3) {
        semanticBoost += 0.05;
      }
      
      // Boost if query words appear in close proximity in the chunk
      const queryWords = this.extractKeywords(queryLower);
      const contentLower = result.chunk.content.toLowerCase();
      let proximityBoost = 0;
      
      for (let i = 0; i < queryWords.length - 1; i++) {
        const word1Pos = contentLower.indexOf(queryWords[i]);
        const word2Pos = contentLower.indexOf(queryWords[i + 1]);
        
        if (word1Pos !== -1 && word2Pos !== -1) {
          const distance = Math.abs(word2Pos - word1Pos);
          if (distance < 50) { // Words within 50 characters
            proximityBoost += 0.05;
          }
        }
      }
      
      semanticBoost += Math.min(proximityBoost, 0.15);
      
      return {
        ...result,
        score: result.score + semanticBoost
      };
    }).sort((a, b) => b.score - a.score);
  }
}

// Export singleton instance
export const vectorStorage = new VectorStorage();