import torch
import numpy as np

def norm_1d(x, sig_p = 1, mean = 0):
    '''
    Realize the normalization function of one-dimensional data
    Parameters:
    x: array_like
    Input one-dimensional data array
    sig_p: float, optional
    Normalization data power value, default to 1
    mean: bool, optional
    Represents whether to do the 0-means normalization

    Returns:
    y: array_like
    Output one-dimensional normalized data array
    Raises:
    AttributeError
    When the input data format is not numpy or tensor, show that 'Data type is not supported'
    '''
    sym_num = x.shape[0]
    if type(x) == torch.Tensor:
        if mean:
            m = x.mean()
            y = x - m
        else:
            y = x
        mean_power = torch.sum((y.abs()**2)) / sym_num  
        scale = (mean_power / sig_p) ** 0.5
    elif type(x) == np.ndarray:
        if mean:
            m = np.mean(x)
            y = x - m
        else:
            y = x
        mean_power = np.sum((np.abs(y)**2)) / sym_num  
        scale = (mean_power / sig_p) ** 0.5
    else:
        raise AttributeError('Data type is not supported')
    
    y = y / scale
    return y


def norm_dual_pol(x, y, sig_p = 1, mean = 0, axis = 0, data_mode = 'tensor'):
    '''
    Realize the normalization function of dual-polar data
    Parameters:
    x: array_like
    The first dimension of the input data array
    y: array_like
    The second dimension of the input data array
    sig_p: float, optional
    Normalization data power value, default to 1
    mean: bool, optional
    Represents whether to do the 0-means normalization
    axis: int, optional
    Specifies which axis of the input data array to normalize along. The default is 0
    data_mode: str, optional
    Input data format (numpy or tensor)

    Returns:
    x1: array_like
    Output the first dimension of the normalized data array
    y1: array_like
    Output the second dimension of the normalized data array
    '''
    sym_num = x.shape[axis]
    if data_mode == 'tensor':
        if mean:
            x_mean = x.mean(dim = axis)
            y_mean = y.mean(dim = axis)
            x1 = x - x_mean
            y1 = y - y_mean
        else:
            x1 = x
            y1 = y
        mean_power = torch.sum((x1.abs()**2 + y1.abs()**2), dim = axis, keepdim = True) / sym_num     
    else:
        if mean:
            x_mean = np.mean(x, axis = axis)
            x_mean = np.mean(x, axis = axis)
            x1 = x - x_mean
            y1 = y - y_mean
        else:
            x1 = x
            y1 = y
        mean_power = np.sum(np.abs(x1) ** 2 + np.abs(y1) ** 2, axis = axis, keepdims = True) / sym_num
    
    scale = (mean_power / sig_p) ** 0.5
        
    x1 = x1 / scale
    y1 = y1 / scale
    return x1, y1

def norm_hd(x, axis = -1, sig_p = 1, mean = 0):  
    '''
    Realize normalization function of high-dimensional data
    Parameters:
    x: array_like
    Input high-dimensional data array
    axis: int, optional
    Specify which dimensions to standardize. The default value is -1, normalization all dimensions
    sig_p: float, optional
    Normalization data power value, default to 1
    mean: bool, optional
    Represents whether to do the 0-means normalization

    Returns:
    y: array_like
    Output normalized data array
    '''
    sym_num = x.shape[axis]
    if type(x) == torch.Tensor:
        if mean:
            x_mean = x.mean(dim = axis)
            y = x - x_mean
        else:
            y = x
        mean_power = torch.sum((y.abs()**2), dim = axis, keepdim = True) / sym_num     
    else:
        if mean:
            x_mean = np.mean(x, axis = axis)
            y = x - x_mean
        else:
            y = x
        mean_power = np.sum(np.abs(x) ** 2, axis = axis, keepdims = True) / sym_num
    scale = (mean_power / sig_p) ** 0.5
        
    y = y / scale
    return y
