$LOAD_PATH.unshift(File.dirname(__FILE__)) unless $LOAD_PATH.include?(File.dirname(__FILE__))

require 'lda-ruby/lda'
require 'lda-ruby/document/document'
require 'lda-ruby/document/data_document'
require 'lda-ruby/document/text_document'
require 'lda-ruby/corpus/corpus'
require 'lda-ruby/corpus/data_corpus'
require 'lda-ruby/corpus/text_corpus'
require 'lda-ruby/corpus/directory_corpus'
require 'lda-ruby/vocabulary'

module Lda
  class Lda
    attr_reader :vocab, :corpus

    def initialize(corpus)
      load_default_settings

      @vocab = nil
      self.corpus = corpus
      @vocab = corpus.vocabulary.to_a if corpus.vocabulary

      @phi = nil
    end

    def load_default_settings
      self.max_iter = 20
      self.convergence = 1e-6
      self.em_max_iter = 100
      self.em_convergence = 1e-4
      self.num_topics = 20
      self.init_alpha = 0.3
      self.est_alpha = 1

      [20, 1e-6, 100, 1e-4, 20, 0.3, 1]
    end

    def load_corpus(filename)
      @corpus = Corpus.new
      @corpus.load_from_file(filename)

      true
    end

    def load_vocabulary(vocab)
      if vocab.is_a?(Array)
        @vocab = Marshal::load(Marshal::dump(vocab)) # deep clone array
      elsif vocab.is_a?(Vocabulary)
        @vocab = vocab.to_a
      else
        @vocab = File.open(vocab, 'r') { |f| f.read.split(/\s+/) }
      end

      true
    end

    #
    # Visualization method for printing out the top +words_per_topic+ words
    # for each topic.
    #
    # See also +top_words+.
    #
    def print_topics(words_per_topic = 10)
      raise 'No vocabulary loaded.' unless @vocab

      beta.each_with_index do |topic, topic_num|
        # Sort the topic array and return the sorted indices of the best scores
        indices = topic.zip((0...@vocab.size).to_a).sort { |x| x[0] }.map { |_i, j| j }.reverse[0...words_per_topic]

        puts "Topic #{topic_num}"
        puts "\t#{indices.map { |i| @vocab[i] }.join("\n\t")}"
        puts ''
      end

      nil
    end

    #
    # After the model has been run and a vocabulary has been loaded, return the
    # +words_per_topic+ top words chosen by the model for each topic.  This is
    # returned as a hash mapping the topic number to an array of top words
    # (in descending order of importance).
    #
    #   topic_number => [w1, w2, ..., w_n]
    #
    # See also +print_topics+.
    #
    def top_word_indices(words_per_topic = 10)
      raise 'No vocabulary loaded.' unless @vocab

      # find the highest scoring words per topic
      topics = {}

      beta.each_with_index do |topic, topic_num|
        topics[topic_num] = topic.zip((0...@vocab.size).to_a).sort { |x| x[0] }.map { |_i, j| j }.reverse[0...words_per_topic]
      end

      topics
    end

    def top_words(words_per_topic = 10)
      output = {}

      topics = top_word_indices(words_per_topic)
      topics.each_pair do |topic_num, words|
        output[topic_num] = words.map { |w| @vocab[w] }
      end

      output
    end

    #
    # Get the phi matrix which can be used to assign probabilities to words
    # belonging to a specific topic in each document.  The return value is a
    # 3D matrix:  num_docs x doc_length x num_topics.  The value is cached
    # after the first call, so if it needs to be recomputed, set the +recompute+
    # value to true.
    #
    def phi(recompute = false)
      if @phi.nil? || recompute
        @phi = self.compute_phi
      end

      @phi
    end

    #
    # Compute the average log probability for each topic for each document in the corpus.
    # This method returns a matrix:  num_docs x num_topics with the average log probability
    # for the topic in the document.
    #
    def compute_topic_document_probability
      outp = []

      @corpus.documents.each_with_index do |doc, idx|
        tops = [0.0] * num_topics
        ttl  = doc.counts.inject(0.0) { |sum, i| sum + i }
        phi[idx].each_with_index do |word_dist, word_idx|
          word_dist.each_with_index do |top_prob, top_idx|
            tops[top_idx] += Math.log(top_prob) * doc.counts[word_idx]
          end
        end
        tops = tops.map { |i| i / ttl }
        outp << tops
      end

      outp
    end

    #
    # String representation displaying current settings.
    #
    def to_s
      outp = ['LDA Settings:']
      outp << '    Initial alpha: %0.6f'.format(init_alpha)
      outp << '      # of topics: %d'.format(num_topics)
      outp << '   Max iterations: %d'.format(max_iter)
      outp << '      Convergence: %0.6f'.format(convergence)
      outp << 'EM max iterations: %d'.format(em_max_iter)
      outp << '   EM convergence: %0.6f'.format(em_convergence)
      outp << '   Estimate alpha: %d'.format(est_alpha)

      outp.join("\n")
    end
  end
end
