#!/usr/bin/env python

from __future__ import print_function
from mpi4py import MPI
import time
import numpy as np

rand = np.random.rand

IMG_H, IMG_W = 400, 400
WIN_H, WIN_W = 5, 5
LAY_H, LAY_W = IMG_H - WIN_H + 1, IMG_W - WIN_W + 1
FILT_NUM = 6

image = rand(IMG_H, IMG_W)
all_weights = rand(FILT_NUM, WIN_H, WIN_W)
biases = rand(FILT_NUM)
layer = rand(FILT_NUM, LAY_H, LAY_W)

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
if rank == 0:
    # print('layer', layer.shape)
    start = time.time()

left = int((FILT_NUM / comm.size) * (comm.rank))
right = int((FILT_NUM / comm.size) * (comm.rank + 1))
print('left', left, 'right', right)
total = 0
for x in range(left, right):
    part = layer[x]
    weights = all_weights[x].reshape(all_weights[x].size)
    bias = biases[x]
    for i in range(LAY_H):
        for j in range(LAY_W):
            sum = 0
            counter = 0
            for k in range(i, i + WIN_H):
                for l in range(j, j + WIN_W):
                    # print('weights', weights)
                    sum += image[k, l] * weights[counter]
                    counter += 1
                    total = total + 1
            part[i, j] = sum
print('total', total)

# print("Hello! I'm rank %d from %d running in total..." % (comm.rank, comm.size))

comm.Barrier()  # wait for everybody to synchronize _here_

if rank == 0:
    end = time.time()
    # print('layer', layer.shape)
    print('time', end - start)