import sys
import math
import numpy as np
import copy
from mpl_toolkits.basemap import Basemap
from matplotlib import pyplot as plt

# class object for each data
class Tornado(object):
        def __init__(self, longitude, latitude, time_in, fujita):
                self.longitude = longitude
                self.latitude = latitude
                self.time_in = time_in
                self.fujita = fujita
        def __repr__(self):
                return '{}: {} {}'.format(self.__class__.__name__,
                                          self.longitude,
                                          self.latitude,
                                          self.time_in,
                                          self.fujita)

#INPUT: a Tornado object
#EFFECT: return the time for soring
def getkey(Tornado):
        return Tornado.time_in

# INPUT: longitude, latitude
# EFFECT: calculates the kernel based on longitude and latitude
def kernel (longitude, latitude):
        CONST_e = math.e
        CONST_pi = math.pi
        # kernel = (2Pi)^(-1) * (e)^(-(x^2 + y^2)/2) where x is longitude and y is latitude
        kernel = math.pow(CONST_e, -(math.pow(longitude, 2) + math.pow(latitude, 2))/2) / (2* CONST_pi)
        return kernel

# INPUT: x position as longitude, y position as latitude, the bandwidth for space h
# EFFECT: calculate the kernel at M(x,y)
def kernel_bandwidth_h(position_x, position_y, bandwidth):
        kernel_h = (1/ math.pow(bandwidth, 2)) * kernel (position_x/bandwidth, position_y/bandwidth)
        return kernel_h

# INPUT: min_value as the minimum value of longitude or latitude,
#        max_value as the maximum value of longitude or latitude,
#        position_value as longitude or latitude
#        pixel_value as the user input m_x or m_y
# EFFECT: Calculate the indice where the point longitude[i], latitude[i] falls in in the matrix_M
def indice_calc (min_value, max_value, position_value, pixel_value):
        indice_value = int(round(pixel_value * (position_value - min_value)/(max_value - min_value)))
        return indice_value

# INPUT: k as the parameter that determines the size of the matrix_N
#        bandwidth h, row number and column number
# EFFECT: calculates the N(row, col) entry of the matrix_N
def N_matrix_entries (k, bandwidth, row, col):
        CONST_e = math.e
        CONST_pi = math.pi
        # N(i, j) = e^(-((i - (k+1))/h)^2 - ((j -(k+1))/h)^2) /(2*pi*h^2)
        entry = math.pow(CONST_e, -math.pow((row - k)/bandwidth, 2) - math.pow((col - k)/bandwidth, 2)) / (2*CONST_pi*math.pow(bandwidth, 2))
        return entry

#INPUT: round_robin as a 3D nparray that stores the matrix of different time
#       matrix as the matrix to be stored, position as the indicator and size as the size of the round_robin
#EFFECT: store the matrix in round_robin and update position accordingly
def screen_shot(round_robin, matrix, position, size, m_x, m_y):
        round_robin[position]= copy.deepcopy(matrix)
        if (position < size - 1):
                position += 1
        else:
                position = 0
        return position

def main() :
        CONST_e = math.e
        x0 = -163.53
        x1 = -64.9
        y0 = 18.25
        y1 = 61.02
        
        k = 20
        bandwidth_h = float(k/4.0)

        # form a smaller matrix N of size 2k+1 x 2k+1 and fill in the values
        N = np.zeros((2*k+1, 2*k+1))

        for row in range (0, 2*k+1):
                for col in range (0, 2*k+1):
                        N[row][col] = N_matrix_entries(k, bandwidth_h, row, col)

        # Normalize N
        N = N/np.sum(N)
        ######################################################################

        m_x = input("Please enter the number of pixels in x direction: ")
        m_y = input("Please enter the number of pixels in y direction: ")
        bandwidth_t = float(input("Please enter the bandwidth for time: "))

        # initialize the matrix of size m_x by m_y
        matrix = np.zeros((m_x, m_y))

        x = x0
        y = y0
        t_old = 0.0

        ################ round robin data structure ###############
        size = 9
        round_robin = np.zeros((size, m_x, m_y))
        #change the screen shot time 
        i = 0
        ############################################################
        t_original = 10
        t_count = 0

        map = Basemap(llcrnrlon=x0,llcrnrlat=y0,urcrnrlon=x1,urcrnrlat=y1,
              projection='cyl',lat_1=33,lat_2=45,lon_0=-95)

        # draw the map of the US
        map.drawcountries()
        map.drawstates()
        map.drawcoastlines()

        #draw parallels
        parallels = np.arange(0, 90,10.)
        map.drawparallels(parallels, labels=[1,0,0,0],fontsize=10)
        #draw meridians
        meridians = np.arange(10.,351.,20.)
        map.drawmeridians(meridians, labels=[0,0,0,1],fontsize=10)
        filename = open('tornado_complete.txt', 'r')
        error = 0
        points = []
        # read in points from the file
        for line in filename:
                data = line.split()
                try: # for invalid inputs
                        longitude = float(data[0])
                        latitude = float(data[1])
                        time_in = float(data[2])
                        fujita = float(data[3])
                except ValueError:
                        error += 1
                        continue
                point = Tornado(longitude, latitude, time_in, fujita)
                points.append(point)
                
        # sort based on the time
        points = sorted(points, key=getkey)
        for ele in points:
                ix = indice_calc(x0, x1, ele.longitude, m_x)
                iy = indice_calc(y0, y1, ele.latitude, m_y)


                ##################################################################################
                #for smaller matrix N, the range of the indices when it is overlapped completely #
                # N[0:2k+1, 0:2k+1]                                                              #
                # When it is not overlapped completely,                                          #
                # we choose the x-range to be the max(0, k-ix) to min(k+m_x-ix+1, 2k+1)          #
                # where as the y-range is the max(0, k-iy) to min(k+m_y-iy+1, 2k+1)              #
                ##################################################################################
                m_idx_left = max(0,(ix-k))
                m_idx_right = min((ix+k+1), m_x)
                m_idy_left = max(0,(iy-k))
                m_idy_right = min((iy+k+1),m_y)
                n_idx_left = max(0, k-ix)
                n_idx_right = min(k+m_x-ix,2*k+1)
                n_idy_left = max(0, k-iy)
                n_idy_right = min(k+m_y-iy,2*k+1)
                matrix = matrix*math.pow(CONST_e, -bandwidth_t*(ele.time_in - t_old));
                matrix[m_idx_left : m_idx_right, m_idy_left : m_idy_right] += N[n_idx_left : n_idx_right, n_idy_left : n_idy_right]*math.pow(CONST_e, ele.fujita);
                t_old = ele.time_in
                
                t_count += 1
                if (t_count % 100 == 0):
                       print t_count, ele.time_in
                if ((int(ele.time_in)) == t_original):
                        i = screen_shot(round_robin, bandwidth_t*matrix, i, size, m_x, m_y)
                        t_original += 10
        screen_shot(round_robin, bandwidth_t*matrix, i, size, m_x, m_y)
        total_intensity_decade = []
        for i in range(0,6):
        	total_intensity_decade.append(np.sum(round_robin[i]))
        	print np.sum(round_robin[i])
        diff_matrix = round_robin[0] #- round_robin[0]
        cax = plt.imshow(diff_matrix.T, cmap = 'seismic', vmin=-0.5, vmax=0.5, interpolation='nearest', aspect='auto',origin='lower', extent=(x0,x1,y0,y1))
        cbar = map.colorbar(cax,location ='right', label='Intensity')
        np.savetxt('matrix_out.txt', matrix)
        plt.show() 

if __name__ == "__main__":
        main()