#
# Copyright (C) 2007-2013 German Aerospace Center (DLR/SC)
#
# Created: 2012-12-17 Martin Siggel <Martin.Siggel@dlr.de>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @file ms_componentSegmentGeom.py
# @brief Implements coordinate transforms on the component segment geometry
#

from numpy import *
from ms_optAlgs import *
from ms_segmentGeometry import *

import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

class Elongation:
	Left, No, Right, LeftRight = range(4)

class ComponentSegment:

	def __init__(self):
		self.segments = []

	def addSegment(self, p1, p2, p3, p4):
		self.segments.append(ComponentSegmentGeometry(p1,p2,p3,p4))
		self.__calcEtaRanges()

	def __calcEtaRanges(self):
		size = len(self.segments)
		# calculate total projected length
		len_tot = 0
		for item in self.segments:
			len_tot = len_tot + item.calcProjectedLeadingEdgeLength(Elongation.No)

		len_tot = len_tot + self.segments[0].calcProjectedLeadingEdgeLength(Elongation.Left) \
						- self.segments[0].calcProjectedLeadingEdgeLength(Elongation.No) \
						+ self.segments[size-1].calcProjectedLeadingEdgeLength(Elongation.Right) \
						- self.segments[size-1].calcProjectedLeadingEdgeLength(Elongation.No)

		#calculate inner eta of first segment
		etastart = (self.segments[0].calcProjectedLeadingEdgeLength(Elongation.Left) \
					- self.segments[0].calcProjectedLeadingEdgeLength(Elongation.No) )/len_tot

		for item in self.segments:
			etastop = etastart + item.calcProjectedLeadingEdgeLength(Elongation.No)/len_tot
			item.setLeadingEdgeEtas(etastart, etastop)
			etastart = etastop

	def draw(self, axis):
		for item in self.segments:
			item.drawSegment(axis)

	def calcPoint(self, eta, xsi):
		for item in self.segments:
			if item.checkCoordValidity(eta,xsi) == True:
				return item.calcCSPoint(eta,xsi)

	def drawPoint(self, axis, eta, xsi):
		point = self.calcPoint(eta, xsi)
		axis.plot([point[0]], [point[1]], [point[2]],'rx')

