#!/usr/bin/env python
'''
wirelessheatmap
Copyright 2008 Gerald Kaszuba
http://geraldkaszuba.com/

This program 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.

This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
'''

import colorsys
import math
import time
import pickle
from ConfigParser import ConfigParser
from optparse import OptionParser

import pyglet
pyglet.options['debug_gl'] = 0

from pyglet.window import key
from pyglet.window import mouse
from pyglet import window
from pyglet import image
from pyglet import font
from pyglet.gl import *

from ap import AccessPoint

class Heat(object):

    def __init__(self):

        self.parse_config()

        self.window = window.Window(
#            640, 480,
#            1024, 768,
            800, 600,
#            caption='Heat', \
#            fullscreen=True,
            )
        self.window.on_mouse_press = self.on_mouse_press
        self.window.on_mouse_drag = self.on_mouse_drag
        self.window.on_mouse_scroll = self.on_mouse_scroll
        self.window.on_key_release = self.on_key_release
        
        self.current_ap = ''
       
        # defaults before image loads
        self.cellsize = 12
        self.img_zoom = 1.

        if self.cfg.image_file:
            self.img = image.load(self.cfg.image_file)
            self.img_zoom = float(self.window.width) / self.img.width
            zoomier = float(self.window.height) / self.img.height
            if zoomier < self.img_zoom:
                self.img_zoom = zoomier

            self.cellsize = self.img.width / 100

        else:
            self.img = None

        self.img_pos = [0., 0.]

        self.font_size = 20
        self.font = font.load('Arial', self.font_size)

        self.load_data()

        self.last_dump_read = 0
                
        self.samples = 10
        self.grid = {}

        self.grid_margin = 100  # pixels surrounding the image

    def parse_config(self):
        self.cfg = ConfigParser()
        self.opt = OptionParser(usage='usage: %prog [options]')

        # Configuration loading and saving
        self.opt.add_option('--config', dest='config_file', \
            help='Read configuration from a file.')
        self.opt.add_option('--save', dest='dest_cfg', \
            help='Generate a config file from your options into a file.')

        # General shit
        self.opt.add_option('--image', dest='image_file')
        self.opt.add_option('--dump', dest='dump_file')
        self.opt.add_option('--store', dest='store_file')

        (opts, args) = self.opt.parse_args()

        # If we want to load from a cfg file, update the defaults then override
        # them with the command line arguments.
        if opts.config_file:
            self.cfg.read(opts.config_file)
            for k, v in self.cfg.items('main'):
                # Since these are saved as strings we have to cast them back.
                # XXX: This needs to be cleaner...
                if v == 'False':
                    v = False
                if v == 'True':
                    v = True
                self.opt.set_default(k, v)
            (opts, args) = self.opt.parse_args()

        if opts.dest_cfg:
            f = open(opts.dest_cfg, 'w')
            f.write('[main]\n')
            for a in self.opt.defaults:
                if a in ('dest_cfg', 'config_file'):
                    continue
                f.write('%s = %s\n' % (a, getattr(opts, a)))
            f.close()
            print '%s written' % opts.dest_cfg
            sys.exit(0)

        self.cfg = opts

        return args

    def draw_text(self, text, line):
        t = font.Text(self.font, text, y=line*self.font_size)
        t.draw()

    def read_dump(self):
        if not self.cfg.dump_file:
            return
        d = open(self.cfg.dump_file).read()
        lines = d.split('\n')
        
        for line in lines:

            bits = line.split(',')

            if len(bits) != 15:
                continue
            if bits[0] == 'BSSID':
                continue

            ap = AccessPoint(*bits)

            self.ap[ap.bssid] = ap
            if ap.bssid not in self.store:
                self.store[ap.bssid] = {}

            # Make the first read AP the currently selected one
            if not self.current_ap:
                self.current_ap = ap.bssid

    def go(self):
        while not self.window.has_exit:

            if self.last_dump_read + 1 < time.time():
                self.read_dump()
                self.last_dump_read = time.time()

            self.window.clear()
            
            if self.current_ap:
                self.calc_heat()

                self.map_matrix()
                self.draw_heat()
                self.draw_strength()
                
                glLoadIdentity()
                self.draw_text('%s %s %f' % (self.current_ap,
                    self.ap[self.current_ap].essid, 
                    self.ap[self.current_ap].power), 0)

            if self.img:
                self.map_matrix()
                glEnable(GL_BLEND)
                glBlendFunc(GL_SRC_ALPHA, GL_ONE)
                glColor4f(1, 1, 1, 0.5)
                self.img.blit(0, 0)
                glDisable(GL_BLEND)

            self.window.flip()
            self.window.dispatch_events()
