from graph import Graph
from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt
import numpy as np
import os

class Map(Graph):
    def __init__(self):
        super().__init__('map')
        self.data = {} # (lat, lon): { 'num': num, 'size': size (gib) }
    
    def process(self, ip, time, request, size, location, log):
        if not location:
            return []

        # since we show pie chart of us states, ensure that the points we plot (if in us) have an associated state
        if location.country.name != 'United States' or len(location.subdivisions) > 0:
            pos = (location.location.latitude, location.location.longitude)
            self.validate(self.data, pos, { 'num': 0, 'size': 0})
            self.data[pos]['num'] += 1
            self.data[pos]['size'] += size

        return []
    
    def draw(self, path):
        print('Generating maps...')

        # https://python-graph-gallery.com/310-basic-map-with-markers/
        plt.figure(figsize=(12.8, 9.6))
        m = Basemap(llcrnrlon=-180, llcrnrlat=-90,urcrnrlon=180,urcrnrlat=90)

        m.drawmapboundary(fill_color='#A6CAE0', linewidth=0)
        m.fillcontinents(color='grey', alpha=0.7)
        m.drawcoastlines(color='grey', linewidth=0.5)
        m.drawcountries(color='grey')

        lat = []
        lon = []
        num = np.zeros(len(self.data))
        size = np.zeros(len(self.data))

        for i, ((_lat, _lon), val) in enumerate(self.data.items()):
            n = val['num']
            s = val['size']
            lat.append(_lat)
            lon.append(_lon)
            num[i] = n
            size[i] = s

        # normalize all values to 0.005 - 1, with 0 values still being 0
        num /= num.max() / 0.99
        size /= size.max() / 0.99
        num[num != 0] += 0.01
        size[size != 0] += 0.01

        plt.title('Map of Downloads by Count', fontsize=14)
        n = m.scatter(lon, lat, s=num * 300, alpha=0.5, edgecolor='black', c='blue', linewidth=0.5, zorder=10)

        plt.savefig(os.path.join(path, 'map_count.png'), bbox_inches='tight', dpi=100)
        n.set_visible(False)

        plt.title('Map of Downloads by File Size', fontsize=14)
        s = m.scatter(lon, lat, s=size * 300, alpha=0.75, edgecolor='black', c='orange', linewidth=0.5, zorder=10)
        plt.savefig(os.path.join(path, 'map_size.png'), bbox_inches='tight', dpi=100)