class ComponentSegmentGeometry:
	def __init__(self, p1, p2, p3, p4, etamin = 0, etamax = 1):
		self.setPoints(p1, p2, p3, p4, etamin, etamax)

	def setPoints(self, p1, p2, p3, p4, etamin, etamax):
		self.__p1 = p1
		self.__p2 = p2
		self.__p3 = p3
		self.__p4 = p4
		self.__etamin = etamin
		self.__etamax = etamax

		sv = p2 - p1
		sh = p4 - p3

		# normal vector of plane
		n = array([0, -sv[1], -sv[2]])

		## calculate extended leading edge and trailing edge points, this has to be done
		#  only once per wing segment.

		#calculate point where le intersects plane
		avo = dot(p4-p1,n) / dot(p2-p1,n)
		if avo > 1:
			self.__p2p = avo*sv + p1
			self.__p4p = p4
		else:
			self.__p2p = p2;
			# check trailing edge
			aho = dot(p2-p3,n) / dot(p4-p3,n)
			assert aho >= 1
			self.__p4p = sh*aho + p3

		# now the inner section, the normal vector should be still the same
		avi = dot(p3-p1,n)/dot(p2-p1,n)
		if avi < 0:
			# leading edge has to be extended
			self.__p3p = p3
			self.__p1p = avi*sv + p1
		else:
			self.__p1p = p1;
			ahi = dot(p1-p3,n)/dot(p4-p3,n)
			self.__p3p = p3 + ahi*sh
			assert ahi <= 0

		#calculate eta values of segment edges, these values define also, when a given cs coordinate is outside the wing segment
		self.__eta1 = dot(p1-self.__p1p,n) / dot(self.__p2p-self.__p1p,n)
		self.__eta2 = dot(p2-self.__p1p,n) / dot(self.__p2p-self.__p1p,n)
		self.__eta3 = dot(p3-self.__p1p,n) / dot(self.__p2p-self.__p1p,n)
		self.__eta4 = dot(p4-self.__p1p,n) / dot(self.__p2p-self.__p1p,n)

	# sets the eta range from inner segment tip to the outer tip
	def setEtaMinMax(self, etamin, etamax):
		self.__etamax = etamax
		self.__etamin = etamin

	def getEtaMinMax(self):
		return self.__etamin, self.__etamax

	# sets the eta range of the leading edge. if e.g. the trailing edge is longer than
	# the leading edge, this makes a difference to setEtaMinMax
	def setLeadingEdgeEtas(self, eta_in, eta_out):
		# we need to scale etamax, etamin accordingly
		etamin = (self.__eta2*eta_in - self.__eta1*eta_out)/(self.__eta2 - self.__eta1)
		etamax = etamin + (eta_out - eta_in)/(self.__eta2 - self.__eta1)
		self.setEtaMinMax(etamin, etamax)

	# only valid with calcsCSPoint (not calcCSPoint2/3)
	def checkCoordValidity(self, eta, xsi):
		if eta < self.__etamin or eta > self.__etamax or xsi < 0 or xsi > 1:
			return False
		else:
			actetamin = (1-xsi)*self.__eta1 + xsi*self.__eta3
			actetamax = (1-xsi)*self.__eta2 + xsi*self.__eta4
			if (eta>= actetamin) and (eta <= actetamax):
				return True
			else:
				return False

	def calcProjectedLeadingEdgeLength(self, elongation):
		p1 = self.__p1;
		p2 = self.__p2;
		if elongation   == Elongation.Left:
			p1 = self.__p1p
		elif elongation == Elongation.Right:
			p2 = self.__p2p
		elif elongation == Elongation.LeftRight:
			p1 = self.__p1p
			p2 = self.__p2p

		# project leading edge into the z-y plane
		vProj = array([0, 1, 1])
		return linalg.norm(vProj*(p2-p1))

	def calcCSPoint(self, eta, xsi):
		eta_ = (eta - self.__etamin)/(self.__etamax - self.__etamin)

		#calculate eta values at given xsi
		eta1p = self.__eta1*(1-xsi) + self.__eta3*xsi
		eta2p = self.__eta2*(1-xsi) + self.__eta4*xsi

		pbeg = self.__p1*(1-xsi) + self.__p3*xsi
		pend = self.__p2*(1-xsi) + self.__p4*xsi

		p = pbeg + (pend-pbeg)*(eta_ - eta1p)/(eta2p-eta1p)

		return p

	def calcCSPoint2(self, eta, xsi):
		eta_ = (eta - self.__etamin)/(self.__etamax - self.__etamin)

		p = self.__p1p + (self.__p2p-self.__p1p)*(eta_);

		a = -self.__p1+self.__p2;
		b = -self.__p1+self.__p3;
		c =  self.__p1-self.__p2-self.__p3+self.__p4;
		d =  self.__p1;

		n = array([0, -a[1], -a[2]])

		# calc some constants
		a1 =  dot(p - d, n);
		a2 = -dot(b, n);
		a3 =  dot(a, n);
		a4 =  dot(c, n);

		# this calculates the intersection curve from the segment with a plane (normal vector n, point p2)
		al = lambda beta:  (a1 + a2*beta)/(a3 + a4*beta);
		# diff( eta(xi). xi), tangent in eta xsi space
		alp = lambda beta: (a2*a3 - a1*a4)/((a3 + a4*beta)**2);

		# 3d intersection curve, parametrized by beta [0,1]
		cu = lambda beta: outer(a,al(beta)) + outer(b,beta) + outer(c,al(beta)*beta) + outer(d,ones(size(beta)));

		# tangent in 3d space
		cup = lambda beta: (outer(a,ones(size(beta))) + outer(c,beta))*outer(ones(3),alp(beta)) + outer(c,al(beta)) + outer(b,ones(size(beta)));

		#norm of tangent curve
		f =  lambda beta: sqrt(sum(cup(beta)**2.,0));

		# we want to integrate int f(x)*dx, we do gaussian method

		# substitution of f to range [-1,1]
		g = lambda x,beta: beta/2.*f((x+1.)*beta/2.);

		# gauss x points 5th order, this is really some dark magic ;)
		x = array([9.06179845938664e-01,
				5.38469310105683e-01,
				0.00000000000000e+00,
				-5.38469310105683e-01,
				-9.06179845938664e-01])

		# gauss weights
		w = array([2.36926885056189e-01,
				4.78628670499366e-01,
				5.68888888888889e-01,
				4.78628670499366e-01,
				2.36926885056189e-01])

		# calculate total length of iso eta curve
		ltot =  dot(g(x,1.),w);

		#now we use Newton Raphson to find beta, such that F(beta) == xi*ltot
		F = lambda beta: dot(g(x,beta),w)/ltot - xsi
		beta = xsi
		diff = F(beta)
		while abs(diff) > 1e-12:
			dir = -diff/(f(beta)/ltot)
			beta = beta + dir
			diff = F(beta)

		return cu(beta)

	# alternative implementation, where xsi resembles the relative coordinate
	# between intersection point of the leading edge and the intersection point
	# of the trailing edge. This intermediate point will then be projected onto
	# the true intersection curve
	def calcCSPoint3(self, eta, xsi):
		debug = True

		eta_ = (eta - self.__etamin)/(self.__etamax - self.__etamin)

		p = self.__p1p + (self.__p2p-self.__p1p)*(eta_);

		a = -self.__p1+self.__p2;
		b = -self.__p1+self.__p3;
		c =  self.__p1-self.__p2-self.__p3+self.__p4;
		d =  self.__p1;

		n = array([0, -a[1], -a[2]])

		# calc some constants
		a1 =  dot(p - d, n);
		a2 = -dot(b, n);
		a3 =  dot(a, n);
		a4 =  dot(c, n);

		# this calculates the intersection curve from the segment with a plane (normal vector n, point p2)
		al = lambda beta:  (a1 + a2*beta)/(a3 + a4*beta);
		# diff( eta(xi). xi), tangent in eta xsi space
		alp = lambda beta: (a2*a3 - a1*a4)/((a3 + a4*beta)**2);

		# 3d intersection curve, parametrized by beta [0,1]
		cu = lambda beta: outer(a,al(beta)) + outer(b,beta) + outer(c,al(beta)*beta) + outer(d,ones(size(beta)));

		# tangent in 3d space
		cup = lambda beta: (outer(a,ones(size(beta))) + outer(c,beta))*outer(ones(3),alp(beta)) + outer(c,al(beta)) + outer(b,ones(size(beta)));

		# calculate intersection with leading and trailing edge
		pbeg = cu(0)[:,0]
		pend = cu(1)[:,0]
		reflen = linalg.norm(pbeg-pend);

		# go along this line
		pact = (1-xsi)*pbeg + xsi*pend

		# project this point onto intersection curve i.e. find beta so that (cu(beta) - pact) * (pbeg-pend) == 0
		# as cu(beta) is not linear, we try to find the solution with newton raphson method
		f = lambda beta: dot(cu(beta)[:,0] - pact, pend - pbeg)/reflen
		fp = lambda beta: dot(cup(beta)[:,0], pend - pbeg)/reflen

		beta =  xsi

		diff = f(beta)
		iter = 0;
		if debug: print 'Iter:', iter, ' Error=', abs(diff), ' @ Beta=' , beta
		while abs(diff) > 1e-12 and iter  < 20:
			iter += 1
			dir = -diff/(fp(beta))
			# maybe we need a linesearch here...
			beta = beta + dir
			diff = f(beta)
			if debug: print 'Iter:', iter, ' Error=', abs(diff), '@ Beta=' , beta

		if iter >= 20:
			print "ERROR: could not project intersection curve onto line"

			if debug == True:
				myb = linspace(-1.,1., 1000)
				val = 0*myb
				for i in range(len(myb)):
					val[i] = f(myb[i]);

				fig = plt.figure()
				ax2 = fig.gca()
				ax2.plot(myb, val)

		# calculate result
		point = cu(beta)

		# here we got for free our segment coordinates also, which are
		# eta_s = al(beta), xsi_s = beta

		return point

	# calculates the tangents in eta and xsi direction at the given point
	def calcCSPointTangents(self, eta, xsi):
		eta_ = (eta - self.__etamin)/(self.__etamax - self.__etamin)
		deta_ = 1. /(self.__etamax - self.__etamin)

		#calculate eta values at given xsi
		eta1p = self.__eta1*(1-xsi) + self.__eta3*xsi
		eta2p = self.__eta2*(1-xsi) + self.__eta4*xsi

		pbeg = self.__p1*(1-xsi) + self.__p3*xsi
		pend = self.__p2*(1-xsi) + self.__p4*xsi

		# calculate derivatives
		deta1p = self.__eta3 - self.__eta1;
		deta2p = self.__eta4 - self.__eta2;
		dpbeg  = self.__p3   - self.__p1;
		dpend  = self.__p4   - self.__p2;

		J = zeros((3,2))
		J[:,0] = (pend-pbeg)/(eta2p-eta1p)*deta_;
		J[:,1] = dpbeg + (dpend-dpbeg)*(eta_ - eta1p)/(eta2p-eta1p) + (pend - pbeg)*(-deta1p/(eta2p-eta1p) - (eta_ - eta1p)/((eta2p-eta1p)**2)*(deta2p - deta1p) );

		return J

	def calcCSPointNormal(self,  eta, xsi):
		J = self.calcCSPointTangents(eta, xsi);
		normal = cross(J[:,1],J[:,0])
		return normal/linalg.norm(normal)

	def __calcCSHessian(self, eta, xsi, p):
		eta_ = (eta - self.__etamin)/(self.__etamax - self.__etamin)
		deta_ = 1. /(self.__etamax - self.__etamin)

		#calculate eta values at given xsi
		eta1p = self.__eta1*(1-xsi) + self.__eta3*xsi
		eta2p = self.__eta2*(1-xsi) + self.__eta4*xsi

		pbeg = self.__p1*(1-xsi) + self.__p3*xsi
		pend = self.__p2*(1-xsi) + self.__p4*xsi

		# calculate derivatives
		deta1p = self.__eta3 - self.__eta1;
		deta2p = self.__eta4 - self.__eta2;
		dpbeg = self.__p3 - self.__p1;
		dpend = self.__p4 - self.__p2;

		# helper variables and their derivatives to xsi
		hv1 =  pend -  pbeg;
		dhv1 = dpend - dpbeg;

		h2 = eta_ -  eta1p;
		dh2 =    - deta1p;

		h3 =  1/(eta2p-eta1p);
		dh3 = -1/(eta2p-eta1p)**2*(deta2p-deta1p);
		d2h3 = 2/(eta2p-eta1p)**3*(deta2p-deta1p)**2;

		# p(eta, xsi)
		p_ = pbeg + hv1*h2*h3;

		# first derivative, dp(eta, xsi)
		J1 = hv1*h3*deta_;
		J2 = dpbeg + dhv1*h2*h3 + hv1*(dh2*h3 + h2 * dh3);

		# second order derivative d2p(eta, xsi), H11 is zero!
		H21 =   (dhv1*h3 + hv1*dh3)*deta_;
		H22 =   dhv1*(dh2*h3 + h2*dh3)*2 + hv1*( 2*dh2*dh3 + h2*d2h3 );

		# finally applying for the object function
		H = zeros((2,2))
		H[0,0] = dot(J1,J1)
		H[0,1] = dot(J1,J2) + dot(p_ - p, H21)
		H[1,1] = dot(J2,J2) + dot(p_ - p, H22)
		H[1,0] = H[0,1];

		return 2.*H

	def projectOnCS(self, p):
		opttype = 'newton'
		# calculate initial guess, project onto leading edge and inner section
		eta = dot(p - self.__p1p, self.__p2p - self.__p1p)/( linalg.norm(self.__p2p - self.__p1p)**2)
		xsi = dot(p - self.__p1, self.__p3 - self.__p1)/( linalg.norm(self.__p3 - self.__p1)**2)
		# scale according to local eta range
		eta = eta*(self.__etamax-self.__etamin) + self.__etamin
		x = array([eta,xsi])

		of    = lambda x: linalg.norm(self.calcCSPoint(x[0], x[1])-p)**2;
		ograd = lambda x: 2.* dot(self.calcCSPointTangents(x[0], x[1]).transpose(), self.calcCSPoint(x[0], x[1])-p);
		#ograd = lambda x: ms_numGrad(of, x, 1e-9)
		ohess = lambda x: self.__calcCSHessian( x[0], x[1], p);
		#ohess = lambda x:  ms_numHess(ograd, x, 1e-9)

		fig2 = plt.figure();

		X, Y = meshgrid(arange(self.__etamin-0.2, self.__etamax+0.2, 0.02), arange(-0.2, 1.2, 0.02))
		Z = zeros(X.shape);

		for i in range(0,size(X,0)):
			for j in range(0,size(X,1)):
				Z[i,j] = of([X[i,j], Y[i,j]])

		plt.imshow(Z,origin='lower', extent=[self.__etamin-0.2, self.__etamax+0.2,-0.2,1.2], aspect=1./1.)
		plt.colorbar();
		plt.contour(X,Y,Z)
		plt.title('Objective function opt:'+opttype)
		plt.xlabel('eta');
		plt.ylabel('xsi');

		if opttype == 'bfgs':
			x_= ms_optQuasiNewton(of,ograd, x, 'bfgs')
		elif opttype == 'sr1':
			x_= ms_optQuasiNewton(of,ograd, x, 'sr1')
		elif opttype == 'gradient':
			x_= ms_optSteepestDescent(of,ograd, x)
		elif opttype == 'cg':
			x_= ms_optCG(of,ograd, x, 'fr')
		else:
			x_= ms_optNewton(of,ograd,ohess,x)

		eta = x_[0]; xsi = x_[1];
		return (eta, xsi)

	def projectOnCS3(self, p):
		segment = SegmentGeometry(self.__p1, self.__p2, self.__p3, self.__p4)
		# get the projection point on the, here we get some numerical uncertainty
		# if we'd knew that p is already on the plane, we could directly use p
		alpha, beta = segment.projectOnSegment(p);
		return self.convertAlphaBetaToEtaXsi(alpha, beta)
		#return self.convertXYZtoEtaXsi(p[:,0])

	def convertAlphaBetaToEtaXsi(self, alpha, beta):
		segment = SegmentGeometry(self.__p1, self.__p2, self.__p3, self.__p4)
		p_proj = segment.getPoint(alpha, beta)[:,0]
		return self.convertXYZtoEtaXsi(p_proj)

	# we must ensure that p_proj already lies on the segment, if unsure use projectOnCS3
	def convertXYZtoEtaXsi(self, p_proj):
		# project leading edge into the z-y plane
		vProj = array([0, 1, 1])
		n = vProj*(self.__p2p-self.__p1p)

		# calc eta koordinate of that point
		eta = (dot(p_proj,n)*(self.__eta2 - self.__eta1) + dot(self.__p2,n)*self.__eta1 - dot(self.__p1,n)*self.__eta2) \
				/ dot(self.__p2 - self.__p1, n);

		# intersection point of plane with leading edge
		p_ = 1./(self.__eta2 - self.__eta1)*( (self.__eta2 - eta)*self.__p1 + (eta - self.__eta1)*self.__p2 )

		a = -self.__p1+self.__p2;
		b = -self.__p1+self.__p3;
		c =  self.__p1-self.__p2-self.__p3+self.__p4;
		d =  self.__p1;
		# calc some constants
		a1 =  dot(p_ - d, n);
		a2 = -dot(b, n);
		a3 =  dot(a, n);
		a4 =  dot(c, n);

		# this calculates the intersection curve from the segment with a plane (normal vector n, point p2)
		al = lambda beta:  (a1 + a2*beta)/(a3 + a4*beta);
		# 3d intersection curve, parametrized by beta [0,1]
		cu = lambda beta: outer(a,al(beta)) + outer(b,beta) + outer(c,al(beta)*beta) + outer(d,ones(size(beta)));

		# now we have to find the xi coordinate, i.e. (pbeg + (pend-pbeg)*xi - p_proj)*(pbeg-pend) == 0
		pbeg = cu(0)[:,0]
		pend = cu(1)[:,0]

		xsi = dot(p_proj - pbeg, pbeg-pend)/dot(pend-pbeg, pbeg-pend);

		return eta, xsi

	def calcCSIsoXsiLine(self, xsi, extentToGeometry = False):
		etamin = self.__etamin
		etamax = self.__etamax

		if not extentToGeometry:
			etamin = xsi * (self.__eta3 - self.__eta1) + self.__eta1
			etamin = etamin * (self.__etamax - self.__etamin) + self.__etamin
			etamax = xsi * (self.__eta4 - self.__eta2) + self.__eta2
			etamax = etamax * (self.__etamax - self.__etamin) + self.__etamin

		P1 = self.calcCSPoint(etamin,xsi);
		P2 = self.calcCSPoint(etamax,xsi);
		return ([P1[0], P2[0]],  [P1[1], P2[1]], [P1[2], P2[2]] )

	def calcCSIsoEtaLine(self, eta, extentToGeometry = False, npoints = 30):
		eta_ = (eta - self.__etamin)/(self.__etamax - self.__etamin)

		# calculate bilinear vectors
		a = -self.__p1 + self.__p2;
		b = -self.__p1 + self.__p3;
		c =  self.__p1 - self.__p2 - self.__p3 + self.__p4;
		d =  self.__p1;

		# leading edge vector
		sv = self.__p2 - self.__p1
		# normal vector of intersection plane
		n = array([0, -sv[1], -sv[2]])

		# calculate eta point on leading edge
		p_ = self.__p1p*(1-eta_) + eta_*self.__p2p;

		a1 =  dot(p_-d,n);
		a2 = -dot(b,n);
		a3 =  dot(a,n);
		a4 =  dot(c,n);

		# this calculates the intersection curve from the segment with a plane (normal vector n, point p2)
		al = lambda beta: (a1 + a2*beta)/(a3 + a4*beta);

		# 3d intersection curve, parameterized by beta [0,1]
		cu = lambda beta: outer(a,al(beta)) + outer(b,beta) + outer(c,al(beta)*beta) + outer(d, ones(size(beta)));

		xsistart = 0 if (eta_ <= self.__eta2) else (eta_ - self.__eta2)/(self.__eta4 - self.__eta2)
		xsistop  = 1 if (eta_ <= self.__eta4) else (eta_ - self.__eta2)/(self.__eta4 - self.__eta2)

		xsistart = xsistart if (eta_ >= self.__eta1) else (eta_ - self.__eta1)/(self.__eta3 - self.__eta1)
		xsistop  = xsistop  if (eta_ >= self.__eta3) else (eta_ - self.__eta1)/(self.__eta3 - self.__eta1)

		if extentToGeometry:
			xsi = linspace(0,1,npoints)
		else:
			xsi = linspace(xsistart,xsistop,npoints)
		points = cu(xsi);

		X = points[0,:];
		Y = points[1,:];
		Z = points[2,:];

		return (X,Y,Z)

	def drawSegment(self, axis, extentToGeometry = False):
		start  = math.ceil (self.__etamin*10.)/10.
		stop =   math.floor(self.__etamax*10.)/10.

		alpha = start
		while alpha <= stop:
				X,Y,Z = self.calcCSIsoEtaLine(alpha, extentToGeometry)
				axis.plot(X,Y,Z,'g')
				alpha = alpha + 0.1

		beta = 0.0
		while beta <= 1.0:
			X,Y,Z = self.calcCSIsoXsiLine(beta, extentToGeometry)
			axis.plot(X,Y,Z,'g');
			beta = beta + 0.1

		lw = 2
		axis.plot([self.__p1[0], self.__p2[0]], [self.__p1[1], self.__p2[1]], [self.__p1[2], self.__p2[2]],'b',linewidth=lw);
		axis.plot([self.__p1[0], self.__p3[0]], [self.__p1[1], self.__p3[1]], [self.__p1[2], self.__p3[2]],'b',linewidth=lw);
		axis.plot([self.__p4[0], self.__p2[0]], [self.__p4[1], self.__p2[1]], [self.__p4[2], self.__p2[2]],'b',linewidth=lw);
		axis.plot([self.__p3[0], self.__p4[0]], [self.__p3[1], self.__p4[1]], [self.__p3[2], self.__p4[2]],'b',linewidth=lw);


