#!/usr/bin/env python
from argparse import ArgumentParser
from sys import exit
import json
import matplotlib.pyplot as plt
import numpy as np
from math import sqrt

json_key_words = ["config", "global"]

def main():
    '''
    Make plots from SARS-CoV-2 simulation in Slovakia
    '''

    parser = ArgumentParser(description='Yay, make Slovak SARS-CoV-2 plots!')
    parser.add_argument('-c', '--config',
                        type=str, default='../../config/plot.json',
                        help='location of configuration file')
    parser.add_argument('-t', '--type',
                        type=str, default='slovakia',
                        help='plot type')
    required_arguments = parser.add_argument_group('required arguments')
    required_arguments.add_argument('-i', '--input-file',
                                    type=str, required=True,
                                    help='input file path')
    args = parser.parse_args()

    try:
        infile = open(args.config, "r")
        config = json.load(infile)
        infile.close()
    except IOError:
        print("ERROR: Configuration file not found!")
        exit(1)

    try:
        infile = open(args.input_file, "r")
        indata = json.load(infile)
        infile.close()
    except IOError:
        print("ERROR: Input file not found!")
        exit(1)

    print('Number of municipalities: ' + str(len(indata)))

    if args.type == 'slovakia':
        plot_slovakia(indata, config)
    elif args.type == 'slovakia_map':
        plot_slovakia_map(indata, config)
    elif args.type == 'by_municipality':
        plot_by_municipality(indata, config)


def plot_slovakia(indata, config):
    '''
        Create plots for the whole Slovakia.
    '''

    outfile_path = '../../plots/slovakia'

    lines = config['type']['slovakia']['lines']
    days = indata['Nitra']['simulation_days']

    height = len(lines)
    width = len(days)
    line_arr = [[0 for x in range(width)] for y in range(height)]

    for index, key in enumerate(indata):
        if (key in json_key_words):
            continue
        place = indata[key]

        percent = round(100.*index/len(indata))
        print(str(percent) + '%, municipality: ' + key)

        for line_index, line in enumerate(lines):
            line_arr[line_index] = np.add(line_arr[line_index], place[line])

    for index, line in enumerate(lines):
        plt.plot(days, line_arr[index], label=line)

    plt.title('Slovakia')
    plt.xlabel('Days')
    plt.ylabel('Number of people')
    plt.yscale(config['type']['by_municipality']['yscale'])
    plt.legend(lines)
    plt.grid()

    outfile_path += '.pdf'
    plt.savefig(outfile_path)

    plt.clf()


def plot_by_municipality(indata, config):
    '''
        Create plots for every municipality in Slovakia.
    '''

    for index, key in enumerate(indata):
        if (key in json_key_words):
            continue
        place = indata[key]
        outfile_path = '../../plots/by_municipality/' + key

        percent = round(100.*index/len(indata))
        print(str(percent) + '%, municipality: ' + key)

        days = place['simulation_days']
        lines = config['by_municipality']['lines']

        for line in lines:
            lines_to_sum = line.split('+')
            line_sum = np.zeros(len(days))
            for lin in lines_to_sum:
                line_sum = np.add(line_sum, place[lin])

            plt.plot(days, line_sum, label=line)

        plt.title(key)
        plt.xlabel('Days')
        plt.ylabel('Number of people')
        plt.yscale(config['type']['by_municipality']['yscale'])
        plt.legend(lines)
        plt.grid()

        outfile_path += '.pdf'
        plt.savefig(outfile_path)

        plt.clf()


def plot_slovakia_map(indata, config):
    '''
        Create dot on the map for every municipality in Slovakia.
        To create video run:
        ffmpeg -framerate 4 -i day_%03d.png -c:v libx264 -r 30 \
               -pix_fmt yuv420p slovakia_map.mp4
    '''

    outfile_path_base = '../../plots/slovakia_map/day_'

    # variable = config['slovakia_map']['variable']
    days = indata['Nitra']['simulation_days']

    lon = np.array([])
    lat = np.array([])
    pop = np.array([])
    number_of_municipalities = 0
    for index, key in enumerate(indata):
        if (key in json_key_words):
            continue
        place = indata[key]

        percent = round(100.*index/len(indata))
        print(str(percent) + '%, municipality: ' + key)

        lon = np.append(lon, place['longitude'])
        lat = np.append(lat, place['latitude'])
        pop = np.append(pop, sqrt(place['size']*0.01) + 1)
        number_of_municipalities += 1

    var = np.empty((0, number_of_municipalities), int)
    for index, day in enumerate(days):
        var_one_day = np.array([])

        percent = round(100.*index/len(days))
        print(str(percent) + '%, day: ' + str(day))

        for key in indata:
            place = indata[key]
            if (key in json_key_words):
                continue
            var_one_day = np.append(
                var_one_day,
                place['infected'][index])

        var = np.append(var, [var_one_day], axis=0)

    print(np.amax(var))


    for index, day in enumerate(days):
        outfile_path = outfile_path_base
        plt.scatter(lon, lat, s=pop, c=var[index],
                    cmap='jet', vmin=0, vmax=20000)

        plt.title('Infected, day: ' + str(day))
        plt.xlabel('Longitude')
        plt.ylabel('Latitude')

        outfile_path += str(day).zfill(3)
        outfile_path += '.png'
        plt.savefig(outfile_path, dpi=300)

        plt.clf()


if __name__ == '__main__':
    main()
