#!/usr/bin/python
#
#  Copyright 2002-2019 Barcelona Supercomputing Center (www.bsc.es)
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

# -*- coding: utf-8 -*-
import os
import time
import json
from pycompss.api.api import compss_barrier

def initialize_variables():
    for i in range(MSIZE):
        A.append([])
        B.append([])
        C.append([])
        for j in range(MSIZE):
            A[i].append("A.%d.%d" % (i, j))
            B[i].append("B.%d.%d" % (i, j))
            C[i].append("C.%d.%d" % (i, j))


def fill_matrices():
    for c in ['A', 'B', 'C']:
        for i in range(MSIZE):
            for j in range(MSIZE):
                tmp = "%s.%d.%d" % (c, i, j)
                f = open(tmp, 'w')
                for _ in range(BSIZE):
                    for jj in range(BSIZE):
                        if c == 'C':
                            f.write('0.0')
                        else:
                            f.write('2.0')
                        if jj < BSIZE - 1:
                            f.write(' ')
                    f.write('\n')
                f.close()


# ## MAIN PROGRAM ## #

if __name__ == "__main__":
    import sys
    from matmul_tasks import multiply

    print('this is the python path')
    print(os.getenv("PYTHONPATH"))

    args = sys.argv[1:]
    MSIZE = int(args[0])
    BSIZE = int(args[1])

    A = []
    B = []
    C = []

    print(f'\nCase with {MSIZE} number of blocks and {BSIZE} elements.\n')

    start = time.time()
    initialize_variables()
    fill_matrices()
    initialize = time.time()
    for i in range(MSIZE):
        for j in range(MSIZE):
            for k in range(MSIZE):
                multiply(A[i][k], B[k][j], C[i][j], BSIZE)
    compss_barrier(no_more_tasks=True)
    multiplication = time.time()

    initialization_time = initialize - start
    multiplication_time = multiplication - initialize
    total_time = multiplication - start

    print("-----------------------------------------")
    print("-------------- RESULTS ------------------")
    print("-----------------------------------------")
    print("Initialization time: %f" % (initialization_time))
    print("Multiplication time: %f" % (multiplication_time))
    print("Total time: %f" % (total_time))
    print("-----------------------------------------")

    result = {
        'num_blocks': MSIZE,
        'elems_per_block': BSIZE,
        'init_time': initialization_time,
        'mult_time': multiplication_time,
        'total_time': total_time
    }
    with open('output.json', 'a') as out:
        json.dump(result, out, indent=4)
        out.write("\n")