import torch
import math


def compute_ls_grid(As, y, sns_vec, m, ks, n_ls, l_eps, dtype):
    """Compute l values for each given k and return a dictionary mapping
    k to a list (in decreasing order) of lambda values.

    Arguments have the same meaning as in gel_paths2. sns_vec is a vector of
    sns_j values as opposed to the matrix computed in gel_paths2.
    """
    ls_grid = {}
    # The bound is given by max{||A_j'@(y - b_0)||/(m*sqrt{n_j}*k)} where b_0 =
    # 1'@y/m. So most things can be precomputed.
    l_max_b_0 = y.mean()
    l_max_unscaled = max(
        (A_j.t() @ (y - l_max_b_0)).norm(p=2) / (m * sns_j)
        for A_j, sns_j in zip(As, sns_vec)
    )
    for k in ks:
        l_max = l_max_unscaled / k
        if n_ls == 1:
            ls_grid[k] = [l_max]
        else:
            l_min = l_max * l_eps
            ls = torch.logspace(
                math.log10(l_min), math.log10(l_max), steps=n_ls, dtype=dtype
            )
            ls = sorted([l.item() for l in ls], reverse=True)
            ls_grid[k] = ls
    return ls_grid