//    Copyright 2011 Brad Block, Pawjaw, LLC. (brad@pawjaw.com)
//
//    This program is free software: you can redistribute it and/or modify
//    it under the terms of the GNU General Public License as published by
//    the Free Software Foundation, either version 3 of the License, or
//    (at your option) any later version.
//
//    This program is distributed in the hope that it will be useful,
//    but WITHOUT ANY WARRANTY; without even the implied warranty of
//    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//    GNU General Public License for more details.
//
//    You should have received a copy of the GNU General Public License
//    along with this program.  If not, see <http://www.gnu.org/licenses/>.

package com.pawjaw.topics;

import com.pawjaw.linalg.Matrix;
import com.pawjaw.utils.OptArg;
import com.pawjaw.utils.OptArg.OptType;
import com.pawjaw.utils.OptArgException;
import com.pawjaw.utils.ProgressCallback;
import com.pawjaw.utils.SortedPairs;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.List;

public class Main {
    private static final int NONE = 1;
    private static final int SYSTEM = 2;
    private static final int CGLDA = 4;
    private static final int CVHDP = 8;

    private int verbosity = 2, model;
    private long start = System.currentTimeMillis();
    private String timing;
    private List<int[]> n_ks;
    private List<Double> avglls;
    private List<double[]> E_n_ks;
    private Integer[] dk_k;
    private double[][] phi_kw, theta_dk;

    public Main(String[] args) {
        List<OptArg> optargs;
        try {
            initOptArgs();
            try {
                if((optargs = OptArg.parseArgs(args)) == null || optargs.isEmpty()) {
                    err("error parsing command line, run with -h flag for help\n");
                    System.exit(1);
                }
                for(OptArg optarg : optargs)
                    if(optarg.short_opt.equals("-h")) {
                        showHelp();
                        System.exit(0);
                    }
                check();
            } catch(OptArgException e) {
                if(verbosity >= 1) {
                    err("error parsing command line, run with -h flag for help\n");
                    if(e.optarg == null)
                        err("  " + e.getMessage() + "\n");
                    else {
                        err("  with option " + e.optarg.label() + " " + e.getMessage() + "\n");
                        err("  " + e.optarg.longHelp() + "\n");
                    }
                }
                System.exit(1);
            }
            if(verbosity == 2)
                showParams();
            runModel();
            writeResults();
            info("done\n");
        } catch(Exception f) {
            if(verbosity >= 1)
                f.printStackTrace(System.err);
        }
    }

