import math
import torch


def cp_init(weights, factors, std=0.02):
    """Initializes directly the weights and factors of a CP decomposition so the reconstruction has the specified std and 0 mean

    Parameters
    ----------
    weights : 1D tensor
    factors : list of 2D factors of size (dim_i, rank)
    std : float, default is 0.02
        the desired standard deviation of the full (reconstructed) tensor

    Notes
    -----
    We assume the given (weights, factors) form a correct CP decomposition, no checks are done here.
    """
    rank = factors[0].shape[1] # We assume we are given a valid CP 
    order = len(factors)
    std_factors = (std/math.sqrt(rank))**(1/order)

    with torch.no_grad():
        weights.fill_(1)
        for i in range(len(factors)):
            factors[i].normal_(0, std_factors)
    return weights, factors