import torch
import time


def cs_sim(image, cs_transform, std, cs_algo):
    """Simulate compressive sensing problem

    Args:
        image (torch.Tensor): image tensor with shape (C, H, W).
        cs_transform (CStransform): a CS transformation instance.
        std (float): standard deviation of additive Gaussian noise
            on the scale [0., 1.].
        cs_algo (algorithms.csalgo.IterativeDenoisingCSRecon):
            an instance of CS reconstruction algorithm.

    Returns:
        recon_image (torch.Tensor): tensor of reconstructed image.
        r_t (torch.Tensor): tensor of the pseudo noisy image (r_t)
            from last iteration of cs_algo.
        std_est (float): estimated standard deviation of noise in r_t
        psnr (np.ndarray): PSNR at each iteration of reconstruction with cs_algo.
            psnr is None if cs_algo does not have the reference ground truth image.
    """
    print('deploy.sim.cs_sim: Simulating compressive sensing...')
    y = cs_transform.Afun(image.view(-1, 1))
    if std > 0:
        y += torch.normal(mean=torch.zeros(y.shape), std=std)

    print('deploy.sim.cs_sim: Begin compressive sensing reconstruction...')
    start_time = time.time()
    recon_image, psnr, r_t, std_est = cs_algo(y, cs_transform.Afun, cs_transform.Atfun,
                                              cs_transform.get_m(), cs_transform.get_n())
    print('deploy.sim.cs_sim: Reconstruction took {:.3f} s'.format(time.time() - start_time))

    return recon_image, r_t, std_est, psnr