#!/usr/bin/env python3
''' quick plan_gen visualizer '''

import sys
import numpy as np
import lxml.etree as etree
import matplotlib

def get_nodes(input):
    tree = etree.parse(input)
    return [node for node in tree.xpath("/network/nodes/node")]

def get_persons(input):
    tree = etree.parse(input)
    return [p for p in tree.xpath("/plans/person")]

if __name__ == '__main__':
    matplotlib.use('TkAgg')
    import matplotlib.pyplot as plt

    # get data
    NB_CLUSTERS = int(sys.argv[1])
    NODES = get_nodes(sys.argv[2])
    PERSONS = get_persons(sys.argv[3])
    PERSONS_XY = ['{}|{}'.format(
                  p.find('plan/act').get('x'),
                  p.find('plan/act').get('y')) for p in PERSONS]
    P_XY_UNIQUE, P_XY_COUNTS = np.unique(PERSONS_XY, return_counts=True)
    
    # plot init
    FIG = plt.figure()
    AX = FIG.add_subplot(111)

    # plot nodes
    NODES_X = [float(n.get('x')) for n in NODES]
    NODES_Y = [float(n.get('y')) for n in NODES]
    MIN_X, MAX_X = min(NODES_X), max(NODES_X)
    MIN_Y, MAX_Y = min(NODES_Y), max(NODES_Y)
    AX.scatter(NODES_X, NODES_Y,
                marker='.', c='grey', linewidth=0.5, s=10)

    # plot persons
    PERSONS_X = [float(coord.split('|')[0]) for coord in P_XY_UNIQUE]
    PERSONS_Y = [float(coord.split('|')[1]) for coord in P_XY_UNIQUE]
    SC = AX.scatter(PERSONS_X, PERSONS_Y,
                     alpha=0.75, s=P_XY_COUNTS*10,
                     c=P_XY_COUNTS, cmap='rainbow')

    # plot map
    img = plt.imread('input/map.png')
    implot = plt.imshow(img, alpha=1, extent=[MIN_X, MAX_X, MIN_Y, MAX_Y])

    # final plot
    X_TICKS = np.arange(MIN_X, MAX_X, (MAX_X - MIN_X) / NB_CLUSTERS)
    Y_TICKS = np.arange(MIN_Y, MAX_Y, (MAX_Y - MIN_Y) / NB_CLUSTERS)
    TICK_LABELS = list(range(NB_CLUSTERS))
    AX.set_title('Number of agents by nodes (total: {} agents)'.format(len(PERSONS)))
    AX.set_xlim(MIN_X, MAX_X)
    AX.set_ylim(MIN_Y, MAX_Y)
    AX.set_xticks(X_TICKS)
    AX.set_yticks(Y_TICKS)
    AX.set_xticklabels(TICK_LABELS)
    AX.set_yticklabels(TICK_LABELS)
    AX.tick_params(labelsize=5)
    AX.grid(True)
    plt.colorbar(SC)
    plt.show()