pyvtraceroute/vtracemap.py

# -*- coding: utf-8 -*-#

# Copyright 2015 Weber Yann
#
# This file is part of pyMapTraceroute.
#
#        pyMapTraceroute is free software: you can redistribute it and/or modify
#        it under the terms of the GNU General Public License as published by
#        the Free Software Foundation, either version 3 of the License, or
#        (at your option) any later version.
#
#        pyMapTraceroute is distributed in the hope that it will be useful,
#        but WITHOUT ANY WARRANTY; without even the implied warranty of
#        MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#        GNU General Public License for more details.
#
#        You should have received a copy of the GNU General Public License
#        along with pyMapTraceroute.  If not, see <http://www.gnu.org/licenses/>.
#

import math, pygame, os, time, multiprocessing, struct, sys

from multiprocessing import Pool
from functools import partial

import vtracert

PID4 = math.pi / 4.0
LOGTPID4 = math.log1p(math.tan(PID4))

## Function runned by the pool when drawing the map
# @param (lon, lat) Coordinates
# @param arg ( (lonmin,lonmax,latmin,latmax), scaleW, scaleH, xlim)
# @return (x,y) or None
def draw__ProcessFun((lon, lat), arg ):
	
	( (lonmin,lonmax,latmin,latmax), scaleW, scaleH, xlim) = arg
	
	if lonmin > lonmax:
		test = (lon < lonmax and lon > lonmin)
	else:
		test = (lon > lonmin and lon < lonmax)
		
	if lat > latmin and lat < latmax and test:
		
		(x,y) = TraceMap._latlon2ortho(lon,lat,scaleW,scaleH,xlim)
		
		x%=scaleW
		y%=scaleH
		return (x,y)
	return None