# Noysim -- Noise simulation tools for Aimsun.
# Copyright (c) 2010-2011 by Bert De Coensel, Ghent University & Griffith University.
#
# Classes for sending and viewing noise levels in real-time

import os
import sys
import socket
import threading
import time
import random
import msvcrt

if not hasattr(sys, 'frozen'):
  import wxversion
  wxversion.select('2.8-msw-unicode') # version of wxPython
import wx
from wx.lib.agw.floatspin import FloatSpin, EVT_FLOATSPIN

import matplotlib
matplotlib.use('WXAgg')
from matplotlib.figure import Figure
from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigCanvas, NavigationToolbar2WxAgg as NavigationToolbar

import numpy
import pylab

USERPYC = True # if set to False, low level sockets are used
if USERPYC:
  try:
    # check if rpyc is installed
    import rpyc
    from rpyc.utils.server import ThreadedServer
    from rpyc.utils.classic import DEFAULT_SERVER_PORT
    from rpyc.utils.registry import UDPRegistryClient
    from rpyc.core import SlaveService
  except:
    # revert to using low level sockets
    USERPYC = False
    raise Exception('rpyc has to be installed')

import version

#---------------------------------------------------------------------------------------------------
# Parameters
#---------------------------------------------------------------------------------------------------

# general parameters
NAME = '%s %s Viewer' % (version.name.capitalize(), version.version)
ABOUT = NAME + '\n\n' + version.copyright.replace(', ', '\n') + '\n' + version.email

