# Eclipse SUMO, Simulation of Urban MObility; see https://eclipse.org/sumo
# Copyright (C) 2016-2017 German Aerospace Center (DLR) and others.
# SUMOPy module
# Copyright (C) 2012-2017 University of Bologna - DICAM
# This program and the accompanying materials
# are made available under the terms of the Eclipse Public License v2.0
# which accompanies this distribution, and is available at
# http://www.eclipse.org/legal/epl-v20.html

# @file    results_mpl.py
# @author  Joerg Schweizer
# @date    
# @version $Id$


import numpy as np
from collections import OrderedDict
import matplotlib as mpl
from matplotlib.patches import Arrow, Circle, Wedge, Polygon, FancyArrow
from matplotlib.collections import PatchCollection
import matplotlib.pyplot as plt

import agilepy.lib_base.classman as cm
import agilepy.lib_base.arrayman as am
from agilepy.lib_base.geometry import *
from agilepy.lib_base.processes import Process


def init_plot():
    plt.close("all")
    fig = plt.figure()
    ax = fig.add_subplot(111)
    return ax


def plot_net(ax, net, edgecolor="gray", edgewidth=2, color_node=None,
             alpha=0.5):
    for shape in net.edges.shapes.get_value():
        x_vec = np.array(shape)[:, 0]
        y_vec = np.array(shape)[:, 1]
        ax.plot(x_vec, y_vec, color=edgecolor, lw=edgewidth, alpha=alpha)

    # do nodes
    if color_node is None:
        is_nodefill = False
        color_node = 'none'
    else:
        is_nodefill = True

    #ax.scatter(x_vec[0], y_vec[0], s=np.pi * (10.0)**2, c=colors, alpha=0.5)
    coords = net.nodes.coords.get_value()
    radii = net.nodes.radii.get_value()
    #ax.scatter(coords[:,0], coords[:,1], s=np.pi * (radii)**2, alpha=0.5)
    #patches = []
    for coord, radius in zip(coords, radii):
        ax.add_patch(Circle(coord, radius=radius,
                            linewidth=edgewidth,
                            edgecolor=edgecolor,
                            facecolor=color_node,
                            fill=is_nodefill,
                            alpha=alpha))


def plot_edgevalues(ax,  ids_result, config_ids_edge, config_values, width_max=10.0):
    edges = config_ids_edge.get_linktab()
    ids_edges = config_ids_edge[ids_result]
    values = config_values[ids_result]
    values_norm = np.array(values, dtype=np.float32) / np.max(values)
    colors = mpl.cm.cool_r(values_norm, 1)
    for shape, val_norm, val, color, _id in zip(edges.shapes[ids_edges], values_norm, values, colors, ids_result):
        x_vec = np.array(shape)[:, 0]
        y_vec = np.array(shape)[:, 1]
        #ax.text(x_vec[-1]+0.5*width_max,y_vec[-1]+0.5*width_max, config_values.format_value(_id) )
        ax.text(x_vec[-1] + 0.5 * width_max, y_vec[-1] +
                0.5 * width_max, "%.1f" % config_values[_id])
        ax.plot(x_vec, y_vec, color=color, linewidth=width_max * val_norm)
    # plt.colorbar() # no


