# Copyright (C) 2013-2023 Karlsruhe Institute of Technology
#
# This file is part of syris.
#
# This library is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library. If not, see <http://www.gnu.org/licenses/>.

"""Show multi-device speedup on a problem of size n x m^k, where n is the number of pixels to
compute, m is the base number of operations per pixel powered to k.
"""
import time
import pyopencl as cl
import matplotlib.pyplot as plt
import numpy as np
import syris
import syris.config as cfg
import syris.gpu.util as gutil
from .util import get_default_parser


def get_kernel():
    return """
        kernel void foo (global float *output,
                         const int stop)
        {
            int idx = get_global_id (0);
            int n = get_global_size (0);
            int i;
            float result = (float) idx;

            for (i = 1; i < stop + 1; i++) {
                result = result * result;
            }

            output[idx] = result;
        }
        """


def parse_args():
    parser = get_default_parser(__doc__)
    parser.add_argument("--n", type=int, default=512 ** 2, help="Number of pixels (default 512^2)")
    parser.add_argument("--m", type=int, default=64, help="Number of pixel operations (default 64)")
    parser.add_argument(
        "--k",
        type=float,
        default=[1],
        nargs="*",
        help="Time complexity, there are n x m^k operations in total "
        "(default 1). The argument might have one (single run) or 3 numbers "
        "(meaning start stop step) when a scan is performed.",
    )
    parser.add_argument(
        "--runs",
        type=int,
        default=1,
        help="Number of runs per one complexity. " "The result speedup is the mean.",
    )
    parser.add_argument("--plot", action="store_true", help="Plot results")
    parser.add_argument(
        "--verbose", action="store_true", help="Print infomation from individual runs"
    )
    parser.add_argument(
        "--output",
        type=str,
        help="Output individual speedups from all "
        "runs as a numpy array in this filename k is the complexity. "
        "It must conatin a formatting string which will hold the complexity, "
        "e.g. speedups{:>04}.npy.",
    )

    args = parser.parse_args()
    m = args.k
    if not (len(m) == 1 or len(m) == 3):
        raise ValueError("--k must contain either 1 or 3 numbers")

    return args


def run(n, m, complexity, prg, verbose=False):
    devices = cfg.OPENCL.devices
    queues = cfg.OPENCL.queues

    stop = int(m ** complexity)
    complexity_fmt = "complexity: {} x {}^{}, pixel operations: {}"
    if verbose:
        print(complexity_fmt.format(n, m, complexity, stop))
    num_items = len(queues)
    events = []

    def process(item, queue):
        data = cl.array.Array(queue, (n,), dtype=cfg.PRECISION.np_float)
        event = prg.foo(queue, (n, 1), None, data.data, np.int32(stop))
        events.append(event)
        # Wait for the event so that the command queue will not be scheduled before the work has
        # been done.
        event.wait()

    start = time.time()
    gutil.qmap(process, list(range(num_items)))
    host_duration = time.time() - start

    if verbose:
        print("-------------------------------")
        print("     All duration: {:.2f} s".format(host_duration))
        print("-------------------------------")

    all_duration = 0
    for i, event in enumerate(events):
        duration = get_duration(event)
        all_duration += duration
        if verbose:
            print("Device {} duration: {:.2f} s".format(i, duration))

    speedup = all_duration / host_duration
    if verbose:
        print("-------------------------------")
        print("    Mean duration: {:.2f} s".format(all_duration / len(devices)))
        print("-------------------------------")
        print("          Speedup: {:.2f} / {}".format(speedup, len(devices)))
        print("-------------------------------")

    return speedup


def main():
    args = parse_args()
    syris.init(profiling=True)
    prg = cl.Program(cfg.OPENCL.ctx, get_kernel()).build()

    if len(args.k) == 1:
        complexities = args.k
    else:
        complexities = np.arange(args.k[0], args.k[1] + args.k[2], args.k[2])
    print("Complexities: {}".format(complexities))

    results = []
    for complexity in complexities:
        runs = []
        for i in range(args.runs):
            print("Run {} / {}".format(i + 1, args.runs))
            runs.append(run(args.n, args.m, complexity, prg, verbose=args.verbose))
        if args.output:
            np.save(args.output.format(complexity), runs)
        results.append(np.mean(runs))

    print()
    print("===============================")
    for i, result in enumerate(results):
        print("Complexity: {:.2f}, speedup: {:.2f}".format(complexities[i], result))
    print("===============================")

    if args.plot:
        plt.figure()
        plt.plot(complexities, results)
        plt.xlabel("Complexity")
        plt.xlabel("Speedup")
        plt.grid()
        plt.show()


def get_duration(event):
    return (event.profile.end - event.profile.start) * 1e-9


if __name__ == "__main__":
    main()
