from __future__ import division, print_function
from six import StringIO
from svgpath2mpl import parse_path
from collections import defaultdict
from src.pred_plot import hour_rounder
import xml.etree.ElementTree as etree
import re
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import requests
import pandas as pd
import datetime

import warnings
warnings.filterwarnings("ignore")

def calling_pred_map_viz(counts_df1, input_day = None, input_hour = None):
    r = "svg/snazzy-image-01.svg"
    tree = etree.parse(r)
    root = tree.getroot()
    path_elems = root.findall('.//{http://www.w3.org/2000/svg}path')

    paths = [parse_path(elem.attrib['d']) for elem in path_elems]
    facecolors = []
    edgecolors = []
    linewidths = []
    for elem in path_elems:
        facecolors.append(dict(item.split(":") for item in elem.attrib.get('style', 'none').split(";")).get("fill", "none"))
        edgecolors.append(dict(item.split(":") for item in elem.attrib.get('style', 'none').split(";")).get("stroke", "none"))
        linewidths.append(dict(item.split(":") for item in elem.attrib.get('style', 'none').split(";")).get("stroke-width", "none").replace("px", ""))

    path_id = defaultdict(int)
    for i, elem in enumerate(path_elems):
        try:
            #print(i, elem.attrib['id'])
            path_id[elem.attrib['id']] = i
        except:
            continue

    counts_df1['total'] = counts_df1['car'] + counts_df1['motorcycle'] + counts_df1['large_vehicle']
    counts_df1.loc[:,'date_time'] = pd.to_datetime(counts_df1.loc[:,'date'] + " "+ counts_df1.loc[:,'time'], format='%Y-%m-%d %H:%M:%S')
    counts_df1.loc[:,'hour'] = counts_df1.loc[:,'date_time'].apply(hour_rounder)
    counts_df1.loc[:,'day_name'] = counts_df1.loc[:,'date_time'].dt.day_name()

    if (input_day != None) & (input_hour != None):
        filtered_day = input_day
        filtered_hour = input_hour
    else:
        filtered_date = counts_df1.iloc[-1]['date']
        filtered_time = counts_df1.iloc[-1]['time']
        filtered_day = counts_df1.iloc[-1]['day_name']
        filtered_hour = counts_df1.iloc[-1]['hour']

    
    # filtered_date = counts_df1.iloc[-1]['date']
    # filtered_time = counts_df1.iloc[-1]['time']
    # filtered_day = counts_df1.iloc[-1]['day_name']
    # filtered_hour = counts_df1.iloc[-1]['hour']


    day_hour_view_group = counts_df1.groupby(by=['view', 'day_name', 'hour'])['total'].mean().reset_index()
    count_max = day_hour_view_group['total'].max()
    count_min = day_hour_view_group['total'].min()


    count_dict = {"woodlands_to_sg" :day_hour_view_group.loc[(day_hour_view_group['view'] == 'Woodlands - to SG') & (day_hour_view_group['day_name'] == filtered_day) & (day_hour_view_group['hour'] == filtered_hour), "total" ].iloc[0],
                  "woodlands_to_jh" :day_hour_view_group.loc[(day_hour_view_group['view'] == 'Woodlands - to Johor') & (day_hour_view_group['day_name'] == filtered_day) & (day_hour_view_group['hour'] == filtered_hour), "total" ].iloc[0],
                  "tuas_to_sg" :day_hour_view_group.loc[(day_hour_view_group['view'] == 'Tuas - to SG') & (day_hour_view_group['day_name'] == filtered_day) & (day_hour_view_group['hour'] == filtered_hour), "total" ].iloc[0],
                  "tuas_to_jh" :day_hour_view_group.loc[(day_hour_view_group['view'] == 'Tuas - to Johor') & (day_hour_view_group['day_name'] == filtered_day) & (day_hour_view_group['hour'] == filtered_hour), "total" ].iloc[0]
                  }

    values = np.array([0., 0.25, 1.])
    values = np.sort(np.array(values))
    values = np.interp(values, (values.min(), values.max()), (0., 1.))
    colors = ["#539f6b", "#ffc835", "#bf0000"]
    cmap = mpl.colors.LinearSegmentedColormap.from_list("custom", list(zip(values, colors)))

    norm = mpl.colors.Normalize(vmin=count_min, vmax=count_max)



    hex_dict = {k: mpl.colors.to_hex(cmap(norm(v))) for k, v in count_dict.items()}
    color_dict = defaultdict(str)

    for k, i in path_id.items():
        #print(k, i)
        color_dict[i] = hex_dict[k]

    for k, i in color_dict.items():
        #print(k,i)
        facecolors[k] = i

    collection = mpl.collections.PathCollection(paths, 
                                          edgecolors=edgecolors, 
                                          linewidths=[int(i)/100 for i in linewidths if i != 'none'],
                                          facecolors=[i.strip() for i in facecolors])



    fig = plt.figure(figsize=(10,10))
    ax = fig.add_subplot(111)
    collection.set_transform(ax.transData)
    ax.add_artist(collection)
    ax.set_xlim([100, 1900])
    ax.set_ylim([1800, 0])
    ax.set_title(filtered_day+ " | " + filtered_hour + " SGT", fontname = 'Georgia')
    return fig