    private void initOptArgs() {
        new OptArg(NONE, OptType.FLAG, "-h", "--help", null, "display help");
        new OptArg(SYSTEM, OptType.INT, "-v", "--verbosity", 2, "verbosity, 0 = quiet, 1 = errors, 2 = info/errors", new Object[] { 0, 1, 2});
        new OptArg(SYSTEM, OptType.INT, "-m", "--topic-model", null, "topic model to use, 0 = CGLDA, 1 = CVHDP", new Object[] { 0, 1 });
        new OptArg(SYSTEM, OptType.STRING, "-it", "--training-matrix", null, "input training count data file in \"row col count\" sparse matrix text format (first row/col = 0, documents in columns)");
        new OptArg(SYSTEM, OptType.STRING, "-ie", "--test-matrix", null, "input test count data file for held-out log-likelihood calculations in \"row col count\" sparse matrix text format  (first row/col = 0, document in columns)");
        new OptArg(CGLDA, OptType.STRING, "-ib", "--beta-matrix", null, "input Dirichlet alpha parameters file to specify CGLDA beta prior parameter for each word (row) and topic (column) in dense matrix text format");
        new OptArg(SYSTEM, OptType.STRING, "-id", "--dict", null, "input dictionary file with one word per line corresponding to each row in data matrices in text format");
        new OptArg(SYSTEM, OptType.STRING, "-opm", "--phi-matrix", null, "output phi topic-word mixture file for each word (row) and topic (column) in dense matrix text format");
        new OptArg(SYSTEM, OptType.STRING, "-otm", "--theta-matrix", null, "output theta document-topic mixture file for each topic (row) and document (column) in dense matrix text format");
        new OptArg(SYSTEM, OptType.STRING, "-ol", "--lls", null, "output held-out log-likelihood file one line per estimate in text format");
        new OptArg(SYSTEM, OptType.STRING, "-on", "--nks", null, "output token topic assignment count file for each estimate (row) and topic (column) in dense matrix text format");
        new OptArg(SYSTEM, OptType.STRING, "-oc", "--timing", null, "output timing file");
        new OptArg(SYSTEM, OptType.STRING, "-op", "--params", null, "output parameter file");
        new OptArg(SYSTEM, OptType.STRING, "-ot", "--top-topic-word", null, "output top topic words file one line for each topic (row), words separated by tabs sorted by descending phi value");
        new OptArg(SYSTEM, OptType.POSITIVEINT, "-tt", "--top-topic-word-count", 20, "number of words to display per topic in top-topic-word file");
        new OptArg(CGLDA|CVHDP, OptType.POSITIVEINT, "-k", "--topics", null, "number of topics (CGLDA) or topic truncation threshold (CVHDP)");
        new OptArg(CVHDP, OptType.POSITIVEINT, "-i", "--iterations", null, "number of variational iterations");
        new OptArg(CGLDA, OptType.POSITIVEINT, "-l", "--samples", null, "number of estimates to use from Gibbs sampling phase");
        new OptArg(CGLDA, OptType.POSITIVEINT, "-q", "--sample-freq", 1, "number of Gibbs iterations between estimate sampling");
        new OptArg(CGLDA|CVHDP, OptType.NONNEGATIVEINT, "-e", "--estimates", 1, "number of log-likelihood calculations to record");
        new OptArg(CGLDA, OptType.NONNEGATIVEINT, "-n", "--burnin", 1, "number of burnin iterations");
        new OptArg(CVHDP, OptType.POSITIVEINT, "-t", "--threads", 1, "number of threads to use for variational inference");
        new OptArg(CVHDP, OptType.BOOLEAN, "-g", "--group-tokens", true, "0 = calculate variational posterior for each token, 1 = for each word");
        new OptArg(CGLDA, OptType.POSITIVEDOUBLE, "-as", "--alpha-sum", 1.0, "uniform Dirichlet alpha parameter, gets divided by number of topics (prior on theta, document-topic mixture)");
        new OptArg(CGLDA, OptType.POSITIVEDOUBLE, "-b", "--beta", 0.01, "uniform Dirichlet alpha parameter (prior on phi, topic-word mixture)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-eb", "--e-beta", 100.0, "expectation of beta (prior on phi, topic-word mixture)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-gb", "--g-beta", 100.0, "geometric expectation of beta (prior on phi, topic-word mixture)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-aa", "--a-alpha", 3.0, "Gamma shape parameter for prior on alpha (prior on theta, document-topic mixture)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-ba", "--b-alpha", 5.0, "Gamma scale parameter for prior on alpha (prior on theta, document-topic mixture)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-ab", "--a-beta", 3.0, "Gamma shape parameter for prior on beta (prior on phi, topic-word mixture)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-bb", "--b-beta", 5.0, "Gamma scale parameter for prior on beta (prior on phi, topic-word mixture)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-ag", "--a-gamma", 6.0, "Gamma shape parameter for prior on gamma (prior on pi, topics)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-bg", "--b-gamma", 8.0, "Gamma scale parameter for prior on gamma (prior on pi, topics)");
        new OptArg(CVHDP, OptType.POSITIVEDOUBLE, "-at", "--a-tau", 40.0, "uniform Dirichlet alpha parameter for prior on tau, gets divided by vocabulary size (prior on phi, topic-word mixture)");
    }

    private void checkCGLDA() throws OptArgException, IOException {
        OptArg optarg;
        String s = null;
        BufferedReader br;
        model = CGLDA;
        if((optarg = OptArg.optarg_flag.get("-l")).arg == null)
            s = "CGLDA requires -l option";
        if(s != null)
            throw new OptArgException(s, optarg);
        if((optarg = OptArg.optarg_flag.get("-ib")).arg != null) {
            info("loading beta file " + optarg.arg + "\n");
            br = new BufferedReader(new FileReader(optarg.getFileToRead()));
            optarg.data = Matrix.read(br);
            br.close();
            info("loaded beta file: " + ((double[][])optarg.data).length + "x" + ((double[][])optarg.data)[0].length + "\n");
        }
    }

    private void checkCVHDP() throws OptArgException {
        OptArg optarg;
        String s = null;
        model = CVHDP;
        if((optarg = OptArg.optarg_flag.get("-i")).arg == null)
            s = "CVHDP requires -s option";
        if(s != null)
            throw new OptArgException(s, optarg);
    }

