import argparse
import ast
import functools
import gc
import operator
import time

import numpy
import cupy


parser = argparse.ArgumentParser()
parser.add_argument('--number', type=int, default=10000)
parser.add_argument('--src-xp', type=str, choices=["numpy", "cupy"], required=True)
parser.add_argument('--dst-xp', type=str, choices=["numpy", "cupy"], required=True)
parser.add_argument('--shape', type=str, default="(1,)")
parser.add_argument('--batch-size', type=int, default=1)
parser.add_argument('--debug', action='store_true')


def to_xp(s):
    if s == 'numpy':
        return numpy
    if s == 'cupy':
        return cupy
    else:
        raise


def setup(args):
    src_xp = to_xp(args.src_xp)
    dst_xp = to_xp(args.dst_xp)
    shape = ast.literal_eval(args.shape)
    size = functools.reduce(operator.mul, shape, 1)
    batch_size = args.batch_size
    number = args.number

    # e.g. a = [src_xp.array([[0, 1], [2, 3]]), src_xp.array([[4, 5], [6, 7]])]
    a = [
        src_xp.array(range(i*size, (i+1)*size), dtype=src_xp.float32).reshape(shape)
        for i in range(batch_size)]

    print("Shape: " + str(shape))
    print("Batch size: " + str(batch_size))
    print("Running " + dst_xp.__name__ + ".array(<List[" + str(src_xp.__name__) + ".ndarray]>) in " + str(number) + " times...")

    return (dst_xp, a)


def time_numpy(number, f, *args, **kwargs):
    start = time.time()
    for i in range(number):
        f(*args, **kwargs)
    end = time.time()
    return end - start


def time_cupy(number, f, *args, **kwargs):
    stream = cupy.cuda.Stream.null
    start = stream.record()
    for i in range(number):
        f(*args, **kwargs)
    end = stream.record()
    stream.synchronize()
    return cupy.cuda.get_elapsed_time(start, end) / 1000


def time_chainerx(number, device, f, *args, **kwargs):
    start = time.time()
    for i in range(number):
        f(*args, **kwargs)
    device.synchronize()
    end = time.time()
    return end - start


def main():
    args = parser.parse_args()
    xp, a = setup(args)
    if xp is numpy:
        print(time_numpy(args.number, xp.array, a))
    elif xp is cupy:
        print(time_cupy(args.number, xp.array, a))
    else:
        raise


if __name__ == '__main__':
    main()