import math
import os
from time import sleep


# Momentum: https://blog.paperspace.com/intro-to-optimization-momentum-rmsprop-adam/
# V_j = coefficient_momentum * V_j - learning_rate * gradient
# W_j = V_j + W_jx  x
# https://www.youtube.com/watch?v=k8fTYJPd3_I
def adjust_learning_rate(optimizer, initial_lr, round_index, total_round):
    """Decay the learning rate based on schedule"""
    lr = initial_lr
    # cosine lr schedule
    lr *= 0.5 * (1. + math.cos(math.pi * round_index / total_round))

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


def post_complete_message_to_sweep_process(args):
    pipe_path = "./tmp/fedml"
    if not os.path.exists(pipe_path):
        os.mkfifo(pipe_path)
    pipe_fd = os.open(pipe_path, os.O_WRONLY)

    with os.fdopen(pipe_fd, 'w') as pipe:
        pipe.write("training is finished! \n%s\n" % (str(args)))


def wait_for_the_training_process():
    pipe_path = "./tmp/fedml"
    if not os.path.exists(pipe_path):
        os.mkfifo(pipe_path)
    pipe_fd = os.open(pipe_path, os.O_RDONLY | os.O_NONBLOCK)
    with os.fdopen(pipe_fd) as pipe:
        while True:
            message = pipe.read()
            if message:
                print("Received: '%s'" % message)
                print("Training is finished. Start the next training with...")
                os.remove(pipe_path)
                return
            sleep(5)
            print("Daemon is alive. Waiting for the training result.")