    private void loadFiles() throws OptArgException, IOException {
        String line;
        OptArg optarg;
        List<String> dict;
        BufferedReader br = new BufferedReader(new FileReader((optarg = OptArg.optarg_flag.get("-it")).getFileToRead()));
        info("loading training matrix from file " +  optarg.arg + "\n");
        optarg.data = new Matrix(br);
        br.close();
        info("loaded training matrix: " + ((Matrix)optarg.data).toSummary() + "\n");
        if((optarg = OptArg.optarg_flag.get("-ie")).arg != null) {
            info("loading test matrix from file " + optarg.arg + "\n");
            br = new BufferedReader(new FileReader(optarg.getFileToRead()));
            optarg.data = new Matrix(br);
            br.close();
            info("loaded test matrix: " + ((Matrix)optarg.data).toSummary() + "\n");

        }
        if((optarg = OptArg.optarg_flag.get("-id")).arg != null) {
            info("loading dictionary from file " + optarg.arg + "\n");
            br = new BufferedReader(new FileReader(optarg.getFileToRead()));
            dict = new ArrayList();
            while((line = br.readLine()) != null)
                dict.add(line);
            optarg.data = dict;
            br.close();
            info("loaded dictionary: " + dict.size() + " elements\n");
        }
    }

    private void check() throws OptArgException, IOException {
        OptArg optarg;
        String s = null;
        verbosity = (optarg = OptArg.optarg_flag.get("-v")).arg == null ? 2 : (Integer)optarg.arg;
        if((optarg = OptArg.optarg_flag.get("-m")).arg == null)
            s = "option required";
        else if((optarg = OptArg.optarg_flag.get("-it")).arg == null)
            s = "option required";
        else if((optarg = OptArg.optarg_flag.get("-k")).arg == null)
            s = "option required";
        else if(OptArg.optarg_flag.get("-ol").arg != null && (optarg = OptArg.optarg_flag.get("-ie")).arg == null)
            s = "-ie option required when -ol option is used";
        else if(OptArg.optarg_flag.get("-ot").arg != null && (optarg = OptArg.optarg_flag.get("-id")).arg == null)
            s = "-id option required when -ot option is used";
        else if(OptArg.optarg_flag.get("-e").arg != null && (optarg = OptArg.optarg_flag.get("-ie")).arg == null)
            s = "-ie option required when -e option is used";
        if(s != null)
            throw new OptArgException(s, optarg);
        switch((Integer)OptArg.optarg_flag.get("-m").arg) {
            case 0: checkCGLDA(); break;
            case 1: checkCVHDP(); break;
        }
        loadFiles();
    }

    private void showHelp() {
        StringBuilder sb = new StringBuilder();
        sb.append("command line options:\n\n");
        for(int i : new int[] { NONE, SYSTEM })
            for(OptArg optarg : OptArg.optargs_applies.get(i)) {
                sb.append("    ");
                sb.append(optarg.longHelp());
                sb.append("\n\n");
            }
        sb.append("  CGLDA:\n\n");
        for(OptArg optarg : OptArg.optargs_applies.get(CGLDA)) {
            sb.append("    ");
            sb.append(optarg.longHelp());
            sb.append("\n\n");
        }
        sb.append("  CVHDP:\n\n");
        for(OptArg optarg : OptArg.optargs_applies.get(CVHDP)) {
            sb.append("    ");
            sb.append(optarg.longHelp());
            sb.append("\n\n");
        }
        info(sb.toString());
    }

    private String paramsToString() {
        StringBuilder sb = new StringBuilder();
        for(OptArg optarg : OptArg.optargs_applies.get(SYSTEM)) {
            sb.append(optarg.label());
            sb.append(": ");
            sb.append(optarg.argOrDefault());
            sb.append("\n");
        }
        for(OptArg optarg : OptArg.optargs_applies.get(model)) {
            sb.append(optarg.label());
            sb.append(": ");
            sb.append(optarg.argOrDefault());
            sb.append("\n");
        }
        return sb.toString();
    }

    private void showParams() {
        StringBuilder sb = new StringBuilder();
        sb.append("parameter settings:\n");
        sb.append(paramsToString());
        info(sb.toString());
    }