#            time.sleep(0.5)
   
    def map_matrix(self):
        glLoadIdentity()
        glTranslatef(self.img_pos[0], self.img_pos[1], 0)
        glScalef(self.img_zoom, self.img_zoom, 1)

    def screen2map(self, x, y):
        x -= self.img_pos[0]
        y -= self.img_pos[1]
        x /= self.img_zoom
        y /= self.img_zoom
        return x, y

    def draw_strength(self):
        if not self.current_ap:
            return
        if self.current_ap not in self.store:
            return
        if not self.store[self.current_ap]:
            return

        max_str = max(self.store[self.current_ap].values())
        min_str = min(self.store[self.current_ap].values())

        if not max_str:
            return
        if max_str == min_str:
            return

        glPointSize(2)
        glBegin(GL_POINTS)
        for pos, strength in self.store[self.current_ap].items():
            s = 0.5 + (strength - min_str) / (max_str - min_str)
            glColor3f(s, s, s)
            glVertex2f(pos[0], pos[1], 0)
        glEnd()

    def calc_heat(self):

        if self.grid:
            return

        ap = self.store[self.current_ap]
        self.grid = {}
        
        gmin = -self.grid_margin
        gmaxx = self.img.width + self.grid_margin
        gmaxy = self.img.height + self.grid_margin

        for x in range(gmin, gmaxx, self.cellsize):
            for y in range(gmin, gmaxy, self.cellsize):

                distance_strengths = []
                for pos, strength in ap.items():
                    dist = math.hypot(x - pos[0], y - pos[1])
                    distance_strengths.append((dist, strength))

                heat = 0
                def srt(a, b):
                    if a[0] > b[0]:
                        return 1
                    return -1

                distance_strengths.sort(srt)

                for d, s in distance_strengths[:self.samples]:
                    heat += s
                heat /= self.samples
                self.grid[x, y] = heat

    def draw_heat(self):
        if not self.grid:
            return
        max_heat = max(self.grid.values())
        min_heat = min(self.grid.values())
        if not max_heat:
            return
        if max_heat == min_heat:
            return

        glBegin(GL_QUADS)

        for pos, heat in self.grid.items():
            x = pos[0] - self.cellsize / 2
            y = pos[1] - self.cellsize / 2
            heat -= min_heat
            heat /= (max_heat - min_heat)
            rgb = self.col(heat)
            if not rgb:
                continue

            glColor3f(*rgb)
            glVertex2f(x, y)
            glVertex2f(x + self.cellsize, y)
            glVertex2f(x + self.cellsize, y + self.cellsize)
            glVertex2f(x, y + self.cellsize)

        glEnd()

    def col(self, val):
        return colorsys.hsv_to_rgb((1 - val) * 2. / 3, 1, val * 0.6)

    def on_mouse_press(self, x, y, button, modifiers):
        if button == mouse.LEFT:
            x, y = self.screen2map(x, y)
            for ap in self.ap.values():
                self.store[ap.bssid][x, y] = ap.power
            self.save_data()
   
    def on_mouse_drag(self, x, y, dx, dy, button, modifiers):
        if button == mouse.RIGHT:
            self.img_pos[0] += dx
            self.img_pos[1] += dy

    def on_mouse_scroll(self, x, y, scroll_x, scroll_y):
        if scroll_y > 0:
            self.img_zoom /= 0.8 * scroll_y
        if scroll_y < 0:
            self.img_zoom *= 0.8 * -scroll_y

    def on_key_release(self, symbol, mods):
        self.grid = {}
        if symbol == key.S:
            pyglet.image.get_buffer_manager().get_color_buffer(). \
                save('screenshot.png')
            print 'Saved as screenshot.png'
        if symbol == key.SPACE:

            bssids = self.ap.keys()
            bssids.sort()

            try:
                pos = bssids.index(self.current_ap)
            except ValueError:
                pos = 0

            try:
                self.current_ap = bssids[pos + 1]
            except IndexError:
                self.current_ap = bssids[0]

        if symbol == key.UP:
            self.samples += 1
            print self.samples
        if symbol == key.DOWN:
            self.samples -= 1
            print self.samples

    def load_data(self):
        # store[essid][x, y] = strength
        # ap[essid] = AccessPoint()
        try:
            self.store, self.ap = pickle.load(open(self.cfg.store_file))
        except Exception:
            self.store = {}
            self.ap = {}
            print 'Creating new store'
            return

        print self.cfg.store_file, 'loaded'
        print len(self.store), 'stores'
        print self.ap

    def save_data(self):
        pickle.dump((self.store, self.ap), open(self.cfg.store_file, 'wb'))


Heat().go()