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

import java.io.*;

public class State implements Serializable {
    public double eta;
    public double[][] stats;
    public int numberOfDocuments;

    public State(double eta, int numberOfTopics, int numberOfTerms) {
        this.eta = eta;
        stats = new double[numberOfTopics][numberOfTerms];
    }

    public double[][] getLambda() {
        return Compute.add(stats, eta);
    }

    public double[][] getELogBeta() {
        return Compute.dirichletExpectation(getLambda());
    }

    public void addToStats(double[][] add) throws Exception {
        Compute.addMutable(stats, add);
    }

    public void blend(double rho, State other) throws Exception {
        Compute.mulMutable(stats, (1.0 - rho));

        double scale;
        if (other.numberOfDocuments == 0 || numberOfDocuments == other.numberOfDocuments) {
            scale = 1.0;
        } else {
            scale = 1.0 * numberOfDocuments / (double)other.numberOfDocuments;
        }

        Compute.addMutable(stats, Compute.mul(other.stats, rho * scale));
    }

    public int getNumberOfTokens() {
        return stats[0].length;
    }

    public int getNumberOfTopics() {
        return stats.length;
    }

    public void saveToFile(File file) throws IOException {
        OutputStream fileOutputStream = new FileOutputStream(file);
        OutputStream buffer = new BufferedOutputStream(fileOutputStream);
        ObjectOutput output = new ObjectOutputStream(buffer);
        output.writeObject(this);

        output.flush();
        output.close();
    }

    public static State loadFromFile(File file) throws IOException, ClassNotFoundException {
        InputStream inputStream = new FileInputStream(file);
        InputStream buffer = new BufferedInputStream(inputStream);
        ObjectInput input = new ObjectInputStream (buffer);

        return (State) input.readObject();
    }
}
