#! python

import os
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 nice_hash as H
from gmisclib import dict_vector as DV
from gmisclib import g2_select as GS

import read_results as RR

def filter(data, n):
	for datum in data:
		if int(datum.get('dof', 0)) <= 0:
			continue
		if int(datum['significant'])==1 and int(datum['N'])==n:
			yield datum

def sign(x):
	if x >= 0:
		return 1
	return -1

MAXPASSES = 30
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, per_predictor_sigmas):
	print 'per_predictor_sigmas=', per_predictor_sigmas
	a = DO.dict_of_lists()
	sf = DO.dict_of_averages()
	h = H.simple()
	for d in dlist:
		target = d['target']
		tgt_std = math.sqrt(float(d['err'])/int(d['dof']))
		for k in d.keys():
			if k.startswith('C:'):
				i, dist, pred = k.split(':')
				if dist != '*':
					x = float(d[k])*per_predictor_sigmas[pred]/tgt_std
					pt = (pred, target)
					key = (int(dist), pt)
					a.add(key, x)
					sf.add(pt, float(d.get('sampled_frac', 1.0)))
					h.add(pt)
	if h.n == 0:
		raise NotEnoughData
	for (k,v) in h.map().items():
		print k, v, COLORMAP[v]
	posns = {}
	for ((dist,pt),vlist) in a.items():
		if dist not in posns:
			posns[dist] = [None] * h.n
		if len(vlist) <= 3:
			continue
		ipred = h.add(pt)
		mean = numpy.average(vlist)
		isf = 1.0/sf.get_avg(pt)
		assert 10000.0 > isf >= 1.0
		stdfac = math.sqrt(min(isf**2, len(vlist)))
		assert isf >= stdfac >= 1.0
		std = math.sqrt(len(vlist)/float(len(vlist)-1)) * numpy.std(vlist) / stdfac
		posns[dist][ipred] = pylab.startplot.errorbar(mean-3*std, mean, mean+3*std,
								dim=abs(mean)<=3*std)

	pylab_starplot.plot(sorted(posns.items()), lcolormap=COLORMAP)


def process(flist, n):
	assert 0 < n < 100
	h = H.simple()
	data, per_predictor_sigmas = RR.read_many(flist)
	for d in data:
		h.add((d['lbl'], d['target']))
	i = 0
	for (lbl, target) in h.classes():
		print 'lbl, target=', lbl, target
		selector = "lbl=='%s' and target=='%s'" % (lbl, target)
		print 'selector=', selector, flist[0]
		s = GS.selector_c(selector, global_values={'math': math, 'os': os})
		d2 = [ d for d in filter(data, n) if s.eval(d) ]
		try:
			plot_d(d2, per_predictor_sigmas)
		except NotEnoughData:
			print '\tNot enough data'
			pass
		else:
			pylab.title("n=%d, %s, %s" % (n, selector, flist[0]))
			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[1:], int(arglist[0]))


