import torch
import numpy as np
from ... comm_tools.unwrap import unwrap 
from ... comm_tools.movmean import movmean
    
def vv_cpe(x, move_mean = 1, *args, **kwargs):
    r"""
        The main function of the viterbi_viterbi algorithm.
        Parameters: x：ndarray
                        Signal sequence of receiver.
                    move_mean: int
                        Whether to take the moving average operation, defaults to 1.
        Return:     y: ndarray
                        The signal sequence of receiver after phase offset compensation.
                    phase：ndarray
                        The phase offset in receiver.
    """
    x_4 = - x ** 4
    if move_mean:
        window_size = kwargs.get('window_size', 100)    # CPE memory
        x_4 = movmean(x_4, window_size)
    if type(x_4) == torch.Tensor:
        phase = torch.angle(x_4)
    else:
        phase = np.angle(x_4)
    
    phase = unwrap(phase) / 4
    y = x * torch.exp(1j * phase)
        
    return phase, y