def plot_edgevalues_patches(ax, ids_result, config_ids_edge, config_values,
                            width_max=10.0, alpha=0.8, printformat='%.2f',
                            color_outline=None, color_fill=None, is_antialiased=True,
                            is_fill=True, is_widthvalue=True,
                            arrowshape='left', length_arrowhead=10.0,
                            headwidthstretch=1.3):

    head_width = headwidthstretch * width_max

    edges = config_ids_edge.get_linktab()
    ids_edges = config_ids_edge[ids_result]
    values = config_values[ids_result]
    values_norm = np.array(values, dtype=np.float32) / np.max(values)

    patches = []
    displacement = float(width_max) / 4.0
    if is_widthvalue:
        linewidths = width_max * values_norm
    else:
        linewidths = width_max * np.ones(len(values_norm), np.float32)

    deltaangle_text = -np.pi / 2.0
    displacement_text = displacement + width_max
    for id_edge, value, value_norm, linewidth in zip(ids_edges, values, values_norm, linewidths):
        shape, angles_perb = get_resultshape(edges, id_edge, displacement)
        x_vec = np.array(shape)[:, 0]
        y_vec = np.array(shape)[:, 1]
        deltax = x_vec[-1] - x_vec[0]
        deltay = y_vec[-1] - y_vec[0]

        if printformat is not '':
            angles_text = np.arctan2(deltay, deltax)

            x_label = x_vec[0] + 0.5 * deltax + displacement_text * \
                np.cos(angles_text + deltaangle_text)
            y_label = y_vec[0] + 0.5 * deltay + displacement_text * \
                np.sin(angles_text + deltaangle_text)

            ax.text(x_label, y_label, printformat %
                    value, rotation=angles_text / (np.pi) * 180)

        if is_widthvalue:
            head_width = headwidthstretch * linewidth
        arrow = FancyArrow(x_vec[0], y_vec[0], deltax, deltay, width=linewidth,
                           antialiased=is_antialiased,
                           edgecolor=color_outline, facecolor=color_fill,
                           head_width=head_width, head_length=length_arrowhead,
                           length_includes_head=True,
                           fill=True, shape=arrowshape)
        patches.append(arrow)

    if is_fill:
        alpha_patch = alpha
        patchcollection = PatchCollection(
            patches, cmap=mpl.cm.jet, alpha=alpha_patch)
        patchcollection.set_array(values)
        ax.add_collection(patchcollection)
        plt.colorbar(patchcollection)
    else:
        for patch in patches:
            ax.add_patch(patch)


def get_resultshape(edges, id_edge, dispacement):
    """
    Return resultshape coords for this edge.
    """
    shape = np.array(edges.shapes[id_edge], np.float32)
    n_vert = len(shape)
    resultshape = np.zeros(shape.shape, np.float32)
    #laneshapes = np.zeros((n_lanes,n_vert,3), np.float32)
    angles_perb = get_angles_perpendicular(shape)
    if edges.types_spread[id_edge] == 0:  # right spread

        resultshape[:, 0] = shape[:, 0] + np.cos(angles_perb) * dispacement
        resultshape[:, 1] = shape[:, 1] + np.sin(angles_perb) * dispacement
        return resultshape, angles_perb
    else:
        return shape.copy(), angles_perb


def show_plot():
    plt.show()


