#!/usr/bin/python
# ./generate-worldmap.py -i mapping/f7.txt -o public_html/maps/testing/7/f7.png
#
import sys
import getopt
import GeoIP
import random
import matplotlib
matplotlib.use('Agg2')
from pylab import *
from matplotlib.numerix import ma
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.toolkits.basemap import Basemap
from matplotlib.colors import LogNorm
from matplotlib.ticker import LogFormatterMathtext as JefFormatter

gi = GeoIP.open("/usr/share/GeoIP/GeoLiteCity.dat", GeoIP.GEOIP_MEMORY_CACHE)
random.seed()

def lookup_client_locations():
    results = []
    f = open(inputfile, 'r')
    for line in f:
        try:
            gir = gi.record_by_addr(line.strip())
        except:
            continue
        if gir != None:
            t = (line.strip(), gir['country_code'], gir['latitude'], gir['longitude'])
            results.append(t)
    f.close()
    return results


def lookup_host_locations():
    results = []
    for h in Host.select():
        if h.private or h.site.private or \
               not h.user_active or not h.admin_active or \
               not h.site.user_active or not h.site.admin_active:
            continue
        try:
            gir = gi.record_by_name(h.name)
        except:
            print "Cannot find location for %s" % (h.name)
            continue
        if gir != None:
            t = (h.name, gir['country_code'], gir['latitude'], gir['longitude'])
            results.append(t)

    return results

def draw_client_density():

    m = Basemap(llcrnrlon=-180.,llcrnrlat=-90,urcrnrlon=180.,urcrnrlat=90.,\
                resolution='l',projection='cyl')

    # plot them as filled circles on the map.
    # first, create a figure.
    dpi=100
    dimx=800/dpi
    dimy=400/dpi
    fig=figure(figsize=(dimx,dimy), dpi=dpi, frameon=True, facecolor='blue',edgecolor='white')
#    ax=fig.add_axes([0.1,0.1,0.7,0.7],axisbg='g')
    ax=fig.add_axes([0.05,0.1,0.8,0.8],axisbg=(0.05,0.65,0.05))
    canvas = FigureCanvas(fig)
    results = lookup_client_locations()
    X,Y,Z = find_client_density(m,results)
#    s = random.sample(results, 40000)
#    for t in s:
#        lat=t[2]
#        lon=t[3]
#        # draw a red dot at the center.
#        xpt, ypt = m(lon, lat)
#        m.plot([xpt],[ypt],'ro', zorder=10)
    # draw coasts and fill continents.
    m.drawcoastlines(linewidth=0.5)
    m.drawcountries(linewidth=0.5)
    m.drawlsmask([100,100,100,0],[100,210,210,255])
#    m.fillcontinents(color='green')
    palette = cm.YlOrRd
#    palette = cm.cool
    m.imshow(Z,palette,extent=(m.xmin,m.xmax,m.ymin,m.ymax), norm=LogNorm(),interpolation='gaussian')
    l,b,w,h = ax.get_position() 
    cax = axes([l+w+0.005, b, 0.03, h])
    colorbar(cax=cax,format=JefFormatter()) # draw colorbar
    figtext(l+w+0.1,0.5,"Client Density: clients per degree lat/lon",va="center",ha="center",rotation="vertical",fontsize=10)
    figtext(0.05,0.05,footer,backgroundcolor='white',fontsize="smaller",va="bottom")
    figtext(0.05,0.95,header,backgroundcolor='white',fontsize="smaller",va="top")
    canvas.print_figure(outputfile, dpi=100,facecolor='white',edgecolor='white')


def find_client_density(m,client_locations,latscale=1.0,lonscale=1.0,lat_smooth=1,lon_smooth=1):
    lat_array=arange(m.ymin,m.ymax+latscale,latscale)
    lon_array=arange(m.xmin,m.xmax+lonscale,lonscale)
    maxlat=len(lat_array)-1
    maxlon=len(lon_array)-1
    Z=zeros((len(lat_array),len(lon_array)),dtype='float')
    for client in client_locations:
      lat=client[2]
      i_lat=int(float((lat-lat_array[0]))/float(latscale)) 
      lon=client[3]
      i_lon=int(float((lon-lon_array[0]))/float(lonscale)) 
      for i in xrange(-int(lat_smooth),int(lat_smooth+1),1):
        for j in xrange(-int(lon_smooth),int(lon_smooth+1),1):
          if ( i_lat+i >= 0 ) and (i_lat+i < maxlat) : 
            if ( i_lon+j >= 0) and ( i_lon+j < maxlon) :
              Z[i_lat+i,i_lon+j]+=1.0
#    Z=Z/(2.0*lat_smooth+1)/(2.0*lon_smooth+1)
    Lon,Lat=meshgrid(lon_array,lat_array)
    X,Y=m(Lon,Lat)
#    Z= Z + 1.0
#    Z=log(Z)
    Z = where(Z <= 0.,1.e10,Z)
    Z = ma.masked_values(Z, 1.e10)
    return X,Y,Z

def usage():
  print "%s -h header -f footer -i inputfile -o outputfile --indir=indir --outdir=outdir -v" % sys.argv[0]    

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "h:f:i:o:v", ["help","verbose","indir=","outdir=","input=","output=","footer=","header="])
    except getopt.GetoptError:
        # print help information and exit:
        usage(args)
        sys.exit(2)
    verbose = False
    global header, footer, inputfile,outputfile,indir,outdir
    inputfile=None
    outputfile=None
    header=""
    footer=""
    indir='./'
    outdir='./'
    for o, a in opts:
        if o in ("-h","--header"):
            header = a
        if o in ("-f","--footer"):
            footer = a
        if o in ("-i","--input"):
            inputfile = a
        if o in ("-o","--output"):
            outputfile = a
        if o == "--indir":
            indir = a
        if o == "--outdir":
            outdir = a
        if o in ("-v","--verbose"):
            verbose = True
        if o in ("--help"):
            usage()
            sys.exit()

    if outputfile is None : outputfile = outdir+'/clientmap.png'
    if inputfile is None : inputfile = indir+'/ips.txt'

    draw_client_density()



if __name__ == "__main__":
    sys.exit(main())