# Library
import time
import pickle
import imageio
import geocoder
import argparse
import datetime
import webbrowser
from dateutil import parser

# Custom
import GeoPlot
from GeoPlot import geo
from TimeSplit import split


print('\n\nRunning mapping utility')
# Allow for input file to break up tasks
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('-i', '--input', type=str, help='The input binary file of tweets', nargs='+', required=True)
arg_parser.add_argument('-o', '--output', type=str, help='The output file directory for files', nargs='+', default=None)
arg_parser.add_argument('-t', '--time', type=bool, help='If times need to be formatted to datetime', default=True)
arg_parser.add_argument('-l', '--location', type=bool, help='If the locations need to convert to lat lon', default=False)
arg_parser.add_argument('-n', '--divisions', type=int, help='The time range of one map in days', default=10)
arg_parser.add_argument('-m', '--scale', type=int, help='The size of the map to use', default=0)
arg_parser.add_argument('-z', '--size', type=int, help='The size of the map points', default=3)
arg_parser.add_argument('-c', '--color', type=str, help='The color of the map points', default='red')
arg_parser.add_argument('-f', '--fps', type=int, help='The frame rate of the gif', default=1)
arg_parser.add_argument('-s', '--save', type=bool, help='Save the file', default=False)
arg_parser.add_argument('-d', '--display', type=bool, help='Display the images', default=False)
args = arg_parser.parse_args()

if args.save and args.output is None:
    print('In order to save you must specify an output directory')
    exit()

# Load input a list of sets of text location and time
with open(' '.join(args.input), 'rb') as file:
    tweets = pickle.load(file)

if args.time:
    # Change time scale and order
    for tweet in tweets:
        tweet['time'] = parser.parse(tweet['time'])

if args.location:
    # Find locations for non lats and lons
    print('Finding locations')
    for tweet in tweets:
        if tweet['latitude'] is None or tweet['longitude'] is None:
            if tweet['location'] is not None:
                loc = geocoder.google(tweet['location'])
                print(loc)
                tweet['latitude'] = loc.latlng[0]
                tweet['longitude'] = loc.latlng[1]
                time.sleep(0.1)
            else:
                tweets.remove(tweet)

# Split the data into even times
start = time.time()
splitter = split(unit=datetime.timedelta(days=args.divisions))
chunks = splitter.split(tweets)
ranges = splitter.ranges()

# Select the map size
if args.scale == 1:
    scale = GeoPlot.RocMap
elif args.scale == 2:
    scale = GeoPlot.NYMap
else:
    scale = GeoPlot.FullMap

# Print and Save
folder = ' '.join(args.output)
save = args.save
display = args.display
pics = []
pos = 0
neg = 0
# Create each map
for i in xrange(len(chunks)):
    temp_map = geo(local=scale, res=GeoPlot.Crude)
    # Only add point of the flu
    for item in chunks[i]:
        if item['tag']:
            pos = pos + 1
            temp_map.add_points(item['latitude'], item['longitude'])
        else:
            neg = neg + 1
    # Names for the files and plots
    step = str(i+1) + ' of ' + str(len(chunks))
    title = 'From (' + str(ranges[i][0].strftime("%y-%m-%d %H:%M")) + ') to (' + str(ranges[i][1].strftime("%y-%m-%d %H:%M")) + '): ' + step
    name = folder + step + '.png'
    print title
    # Prepare to make gif
    if save:
        pics.append(name)
    # Output as requested
    temp_map.output(display=display, save=save, path=name, title=title, size=args.size, color=args.color)

print('\n\n')
print(str(pos) + ' positive cases plotted')
print(str(neg) + ' negative cases ignored')

# Save as a gif
if save:
    imgs = []
    for pic in pics:
        imgs.append(imageio.imread(pic))
    imageio.mimsave(folder + 'animation.gif', imgs, fps=args.fps)
    if display:
        webbrowser.open(folder+'animation.gif')

print('\n\nTime elapsed: ' + str(time.time() - start) + ' seconds')

