import gpkavg
import avio
import sets
import re
import chunkio
import die


def process_PNP(w, mlist, Kc):
	f = "%s/PNP.av" % w
	data, comments = avio.read(open(f, 'r'))
	o = {}
	for m in mlist:
		Ctmp = []
		Ktmp = []
		for datum in data:
			if m.check(datum):
				Ctmp.append( float(datum['Ch']) )
				Ktmp.append( float(datum['K']) )
		Ch0, Chsig = gpkavg.avg(Ctmp, None, 0.1)
		# Basically, we adjust the chance probability and re-compute K.
		Ch1 = Ch0 + Kc[m.Kc_name()]*(1.0-Ch0)

		K, Ksig = gpkavg.avg(Ktmp, None, 0.1)
		Pcor = Ch0 + K * (1-Ch0)
		Knew = (Pcor - Ch1) / (1.0 - Ch1)
		o[m.name() ] = Knew
	return o


def process_PNPwsc(w, mlist, loose=1, Kczero=False):
	f = "%s/PNP.av" % w
	data, comments = avio.read(open(f, 'r'))
	o = {}
	for m in mlist:
		Ctmp = []; CStmp = []
		Ktmp = []; KStmp = []
		Ptmp = []; PStmp = []
		for datum in data:
			if loose and not ('Ch' in datum and 'K' in datum):
				loose -= 1
				continue
			if m.check(datum):
				if not 'Ch' in datum or not 'K' in datum:
					die.warn('Missing K or Ch: %s'
							% str(datum))
					continue
				if int(datum['scr'])==0:
					Ctmp.append( float(datum['Ch']) )
					Ktmp.append( float(datum['K']) )
					Ptmp.append( float(datum['Perfect']) )
				else:
					CStmp.append( float(datum['Ch']) )
					KStmp.append( float(datum['K']) )
					PStmp.append( float(datum['Perfect']) )
		if len(Ctmp) == 0:
			die.warn('No data for m=%s' % str(m))
			continue

		Ch0, Chsig = gpkavg.avg(Ctmp + CStmp, None, 0.1)
		Pc0, Pcsig = gpkavg.avg(Ptmp + PStmp, None, 0.1)

		if Kczero:
			Kc = 0.0
		else:
			if len(KStmp) == 0:
				die.warn('No data for m=%s' % str(m))
				continue
			Kc, KCsig = gpkavg.avg(KStmp, None, 0.1)

		# Basically, we adjust the chance probability and re-compute K.
		Ch1 = Ch0 + Kc*(Pc0-Ch0)

		Kold, Ksig = gpkavg.avg(Ktmp, None, 0.1)
		Pcor = Ch0 + Kold * (Pc0-Ch0)
		Knew = (Pcor - Ch1) / (Pc0-Ch1)
		o[m.name() ] = Knew
	return o




def get_datatypes(fn, pattern= 'D-([a-zA-Z0-9_:,]+)$', key='d'):
	return get_X(fn, pattern, key)


def get_X(fn, pattern, key):
	if pattern is not None:
		return get_X_pat(fn, pattern, key)
	return get_X_simple(fn, key)


def get_X_pat(fn, pattern, key):
	gdt = re.compile(pattern)
	dtl = sets.Set()
	data, comment = avio.read(open(fn, 'r'))
	for datum in data:
		dt = gdt.search(datum[key])
		if dt:
			dtl.add( dt.group(1) )
	return dtl


def get_X_simple(fn, key):
	dtl = sets.Set()
	data, comments = avio.read(open(fn, 'r'))
	for datum in data:
		dtl.add( datum[key] )
	return dtl



def get_styles(fn, pattern='T([a-zA-Z0-9_]+)N[0-9]D-', key='d'):
	return get_X(fn, pattern, key)




def get_dialects(fn, key='dialect'):
	return get_X_simple(fn, key)



def get_Kc(fn):
	"""Read in a data file of the raw K-values
	produced by chance.
	"""
	return chunkio.datachunk(open(fn, 'r')).read_dict(float)


class zero_Kc:
	def __init__(self):
		pass
	def values(self):
		return [ 0.0 ]
	def has_key(self, key):
		return True
	def get(self, key, dflt=0.0):
		return 0.0
	def clear(self):
		pass
	def setdefault(self, k, a):
		return 0.0
	def copy(self):
		return self
	def __contains__(self, x):
		return True
	def __getitem__(self, k):
		return 0.0
