#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
Find features, that are most descriptive for each cluster obtained by Cluto.
"""

__author__ = 'twalkow'


import os
import itertools
import csv
import tempfile,shutil
import scipy
from copy import deepcopy
from scipy import stats
import numpy
from scipy.sparse import lil_matrix
import ltlearn.transformations.selection.weka.wekaattributeselection as was
from ltcore.io.matrix import load_matrix, save_matrix
from ltcore import SparseMatrix
from sklearn.externals import joblib
import json


SCIPY_METHODS = ['MannWhitney', 'WilcoxonRankSum', 'KolmogorovSmirnov', 'KruskalWallis', 'TTest']
WEKA_METHODS = ['InfoGainAttributeEval', 'CfsSubsetEval', 'ConsistencySubsetEval',
                'GainRatioAttributeEval', 'ChiSquaredAttributeEval']


weka_jar = '/usr/share/java/weka-3.6.6.jar'
stats_file = 'descriptive_features.csv'

class Object(object):
	pass
	
def main():
	options={};
	run('dane3','result',options);
	
def run(inputFile, output_dir, options):
    
	attr_eval = 'InfoGainAttributeEval'
	#attr_eval = 'KolmogorovSmirnov'
	if 'method' in options:
			attr_eval=options['method'];
	
	no_features_to_select = 100
	if 'no_features_to_select' in options:
			no_features_to_select=options['no_features_to_select'];
		
	
	if not os.path.exists(output_dir):
		os.mkdir(output_dir)

	temp_folder = tempfile.mkdtemp()
	
	orginal_rowlabels=joblib.load(inputFile+"/rowlabels.pkl");
	cluster_file=inputFile+"/result.clustering";
	clusters=save_cluster_contents(cluster_file, orginal_rowlabels, output_dir)
	
	
	input_matrix=Object();
	input_matrix.class_labels=[0]*len(orginal_rowlabels);
	input_matrix.contexts=orginal_rowlabels;
	input_matrix.column_labels=joblib.load(inputFile+"/columns.pkl");
	input_matrix.boundaries=False;
	input_matrix.attributes=False;
	input_matrix.documents_ids=orginal_rowlabels;
	input_matrix.data=scipy.sparse.lil.lil_matrix(joblib.load(inputFile+"/data.pkl"));
	
	if len(input_matrix.column_labels)<no_features_to_select:
	    no_features_to_select=len(input_matrix.column_labels);
	
	fs = FeatureSelection(attr_eval, no_features_to_select, weka_jar,temp_folder)
	
	switch_labels_to_cluster_ids(cluster_file, input_matrix)
	
	cluster_matrices = split_matrix(input_matrix)
	
	# Run feature selection and generate descriptive_features.csv
	feature_dict = {}
	for cluster_matrix in cluster_matrices:
		cluster_matrix = load_matrix(cluster_matrix.name)
		cluster_id = cluster_matrix.head_word
		feature_dict[int(cluster_id)] = fs.select_features(cluster_matrix)
	save_features_with_stats(feature_dict, os.path.join(output_dir, stats_file))
	#joblib.dump(feature_dict,output_dir+'/feature_dict.pkl');
	#joblib.dump(orginal_rowlabels,output_dir+'/orginal_rowlabels.pkl');
	unify2JSON(feature_dict,clusters,output_dir+"/matrix.json");
	shutil.rmtree(temp_folder)	
	

def unify2JSON(feature_dict,clusters,out_file):
	common_features=[];
	max=50;
	i=0;
	index=[0]*len(feature_dict);
	
	
	while len(common_features)<max:
		n=i%len(feature_dict);
		i=i+1;
		if i>len(feature_dict[n]*len(feature_dict)):
			break;	
		if (index[n]>=len(feature_dict[n])):
			continue;
		
		name=feature_dict[n][index[n]][0];
		if not(name in common_features):
			common_features.append(name);
		index[n]=index[n]+1;	
		
	arr=[[0 for i in range(len(common_features))] for j in range(len(feature_dict))]
	
	for i in range(len(common_features)):
		name=common_features[i];
		for j in range(len(feature_dict)):
			tab=feature_dict[j];
			for k in range(len(tab)):
				if (tab[k][0]==name):
					break;
				
			arr[j][i]=1-1.0*k/len(tab);		
	
	
	rows=[];
	for j in clusters:
		rows.append(clusters[j]);
	obj={};
	obj["matrix"]=arr;
	obj["rows"]=rows;
	obj["columns"]=common_features;
	obj["x"]=len(common_features);
	obj["y"]=len(feature_dict);
	with open(out_file, 'w') as outfile:
		json.dump(obj, outfile)
	
def save_cluster_contents(cluster_file, doc_ids, out_dir):
	"""
	Generate a csv file showing which documents were assigned to which cluster.

	:param cluster_file: Cluto output file: input_file_name.clustering.no_clusters
	:param doc_ids: list with document_ids of clustered documents
	:param out_dir: where to save output
	"""

	with open(cluster_file, 'rb') as f:
		clusters = f.readlines()
	assert len(clusters) == len(doc_ids)

	cluster_contexts = {}

	for i in range(len(clusters)):
		if int(clusters[i].split(' ')[0]) not in cluster_contexts:
			cluster_contexts[int(clusters[i].split(' ')[0])] = [doc_ids[i].strip('\n')]
		else:
			cluster_contexts[int(clusters[i].split(' ')[0])].append(doc_ids[i].strip('\n'))

	res_dict = {}
	for cluster in cluster_contexts:
		res_dict[cluster] = [context.split(';')[-1] for context in cluster_contexts[cluster]]
	save_results(res_dict, os.path.join(out_dir, 'cluster_contents.csv'))
	return res_dict;

def save_features_with_stats(results, filename):
	"""
	Generate a csv file showing features that were characteristic for each cluster along with means
	and standard deviations of those features in this cluster and in all the other clusters.

	results = {cluster_id:[(feature, mean in that cluster, std in that cluster, mean in all the other clusters,
							std in all the other clusters)]}
	"""

	all_res = []
	for r in sorted(results):
		header = ('feature',
				  'mean in cluster %s' % str(r),
				  'standard deviation in cluster %s' % str(r),
				  'mean outside of cluster %s' % str(r),
				  'standard deviation outside of cluster %s' % str(r))
		stats = []
		for res in results[r]:
			feature = res[0]
			mean_in_class = res[1]
			std_in_class = res[2]
			mean_out_of_class = res[3]
			std_out_of_class = res[4]
			stats.append((feature, mean_in_class, std_in_class, mean_out_of_class, std_out_of_class))
		all_res.append(tuple([(r, '', '', '', ''), header] + stats))

	res_rows = itertools.izip_longest(*all_res)
	with open(filename, 'wb') as res_file:
		csv_writer = csv.writer(res_file, delimiter=';')
		for row in res_rows:
			row_list = []
			for r in row:
				if r is not None:
					row_list.extend(list(r))
				else:
					row_list.extend(['', '', '', '', ''])
			csv_writer.writerow(row_list)


def split_matrix(matrix):
		"""
		This splitter splits data from a lexcsd matrix into a number of matrices (one for each possible class).
		The outcome will be a list of matrices, each corresponding to one class, with class_labels true or false.
		e.g. if input matrix has class_labels looking like this:
		class1:class2:class3
		class1:class4
		class2:class3
		Matrix for class2 will have class_labels looking like this:
		true
		false
		true

		@type matrix: lexcsd matrix
		@rtype: list(tempfile.NamedTemporaryFile)
		"""

		matrix = deepcopy(matrix)

		class_labels = matrix.class_labels
		classes = []
		for line in class_labels:
			classes.extend(line.split(':'))
		classes = set(classes)

		#creating dict: {tag: [true, false, false ...]}
		label_list_dict = {}
		for tag in classes:
			new_label_list = []
			for line in class_labels:
				if tag in line.split(':'):
					new_label_list.append('true')
				else:
					new_label_list.append('false')
			label_list_dict[tag] = new_label_list

		matrix.possible_classes = [u'true', u'false']

		#saving the matrices
		temp_files = []
		for tag in classes:
			out_file = tempfile.NamedTemporaryFile(prefix='%s__' % tag, suffix='.tar.bz2')
			matrix.class_labels = label_list_dict[tag]
			matrix.head_word = tag
			temp_files.append(out_file)
			save_matrix(out_file.name, matrix)

		return temp_files


def switch_labels_to_cluster_ids(cluster_file, matrix):
	"""
	Switching labels of the documents used for clustering to ids of clusters, to which these documents were assigned.
	e.g.: sienkiewicz sienkiewicz nalkowska nalkowska zeromski -> 1 1 2 1 2

	:param cluster_file: Cluto output file: input_file_name.clustering.no_clusters
	:param matrix: LexCSD matrix
	"""

	with open(cluster_file, 'rb') as f:
		clusters = [cluster_id.strip('\n') for cluster_id in f.readlines()]
	for i, cluster_id in enumerate(clusters):
		matrix.class_labels[i] = cluster_id
	matrix.possible_classes = set([cluster_id for _, cluster_id in enumerate(clusters)])


def complete_with_zeros(sample, length):
	"""
	Extend a list with zeros, so that it has specified length in the end.
	"""

	return numpy.append(sample, numpy.zeros(length - len(sample)), 0)


def split_matrix_by_classes(matrix):
	"""
	Splits matrix into smaller matrices, one for each of the possible classes.
	"""

	oneclass_matrices = []

	cls_dict = {}
	for i, cls in enumerate(matrix.class_labels):
		for c in cls.split(':'):
			if c in cls_dict:
				cls_dict[c].append(i)
			else:
				cls_dict[c] = [i]

	for cls in cls_dict:
		matrix.head_word = cls
		oneclass_matrices.append(cut_matrix_by_indices(matrix, cls_dict[cls]))

	return oneclass_matrices


def cut_matrix_by_indices(matrix, indices):
	"""
	Returns SparseMatrix with rows at given indices from given matrix.
	"""

	documents_ids = []
	class_labels = []
	contexts = []
	data = lil_matrix((len(indices), matrix.data.shape[1]))
	row = 0
	for i in indices:
		documents_ids.append(matrix.documents_ids[i])
		class_labels.append(matrix.class_labels[i])
		contexts.append(matrix.contexts[i])
		data.data[row] = matrix.data.data[i]
		data.rows[row] = matrix.data.rows[i]
		row += 1

	new_matrix = SparseMatrix(matrix.head_word, class_labels, matrix.column_labels, contexts, data, matrix.attributes,
							  matrix.boundaries, list(set(class_labels)), documents_ids)
	return new_matrix


def get_column(matrix, col_name):
	"""
	Return data from a specific column.
	"""

	mat = matrix.data.tocsc()
	for i in range(len(matrix.column_labels)):
		if matrix.column_labels[i] == col_name:
			return mat.data[mat.indptr[i]:mat.indptr[i+1]]
	return []


def byThirdTerm(feat_tuple):
	return feat_tuple[2]


def save_results(results, filename):
	"""
	Save a dict to a csv file.

	results = {cluster_id:[doc1, doc2, doc3...]}
	"""
	all_res = []
	for r in sorted(results):
		all_res.append(tuple([r] + results[r]))

	res_rows = itertools.izip_longest(*all_res)
	with open(filename, 'wb') as res_file:
		csv_writer = csv.writer(res_file, delimiter=';')
		for row in res_rows:
			csv_writer.writerow(row)


class FeatureSelection(object):
	"""
	A class for feature selection with Weka or Scipy.
	"""

	def __init__(self, attr_eval, no_features_to_select, weka_jar,tmp_dir):
		self.attr_eval = attr_eval
		self.no_feats = no_features_to_select
		self.weka_jar = weka_jar
		self.tmp_dir= tmp_dir
		if self.attr_eval in WEKA_METHODS:
			if 'Subset' in self.attr_eval:
				self.search_method = 'GreedyStepwise'
			else:
				self.search_method = 'Ranker'

	def select_features(self, matrix):
		"""
		Select features differing classes from the matrix. Matrix has to have only two classes inside.

		:return {cluster_id:[(feature, mean in that cluster, std in that cluster, mean in all the other clusters,
							std in all the other clusters)]}
		"""

		if self.attr_eval in WEKA_METHODS:
			return self.weka_feature_selection(matrix)
		elif self.attr_eval in SCIPY_METHODS:
			return self.scipy_feature_selection(matrix)
		else:
			print 'Unrecognized feature selection method: %s. Check config file.' % self.attr_eval

	def scipy_feature_selection(self, matrix):
		"""
		Select features with Scipy.
		"""

		subclasses = split_matrix_by_classes(matrix)
		results = []
		for col_name in subclasses[0].column_labels:
			sample1 = get_column(subclasses[0], col_name)
			sample1 = complete_with_zeros(sample1, subclasses[0].data.shape[0])
			sample2 = get_column(subclasses[1], col_name)
			sample2 = complete_with_zeros(sample2, subclasses[1].data.shape[0])
			try:
				if self.attr_eval == 'MannWhitney':
					res = stats.mannwhitneyu(sample1, sample2)
				elif self.attr_eval == 'WilcoxonRankSum':
					res = stats.ranksums(sample1, sample2)
				elif self.attr_eval == 'KolmogorovSmirnov':
					res = stats.ks_2samp(sample1, sample2)
				elif self.attr_eval == 'KruskalWallis':
					res = stats.kruskal(sample1, sample2)
				elif self.attr_eval == 'TTest':
					res = stats.ttest_ind(sample1, sample2)
				else:
					print('No correlation test named %s' % self.attr_eval)
					res = None
				if res is not None and res[1] < 0.05:   # res[1] --> p-value
					if subclasses[0].head_word == 'true':
						mean_in = numpy.mean(sample1)
						std_in = numpy.std(sample1)
						mean_out = numpy.mean(sample2)
						std_out = numpy.std(sample2)
					else:
						mean_in = numpy.mean(sample2)
						std_in = numpy.std(sample2)
						mean_out = numpy.mean(sample1)
						std_out = numpy.std(sample1)
					results.append((col_name, res[0], res[1], mean_in, std_in, mean_out, std_out))
			except ValueError:
				print '%s error for %s' % (self.attr_eval, col_name), set(sample1), set(sample2)

		return [(r[0].encode('utf-8'), r[3], r[4], r[5], r[6]) for r in sorted(results, key=byThirdTerm)]

	def weka_feature_selection(self, matrix):
		"""
		Select features with Weka.
		"""

		print 'Selecting features for %s' % matrix.head_word
		tmp_mat = SparseMatrix(matrix.head_word, matrix.class_labels, matrix.column_labels, matrix.contexts,
							   matrix.data, boundaries=matrix.boundaries, possible_classes=matrix.possible_classes)

		w = was.AttributeSelection()
		w._weka_jar = self.weka_jar
		w._work_dir = self.tmp_dir
		w.set_attribute_evaluator(self.attr_eval)
		w.set_search_method(self.search_method)

		if self.no_feats is not None and 'Subset' not in self.attr_eval:
			w.set_search_method_options('-N %s' % self.no_feats)
		elif 'Subset' not in self.attr_eval:
			print 'WARNING: number of features to select is not given in config file'
		
		w.make(tmp_mat)
		feats = w._selected_features[:-1]

		subclasses = split_matrix_by_classes(matrix)
		results = []
		for feature in feats:
			sample1 = get_column(subclasses[0], feature)
			sample1 = complete_with_zeros(sample1, subclasses[0].data.shape[0])
			sample2 = get_column(subclasses[1], feature)
			sample2 = complete_with_zeros(sample2, subclasses[1].data.shape[0])
			if subclasses[0].head_word == 'true':
				mean_in = numpy.mean(sample1)
				std_in = numpy.std(sample1)
				mean_out = numpy.mean(sample2)
				std_out = numpy.std(sample2)
			else:
				mean_in = numpy.mean(sample2)
				std_in = numpy.std(sample2)
				mean_out = numpy.mean(sample1)
				std_out = numpy.std(sample1)
			results.append((feature.encode('utf-8'), mean_in, std_in, mean_out, std_out))

		return results


if __name__ == '__main__':
	main()
