import numpy as np
import torch
import matplotlib.pyplot as plt

def main(x, tx):
    r"""
        The main function of the coarse estimation algorithm.
        Parameters: x：ndarray
                       Signal sequence of receiver.
                    tx：ndarray
                       Signal sequence of transmitter.
        Return:     y: ndarray
                       The signal sequence of receiver after phase offset compensation.
                    phi：ndarray
                       The phase offset in receiver.
    """

    """ 
    Rn = r * exp(j(an + 2 * pi * n * Δv * Ts + θn))，
    Tn = r * exp(j(an))。
    Tn * Rn.conj = r ^ 2 * exp(-j(2 * pi * n * Δv * Ts + θn))。
    """
    if type(x) == np.ndarray:
        sig_complex = np.mean(tx * np.conj(x))
        phi = - np.angle(sig_complex)
        y = x * np.exp(-1j * phi)
    else:
        sig_complex = torch.mean(tx * x.conj())
        phi = - torch.angle(sig_complex)
        y = x * torch.exp(-1j * phi)
    
    return phi, y