class Resultplotter(Process):

    def __init__(self, results, name='Plot results with Matplotlib',
                 info="Creates plots of different results using matplotlib",
                 logger=None):

        self._init_common('resultplotter', parent=results, name=name,
                          info=info, logger=logger)

        print 'Resultplotter.__init__', results, self.parent
        attrsman = self.get_attrsman()

        # edgeresultes....
        attrnames_edgeresults = {}
        edgeresultattrconfigs = self.parent.edgeresults.get_group_attrs(
            'results')
        edgeresultattrnames = edgeresultattrconfigs.keys()
        edgeresultattrnames.sort()
        for attrname in edgeresultattrnames:
            attrconfig = edgeresultattrconfigs[attrname]

            attrnames_edgeresults[
                attrconfig.format_symbol()] = attrconfig.attrname

        #attrnames_edgeresults = {'Entered':'entered'}
        self.edgeattrname = attrsman.add(cm.AttrConf('edgeattrname', 'entered',
                                                     choices=attrnames_edgeresults,
                                                     groupnames=['options'],
                                                     name='Edge Quantity',
                                                     info='The edge related quantity to be plotted.',
                                                     ))

        self.is_show_network = attrsman.add(cm.AttrConf('is_show_network', False,
                                                        groupnames=['options'],
                                                        name='Show network',
                                                        info='Shows a schematic network in the background.',
                                                        ))

        self.is_show_title = attrsman.add(cm.AttrConf('is_show_title', True,
                                                      groupnames=['options'],
                                                      name='Show tile',
                                                      info='Shows title and unit.',
                                                      ))

        self.resultwidth = attrsman.add(cm.AttrConf('resultwidth', 10.0,
                                                    groupnames=['options'],
                                                    name='Result width',
                                                    unit='m',
                                                    info='Maximum width of graphical resuls on map.',
                                                    ))

        self.length_arrowhead = attrsman.add(cm.AttrConf('length_arrowhead', 10.0,
                                                         groupnames=[
                                                             'options'],
                                                         name='Arrow length',
                                                         unit='m',
                                                         info='Length of arrowhead on result map.',
                                                         ))

        self.is_widthvalue = attrsman.add(cm.AttrConf('is_widthvalue', False,
                                                      groupnames=['options'],
                                                      name='Value width?',
                                                      info='If True, the arrow width of the graphical representation is proportional to the result value.',
                                                      ))

        self.is_colorvalue = attrsman.add(cm.AttrConf('is_colorvalue', True,
                                                      groupnames=['options'],
                                                      name='Value color?',
                                                      info='If True, the arrows of the graphical representation are filled with a colour representing the result value.',
                                                      ))

        self.color_outline = attrsman.add(cm.AttrConf('color_outline', np.array([0.0, 0.0, 0.0, 0.95], dtype=np.float32),
                                                      groupnames=['options'],
                                                      perm='wr',
                                                      metatype='color',
                                                      name='Outline color',
                                                      info='Outline color of result arrows in graphical representation. Only valid if no color-fill is chosen.',
                                                      ))

        self.color_fill = attrsman.add(cm.AttrConf('color_fill', np.array([0.3, 0.3, 1.0, 0.95], dtype=np.float32),
                                                   groupnames=['options'],
                                                   perm='wr',
                                                   metatype='color',
                                                   name='Fill color',
                                                   info='Fill color of result arrows in graphical representation. Only valid if no color-fill is chosen.',
                                                   ))

        self.alpha_results = attrsman.add(cm.AttrConf('alpha_results', 0.8,
                                                      groupnames=['options'],
                                                      name='Result transparency',
                                                      info='Transparency of result arrows in graphical representation.',
                                                      ))

        self.printformat = attrsman.add(cm.AttrConf('printformat', '%.1f',
                                                    choices=OrderedDict([
                                                        ('Show no values', ''),
                                                        ('x', '%.d'),
                                                        ('x.x', '%.1f'),
                                                        ('x.xx', '%.2f'),
                                                        ('x.xxx', '%.3f'),
                                                        ('x.xxxx', '%.4f'),
                                                    ]),
                                                    groupnames=['options'],
                                                    name='Value formatting',
                                                    info='Value formatting of displayed text in graphical representation.',
                                                    ))

        self.is_grid = attrsman.add(cm.AttrConf('is_grid', False,
                                                groupnames=['options'],
                                                name='Show grid?',
                                                info='If True, shows a grid on the graphical representation.',
                                                ))

        self.axis = None

    def show(self):
        # print 'show',self.edgeattrname
        # if self.axis == None:
        axis = init_plot()
        # else:
        net = self.parent.get_scenario().net
        if self.is_show_network:
            plot_net(axis, net)

        if (self.edgeattrname is not ""):
            resultattrconf = getattr(
                self.parent.edgeresults, self.edgeattrname)
            ids = self.parent.edgeresults.get_ids()
            if len(ids > 0):
                title = resultattrconf.get_info(
                ) + resultattrconf.format_unit(show_parentesis=True)  # format_symbol()
                #plot_edgevalues(axis, ids, self.parent.edgeresults.ids_edge, resultattrconf)
                plot_edgevalues_patches(axis, ids,
                                        self.parent.edgeresults.ids_edge,
                                        resultattrconf,
                                        width_max=self.resultwidth,
                                        alpha=self.alpha_results,
                                        printformat=self.printformat,
                                        color_outline=self.color_outline,
                                        color_fill=self.color_fill,
                                        is_antialiased=True,
                                        is_fill=self.is_colorvalue,
                                        is_widthvalue=self.is_widthvalue,
                                        length_arrowhead=self.length_arrowhead,
                                        )

                axis.set_title(title)

        axis.axis('equal')
        # ax.legend(loc='best',shadow=True)

        axis.grid(self.is_grid)
        axis.set_xlabel('West-East [m]')
        axis.set_ylabel('South-North [m]')
        show_plot()

    def do(self):
        # print 'do',self.edgeattrname
        self.show()

    def get_scenario(self):
        return self._scenario
