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


import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

public class Lda {

    private int numberOfTopics;
    private int numberOfTerms;
    private double alpha;
    private double eta;
    private State state;
    private double[][] expELogBeta;
    private int iterations;
    private double gammaThreshold;
    private int numberOfUpdates;
    private int chunkSize;

    public Lda(Corpus corpus, int numberOfTopics, int iterations, double gammaThreshold, int chunkSize) throws Exception {
        this(numberOfTopics, iterations, gammaThreshold, chunkSize);

        numberOfTerms = corpus.getNumberOfTokens();
        alpha = 1.0 / numberOfTopics;
        eta = 1.0 / numberOfTopics;

        state = new State(eta, numberOfTopics, numberOfTerms);
        state.stats = Compute.gammaDistribution(100.0, 1.0 / 100.0, numberOfTopics, numberOfTerms);
        sync();

        update(corpus);
    }

    public Lda(State state, int iterations, double gammaThreshold, int chunkSize) {
        this(state.getNumberOfTopics(), iterations, gammaThreshold, chunkSize);

        numberOfTerms = state.getNumberOfTokens();
        alpha = 1.0 / numberOfTopics;
        eta = 1.0 / numberOfTopics;

        this.state = state;
        sync();
    }

    private Lda(int numberOfTopics, int iterations, double gammaThreshold, int chunkSize) {
        this.numberOfTopics = numberOfTopics;
        this.iterations = iterations;
        this.gammaThreshold = gammaThreshold;
        this.chunkSize = chunkSize;
    }

    public List<DocumentTopicDistance> findTopic(Document document, double eps) throws Exception {
        Corpus corpus = new Corpus();
        corpus.putDocument(document);

        InferenceOutput inferenceOutput = inference(corpus, false);
        double[][] topicDistance = Compute.div(inferenceOutput.gamma, Compute.sum(inferenceOutput.gamma));

        ArrayList<DocumentTopicDistance> documentTopicDistances = new ArrayList<DocumentTopicDistance>();

        for (int i = 0; i < topicDistance[0].length; i++) {
            double distance = topicDistance[0][i];

            if (distance > eps) {
                documentTopicDistances.add(new DocumentTopicDistance(topicDistance[0][i], getTopic(i)));
            }
        }

        Collections.sort(documentTopicDistances, new Comparator<DocumentTopicDistance>() {
            @Override
            public int compare(DocumentTopicDistance o1, DocumentTopicDistance o2) {
                return o2.getDistance().compareTo(o1.getDistance());
            }
        });

        return documentTopicDistances;
    }

    public Topic findBestTopic(Document document) throws Exception {
        return findTopic(document, 0.001).get(0).getTopic();
    }

    public Topic getTopic(int topicId) {
        return new Topic(topicId, new double[][] { state.getLambda()[topicId] });
    }

    public int getNumberOfTopics() {
        return numberOfTopics;
    }

    public State getState() {
        return state;
    }

    private void sync() {
        expELogBeta = Compute.exp(state.getELogBeta());
    }

    private void update(Corpus corpus) throws Exception {
        int corpusLength = corpus.documentCount();
        state.numberOfDocuments += corpusLength;

        for (Corpus chunk : corpus.getChunks(chunkSize)) {
            State chunkState = new State(eta, numberOfTopics, numberOfTerms);
            eStep(chunk, chunkState);
            mStep(rho(), chunkState);

            /*System.out.print("Trained ");
            System.out.print(numberOfUpdates);
            System.out.print(" documents from ");
            System.out.println(state.numberOfDocuments);*/
        }
    }

    private double rho() {
        //            offset                                decay
        return Math.pow(1.0 + numberOfUpdates / chunkSize, - 0.5);
    }

    // complete
    private void eStep(Corpus chunk, State chunkState) throws Exception {
        InferenceOutput inferenceOutput = inference(chunk, true);

        chunkState.addToStats(inferenceOutput.stats);
        chunkState.numberOfDocuments += inferenceOutput.gamma.length;
    }

    private void mStep(double rho, State chunkState) throws Exception {
        state.blend(rho, chunkState);
        sync();

        numberOfUpdates += chunkState.numberOfDocuments;
    }

    private InferenceOutput inference(Corpus chunk, boolean collectStats) throws Exception {
        double[][] gamma = Compute.gammaDistribution(100.0, 1.0 / 100.0, chunk.documentCount(), numberOfTopics);
        double[][] eLogTheta = Compute.dirichletExpectation(gamma);
        double[][] expELogTheta = Compute.exp(eLogTheta);

        double[][] stats = null;
        if (collectStats) {
            stats = new double[expELogBeta.length][expELogBeta[0].length];
        }

        int id = 0;
        for (Document doc : chunk) {
            int[] termIds = doc.getTermIds();
            double[][] frequencies = new double[][] { doc.getFrequencies() };

            double[][] gammaForDocument = new double[][] { gamma[id] };
            double[][] expELogThetaForDocument = new double[][] { expELogTheta[id] };
            double[][] expELogBetaForDocument = Compute.getColumns(expELogBeta, termIds);

            double[][] phiNorm = Compute.matrixMultiplication(expELogThetaForDocument, expELogBetaForDocument);
            Compute.addMutable(phiNorm, 1e-100);

            int i = 0;
            for (i = 0; i < iterations; i++) {
                double[][] lastGamma = Compute.copy(gammaForDocument);

                gammaForDocument = Compute.matrixMultiplication(Compute.div(frequencies, phiNorm), Compute.transpose(expELogBetaForDocument));
                Compute.addMutable(Compute.mulMutable(gammaForDocument, expELogThetaForDocument), alpha);

                double[][] eLogThetaForDocument = Compute.dirichletExpectation(gammaForDocument);
                expELogThetaForDocument = Compute.exp(eLogThetaForDocument);

                phiNorm = Compute.matrixMultiplication(expELogThetaForDocument, expELogBetaForDocument);

                double meanChange = Compute.mean(Compute.subAbsolute(gammaForDocument, lastGamma));
                if (meanChange < gammaThreshold) {
                    break;
                }
            }

            gamma[id] = gammaForDocument[0];

            if (collectStats) {
                double[][] sst = Compute.outer(expELogThetaForDocument[0], Compute.div(frequencies, phiNorm)[0]);
                Compute.addMutableForColumns(stats, sst, termIds);
            }
        }

        if (collectStats) {
            Compute.mulMutable(stats, expELogBeta);
        }

        return new InferenceOutput(gamma, stats);
    }

    class InferenceOutput {
        public double[][] gamma;
        public double[][] stats;

        public InferenceOutput(double[][] gamma, double[][] stats) {
            this.gamma = gamma;
            this.stats = stats;
        }
    }
}