# communication with level viewer
RPYCTHREAD = None # global level thread variable (needed to circumvent the rpyc service factory)
HOST = 'localhost'
PORT = 50007
TIMEOUT = 0.01
SLEEP = 0.001
BUFSIZE = 4096

# timing parameters
REDRAWTIME = 100 # number of milliseconds between redraws
FLASHTIME = 1500 # duration of messages on the status bar, in milliseconds

# visualisation parameters
DPI = 100 # dots per inch for plotting and saving
FIGSIZE = (3.0, 3.0) # size of plotting canvas in inches (defaults to 300x300 pixels)
FONTSIZE = 8 # size of font of labels
BGCOLOR = 'black'
GRIDCOLOR = 'gray'
LINECOLOR = 'yellow'
LINEWIDTH = 1

# axes parameters
SPININC = 5.0 # increment of spin controls
XMIN = 10.0 # minimal x-axis range width
XWIDTH = 30.0 # initial value of x-axis range width
YMIN = (0.0, 10.0) # minimal y-axis low and height
YRANGE = (30.0, 60.0) # initial values of y-axis low and height
MARGIN = 1.0 # margin for auto range of levels

# test parameters
TESTDT = 0.5 # simulation timestep in seconds
TESTSLEEP = 0.2 # time between level updates
TESTLOCS = ['(1.00,2.00,3.00)', '(4.00,5.00,6.00)'] # locations of test receivers
randomLevel = lambda: 40.0 + 30.0*random.random() # function that generates a random sound level

#---------------------------------------------------------------------------------------------------
# Communication from plugin to viewer
#---------------------------------------------------------------------------------------------------

