#!/usr/bin/env python3
''' Python generator for MATSim plans. '''

import sys
import lxml.etree as etree
import numpy as np
import joblib as jl
import plan_gen.plan_gen as pg

NB_CORES = 4

def run_rand_person(i):
    ''' parallelisation '''
    np.random.seed(SEEDS[i])
    return [pg.rand_person(NODES, CLUSTERS, H_DENSITIES, W_DENSITIES)
                for _ in range(int(NB_PERSONS/NB_CORES))]

if __name__ == '__main__':

    # command line arguments
    if len(sys.argv) != 2:
        print('usage:', sys.argv[0], '<params>')
        sys.exit(-1)
    DICT_PARAMS = pg.parse_params(sys.argv[1])

    NB_CLUSTERS = DICT_PARAMS['nc']
    NB_PERSONS = DICT_PARAMS['np']
    INPUT_NETWORK = DICT_PARAMS['nw']
    H_CENTERS = DICT_PARAMS['hc'] if 'hc' in DICT_PARAMS else None
    W_CENTERS = DICT_PARAMS['wc'] if 'wc' in DICT_PARAMS else None
    H_RADIUS = DICT_PARAMS['hr'] if 'hr' in DICT_PARAMS else None
    W_RADIUS = DICT_PARAMS['wr'] if 'wr' in DICT_PARAMS else None

    # prepare data
    NODES = pg.get_nodes(INPUT_NETWORK)
    CLUSTERS = pg.make_clusters(NB_CLUSTERS, NODES)
    H_DENSITIES = pg.make_densities(NB_CLUSTERS, H_CENTERS, H_RADIUS)
    H_DENSITIES = pg.clean_densities(H_DENSITIES, CLUSTERS)
    W_DENSITIES = pg.make_densities(NB_CLUSTERS, W_CENTERS, W_RADIUS)
    W_DENSITIES = pg.clean_densities(W_DENSITIES, CLUSTERS)

    # make xml
    SEEDS = np.random.random_integers(0, 1e9, NB_CORES)
    PERSONS = jl.Parallel(n_jobs=NB_CORES)(jl.delayed(run_rand_person)(i) for i in range(NB_CORES))
    PERSONS = sum(PERSONS,[])
    PLANS = pg.make_plans(PERSONS)

    # print XML
    print('<?xml version="1.0" ?>')
    print('<!DOCTYPE plans SYSTEM "http://www.matsim.org/files/dtd/plans_v4.dtd">')
    print(etree.tostring(PLANS, pretty_print=True).decode('utf-8'))