#!/usr/bin/env python

"""
Generic LPU demo

Notes
-----
Generate input file and LPU configuration by running

cd data
python gen_generic_lpu.py
"""

import argparse
import itertools

import networkx as nx

import neurokernel.core as core
import neurokernel.base as base
from neurokernel.tools.zmq import get_random_port

from neurokernel.LPU.LPU import LPU

dt = 1e-4
dur = 1.0
steps = int(dur/dt)

parser = argparse.ArgumentParser()
parser.add_argument('--debug', default=False,
                    dest='debug', action='store_true',
                    help='Write connectivity structures and inter-LPU routed data in debug folder')
parser.add_argument('-l', '--log', default='none', type=str,
                    help='Log output to screen [file, screen, both, or none; default:none]')
parser.add_argument('-s', '--steps', default=steps, type=int,
                    help='Number of steps [default: %s]' % steps)
parser.add_argument('-d', '--port_data', default=None, type=int,
                    help='Data port [default: randomly selected]')
parser.add_argument('-c', '--port_ctrl', default=None, type=int,
                    help='Control port [default: randomly selected]')
parser.add_argument('-g', '--gpu_dev', default=0, type=int,
                    help='GPU device number [default: 0]')
args = parser.parse_args()

file_name = None
screen = False
if args.log.lower() in ['file', 'both']:
    file_name = 'neurokernel.log'
if args.log.lower() in ['screen', 'both']:
    screen = True
logger = base.setup_logger(file_name=file_name, screen=screen)

if args.port_data is None and args.port_ctrl is None:
    port_data = get_random_port()
    port_ctrl = get_random_port()
else:
    port_data = args.port_data
    port_ctrl = args.port_ctrl

port_time = get_random_port()
man = core.Manager(port_data, port_ctrl, port_time)
man.add_brok()

(n_dict, s_dict) = LPU.lpu_parser('./data/generic_lpu.gexf.gz')

ge = LPU(dt, n_dict, s_dict,
         input_file='./data/generic_input.h5',
         output_file='generic_output.h5', port_ctrl=port_ctrl,
         port_data=port_data,
         device=args.gpu_dev, id='ge',
         debug=args.debug)
man.add_mod(ge)

man.start(steps=args.steps)
man.stop()

@@ -200,22 +200,20 @@
  as nx%0A%0A
-import
+from
  neuroke
@@ -221,20 +221,41 @@
 nel.
-core as core
+tools.logging import setup_logger
 %0Aimp
@@ -274,20 +274,25 @@
 nel.
-base as base
+core_gpu as core%0A
 %0Afro
@@ -309,77 +309,59 @@
 nel.
-tools.zmq import get_random_port%0A%0Afrom neurokernel.LPU.LPU import LPU
+LPU.LPU import LPU%0A%0Aimport neurokernel.mpi_relaunch
 %0A%0Adt
@@ -1512,13 +1512,8 @@
 r = 
-base.
 setu
@@ -1562,293 +1562,26 @@
 n)%0A%0A
-if args.port_data is None and args.port_ctrl is None:%0A    port_data = get_random_port()%0A    port_ctrl = get_random_port()%0Aelse:%0A    port_data = args.port_data%0A    port_ctrl = args.port_ctrl%0A%0Aport_time = get_random_port()%0Aman = core.Manager(port_data, port_ctrl, port_time)%0Aman.add_brok
+man = core.Manager
 ()%0A%0A
@@ -1649,17 +1649,27 @@
 ')%0A%0A
-ge = LPU(
+man.add(LPU, 'ge', 
 dt, 
@@ -1692,17 +1692,16 @@
         
- 
 input_fi
@@ -1738,17 +1738,16 @@
         
- 
 output_f
@@ -1775,60 +1775,9 @@
 5', 
-port_ctrl=port_ctrl,%0A         port_data=port_data,%0A 
+%0A
     
@@ -1804,19 +1804,9 @@
 dev,
- id='ge',%0A 
+%0A
     
@@ -1831,24 +1831,20 @@
 ug)%0A
+%0A
 man.
-add_mod(ge)%0A
+spawn()
 %0Aman
@@ -1876,11 +1876,11 @@
 man.
-stop
+wait
 ()%0A

Update generic demo to work with MPI-based code.