class LevelBuffer(object):
  """ base interface for sending levels to the viewer, implementing the one-way communication protocol
      types of messages:
      - command: 'clear'
      - levels: 't;loc:level;loc:level'
  """
  def __init__(self, host = HOST, port = PORT, active = True, sleep = 0, verbose = False):
    object.__init__(self)
    self.host = host
    self.port = port
    self.queue = [] # queue of messages to send
    self.active = active # if False, nothing is sent
    self.sleep = sleep/1000.0 # time to sleep (in seconds) after sending levels (to slow down a simulation)
    self.verbose = verbose # if True, debug code is printed

  def sendLevels(self, t, levels):
    """ send a series of levels at a particular time at different locations (dict of location:level) """
    if self.active:
      message = ('%.2f;' % t) + ';'.join([('%s:%.2f' % (str(loc), level)) for loc, level in levels.iteritems()])
      self.queue.append(message)
      self.flush()
      if self.sleep > 0.0:
        time.sleep(self.sleep)

  def sendClear(self):
    """ send a 'clear' message """
    if self.active:
      message = 'clear'
      self.queue.append(message)
      self.flush()

  def send(self, message):
    """ should send a single message string to the viewer (raise an error if not succesful) """
    raise NotImplementedError

  def flush(self):
    """ try to send all message strings in the queue to the viewer """
    while (len(self.queue) > 0) and (self.active == True):
      message = self.queue[0]
      try:
        if self.verbose:
          print 'trying to send message "%s"' % message
        self.send(message)
        # remove message from queue
        del self.queue[0]
        if self.verbose:
          print 'sending succesful'
      except:
        if self.verbose:
          print 'sending failed - aborting - length of queue: %d' % len(self.queue)
        break

class SocketLevelBuffer(LevelBuffer):
  """ implement the level buffer using low level sockets """
  def __init__(self, *args, **kwargs):
    LevelBuffer.__init__(self, *args, **kwargs)

  def send(self, message):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((self.host, self.port))
    s.sendall(message)
    s.close()

class RPyCLevelBuffer(LevelBuffer):
  """ implement the level buffer using Remote Python Calls (RPyC) """
  def __init__(self, *args, **kwargs):
    LevelBuffer.__init__(self, *args, **kwargs)

  def send(self, message):
    conn = rpyc.classic.connect('localhost')
    conn.root.processMessage(message)
    conn.close()

def createLevelBuffer(*args, **kwargs):
  """ create a level buffer according to the defined protocol """
  if USERPYC:
    return RPyCLevelBuffer(*args, **kwargs)
  else:
    return SocketLevelBuffer(*args, **kwargs)

#---------------------------------------------------------------------------------------------------
# Viewer thread for receiving levels
#---------------------------------------------------------------------------------------------------

VIEWERLOCK = threading.Lock()

class BaseLevelThread(threading.Thread):
  """ base interface for a thread for receiving levels """
  def __init__(self):
    threading.Thread.__init__(self)
    self.active = True # set this to false for the thread to stop
    self.clear()

  def clear(self):
    """ clear all data """
    VIEWERLOCK.acquire()
    self.data = {} # dict with received levels, for each receiver location
    self.times = [] # list with times
    VIEWERLOCK.release()

  def locations(self):
    """ return the receiver locations """
    VIEWERLOCK.acquire()
    result = self.data.keys()[:]
    VIEWERLOCK.release()
    return result

  def levels(self, loc):
    """ return the times and levels at the given location """
    VIEWERLOCK.acquire()
    result = (numpy.asarray(self.times).copy(), numpy.asarray(self.data[loc]).copy())
    VIEWERLOCK.release()
    return result

class DummyLevelThread(BaseLevelThread):
  """ dummy interface for receiving levels, which adds levels at regular instances in time """
  def __init__(self, dt = TESTDT, sleep = TESTSLEEP, locs = TESTLOCS):
    BaseLevelThread.__init__(self)
    self.dt = dt
    self.sleep = sleep
    self.locs = locs

  def run(self):
    """ instantiate the server """
    print 'thread started...'
    t = 0.0
    while self.active:
      t += self.dt
      VIEWERLOCK.acquire()
      self.times.append(t)
      for loc in self.locs:
        if not loc in self.data:
          self.data[loc] = []
        level = randomLevel()
        self.data[loc].append(level)
        print 'level received succesfully: time %.2fs, %s, %.2f dB' % (t, loc,level)
      VIEWERLOCK.release()
      time.sleep(self.sleep)

class ViewerLevelThread(BaseLevelThread):
  """ interface for receiving levels, as a thread that runs a server which listens to new levels """
  def __init__(self, frame = None, host = HOST, port = PORT, verbose = False):
    BaseLevelThread.__init__(self)
    self.frame = frame # frame to which the thread is connected
    self.host = host
    self.port = port
    self.verbose = verbose # if True, debug code is printed

  def processMessage(self, message):
    """ process an incoming message """
    if message == '':
      pass
    elif message == 'clear':
      self.clear()
      # clear the frame if applicable
      if self.frame != None:
        self.frame.clear_choices()
        self.frame.clear_plot()
      if self.verbose:
        print 'levels cleared'
    else:
      # parse the incoming message
      tokens = message.split(';')
      t = float(tokens[0])
      levels = []
      for token in tokens[1:]:
        loc, level = token.split(':')
        level = float(level)
        levels.append((loc, level))
      # when parsing is succesful, update the data
      if (len(self.times) > 0) and (t < self.times[-1]):
        if self.verbose:
          print 'discarding non-chronological levels: %s' % message
      else:
        VIEWERLOCK.acquire()
        self.times.append(t)
        for loc, level in levels:
          if not loc in self.data:
            self.data[loc] = []
          self.data[loc].append(level)
          if self.verbose:
            print 'level received succesfully: time %.2fs, %s, %.2f dB' % (t, loc,level)
        VIEWERLOCK.release()

class SocketViewerLevelThread(ViewerLevelThread):
  """ implementation of viewer level thread using low level sockets """
  def __init__(self, *args, **kwargs):
    ViewerLevelThread.__init__(self, *args, **kwargs)

  def run(self):
    """ instantiate the server """
    if self.verbose:
      print 'thread started...'
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind((self.host, self.port))
    s.listen(1)
    while self.active:
      # wait for a connection from the plugin
      try:
        s.settimeout(TIMEOUT)
        conn, addr = s.accept()
        s.settimeout(None)
      except:
        time.sleep(SLEEP)
        continue
      # when there is a connection, fetch the message
      if self.verbose:
        print 'connection established'
      data = ''
      try:
        while True:
          temp = conn.recv(BUFSIZE)
          if not temp:
            break
          data += temp
        conn.close()
      except:
        if self.verbose:
          print 'socket error, so skipping message'
      # update the levels
      try:
        self.processMessage(data)
      except:
        if self.verbose:
          print 'error with received message: "%s"' % data
    s.close()

if USERPYC:
  class RPyCViewerService(SlaveService):
    """ service for managing received messages using Remote Python Calls (RPyC) """
    def __init__(self, conn):
      SlaveService.__init__(self, conn)

    def exposed_processMessage(self, message):
      """ send a message to the parent thread for processing """
      global RPYCTHREAD
      RPYCTHREAD.processMessage(message)

class RPyCViewerLevelThread(ViewerLevelThread):
  """ implementation of viewer level thread using Remote Python Calls (RPyC) """
  def __init__(self, *args, **kwargs):
    ViewerLevelThread.__init__(self, *args, **kwargs)

  def run(self):
    """ instantiate the server """
    if self.verbose:
      print 'thread started...'
    global RPYCTHREAD
    RPYCTHREAD = self
    self.server = ThreadedServer(RPyCViewerService, port = DEFAULT_SERVER_PORT, auto_register = False, registrar = UDPRegistryClient())
    self.server.start()

  def join(self):
    self.server.close()
    ViewerLevelThread.join(self)

