import torch
import math


def generate_A_with_L_mu(n, d, L, mu=-1):
    """
    Generate a data matrix A for
        f(x) = \frac{1}{n} || A x - b ||_2^2
    with L-smoothnes and mu-convexity.

    The L-smoothness is the largest eigenvalue of the hessian of f
        hessian(f)(x) = \frac{2}{n} A^T A
    """
    assert mu <= L

    # Generate unitary matrix U and V
    dummy_matrix = torch.randn(n, d)
    U, _, V = torch.linalg.svd(dummy_matrix)

    # Construct matrix S such that S.T @ S has largest elements L
    # and smallest elements mu.
    smallest = math.sqrt(abs(mu) * n / 2)
    largest = math.sqrt(L * n / 2)
    diag = torch.linspace(start=smallest, end=largest, steps=min(n, d))
    S = torch.zeros(n, d)
    S[list(range(d)), list(range(d))] = diag

    # Reconstruct A
    return U @ S @ V.T