# -*- coding: utf-8 -*-
import subprocess, codecs, os, sys, re



gather_classes = set(['AJ0', 'AJC', 'AJS', 'AV0', 'AVP', 'AVQ', 'NN0', 'NN1', 'NN2', 'NP0', 'VM0', 'VVB', 'VVD', 'VVG', ',VVI', 'VVN', 'VVZ'])
morpha_path = '/home/boban/software/morph/morpha.ix86_linux'
morpha_verbstem = '/home/boban/software/morph/verbstem.list'
pat = re.compile(r'<w ([^>]+)>([^<>$]+)')
all_bases = {} # base, word, category -> count

def write_data(filename ):
	file = codecs.open(filename,'w', 'utf8')
	file.write( "base; word; category; count\n" )
	for t in all_bases:
		count = all_bases[t]
		line = "%s;%s;%s;%d\n" % (t[0],t[1], t[2], count)
		file.write(line)
		
	file.close()		
			
def get_base(word, cat):
	to_analyse = word + u"_" + cat + u" "
	p = subprocess.Popen([morpha_path, '-f', morpha_verbstem], stdout=subprocess.PIPE, stdin=subprocess.PIPE)
	p.stdin.write(to_analyse.encode('utf') )
	p.stdin.close()
	return p.stdout.read().lower().strip()

def filter_forms(p):
	if p[0] in gather_classes:
		return True
	return False

def process_file(filename):
	for line in codecs.open(filename, 'r', 'utf8'):
		line = line.strip()
		if len(line) == 0: continue
		m = pat.findall(line)
		if len(m):
			m = filter(filter_forms, m)
			if len(m):
				for c,w in m:
					w = w.strip()
					categoires = c.strip().split('-')
					for c in categoires:					
						base = get_base(w,c)
						triple = (base,w,c)
						if triple in all_bases:
							all_bases[triple] = all_bases[triple] + 1
						else:
							all_bases[triple] = 1
	
def process_single(param, dirname, names):
	for name in names:
		filename = os.path.join(dirname, name)
		if os.path.isfile(filename):
			print "Processing", filename
			process_file(filename)

def load_forms(files):
	all_bases.clear()
	for file in files:
		f = codecs.open(file, 'r', 'utf8')
		for line in f:
			line = line.strip()
			if len(line):
				t = line.split(';')
				all_bases[ (t[0],t[1],t[2]) ] = t[3]

def accept_word(word):
	split = word.split('-')
	for s in split:
		if not s.isalpha():
			return False	
	return True
			
def extract_from_db(outfile):
	import MySQLdb
	import re
	p = re.compile(r'(\W|[-]|\s)+$', re.UNICODE)
	connection = MySQLdb.connect( user = 'plwn', passwd = 'plwn', host = 'b4', db = 'bnc', port = 3306)
	connection.query("SET CHARACTER SET 'utf8'")
	cursor = connection.cursor()
	#query = 'SELECT base,orth,count FROM unigrams WHERE count > 10 AND category LIKE \'V%\' ORDER BY base '
	query = 'SELECT base,orth,count,category FROM unigrams WHERE count > 10 ORDER BY base '
	cursor.execute(query)
	
	data = {} # base -> count, set of orth
	
	for r in cursor:
		base,word,count,category = r
		word = unicode( word.decode('utf8').strip() )
		base = unicode(base.decode('utf8').strip() )
		category = unicode(category.decode('utf8').strip() )

		if len(word) == 0 or len(base) == 0:
			continue
		
		if category[0] == "N":
			bc = (base,"N")
		elif category[0] == "V":
			bc = (base,"V")
		else:
			bc = (base, 'A')
			
		#if p.match(word) != None and p.match(base) != None:
		if  accept_word(word) and  accept_word(base):
			if bc in data:
				count_total, orths = data[bc]
				orths.add(word)
				data[bc] = (count + count_total, orths)
			else:
				pair = (count, set([word]) )
				data[bc] = pair
	
	print "Gathered %d words\nSaving to file %s" % (len(data), outfile)
	
	file = codecs.open(outfile, 'w', 'utf8')
	for bc in sorted(data):
		count, orths = data[bc]
		file.write("%s;%s;%d" % (bc[0],bc[1],count) )
		for o in sorted(orths):
			file.write(';' + o)
		file.write('\n')
	file.close()			
				
	
if __name__ == '__main__':
	if len(sys.argv) == 2:
		extract_from_db(sys.argv[1])
		sys.exit(0)		
	if len(sys.argv) != 3:
		print """Usage:
		python base_extractor.py directory_or_file file_out
		"""
		sys.exit(-1)
	if sys.argv[1] == 'clear':
		create_table()
		print "Done..."
	else:
		print "DICT only"
		dirname = sys.argv[1]
		if os.path.isdir(dirname):
			os.path.walk(dirname, process_single, None)
		else:
			process_file(dirname)
		print "Writing data"
		write_data(sys.argv[2])
	
	