## Handle a map for tracerouting
# Give access to functions for coordinates convertion
class TraceMap(object):
	
	def __init__(self, vtr, lonmin = -179, lonmax = 180, latmin = -100, latmax = 100):
		
		if not isinstance(vtr, vtracert.Vtracert):
			raise TypeError('Vtracert waited')
		
		self.vtr = vtr
		
		#Loading datas
		self.lonlat = TraceMap.lonlatFromFile(self.vtr.mapfile)
		
		self.xlimits = None
		self.scaledWidth = 0 #Store the screen width if we will display all the map
		self.lonlatminmax = (lonmin,lonmax,latmin,latmax)
		self.xlimits = self.genXlimits(lonmin,lonmax,latmin,latmax)
		
		#Pool used to draw the map (convert lon, lat into x,y)
		self.pdraw = Pool(multiprocessing.cpu_count())
		
		pass
	
	def __destroy__(self):
		self.pdraw.close()
	
	## Draw the map
	# Automatically choose the fastest method to draw (with or without the pool)
	# @param animated Use more time but its "fancy"
	def draw(self,animated = True):
		(a,_,b,_) = self.xlimits
		if a + b > 400000:
			self.__poolDraw() #Use the pool
		else:
			self.__soloDraw(False)
		pass
	
	## Draw the map using the TraceMap::pdraw pool
	# @param animated Useless
	# @see TraceMap::draw(), TraceMap::__soloDraw(), __drawP()
	def __poolDraw(self, animated = False):
		
		atime = time.time() #timer
		
		self.vtr.surf['map'].fill(self.vtr.col['bg'])
		
		#Array used because fastest
		sufarr = pygame.surfarray.array2d(self.vtr.surf['map'])
		
		(lonmin,lonmax,latmin,latmax) = self.lonlatminmax
		
		w = self.vtr.surf['map'].get_width()
		h = self.vtr.surf['map'].get_height()
		
		#Pass constants arguments to function
		pDrawP = partial(draw__ProcessFun, arg = ( self.lonlatminmax, w, h, self.xlimits ) )
		
		toUpdate = self.pdraw.map(pDrawP, self.lonlat)
		#Cleaning with set
		toUpdate = set(toUpdate).difference(set([None]))
		
		for x,y in toUpdate:
			sufarr[x][y] = self.vtr.col['map']
		
		#Blit array back on surface
		pygame.surfarray.blit_array(self.vtr.surf['map'], sufarr)
		
		self.vtr.blitScreen('map')
		
		btime = time.time()#timer
		print "Map draw spent ",(btime-atime)*1000,"miliseconds"
		
		pass
	
	## Draw the map using a single process
	# @param animated If True the screen is updated while drawing (small time consumption)
	# @see TraceMap::draw(), TraceMap::__soloDraw()
	def __soloDraw(self, animated = True):
		
		atime = time.time()
		
		self.vtr.surf['map'].fill(self.vtr.col['bg'])
		
		#Array used because fastest
		sufarr = pygame.surfarray.array2d(self.vtr.surf['map'])

		cnt = 1 #Counter
		
		(lonmin,lonmax,latmin,latmax) = self.lonlatminmax

		#Refresh if animated
		lll = len(self.lonlat)
		refresh = [ lll/20, lll/10, lll/2, (lll*3)/4 ]
		
		
		for (lon,lat) in self.lonlat:
			
			if lonmin > lonmax:
				test = (lon < lonmax and lon > lonmin)
			else:
				test = (lon > lonmin and lon < lonmax)
				
			if test and lat > latmin and lat < latmax:
				(x,y) = self.latlon2ortho(lon,lat)
				
				x%=len(sufarr)
				y%=len(sufarr[x])
				
				sufarr[x][y] = self.vtr.col['map']
				
				
				if animated and self.vtr.surf['screen'] != None and cnt in refresh:
					self.vtr.surf['map'].fill((0,0,0))
					pygame.surfarray.blit_array(self.vtr.surf['map'], sufarr)
					self.vtr.blitScreen('map')
					
					pygame.display.flip()
				
			cnt +=1
		
		self.vtr.surf['map'].fill((0,0,0))
		#blit array back
		pygame.surfarray.blit_array(self.vtr.surf['map'], sufarr)
		self.vtr.blitScreen('map')
		
		if animated and self.vtr.surf['screen'] != None:
			pygame.display.flip()
			
		btime = time.time()
		print "Map OLDdraw spent ",(btime-atime)*1000,"miliseconds"
	
	## Generate limits from lon and lat to convert coordinates
	# @param lonmin Longitude mini
	# @param lonmax Longitude max
	# @param latmin Latitude mini
	# @param latmax Latitude max
	def genXlimits(self, lonmin,lonmax,latmin,latmax):
	
		self.lonlatminmax = (lonmin,lonmax,latmin,latmax)
	
		(xs,ys) = self.latlon2ortho(lonmax, latmax, True)
		(xl,yl) = self.latlon2ortho(lonmin, latmin, True)
		
		if xs<xl:
			x1 = xs
			x2 = xl
		else:
			x1 = xl
			x2 = xs
		
		if ys<yl:
			y1 = ys
			y2 = yl
		else:
			y1 = yl
			y2 = ys
		
		self.scaledWidth = int(self.vtr.width * (float(self.vtr.width * 707) / (x2-x1)))
		self.xlimits = (x2-x1,x1,y2-y1,y1)
		return self.xlimits
	
	## Zoom on the map
	# @param (cx,cy) Where to center the zoom
	# @param zoom If True zoom if False zoom out and can be a direction str 'up' 'down' 'left' 'right' to move on the map
	def zoom(self,(cx,cy), zoom = True):
		
		width = self.vtr.surf['map'].get_width()
		height = self.vtr.surf['map'].get_height()
		
		zoomV = 0.3 #Zoom 'power'
		moveD = 50 #Move 'power'
		
		if not isinstance(zoom, str):
			#New height and width	
			if zoom:
				nw = (width * (1-zoomV))/2
				nh = (height * (1-zoomV))/2
			else:
				nw = (width * (1+zoomV))/2
				nh = (height * (1+zoomV))/2

			#New limits virtual coordinates
			xmin = int(cx - nw)
			xmax = int(cx + nw)
			
			ymin = int(cy - nh)
			ymax =  int(cy + nh)
		else:
			#Not zooming but moving
			xmin = 0
			xmax = self.vtr.surf['map'].get_width()
			ymin = 0
			ymax = self.vtr.surf['map'].get_height()
				
			if zoom == 'up':
				ymin -= moveD
				ymax -= moveD
			elif zoom == 'down':
				ymin += moveD
				ymax += moveD
			elif zoom == 'right':
				xmin += moveD
				xmax += moveD
			else:
				xmin -= moveD
				xmax -= moveD
		
		#Calculating new lon and lat limits from new limits
		savlimits = self.lonlatminmax
		
		(lonmin,latmin) = self.ortho2latlon(xmin,ymin)
		(lonmax,latmax) = self.ortho2latlon(xmax,ymax)
		
		if latmin > latmax:
			foo = latmin
			latmin = latmax
			latmax = foo

		if lonmax-lonmin > 360 or latmax - latmin > 180:
			#Less zoom is useless, recenter the map
			lonmin = -179
			lonmax = 179
			latmin = -100
			latmax = 100
		
		if (lonmin,lonmax,latmin,latmax) != savlimits: #only redraw if needed
			
			self.genXlimits(lonmin,lonmax,latmin,latmax)
			
			screen = self.vtr.surf['screen']
			screen.blit(self.vtr.surf['map'],(0,0))
			#zoombox draw
			pygame.gfxdraw.line(screen,xmin,ymin,xmin,ymax,(255,255,100))
			pygame.gfxdraw.line(screen,xmax,ymin,xmax,ymax,(255,255,100))
			pygame.gfxdraw.line(screen,xmin,ymin,xmax,ymin,(255,255,100))
			pygame.gfxdraw.line(screen,xmin,ymax,xmax,ymax,(255,255,100))
			pygame.display.flip()
			
			#Redraw map
			self.draw()
			pass
	
	## Load coordinates from a file
	# The file stores point like : lon lat
	# @param coordFile The filename
	# @return A list of tuple
	@classmethod
	def lonlatFromFile(c, coordFile):
	
		res = []
		
		fd = file(coordFile,"r")

		buff = ''
		lat = 0.0
		lon = 0.0

		#cread = fd.readline()

		for line in fd:
			line = line.strip()
			(lon,lat) = line.rsplit(' ')
			lon = float(lon)
			lat = float(lat)
			res.append((lon,lat))
		
		fd.close()
		return res
	
	## Load packed data from file
	# Packed file use less disk space but take more time to load.
	# Not used by default
	# @param coordFile The filename with packed datas
	# @return a list of (lon,lat)
	@classmethod
	def lonlatFromPackFile(c, coordFile):
		
		res = []
		
		fd = file(coordFile,"rb")
		
		while True:
			"""
			r = fd.read(4)
			if len(r) < 4:
				break
			lon = struct.unpack('f',r)
			r = fd.read(4)
			if len(r) < 4:
				break
			lat = struct.unpack('f',r)
			"""
			r = fd.read(8)
			if len(r) < 8:
				break
			lon = struct.unpack('f',r[:4])
			lat = struct.unpack('f',r[4:])
			res.append((lon,lat))
		
		fd.close()
		return res
			
	## Store data packed
	# Packed file use less disk space but take more time to load.
	# Not used by default
	# @param coordFile The filename where we want to store data
	@classmethod
	def lonlatToPackFile(c, coordFile, lonlat):
		
		fd = file(coordFile, 'w+')
		
		buff = [ struct.pack('f', lon)+struct.pack('f', lat) for (lon, lat) in lonlat ]
		buff = b''.join(buff)
		
		print len(buff)
		
		fd.write(buff)
		fd.close()
		pass
			
	
	## Convert longitude and latitude to screen coordinates
	# @param lon Longitude
	# @param lat Lattitude
	# @param nolimit If true don't scale
	# @return (x,y)
	# @see TraceMap::_latlon2ortho(), TraceMap::ortho2latlon()
	def latlon2ortho(self, lon,lat, nolimit = False):
		scaleW = self.vtr.surf['map'].get_width()
		scaleH = self.vtr.surf['map'].get_height()
		if nolimit:
			xlim = None
		else:
			xlim = self.xlimits
		
		return TraceMap._latlon2ortho(lon,lat,scaleW,scaleH,xlim)
	
	## Body of TraceMap::latlon2ortho()
	# Used like classmethode by __drawP()
	# @param lon Longitude
	# @param lat Lattitude
	# @param scaleW Map width
	# @param scaleH Map height
	# @param xlimits If none dont scale
	# @return (x,y)
	# @see TraceMap::latlon2ortho(), TraceMap::ortho2latlon()
	@classmethod
	def _latlon2ortho(c, lon,lat , scaleW, scaleH, xlimits = None):
		
		width = scaleW * 707
		height = scaleH * (707 * (scaleW/float(scaleH)))
	
		lat=math.radians(lat)
		ratio = int(width/(math.pi*2.0))
		x = int(width*((lon+180)/360.0))
		y = int(height/2.0-LOGTPID4-lat/2.0*ratio)
		
		#Scaling with xlimits infos
		if xlimits != None :
			(dx,xs,dy,ys) = xlimits
			x = int(scaleW * float(x-xs) / dx)
			y = int(scaleH * float(y-ys) / dy)
		
		return (x,y)
	
	## Convert coordinates to longitude and latitude
	# @param x X
	# @param y Y
	# @return (lon,lat)
	# @see TraceMap::latlon2ortho()
	def ortho2latlon(self, x,y):
		
		scaleW = self.vtr.surf['map'].get_width()
		scaleH = self.vtr.surf['map'].get_height()
		width = scaleW * 707
		height = scaleH * (707 * (scaleW/float(scaleH)))
		
		
		ratio = int(width/(math.pi*2.0))
		
		(dx,xs,dy,ys) = self.xlimits
		
		x = (x * dx / float(scaleW))+xs
		y = (y * dy / float(scaleH))+ys
		
		lon = 180 - (360 * (x/float(width))) * -1
		lat = math.degrees( (y- float(height) / 2 + LOGTPID4)/ratio ) * -2
		
		
		lon -= 360
		
		return (lon,lat)