import numpy as np
from scipy.stats import halfnorm, truncnorm


def set_worker_speed(args, logger):
    n_workers = args.n_workers
    tau_max = args.tau_max
    tau_min = args.tau_min
    
    if args.comput_dist == "uniform":
        worker_speeds = np.random.randint(low=tau_min, high=tau_max, size=n_workers)
        
    elif args.comput_dist == "normal":
        # tau_min, tau_max = 1, 5
        tau_min, tau_max = args.tau_min, args.tau_max
        a = (tau_min-args.tau_mean)/args.tau_std
        b = (tau_max-args.tau_mean)/args.tau_std
        x = truncnorm.rvs(a, b, loc=args.tau_mean, scale=args.tau_std, size=n_workers)
        worker_speeds = [max(round(item), tau_min) for item in x]
        
    elif args.comput_dist == "halfnorm":
        x = halfnorm.rvs(loc=args.tau_mean, scale=args.tau_std, size=n_workers)
        worker_speeds = [min(round(item), tau_max) for item in x]
    
    else:
        raise NotImplementedError(f"Computational distribution {args.comput_dist} not implemented yet")
    
    logger.info(f"Worker speeds: {worker_speeds}")

    return worker_speeds
