import os
import torch
import random
import numpy as np

from torch.optim.lr_scheduler import LambdaLR

def load_model(args, save_file):
    path = os.path.join(
        save_file,
        "checkpoint-{}".format(args.max_steps),
        "model.pt",
    )
    assert os.path.exists(path), "Path is not exist"
    model = torch.load(path)
    return model, path

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    # torch.backends.cudnn.benchmark = False
    # torch.backends.cudnn.deterministic = True
    # torch.backends.cudnn.enabled = False
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

def get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps, num_training_steps, last_epoch=-1
):
    """Create a schedule with a learning rate that decreases linearly after
    linearly increasing during a warmup period.
    """

    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(
            0.0,
            float(num_training_steps - current_step)
            / float(max(1, num_training_steps - num_warmup_steps)),
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)

def get_constant_schedule(optimizer, last_epoch=-1):
    """Create a schedule with a constant learning rate."""
    return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)
