#!/usr/bin/env python
"""
Tool to combine a collection of base topic models, generated by NMF, to produce a single ensemble topic model 
using the weighted term coassociation approach (WTCA). This will produce both the ensemble rankings and document partition.

Sample usage:
python ensemble-weighted-coassoc.py -k 5 -m wikipedia2016-w2v-cbow-d100.bin -t 10 data/bbc.pkl data/bbc/nmf_k05/*partition* data/bbc/nmf_k05/*rank* -o results/bbc

"""
import os, sys, random, gensim, textutil
import numpy as np
from pathlib import Path
import logging as log
from optparse import OptionParser
import unsupervised.nmf, unsupervised.rankings, unsupervised.util
from unsupervised.dual import WeightedCoassociationModel

# --------------------------------------------------------------

def main():
	parser = OptionParser(usage="usage: %prog [options] dataset_file partition_file1 partition_file2 rankings_file1 rankings_file2...")
	parser.add_option("-k", action="store", type="int", dest="k", help="number of topics in output model", default=5)
	parser.add_option("-m", "--model", action="store", type="string", dest="model_path", help="path to an embedding file", default=None)
	parser.add_option("-t", "--top", action="store", type="int", dest="top_terms", help="number of top terms to use from base models", default=10)
	parser.add_option("-o","--outdir", action="store", type="string", dest="dir_out", help="output directory (default is current directory)", default=None)
	parser.add_option("-v", "--verbose", action="store_true", dest="verbose", help="display topic descriptors")
	parser.add_option('-d','--debug',type="int",help="Level of log output; 0 is less, 5 is all", default=3)
	(options, args) = parser.parse_args()

	if( len(args) < 3 ):
		parser.error( "Must specify dataset file and at least two partition/ranking files" )
	log_level = max(50 - (options.debug * 10), 10)
	log.basicConfig(level=log_level, format='%(message)s')
	k = options.k
	top = options.top_terms
	model_path = options.model_path

	# Output directory for results
	if options.dir_out is None:
		dir_out = os.getcwd()
	else:
		dir_out = options.dir_out	
		if not os.path.exists(dir_out):
			os.makedirs(dir_out)	

	# Load the preprocessed dataset
	log.info( "Loading data from %s ..." % args[0] )
	(X, all_terms, all_doc_ids, _) = textutil.load_corpus( args[0] )
	num_documents =  len(all_doc_ids)
	log.info( "Read dataset with %d documents, %d terms" % (  num_documents, len(all_terms) ) )

	# Load the embedding model
	if "-ft" in model_path:
		log.info("Loading FastText embedding from %s ..." % model_path )
		embedding = gensim.models.FastText.load(model_path)
		vocab = set(embedding.wv.vocab.keys() )
	else:
		log.info("Loading Word2vec embedding from %s ..." % model_path )
		embedding = gensim.models.KeyedVectors.load_word2vec_format(model_path, binary=True)	
		vocab = set(embedding.vocab.keys())
	log.info("Embedding has vocabulary of size %d" % len(vocab) )

	# Find all of the input files
	ranking_filepaths, partition_filepaths = [], []
	for in_path in sorted(args[1:]):
		in_path = Path( in_path )
		if "ranks" in in_path.stem:
			ranking_filepaths.append( in_path )	
		elif "partition" in in_path.stem:
			partition_filepaths.append( in_path )	
		else:
			log.warning("Ignoring input file of unknown type %s" % in_path )
	# Must have equal number of partitions and rankings
	if len(ranking_filepaths) != len(partition_filepaths):
		log.error("Cannot process input files - unequal number of partition and ranking files specified")
		sys.exit(1)

	# Process each specified base topic model
	model = WeightedCoassociationModel( k, num_documents, embedding, top)
	log.info("Processing %d base topic models. Building representation using top %d terms ..." % (len(ranking_filepaths), top) )
	factors = []
	for base_idx in range( len(ranking_filepaths) ):
		# load the ranks file 
		log.debug("%d: Loading %s" % ( base_idx+1, ranking_filepaths[base_idx]) )
		rankings, _ = unsupervised.util.load_term_rankings( ranking_filepaths[base_idx] )
		trunc_rankings = unsupervised.rankings.truncate_term_rankings( rankings, top )
		# load the corresponding partition file
		log.debug("%d: Loading %s" % ( base_idx+1, partition_filepaths[base_idx]) )
		partition, partition_doc_ids = unsupervised.util.load_partition( partition_filepaths[base_idx] )
		# add them to the current model
		model.add( trunc_rankings, partition )

	# Apply the final topic modelling process
	log.info("Applying dual co-association analysis...")
	model.apply()
	log.info("Generated %d term clusters" % len(model.term_clusters) )

	# Print out the top terms?
	if options.verbose:
		tab = unsupervised.rankings.DescriptorTable( model.term_clusters, top )
		log.info( tab.format() )

	# Write term rankings
	ranks_out_path = os.path.join( dir_out, "ranks_weighted_k%02d.pkl"  % k )
	log.info( "Writing ensemble term rankings to %s" % ranks_out_path )
	unsupervised.util.save_term_rankings( ranks_out_path, model.term_clusters )

	# Write document partition
	doc_partition_out_path = os.path.join( dir_out, "partition_weighted_k%02d.pkl"  % k )
	log.info( "Writing ensemble document partition to %s" % doc_partition_out_path )
	unsupervised.util.save_partition( doc_partition_out_path, model.partition, all_doc_ids )

# --------------------------------------------------------------

if __name__ == "__main__":
	main()