def createViewerLevelThread(*args, **kwargs):
  """ create a viewer level thread according to the defined protocol """
  if USERPYC:
    return RPyCViewerLevelThread(*args, **kwargs)
  else:
    return SocketViewerLevelThread(*args, **kwargs)

#---------------------------------------------------------------------------------------------------
# Utility GUI controls
#---------------------------------------------------------------------------------------------------

class XAxisRangeBox(wx.Panel):
  """ panel for adjusting x-axis range """
  def __init__(self, parent, ID, minvalue = XMIN, initvalue = XWIDTH, increment = SPININC):
    wx.Panel.__init__(self, parent, ID)
    self.minvalue = minvalue
    self.value = initvalue # initial x-axis range width (in sliding mode)
    # controls
    self.radio_full = wx.RadioButton(self, -1, label = 'Full range', style = wx.RB_GROUP)
    self.radio_slide = wx.RadioButton(self, -1, label = 'Sliding')
    self.slide_width = FloatSpin(self, -1, size = (50, -1), digits = 0, value = self.value, min_val = minvalue, increment = increment)
    self.slide_width.GetTextCtrl().SetEditable(False)
    # event bindings
    self.Bind(wx.EVT_UPDATE_UI, self.on_update_radio_buttons, self.radio_full)
    self.Bind(EVT_FLOATSPIN, self.on_float_spin, self.slide_width)
    # layout
    box = wx.StaticBox(self, -1, 'X-axis')
    sizer = wx.StaticBoxSizer(box, wx.VERTICAL)
    slide_box = wx.BoxSizer(wx.HORIZONTAL)
    slide_box.Add(self.radio_slide, flag=wx.ALIGN_CENTER_VERTICAL)
    slide_box.Add(self.slide_width, flag=wx.ALIGN_CENTER_VERTICAL)
    sizer.Add(self.radio_full, 0, wx.ALL, 10)
    sizer.Add(slide_box, 0, wx.ALL, 10)
    self.SetSizer(sizer)
    sizer.Fit(self)

  def on_update_radio_buttons(self, event):
    """ called when the radio buttons are toggled """
    self.slide_width.Enable(self.radio_slide.GetValue())

  def on_float_spin(self, event):
    """ called when the sliding mode spinbox is changed """
    self.value = self.slide_width.GetValue()

  def is_full(self):
    """ return True if full range is checked """
    return self.radio_full.GetValue()

class YAxisRangeBox(wx.Panel):
  """ panel for adjusting y-axis range """
  def __init__(self, parent, ID, minvalue = YMIN, initvalue = YRANGE, increment = SPININC):
    wx.Panel.__init__(self, parent, ID)
    self.value = initvalue # initial y-axis range (in manual mode), i.e. (min, max-min)
    # controls
    self.radio_auto = wx.RadioButton(self, -1, label = 'Auto', style = wx.RB_GROUP)
    self.radio_manual = wx.RadioButton(self, -1, label = 'Manual')
    self.manual_min = FloatSpin(self, -1, size = (50, -1), digits = 0, value = self.value[0], min_val = minvalue[0], increment = increment)
    self.manual_min.GetTextCtrl().SetEditable(False)
    self.manual_width = FloatSpin(self, -1, size = (50, -1), digits = 0, value = self.value[1], min_val = minvalue[1], increment = increment)
    self.manual_width.GetTextCtrl().SetEditable(False)
    # event bindings
    self.Bind(wx.EVT_UPDATE_UI, self.on_update_radio_buttons, self.radio_auto)
    self.Bind(EVT_FLOATSPIN, self.on_float_spin, self.manual_min)
    self.Bind(EVT_FLOATSPIN, self.on_float_spin, self.manual_width)
    # layout
    box = wx.StaticBox(self, -1, 'Y-axis')
    sizer = wx.StaticBoxSizer(box, wx.VERTICAL)
    manual_box = wx.BoxSizer(wx.HORIZONTAL)
    manual_box.Add(self.radio_manual, flag=wx.ALIGN_CENTER_VERTICAL)
    manual_box.Add(self.manual_min, flag=wx.ALIGN_CENTER_VERTICAL)
    manual_box.Add(self.manual_width, flag=wx.ALIGN_CENTER_VERTICAL)
    sizer.Add(self.radio_auto, 0, wx.ALL, 10)
    sizer.Add(manual_box, 0, wx.ALL, 10)
    self.SetSizer(sizer)
    sizer.Fit(self)

  def on_update_radio_buttons(self, event):
    """ called when the radio buttons are toggled """
    toggle = self.radio_manual.GetValue()
    self.manual_min.Enable(toggle)
    self.manual_width.Enable(toggle)

  def on_float_spin(self, event):
    """ called when one of the manual mode spinboxes is changed """
    self.value = (self.manual_min.GetValue(), self.manual_width.GetValue())

  def is_auto(self):
    """ return True if auto range is checked """
    return self.radio_auto.GetValue()

#---------------------------------------------------------------------------------------------------
# Viewer frame class
#---------------------------------------------------------------------------------------------------

