#! python

"""Produce nice star plots of how things depend on the language."""

import sys
import math

import numpy

import pylab
GUI = False
ODir = 'figs'

from gmisclib import die
from gmisclib import dictops as DO
from gmisclib import pylab_starplot
from gmisclib import dict_vector as DV

sys.path.insert(0, '../lib')
import read_results as RR



COLORMAP = [(1.0, 0.0, 0.0), (0.8, 0.8, 0.0), (0.0, 1.0, 0.0), (0.0, 0.8, 0.8),
		(0.0, 0.0, 1.0), (0.0, 0.0, 0.0)
		]


class NotEnoughData(Exception):
	def __init__(self, *s):
		Exception.__init__(self, *s)


def plot_d(dlist, ept, epc):
	per_language = DO.dict_of_lists()
	for d in dlist:
		per_language.add(d['language'], d)
	ebl = []
	for (lang, l) in per_language.items():
		norm_eff_s1 = []
		norm_eff_s4 = []
		norm_eff_s = []
		for d in l:
			assert  ept[d['target']] > 0.0
			assert  epc[d['lbl']] > 0.0
			tmp =  float(d['effect_size']) / ( ept[d['target']] * epc[d['lbl']] )
			norm_eff_s.append(tmp)
			if int(d['N']) == 1:
				norm_eff_s1.append(tmp)
			if int(d['N']) == 4:
				norm_eff_s4.append(tmp)
		mean = numpy.mean(norm_eff_s)
		std = numpy.std(norm_eff_s) * math.sqrt(len(norm_eff_s)/float(len(norm_eff_s)-1))
		ebl.append( pylab_starplot.errorbar(mean-3*std, mean, mean+3*std, sortorder=lang) )
	assert len(ebl) == 5
	pylab_starplot.star(0.0, 0.0, sorted(ebl), lcolormap=COLORMAP)


def process(flist):
	h = DO.dict_of_lists()
	hl = DO.dict_of_lists()
	tepc = DO.dict_of_lists()
	tept = DO.dict_of_lists()
	data = RR.read_many(flist)
	all = []
	for d in data:
		if d['lbl'] in ['initial', 'final']:
			continue
		h.add((d['lbl'], d['target']), d)
		hl.add(d['language'], d)
		if 'effect_size' not in d:
			assert int(d['significant'])==0
			d['effect_size'] = 0
		tepc.add(d['lbl'], float(d['effect_size']))
		tept.add(d['target'], float(d['effect_size']))
		all.append(float(d['effect_size']))
	# print 'tepc=', tepc
	allmean = numpy.mean(all)
	ept = dict([(k, numpy.mean(v)/allmean) for (k,v) in tept.items()])
	epc = dict([(k, numpy.mean(v)/allmean) for (k,v) in tepc.items()])
	i = 0
	for ((lbl, target), d) in h.items():
		try:
			plot_d(d, ept, epc)
		except NotEnoughData:
			print '\tNot enough data'
			pass
		else:
			pylab.title("n=%s, %s" % (lbl, target))
			if GUI:
				pylab.show()
			else:
				pylab.savefig('%s/%d.png' % (ODir, i))
				pylab.cla()
		i += 1


if __name__ == '__main__':
	arglist = sys.argv[1:]
	while arglist and arglist[0].startswith('-'):
		arg = arglist.pop(0)
		if arg == '--':
			break
		elif arg == '-o':
			ODir = arglist.pop(0)
		elif arg == '-gui':
			GUI = True
		else:
			die.die("Unrecognized arg: %s" % arg)
	process(arglist)