    private void runModel() throws InterruptedException {
        CGLDA cglda;
        CVHDP cvhdp;
        Integer topics = (Integer)OptArg.optarg_flag.get("-k").arg;
        Matrix trainingm = (Matrix)OptArg.optarg_flag.get("-it").data, testm = (Matrix)OptArg.optarg_flag.get("-ie").data;
        double[][] beta_wk = (double[][])OptArg.optarg_flag.get("-ib").data;
        switch(model) {
            case CGLDA:
                info("running CGLDA\n");
                if(testm != null & beta_wk != null)
                    cglda = new CGLDA(trainingm, testm, topics,
                            (Double)OptArg.optarg_flag.get("-as").argOrDefault() / (double)topics,
                            beta_wk,
                            (Integer)OptArg.optarg_flag.get("-n").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-q").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-l").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-e").argOrDefault(),
                            new Notifier());
                else if(testm != null)
                    cglda = new CGLDA(trainingm, testm, topics,
                            (Double)OptArg.optarg_flag.get("-as").argOrDefault() / (double)topics,
                            (Double)OptArg.optarg_flag.get("-b").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-n").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-q").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-l").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-e").argOrDefault(),
                            new Notifier());
                else if(beta_wk != null)
                    cglda = new CGLDA(trainingm, topics,
                            (Double)OptArg.optarg_flag.get("-as").argOrDefault() / (double)topics,
                            beta_wk,
                            (Integer)OptArg.optarg_flag.get("-n").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-q").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-l").argOrDefault(),
                            new Notifier());
                else
                    cglda = new CGLDA(trainingm, topics,
                            (Double)OptArg.optarg_flag.get("-as").argOrDefault() / (double)topics,
                            (Double)OptArg.optarg_flag.get("-b").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-n").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-q").argOrDefault(),
                            (Integer)OptArg.optarg_flag.get("-l").argOrDefault(),
                            new Notifier());
                phi_kw = cglda.phi_kw;
                theta_dk = cglda.theta_dk;
                avglls = cglda.avglls;
                n_ks = cglda.n_ks;
                timing = cglda.timing;
                dk_k = cglda.dk_k;
                break;
            case CVHDP:
                info("running CVHDP\n");
                cvhdp = new CVHDP(trainingm, testm, topics,
                        (Integer)OptArg.optarg_flag.get("-t").argOrDefault(),
                        (Integer)OptArg.optarg_flag.get("-i").argOrDefault(),
                        (Boolean)OptArg.optarg_flag.get("-g").argOrDefault(),
                        (Integer)OptArg.optarg_flag.get("-e").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-eb").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-gb").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-aa").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-ba").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-ab").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-bb").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-ag").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-bg").argOrDefault(),
                        (Double)OptArg.optarg_flag.get("-at").argOrDefault(),
                        new Notifier());
                phi_kw = cvhdp.phi_kw;
                theta_dk = cvhdp.theta_dk;
                avglls = cvhdp.avglls;
                E_n_ks = cvhdp.E_n_ks;
                timing = cvhdp.timing;
                dk_k = cvhdp.dk_k;
                break;
        }
    }

    private void closeOrFlush(PrintWriter pw, OptArg optarg) {
        if(!((String)optarg.arg).equals("-"))
            pw.close();
        else {
            pw.println();
            pw.flush();
        }
    }

    private void writeResults() throws IOException {
        int i = 0, I;
        OptArg optarg;
        PrintWriter pw;
        SortedPairs pairs;
        List<String> dict;
        double[] phik_w;
        double[][] array = null;
        if((optarg = OptArg.optarg_flag.get("-opm")).arg != null) {
            info("writing phi matrix to " + optarg.arg + "\n");
            pw = ((String)optarg.arg).equals("-") ? new PrintWriter(System.out) : new PrintWriter(new FileWriter(optarg.getFileToWrite()));
            pw.write("# phi matrix (topic-word multinomial mixture)\n# each row corresponds to a topic, each column corresponds to a word\n");
            Matrix.write(phi_kw, pw);
            closeOrFlush(pw, optarg);
        }
        if((optarg = OptArg.optarg_flag.get("-otm")).arg != null) {
            info("writing theta matrix to " + optarg.arg + "\n");
            pw = ((String)optarg.arg).equals("-") ? new PrintWriter(System.out) : new PrintWriter(new FileWriter(optarg.getFileToWrite()));
            pw.write("# theta matrix (document-topic multinomial mixture)\n# each row corresponds to a document, each column corresponds to a topic\n");
            Matrix.write(theta_dk, pw);
            closeOrFlush(pw, optarg);
        }
        if((optarg = OptArg.optarg_flag.get("-ol")).arg != null && avglls != null && !avglls.isEmpty()) {
            info("writing average log-likelihood results to " + optarg.arg + "\n");
            pw = ((String)optarg.arg).equals("-") ? new PrintWriter(System.out) : new PrintWriter(new FileWriter(optarg.getFileToWrite()));
            pw.write("# average log-likelihood\n# each row contains the log-likelihood of the held-out data set\n");
            for(double avgll : avglls)
                pw.println(avgll);
            closeOrFlush(pw, optarg);
        }
        if((optarg = OptArg.optarg_flag.get("-oc")).arg != null && timing != null && !timing.isEmpty()) {
            info("writing timing information to " + optarg.arg + "\n");
            pw = ((String)optarg.arg).equals("-") ? new PrintWriter(System.out) : new PrintWriter(new FileWriter(optarg.getFileToWrite()));
            pw.write("# model running time information\n");
            pw.write(timing);
            closeOrFlush(pw, optarg);
        }
        if((optarg = OptArg.optarg_flag.get("-op")).arg != null) {
            info("writing parameter information to " + optarg.arg + "\n");
            pw = ((String)optarg.arg).equals("-") ? new PrintWriter(System.out) : new PrintWriter(new FileWriter(optarg.getFileToWrite()));
            pw.write("# experiment parameter settings\n");
            pw.write(paramsToString());
            closeOrFlush(pw, optarg);
        }
        if((optarg = OptArg.optarg_flag.get("-on")).arg != null &&
                ((model == CGLDA && n_ks != null && !n_ks.isEmpty()) ||
                (model == CVHDP && E_n_ks != null && !E_n_ks.isEmpty()))) {
            info("writing token topic assignment counts to " + optarg.arg + "\n");
            pw = ((String)optarg.arg).equals("-") ? new PrintWriter(System.out) : new PrintWriter(new FileWriter(optarg.getFileToWrite()));
            pw.write("# token topic assignment counts\n# each row corresponds to an estimate, each column corresponds to a topic\n");
            switch(model) {
                case CGLDA:
                    array = new double[n_ks.size()][];
                    for(int[] n_k : n_ks)
                        array[i++] = Matrix.asDoubles(n_k);
                    break;
                case CVHDP:
                    array = new double[E_n_ks.size()][];
                    for(double[] E_n_k : E_n_ks)
                        array[i++] = E_n_k;
                    break;
            }
            Matrix.write(array, pw);
            closeOrFlush(pw, optarg);
        }
        if((optarg = OptArg.optarg_flag.get("-ot")).arg != null && (dict = (List)OptArg.optarg_flag.get("-id").data) != null) {
            info("writing top topic words to " + optarg.arg + "\n");
            pw = ((String)optarg.arg).equals("-") ? new PrintWriter(System.out) : new PrintWriter(new FileWriter(optarg.getFileToWrite()));
            pw.write("# top word for each topic according to phi matrix\n# each row corresponds to a topic\n");
            I = (Integer)OptArg.optarg_flag.get("-tt").argOrDefault();
            for(int dk : dk_k) {
                pairs = new SortedPairs((phik_w = phi_kw[dk]).length);
                i = 0;
                for(double phi : phik_w)
                    pairs.add(new double[] { phi, i++ });
                i = 0;
                for(double[] pair : pairs.sortDescending()) {
                    if(i == I)
                        break;
                    else if(i++ > 0)
                        pw.write(", ");
                    pw.write(dict.get((int)pair[1]));
                }
                pw.println();
            }
            closeOrFlush(pw, optarg);
        }
    }

    private void info(String msg) {
        if(verbosity >= 2)
            System.out.print("[" + (System.currentTimeMillis() - start) + "] " + msg);
    }

    private void err(String msg) {
        if(verbosity >= 1)
            System.err.print(msg);
    }

    private class Notifier implements ProgressCallback {
        public void notify(int count, int total) {
            if(count % (int)Math.max(total / 20, 1) == 0)
                info("  " + count + " / " + total + "\n");
        }
    }

    public static void main(String[] args) {
        new Main(args);
    }
}
