package cn.edu.bjut.textmining.chapter3;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.TextField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.MultiTerms;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.Term;
import org.apache.lucene.store.ByteBuffersDirectory;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.BytesRef;

import java.io.IOException;
import java.io.UncheckedIOException;

public class LuceneTfIdfDemo {
	public static void main(String[] args) {
		try (Directory directory = new ByteBuffersDirectory(); 
				Analyzer analyzer = new StandardAnalyzer()) {
			String doc1 = "Java Lucene is powerful.";
			String doc2 = "Lucene is a search library.";

			IndexWriterConfig config = new IndexWriterConfig(analyzer);
			try (IndexWriter writer = new IndexWriter(directory, config)) {
				addDocument(writer, doc1);
				addDocument(writer, doc2);
			}

			try (IndexReader reader = DirectoryReader.open(directory)) {
				int totalDocs = reader.numDocs();
				String targetTerm = "lucene";
				int docId = 0;
				Term term = new Term("content", targetTerm);
				long docFreq = reader.docFreq(term);
				double idf = Math.log((double) totalDocs / (docFreq + 1)) + 1.0;

				int tf = 0;
				PostingsEnum postings = MultiTerms.getTermPostingsEnum(reader, 
						"content", new BytesRef(targetTerm));
				if (postings != null && postings.advance(docId) == docId) {
					tf = postings.freq();
				}

				System.out.println("Total Documents: " + totalDocs);
				System.out.println("Term: " + targetTerm);
				System.out.println("TF (Frequency in doc0): " + tf);
				System.out.println("DocFreq (Docs containing term): " + docFreq);
				System.out.printf("IDF Value: %.4f%n", idf);
				System.out.printf("TF-IDF Score: %.4f%n", tf * idf);
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	private static void addDocument(IndexWriter writer, String content) {
		try {
			Document document = new Document();
			document.add(new TextField("content", content, Field.Store.YES));
			writer.addDocument(document);
		} catch (IOException e) {
			throw new UncheckedIOException("写入 Lucene 文档失败", e);
		}
	}
}