#!/usr/bin/env python

import numpy as np
import sys

from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

tmax = 1000
nxs = range(96, 480+1, 32)	# nx**2, 72.0 KiB ~ 1.76 MiB
#nxs = [96, 128]
if rank == 0:
	dts = np.zeros(len(nxs))
elif rank == 1:
	from datetime import datetime

for i, nx in enumerate(nxs):
	#print('nx = %d, rank = %d' % (nx, rank))
	if rank == 0:
		dts[i] = comm.recv(source=1, tag=10)	# source, tag
		print('[%d] nx = %d, dt = %g' % (rank, nx, dts[i]))

	elif rank == 1:
		arr1 = np.zeros((nx, nx), dtype=np.float32)
		arr2 = np.zeros((nx, nx), dtype=np.float32)
		t0 = datetime.now()
		for tstep in xrange(1, tmax+1):
			comm.Recv(arr1, source=2, tag=21)
			comm.Recv(arr2, source=2, tag=22)
		dt0 = datetime.now() - t0
		dt = dt0.seconds + dt0.microseconds * 1e-6
		#print('[%d] nx = %d, dt = %g' % (rank, nx, dt))
		comm.send(dt, dest=0, tag=10)	# data, dest, tag

	elif rank == 2:
		arr1 = np.random.rand(nx,nx).astype(np.float32)
		arr2 = arr1.copy()
		for tstep in xrange(1, tmax+1):
			comm.Send(arr1, dest=1, tag=21)
			comm.Send(arr2, dest=1, tag=22)
		

if rank == 0:
	np.save('330.npy', dts)
	'''
	import matplotlib.pyplot as plt
	plt.ion()
	plt.plot(dts, linestyle='None', marker='p', markersize=4)
	plt.show()
	'''

"""
if rank == 0:
	# Save as h5
	import pyopencl as cl
	platforms = cl.get_platforms()
	gpu_devices = platforms[0].get_devices(device_type=cl.device_type.GPU)
	gpu_name = gpu_devices[0].get_info(cl.device_info.NAME)
	print gpu_name

	import h5py as h5
	h5_path = './capability_fdtd3d.h5'
	f = h5.File(h5_path, 'a')
	if 'tmax' not in f.attrs.keys():
		f.attrs['tmax'] = tmax
	elif f.attrs['tmax'] != tmax:
		print('The \'tmax\' value is not matched in h5 file\'s attrs')
		sys.exit()

	if 'gpu' not in f.keys():
		f.create_group('gpu')
	if gpu_name not in f['gpu'].keys():
		f['gpu'].create_group(gpu_name)
	f['gpu'][gpu_name].create_dataset('nx', data=np.array(nxs))
	f['gpu'][gpu_name].create_dataset('dt', data=dts)
	f.close()
"""
