# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
import argparse
import os
import sys
import subprocess


def add_poprun_arguments(parser):
    # shared application arguments
    parser.add_argument("--config", default=None, type=str, required=True)
    parser.add_argument("--num_epochs", default=None, type=int, required=True)

    parser.add_argument("--init_lr", default=None, type=float, required=True)
    parser.add_argument("--abs_end_lr", default=None, type=float, required=True)
    parser.add_argument("--warmup_epochs", default=None, type=int, required=True)
    parser.add_argument("--label_smoothing", default=None, type=float, required=True)
    parser.add_argument("--weight_decay", default=None, type=float, required=False)
    parser.add_argument("--optimizer", default=None, type=str, choices=["sgd", "lars"], required=False)
    parser.add_argument("--micro_batch_size", default=None, type=int, required=True)
    parser.add_argument("--grad_accum", default=None, type=int, required=True)
    parser.add_argument("--num_replicas", default=None, type=int, required=True)
    return parser


if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser = add_poprun_arguments(parser)
    args = parser.parse_args()

    hosts = os.environ.get("POPRUN_HOSTS")
    vipu_host = os.environ.get("IPUOF_VIPU_API_HOST")
    vipu_partition = os.environ.get("IPUOF_VIPU_API_PARTITION_ID")
    if None in [hosts, vipu_host, vipu_partition]:
        raise ValueError(
            f"The following environment variables must be defined: "
            f"POPRUN_HOSTS={hosts}\n"
            f"IPUOF_VIPU_API_HOST={vipu_host}\n"
            f"IPUOF_VIPU_API_PARTITION_ID={vipu_partition}"
        )

    user = os.environ["USER"]
    exec_cache = os.environ.get("TF_POPLAR_EXEC_CACHE") or os.path.join("/home", user, "exec_cache")

    poprun_command = [
        "poprun",
        "-vv",
        "--host",
        hosts,
        "--only-output-from-instance",
        "0",
        "--mpi-global-args",
        "--mca oob_tcp_if_include eno1 --mca btl_tcp_if_include eno1",
        "--update-partition",
        "yes",
        "--reset-partition",
        "no",
        "--vipu-server-timeout",
        "600",
        "--vipu-server-host",
        vipu_host,
        "--vipu-partition",
        vipu_partition,
        "--executable-cache-path",
        exec_cache,
        "--num-instances",
        64,
        "--num-replicas",
        64,
        "--ipus-per-replica",
        1,
    ]

    training_command = [
        *poprun_command,
        "python3",
        "train.py",
        "--config",
        args.config,
        "--num-epochs",
        args.num_epochs,
        "--wandb",
        "True",
        "--target-accuracy",
        0.759,
        "--ckpts-per-epoch",
        1 / 4,
        "--first-ckpt-epoch",
        args.num_epochs % 4,
        "--sweep",
        "True",
    ]

    global_batch_size = args.micro_batch_size * args.num_replicas * args.grad_accum
    end_learning_rate_ratio = args.abs_end_lr * (1 / (args.init_lr * global_batch_size))

    # global batch size
    training_command += ["--global-batch-size", global_batch_size]
    # label smoothing
    training_command += ["--label-smoothing", args.label_smoothing]
    # l2 regularization
    training_command += ["--l2-regularization", args.weight_decay] if args.weight_decay is not None else []
    # lr schedule params
    lr_schedule_params = [
        f'"initial_learning_rate": {args.init_lr}',
        f'"end_learning_rate_ratio": {end_learning_rate_ratio}',
        f'"epochs_to_total_decay": {args.num_epochs - args.warmup_epochs}',
        f'"power": 2',
    ]
    training_command += ["--lr-schedule-params", "{" + ",".join(lr_schedule_params) + "}"]
    # warmup params
    lr_warmup_params = [
        f'"warmup_mode": "shift"',
        f'"warmup_epochs": {args.warmup_epochs}',
    ]
    training_command += ["--lr-warmup-params", "{" + ",".join(lr_warmup_params) + "}"]
    # optimizer params
    optimizer_params = [f'"momentum": 0.9']
    optimizer_params += [f'"weight_decay": {args.weight_decay}'] if args.weight_decay is not None else []
    if args.optimizer == "lars":
        optimizer_params += [f'"eeta": 0.001', f'"epsilon": 0']
    training_command += ["--optimizer-params", "{" + ",".join(optimizer_params) + "}"]

    training_command = [str(command) for command in training_command]
    print(" ".join(training_command))

    # run training
    p = subprocess.Popen(training_command, stderr=sys.stderr, stdout=sys.stdout)
    p.wait()