import math
import torch


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

    Parameters
    ----------
    cp_tensor : CPTensor
    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 = cp_tensor.rank # We assume we are given a valid CP 
    order = cp_tensor.orders
    std_factors = (std/math.sqrt(rank))**(1/order)

    with torch.no_grad():
        cp_tensor.weights.fill_(1)
        for factor in cp_tensor.factors:
            factor.normal_(0, std_factors)
    return cp_tensor