import subprocess
import os

from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE


def run_gpt(cmd):
    args = list(cmd.split(" "))
    p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    outs, errs = p.communicate()
    outs = list(str((outs).decode("utf-8")).splitlines())
    success = False
    runtime = 0
    num_params = 0
    for out in outs:
        out = str(out)
        if "Average Iteration Time:" in str(out):
            slicey = out[out.find(":") + 2 :]
            try:
                runtime = float(slicey)
            except:
                print(slicey)
                quit()
        if "Number of Parameters:" in str(out):
            slicey = out[out.find(":") + 2 :]
            try:
                num_params = int(slicey)
            except:
                print(slicey)
                quit()
        if str(out) == str(TEST_SUCCESS_MESSAGE):
            success = True
    return runtime, round(float(int(num_params)) / 10.0 ** 9, 3), success, errs


def plot(runtimes):
    import matplotlib.pyplot as plt

    for distributed_setting in runtimes.keys():
        plt.scatter(
            runtimes[distributed_setting].keys(),
            runtimes[distributed_setting].values(),
            label=distributed_setting,
        )
    plt.legend()
    plt.xlabel("Parameters (Billions)")
    plt.ylabel("Training Iteration time (s)")
    plt.title(str("GPT Scaling w/ Offloading"))
    plt.savefig("offload_gpt_scaling.png")
    plt.close()
    if not os.path.exists("/my_workspace/"):
        os.system("mkdir /my_workspace/")
    os.system("cp *.png /my_workspace/")


def main():
    runtimes = {}
    nlist = (
        list(range(2000, 10000, 2000))
        + list(range(10000, 50000, 5000))
        + list(range(50000, 100000, 10000))
    )
    print("N-List:", nlist)
    for data_parr, tens_parr, pipe_parr in [(8, 1, 1), (4, 2, 1), (2, 1, 4), (1, 2, 4)]:
        for offload in [True, False]:
            dist_setting = (
                "ddp="
                + str(data_parr)
                + ", tensor_parr="
                + str(tens_parr)
                + ", pipe_parr="
                + str(pipe_parr)
                + ", offload="
                + str(offload)
            )
            runtimes[dist_setting] = {}
            print("Beginning Testing for", dist_setting)
            for n in nlist:
                cmd = "python3 -m torch.distributed.launch --nproc_per_node=8 run_gpt_minimal_test.py"
                cmd += (
                    " --micro-batch-size 1 --num-layers "
                    + str(n)
                    + " --hidden-size 128 --num-attention-heads 16"
                )
                cmd += (
                    " --max-position-embeddings 128 --seq-length 128 --tensor-model-parallel-size "
                    + str(tens_parr)
                )
                cmd += (
                    " --pipeline-model-parallel-size "
                    + str(pipe_parr)
                    + (" --cpu-offload" if offload else "")
                )
                print(cmd)
                runtime, bill_params, success, errs = run_gpt(cmd)
                if success:
                    runtimes[dist_setting][bill_params] = runtime
                    print(
                        str(runtime) + "s per training iter for",
                        str(bill_params) + "B parameter GPT-2",
                    )
                    if n >= 10000:
                        plot(runtimes)
                else:
                    print("GPT-2 w/", n, "layers failed using", dist_setting)
                    print("Moving on to the next distributed setting...")
                    print("#" * (25))
                    print()
                    plot(runtimes)
                    break
    print(runtimes)
    plot(runtimes)


if __name__ == "__main__":
    main()