#!/usr/bin/env python3

import matplotlib.pyplot as plt
import numpy as np
import subprocess
import os

RESULT_FILENAME = "result.png"
RESULT_SPEEDUP_FILENAME = "result_speedup.png"
RESULT_APPEND_FILENAME = "result_and_speedup.png"


def plot(title1, xs1, ys1, title2, xs2, ys2, title3):
    fig, ax = plt.subplots()
    ax.plot(xs1, ys1, label=title1)
    ax.plot(xs2, ys2, label=title2)
    ax.set(xlabel="Number of tasks", ylabel='Time (in seconds)',
           title=title3)
    ax.grid()
    plt.tight_layout()
    plt.legend()
    fig.savefig(RESULT_FILENAME)
    plt.show()


def plot_speedup(title1, xs1, ys1, title2, xs2, ys2, title3):
    fig, ax = plt.subplots()
    ax.plot(xs1, ys1, label=title1)
    ax.plot(xs2, ys2, label=title2)
    ax.plot(xs1, xs1, label="Ideal Speedup")
    ax.set(xlabel="Number of tasks", ylabel='Speedup',
           title=title3)
    ax.grid()
    plt.tight_layout()
    plt.legend()
    fig.savefig(RESULT_SPEEDUP_FILENAME)
    plt.show()


TITLES = [
    [
        "Simple Cellular Automaton\nCompute 15 generations with 300'000'000 cells",
        "Simple Cellular Automaton\nSpeedup of computing 15 generations\n with 300'000'000 cells"
    ],
    [
        "Game of Life\nCompute 15 generations with 900'000'000 cells",
        "Game of Life\nSpeedup of computing 15 generations\n with 900'000'000 cells"
    ],
    [
        "Lattice-Boltzmann Method\nCompute 15 generations with 27'000'000 cells",
        "Lattice-Boltzmann Method\nSpeedup of computing 15 generations\n with 27'000'000 cells"
    ]
]


def get_data(filename):
    results = np.genfromtxt(filename, delimiter=';')
    results_ordered = {}
    for result in results:
        key = int(result[0])
        if key in results_ordered:
            results_ordered[key].append(result[-1])
        else:
            results_ordered[int(result[0])] = [result[-1]]
    xs = []
    ys = []
    stds = []
    speedups = []
    for key in results_ordered.keys():
        xs.append(key)
        ys.append(sum(results_ordered[key]) / len(results_ordered[key]))
        stds.append(np.std(results_ordered[key]))
        speedups.append(ys[0] / ys[-1])
    return xs, ys, stds, speedups


def to_md_array(filename, xs, ys, stds, speedups, gpu):
    with open(filename, "w") as fh:
        if gpu:
            fh.write(
                "| Number of tasks | Number of GPUs | Average [s] | Standard Derivation [s] | Speedup | Number of measures |\n")
            fh.write("|:---:|:---:|:---:|:---:|:---:|:---:|\n")
        else:
            fh.write("| Number of tasks | Average [s] | Standard Derivation [s] | Speedup | Number of measures |\n")
            fh.write("|:---:|:---:|:---:|:---:|:---:|\n")
        for i in range(len(xs)):
            if gpu:
                fh.write(
                    f"| {xs[i]} | {min(xs[i], 8)} | {round(ys[i], 3)} [s] | ± {round(stds[i], 3)} [s] | x{round(speedups[i], 1)} | 15 |\n")
            else:
                fh.write(
                    f"| {xs[i]} | {round(ys[i], 3)} [s] | ± {round(stds[i], 3)} [s] | x{round(speedups[i], 1)} | 15 |\n")


def main():
    version1, version2 = "OpenCL", "CUDA"
    xs1, ys1, stds1, speedups1 = get_data("lattice_boltzmann/results/lbm-opencl-2021-07-25.csv")
    xs2, ys2, stds2, speedups2 = get_data("lattice_boltzmann/results/lbm-cuda-2021-07-23.csv")
    plot(version1, xs1, ys1, version2, xs2, ys2, TITLES[2][0])
    plot_speedup(f"{version1} Speedup", xs1, speedups1, f"{version2} Speedup", xs2, speedups2, TITLES[2][1])
    to_md_array("results_opencl.md", xs1, ys1, stds1, speedups1, gpu=True)
    to_md_array("results_cuda.md", xs2, ys2, stds2, speedups2, gpu=True)
    # Need ImageMagick
    subprocess.run(["docker", "run", "-v", f"{os.getcwd()}:/imgs", "dpokidov/imagemagick", "+append",
                    f"/imgs/{RESULT_FILENAME}", f"/imgs/{RESULT_SPEEDUP_FILENAME}", f"/imgs/{RESULT_APPEND_FILENAME}"],
                   capture_output=False)


if __name__ == '__main__':
    main()