def get_optimal_step_size_and_momentum_parameters(sketch, fixed, momentum, n, d, m):
    """
        As defined in:
         1. Jonathan Lacotte, Mert Pilanci, Optimal Randomized First-Order Methods for Least-Squares Problems (https://arxiv.org/abs/2002.09488)
         2. Jonathan Lacotte, Mert Pilanci, Faster Least Squares Optimization (https://arxiv.org/pdf/1911.02675.pdf)
        Args:
            sketch: Sketching matrix used
            fixed: Whether sketching matrix is redrawn at every turn
            momentum: Whether heavy-ball is used or not
            n: Number of rows of data matrix
            d: Number of columns of data matrix
            m: Sketch size

        Returns: step size and momentum parameter

        """
    if sketch == 'SRHT':
        if fixed:
            if momentum:
                alpha = m / d
                step_size = (2 * (alpha - 1)) / (alpha + ((alpha ** 2 - alpha) ** 0.5))
                # beta = ((alpha ** 0.5) - ((alpha - 1) ** 0.5)) / (alpha ** 0.5) + ((alpha - 1) ** 0.5)
                xi, gamma = m / n, d / n
                lambda_h = (((1 - gamma) * xi) ** 0.5 - ((1 - xi) * gamma) ** 0.5) ** 2
                Lambda_h = (((1 - gamma) * xi) ** 0.5 + ((1 - xi) * gamma) ** 0.5) ** 2
                # step_size = 4 / ((1 / Lambda_h)**0.5 + (1 / lambda_h)**0.5)**2
                beta = ((Lambda_h ** 0.5 - lambda_h ** 0.5) / (Lambda_h ** 0.5 + lambda_h ** 0.5)) ** 2
            else:
                beta = 0
                step_size = 1 - d/m
        else:
            beta = 0
            xi, gamma = m/n, d/n
            step_size = (xi - gamma)**2 / (xi**2 + gamma - (2 * gamma * xi))
    else:
        if fixed:
            if momentum:
                beta = d / m
                step_size = (1 - beta) ** 2
            else:
                beta = 0
                step_size = (4 * m * d) / (m**2 + d**2)
        else:
            beta = 0
            step_size = ((m - d) * (m - d - 3)) / (m * (m - 1))
    return beta, step_size