package cz.vutbr.feec.utko.xonder02.lda;

import java.util.Iterator;
import java.util.LinkedList;

public class Corpus implements Iterable<Document> {
    private LinkedList<Document> documents = new LinkedList<Document>();

    public void putDocument(Document document) {
        documents.add(document);
    }

    public int documentCount() {
        return documents.size();
    }

    public Iterator<Document> iterator() {
        return documents.iterator();
    }

    public void merge(Corpus corpus) {
        for (Document document : corpus) {
            documents.add(document);
        }
    }

    public Corpus[] getChunks(int length) {
        int chunkCount = (int) Math.ceil(documents.size() / (double)length);
        Corpus[] chunks = new Corpus[chunkCount];

        int i = 0;
        for (Document doc : documents) {
            int chunkId = i++ / length;

            if (chunks[chunkId] == null) {
                chunks[chunkId] = new Corpus();
            }

            chunks[chunkId].putDocument(doc);
        }

        return chunks;
    }

    public int getNumberOfTokens() {
        int maxTermId = 0;

        for (Document doc : documents) {
            for (Term term : doc) {
                if (term.id > maxTermId) {
                    maxTermId = term.id;
                }
            }
        }

        return maxTermId + 1;
    }

    public double averageNumberOfTokensPerDocument() {
        double sum = 0;
        for (Document doc : documents) {
            sum += doc.getTermCount();
        }
        return sum / documentCount();
    }

    public static Corpus generate(String[][] documentTokens, Dictionary dictionary, boolean removeNotExistsTokens) {
        Corpus corpus = new Corpus();

        for (String[] documentToken : documentTokens) {
            Document document = Document.generate(documentToken, dictionary, removeNotExistsTokens);
            corpus.putDocument(document);
        }

        return corpus;
    }
}
