#python
'''this scripts reads in all files for a specific language and creates a text file: 
category 1: values values values
category 2: values values values
This file can be used for making histograms using histo_any.py

It currently reads the data stored in /proj/aesop/predict_next/SpeechProsody2/r_prose or ../r_poetry. This can be changed by changing the local variable PRED. 

Flags: 
- -l - language (default: all)
-poetry/-prose - default: all
- cat  - <name> category (ttype, lbl, N, target, selector)
- out - <expression> output values (which values should be added to the histogram)
- select-  python expression for selection
- d - directory for the output


Example:

To save the data for effect_size2 in Chinese poetry for different types of 'ttype',  into the folder /proj/aesop/predict_next/SpeechProsody2/temp_txt/lbl_V only using the data where label=V

read_save.py -l Chinese -poetry -cat ttype -out effect_size2 -select "lbl=='V'" -d lbl_V
'''

import sys
import os

from gmisclib import avio
from gmisclib import g2_select as select

PRED='/proj/aesop/predict_next/Jpaper'

def read_arglist(arglist):
	lang='all'
	p='all'
	sel='all'
	outdir=''
	for n in range(0, len(arglist)):
		if arglist[n]=='-l':
			lang=arglist[n+1]
		elif arglist[n]=='-poetry':
			p='poetry'
		elif arglist[n]=='-prose':
			p='prose'
		elif arglist[n]=='-cat':
			cat=arglist[n+1]
		elif arglist[n]=='-out':
			out=arglist[n+1]
		elif arglist[n]=='-select':
			sel=arglist[n+1]
		elif arglist[n]=='-d':
			outdir=arglist[n+1]
	return (lang, p, cat, out, sel, outdir)
		
def list_files(lang, p): 
	flist=[]
	folders=[]
	for d in os.listdir(PRED):
		if 'r_' in d and (p=='all' or p in d):
			print d
			folders.append(os.path.join(PRED, d))
	for folder in folders:
		for fname in os.listdir(folder):
			if lang=='all' or lang in fname:
				flist.append(os.path.join(folder, fname))
	return flist		
			
def read_data(files):
	print 'reading data'
	data=[]
	for f in files:
			d=[]
			fd, c = avio.read(open(f, 'r'))
			for line in fd:
				if not 'np2' in line.keys():
					continue
				for c in line.keys():
					try:
						line[c]=float(line[c])
					except ValueError:
						pass
				d.append(line)				
			data.extend(d)
	return data


def split_data(data, cat, out, sel):
	print 'selecting data'
	if not sel=='all':
		l=len(data)
		data=select.filterlist(sel, data, verbose=False)
		print '%s selects %d lines out of %d' % (sel, len(data), l)
	ddic={}
	for line in data:
		if ' ' in str(line[cat]):
			line[cat]=line[cat].split(' ')[0]
		if not line[cat] in ddic.keys():
			ddic[line[cat]]=[]
		val=eval(out, line)
		ddic[line[cat]].append(str(val))
	return ddic

def print_d(d, f):
	print 'printing output'
	for c in d.keys():
		print '%s: %d values' % (c, len(d[c]))
		print>>f, c, ' '.join(d[c])
	
def get_output(sel, lang, p, cat, out, outdir):
	if outdir=='':
		output=open(os.path.join(PRED, 'temp_txt', '_'.join([lang, p, cat, out])), 'w')
	else:
		try: 
			output=open(os.path.join(PRED, 'temp_txt', outdir, '_'.join([lang, p, cat, out])), 'w')
		except IOError:
			os.mkdir(os.path.join(PRED, 'temp_txt', outdir))
			output=open(os.path.join(PRED, 'temp_txt', outdir, '_'.join([lang, p, cat, out])), 'w')
	return output	


if __name__ == '__main__':
	arglist=sys.argv[1:]
	(lang, p, cat, out, sel, outdir) = read_arglist(arglist)
	print 'Language=%s, poetry/prose=%s, selection: %s, splitting by %s, returning %s' % (lang, p, sel, cat, out)
	output=get_output(sel, lang, p, cat, out, outdir)
	files=list_files(lang, p)
	print 'Processing %d files' % (len(files))
	ddic=split_data(read_data(files), cat, out, sel)
	print_d(ddic, output)
	
