package cn.edu.bjut.chapter2;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class TFIDFAnalyzer {
	private List<Document> documents;

	public TFIDFAnalyzer(List<Document> documents) {
		this.documents = documents;
	}

	public Map<String, Double> calculateIDF() {
		Map<String, Double> idf = new HashMap<String, Double>();
		int totalDocuments = documents.size();

		for (Document document : documents) {
			Set<String> tokens = document.getTermFrequency().keySet();
			for (String token : tokens) {
				idf.put(token, idf.getOrDefault(token, 0.0) + 1);
			}
		}

		for (String key : idf.keySet()) {
			idf.put(key, Math.log(totalDocuments / idf.get(key))); 
		}
		return idf;
	}

	public Map<String, Double> calculateTFIDF() {
		Map<String, Double> tfidf = new HashMap<String, Double>();
		Map<String, Double> idf = calculateIDF();

		for (Document document : documents) {
			Map<String, Double> tf = document.getTermFrequency();
			for (String term : tf.keySet()) {
				double tfidfValue = tf.get(term) * idf.getOrDefault(term, 0.0);
				tfidf.put(term, tfidf.getOrDefault(term, 0.0) + tfidfValue);
			}
		}
		return tfidf;
	}

	public static void main(String[] args) {
		List<Document> documents = new ArrayList<Document>();
		documents.add(new Document("A rose is very beautiful."));
		documents.add(new Document("A rose is a flower, which is a rose."));

		TFIDFAnalyzer tfidfAnalyzer = new TFIDFAnalyzer(documents);
		Map<String, Double> tfidf = tfidfAnalyzer.calculateTFIDF();
		for (Map.Entry<String, Double> entry : tfidf.entrySet()) {
			System.out.println(entry.getKey() + ": " + entry.getValue());
		}
	}
}