#!/bin/env python

from mpi4py import MPI 
from numpy import * 
import sys

NRA = 2880 
NCA = 2000 
NCB = 45 
MASTER = 0
FROM_MASTER = 1
FROM_WORKER = 2
filename = "output_py.dat"

comm = MPI.COMM_WORLD
numtasks = comm.Get_size()
taskid = comm.Get_rank()

if taskid == MASTER:
  rows = NRA
else:
  averow = NRA/(numtasks-1)
  extra = NRA%(numtasks-1)
  if taskid <= extra:
    rows = averow + 1
  else:
    rows = averow

a = zeros((rows,NCA))
b = zeros((NCA,NCB))
c = zeros((rows,NCB))

offset = array([0], dtype=int)
rows_buf = array([0], dtype=int)
my_compute_time = array([0], dtype=float64)
compute_time = array([0], dtype=float64)

b_buf = [b, NCA*NCB, MPI.DOUBLE]

if taskid == MASTER:
  print "Number of worker tasks = %d" %(numtasks-1)
  sys.stdout.flush()

  for i in range(0, NRA):
    for j in range(0, NCA):
      a[i][j] = i+j

  for i in range(0, NCA):
    for j in range(0, NCB):
      b[i][j] = i*j

  averow = NRA/(numtasks-1)
  extra = NRA%(numtasks-1)
  val = 0
  offset[0] = val
  mtype = FROM_MASTER

  for dest in range(1, numtasks):
    if dest <= extra:
      rows = averow + 1
    else:
      rows = averow
    print "   sending %d rows to task %d" %(rows, dest)
    comm.Send([offset, MPI.INT], dest, mtype)
    a_buf = [a[offset[0]], rows*NCA, MPI.DOUBLE]
    comm.Send(a_buf, dest, mtype)
    comm.Send(b_buf, dest, mtype)
    val = val + rows
    offset[0] = val
    del a_buf
else:
  mtype = FROM_MASTER
  a_buf = [a, rows*NCA, MPI.DOUBLE]
  comm.Recv([offset, MPI.INT], MASTER, mtype)
  comm.Recv(a_buf, MASTER, mtype)
  comm.Recv(b_buf, MASTER, mtype)
  del a_buf

del b_buf

comm.Barrier()

if taskid == MASTER:
  mtype = FROM_WORKER
  for source in range(1, numtasks):
    comm.Recv([offset, MPI.INT], source, mtype)
    comm.Recv([rows_buf, MPI.INT], source, mtype)
    c_buf = [c[offset[0]], rows*NCB, MPI.DOUBLE]
    comm.Recv(c_buf, source, mtype)
    del c_buf
else:
  stime = MPI.Wtime()
  c = dot(a, b)
  etime = MPI.Wtime()
  my_compute_time[0] = etime - stime
  rows_buf[0] = rows
  mtype = FROM_WORKER
  comm.Send([offset, MPI.INT], MASTER, mtype)
  comm.Send([rows_buf, MPI.INT], MASTER, mtype)
  c_buf = [c, rows*NCB, MPI.DOUBLE]
  comm.Send(c_buf, MASTER, mtype)
  del c_buf

comm.Reduce([my_compute_time, MPI.DOUBLE], [compute_time, MPI.DOUBLE], 
            MPI.MAX, MASTER)

if taskid == MASTER:
  for i in range(0, NRA):
    for j in range(0, NCB):
      print "%6.2f   " %(c[i][j]),
    print ""

  FILE = open(filename, "w")
  FILE.write("Computation time : %f\n" %(compute_time[0]))
  for i in range(0, NRA):
    for j in range(0, NCB):
      FILE.write("%.2f " %(c[i][j]))
    FILE.write("\n")
  FILE.close()

MPI.Finalize()