class ViewerFrame(wx.Frame):
  """ main frame of the viewer application """
  def __init__(self, test = False):
    wx.Frame.__init__(self, None, -1, NAME)
    self.paused = False
    self.locations = []
    # creation of controls
    self.create_menu()
    self.create_status_bar()
    self.create_main_panel()
    # timer for redrawing
    self.redraw_timer = wx.Timer(self)
    self.Bind(wx.EVT_TIMER, self.on_redraw_timer, self.redraw_timer)
    self.redraw_timer.Start(REDRAWTIME)
    # handle closing the frame
    self.Bind(wx.EVT_CLOSE, self.on_exit, self)
    # manage window style (always on top or not)
    self.wstyle = self.GetWindowStyle()
    self.SetWindowStyle(self.wstyle | wx.STAY_ON_TOP)
    # coordination with data server
    if test:
      self.thread = DummyLevelThread()
    else:
      self.thread = createViewerLevelThread(frame = self)
    self.thread.start()

  def create_menu(self):
    """ construction of menu bar """
    self.menubar = wx.MenuBar()
    # File menu
    menu_file = wx.Menu()
    m_expt = menu_file.Append(-1, '&Save plot\tCtrl-S')
    self.Bind(wx.EVT_MENU, self.on_save_plot, m_expt)
    menu_file.AppendSeparator()
    m_exit = menu_file.Append(-1, 'E&xit\tCtrl-X')
    self.Bind(wx.EVT_MENU, self.on_exit, m_exit)
    # View menu
    menu_view = wx.Menu()
    self.m_ontop = menu_view.Append(-1, '&Stay on top', kind = wx.ITEM_CHECK)
    self.m_ontop.Check(True)
    self.Bind(wx.EVT_MENU, self.on_ontop, self.m_ontop)
    # Help menu
    menu_help = wx.Menu()
    m_about = menu_help.Append(-1, '&About...')
    self.Bind(wx.EVT_MENU, self.on_about, m_about)
    # construction of menu bar
    self.menubar.Append(menu_file, '&File')
    self.menubar.Append(menu_view, '&View')
    self.menubar.Append(menu_help, '&Help')
    self.SetMenuBar(self.menubar)

  def create_status_bar(self):
    """ construction of status bar """
    self.statusbar = self.CreateStatusBar()
    self.statusbar.SetFieldsCount(2)
    self.statusbar.SetStatusWidths([50, -1])

  def create_main_panel(self):
    """ construction of the main controls """
    self.panel = wx.Panel(self)
    # contruct plotting area
    self.fig = Figure(FIGSIZE, dpi = DPI)
    # construct axes
    self.axes = self.fig.add_subplot(111)
    self.axes.set_axis_bgcolor(BGCOLOR)
    # adjust font size of axes labels
    pylab.setp(self.axes.get_xticklabels(), fontsize = FONTSIZE)
    pylab.setp(self.axes.get_yticklabels(), fontsize = FONTSIZE)
    # construct canvas with plotting area
    self.plot_data = self.axes.plot([], linewidth = LINEWIDTH, color = LINECOLOR)[0]
    self.canvas = FigCanvas(self.panel, -1, self.fig)
    # construct location choice box
    self.location_txt = wx.StaticText(self.panel, -1, label = ' Select location:')
    self.location_box = wx.Choice(self.panel, -1, choices = [], size = (150,-1))
    self.location_box.Enable(False)
    self.Bind(wx.EVT_CHOICE, lambda event: self.draw_plot(), self.location_box)
    # layout location choice box
    self.hbox0 = wx.BoxSizer(wx.HORIZONTAL)
    self.hbox0.Add(self.location_txt, border=5, flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
    self.hbox0.Add(self.location_box, border=5, flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
    # construct buttons
    self.pause_button = wx.Button(self.panel, -1, 'Pause')
    self.Bind(wx.EVT_BUTTON, self.on_pause_button, self.pause_button)
    self.Bind(wx.EVT_UPDATE_UI, self.on_update_pause_button, self.pause_button)
    self.clear_button = wx.Button(self.panel, -1, 'Clear')
    self.Bind(wx.EVT_BUTTON, self.on_clear_button, self.clear_button)
    self.cb_grid = wx.CheckBox(self.panel, -1, 'Show grid', style=wx.ALIGN_RIGHT)
    self.Bind(wx.EVT_CHECKBOX, lambda event: self.draw_plot(), self.cb_grid)
    self.cb_grid.SetValue(True)
    self.cb_xlab = wx.CheckBox(self.panel, -1, 'X-labels', style=wx.ALIGN_RIGHT)
    self.Bind(wx.EVT_CHECKBOX, lambda event: self.draw_plot(), self.cb_xlab)
    self.cb_xlab.SetValue(True)
    # layout buttons (add space using self.hbox1.AddSpacer(5))
    self.hbox1 = wx.BoxSizer(wx.HORIZONTAL)
    self.hbox1.Add(self.pause_button, border=5, flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
    self.hbox1.Add(self.clear_button, border=5, flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
    self.hbox1.Add(self.cb_grid, border=5, flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
    self.hbox1.Add(self.cb_xlab, border=5, flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
    # construct axis controls
    self.xrange_control = XAxisRangeBox(self.panel, -1)
    self.yrange_control = YAxisRangeBox(self.panel, -1)
    # layout axis controls
    self.hbox2 = wx.BoxSizer(wx.HORIZONTAL)
    self.hbox2.Add(self.xrange_control, border=5, flag=wx.ALL)
    self.hbox2.Add(self.yrange_control, border=5, flag=wx.ALL)
    # finally, create layout of viewer frame
    self.vbox = wx.BoxSizer(wx.VERTICAL)
    self.vbox.Add(self.canvas, 1, flag=wx.LEFT | wx.TOP | wx.GROW)
    self.vbox.Add(self.hbox0, 0, flag=wx.ALIGN_LEFT | wx.TOP)
    self.vbox.Add(self.hbox1, 0, flag=wx.ALIGN_LEFT | wx.TOP)
    self.vbox.Add(self.hbox2, 0, flag=wx.ALIGN_LEFT | wx.TOP)
    self.panel.SetSizer(self.vbox)
    self.vbox.Fit(self)

  def draw_plot(self):
    """ redraw the plot and update the gui if necessary """
    if not self.paused:
      # check if data is available
      if len(self.locations) == 0:
        self.locations = sorted(self.thread.locations())
        if len(self.locations) > 0:
          self.location_box.AppendItems(self.locations)
          self.location_box.SetSelection(0)
          self.location_box.Enable(True)
          self.flash_status_message('Connection established')
      if len(self.locations) > 0:
        # fetch data at selected receiver location
        loc = self.locations[self.location_box.GetSelection()]
        times, levels = self.thread.levels(loc)
        if (len(times) == len(levels)):
          # calculate x-axis limits
          if self.xrange_control.is_full():
            # show the full range for the x-axis
            xmin = times[0]
            xmax = max(times[0] + self.xrange_control.minvalue, times[-1])
          else:
            # show a sliding window
            xmax = times[-1]
            xmin = xmax - self.xrange_control.value
          # calculate y-axis limits
          if self.yrange_control.is_auto():
            # find the min and max values of the data and add a minimal margin
            ymin = round(min(levels), 0) - MARGIN
            ymax = round(max(levels), 0) + MARGIN
          else:
            # use manual interval
            ymin = self.yrange_control.value[0]
            ymax = ymin + self.yrange_control.value[1]
          # set axis limits
          self.axes.set_xbound(lower = xmin, upper = xmax)
          self.axes.set_ybound(lower = ymin, upper = ymax)
          # finally, plot the data and redraw the plot
          self.plot_data.set_xdata(numpy.array(times))
          self.plot_data.set_ydata(numpy.array(levels))
      # draw grid
      if self.cb_grid.IsChecked():
        self.axes.grid(True, color = GRIDCOLOR)
      else:
        self.axes.grid(False)
      # draw axis labels
      pylab.setp(self.axes.get_xticklabels(), visible = self.cb_xlab.IsChecked())
      self.canvas.draw()

  def clear_plot(self):
    """ clear the data on the plot """
    self.plot_data.set_xdata([])
    self.plot_data.set_ydata([])
    self.canvas.draw()

  def on_redraw_timer(self, event):
    """ redraw the plot """
    self.draw_plot()

  def on_pause_button(self, event):
    """ called when the pause button is clicked """
    self.paused = not self.paused
    if self.paused:
      self.statusbar.SetStatusText('Paused', 0)
    else:
      self.statusbar.SetStatusText('', 0)

  def on_update_pause_button(self, event):
    """ called when the pause button is to be updated """
    label = 'Resume' if self.paused else 'Pause'
    self.pause_button.SetLabel(label)

  def on_clear_button(self, event):
    """ called when the clear butten is clicked """
    self.thread.clear()
    self.clear_choices()
    self.clear_plot()

  def clear_choices(self):
    """ clear the choices box """
    self.locations = []
    self.location_box.Clear()
    self.location_box.Enable(False)
    self.flash_status_message('Cleared')

  def on_save_plot(self, event):
    """ show a window for saving a screenshot """
    dlg = wx.FileDialog(self, message = 'Save plot as...', defaultDir = os.getcwd(), defaultFile = 'plot.png', wildcard = 'PNG (*.png)|*.png', style = wx.SAVE)
    if dlg.ShowModal() == wx.ID_OK:
      path = dlg.GetPath()
      self.canvas.print_figure(path, dpi = DPI)
      self.flash_status_message('Saved to %s' % path)

  def stop_thread(self):
    """ stop the level thread """
    self.thread.active = False
    self.thread.join()

  def on_exit(self, event):
    """ called when the viewer is closed """
    self.stop_thread()
    self.Destroy()

  def on_ontop(self, event):
    """ toggles the stay on top modus """
    if self.m_ontop.IsChecked():
      self.SetWindowStyle(self.wstyle | wx.STAY_ON_TOP)
    else:
      self.SetWindowStyle(self.wstyle)

  def on_about(self, event):
    """ show an about box """
    wx.MessageBox(ABOUT, 'About ' + NAME)

  def flash_status_message(self, message):
    """ flash a message on the status bar """
    try:
      self.statusbar.SetStatusText(message, 1)
      self.timeroff = wx.Timer(self)
      self.Bind(wx.EVT_TIMER, lambda event: self.statusbar.SetStatusText('', 1), self.timeroff)
      self.timeroff.Start(FLASHTIME, oneShot = True)
    except:
      pass

#---------------------------------------------------------------------------------------------------
# Test code
#---------------------------------------------------------------------------------------------------

if __name__ == '__main__':

  if len(sys.argv) <= 1:
    # no command line argument, so run the viewer application
    app = wx.PySimpleApp()
    app.frame = ViewerFrame()
    app.frame.Show()
    app.MainLoop()

  if (len(sys.argv) == 2) and (sys.argv[1] == 'test'):
    # run the viewer in test mode, i.e. generating its own levels for display
    app = wx.PySimpleApp()
    app.frame = ViewerFrame(test = True)
    app.frame.Show()
    app.MainLoop()

  if (len(sys.argv) == 2) and (sys.argv[1] == 'command'):
    # run the viewer in command line mode, i.e. only receiving levels and printing them to the console
    print 'Running viewer in command line mode - press any key to stop...'
    thread = createViewerLevelThread(frame = None, verbose = True)
    thread.start()
    # wait until a key is pressed
    stop = False
    while not stop:
      if msvcrt.kbhit():
        c = msvcrt.getch()
        stop = True
      time.sleep(0.1)
    # stop the thread
    thread.active = False
    thread.join()

  if (len(sys.argv) == 2) and (sys.argv[1] == 'dummy'):
    # run a dummy Aimsun/Noysim2 client that sends random levels (for use with viewer in normal or command line mode)
    print 'Running dummy Aimsun/Noysim2 client - press any key to stop...'
    client = createLevelBuffer(verbose = True, sleep = 1000*TESTSLEEP)
    client.sendClear()
    stop = False
    (t, dt) = (0.0, TESTDT)
    while not stop:
      t += dt
      client.sendLevels(t = t, levels = dict([(loc, randomLevel()) for loc in TESTLOCS]))
      if msvcrt.kbhit():
        c = msvcrt.getch()
        stop = True


# Copyright 2015 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software distributed under the License
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
# or implied. See the License for the specific language governing permissions and limitations under
# the License.

from __future__ import absolute_import
from __future__ import unicode_literals
import collections
import pandas
import sys
import unittest

import google.datalab.bigquery
import google.datalab.utils

class TestCases(unittest.TestCase):

  def test_schema_from_dataframe(self):
    df = TestCases._create_data_frame()
    result = google.datalab.bigquery.Schema.from_data(df)
    self.assertEqual(google.datalab.bigquery.Schema.from_data(TestCases._create_inferred_schema()),
                     result)

  def test_schema_from_data(self):
    variant1 = [
      3,
      2.0,
      True,
      ['cow', 'horse', [0, []]]
    ]
    variant2 = collections.OrderedDict()
    variant2['Column1'] = 3
    variant2['Column2'] = 2.0
    variant2['Column3'] = True
    variant2['Column4'] = collections.OrderedDict()
    variant2['Column4']['Column1'] = 'cow'
    variant2['Column4']['Column2'] = 'horse'
    variant2['Column4']['Column3'] = collections.OrderedDict()
    variant2['Column4']['Column3']['Column1'] = 0
    variant2['Column4']['Column3']['Column2'] = collections.OrderedDict()

    master = [
      {'name': 'Column1', 'type': 'INTEGER'},
      {'name': 'Column2', 'type': 'FLOAT'},
      {'name': 'Column3', 'type': 'BOOLEAN'},
      {'name': 'Column4', 'type': 'RECORD', 'fields': [
          {'name': 'Column1', 'type': 'STRING'},
          {'name': 'Column2', 'type': 'STRING'},
          {'name': 'Column3', 'type': 'RECORD', 'fields': [
              {'name': 'Column1', 'type': 'INTEGER'},
              {'name': 'Column2', 'type': 'RECORD', 'fields': []}
          ]}
      ]}
    ]

    schema_master = google.datalab.bigquery.Schema(master)

    with self.assertRaises(Exception) as error1:
      google.datalab.bigquery.Schema.from_data(variant1)
    if sys.version_info[0] == 3:
      self.assertEquals('Cannot create a schema from heterogeneous list [3, 2.0, True, ' +
                        '[\'cow\', \'horse\', [0, []]]]; perhaps you meant to use ' +
                        'Schema.from_record?', str(error1.exception))
    else:
      self.assertEquals('Cannot create a schema from heterogeneous list [3, 2.0, True, ' +
                        '[u\'cow\', u\'horse\', [0, []]]]; perhaps you meant to use ' +
                        'Schema.from_record?', str(error1.exception))
    schema3 = google.datalab.bigquery.Schema.from_data([variant1])
    schema4 = google.datalab.bigquery.Schema.from_data([variant2])
    schema5 = google.datalab.bigquery.Schema.from_data(master)
    schema6 = google.datalab.bigquery.Schema.from_record(variant1)
    schema7 = google.datalab.bigquery.Schema.from_record(variant2)

    self.assertEquals(schema_master, schema3, 'schema inferred from list of lists with from_data')
    self.assertEquals(schema_master, schema4, 'schema inferred from list of dicts with from_data')
    self.assertEquals(schema_master, schema5, 'schema inferred from BQ schema list with from_data')
    self.assertEquals(schema_master, schema6, 'schema inferred from list with from_record')
    self.assertEquals(schema_master, schema7, 'schema inferred from dict with from_record')

  @staticmethod
  def _create_data_frame():
    data = {
      'some': [
        0, 1, 2, 3
      ],
      'column': [
        'r0', 'r1', 'r2', 'r3'
      ],
      'headers': [
        10.0, 10.0, 10.0, 10.0
      ]
    }
    return pandas.DataFrame(data)

  @staticmethod
  def _create_inferred_schema(extra_field=None):
    schema = [
      {'name': 'some', 'type': 'INTEGER'},
      {'name': 'column', 'type': 'STRING'},
      {'name': 'headers', 'type': 'FLOAT'},
    ]
    if extra_field:
      schema.append({'name': extra_field, 'type': 'INTEGER'})
    return schema