import numpy as np
from numpy.fft import fftn, fftshift, ifftshift



def sigma(field, box_size, assume_invariance=False):
    """
    Compute the first kind of the biased sigma estimator, given a field, using configuration space with periodic boundary condition method. 
    
    Parameters
    ----------
    field : three-dimensiona ndarray
              Input three-dimensional ndarray corresponding 
              to the cosmological field.
    box_size : float64
               The physical size of the box in the desired
               physical units (usually in Mpc).
    assume_invariance : bool, default=False
                        If true, compute bias of first kind estimator
                        by assuming the 'field' input is the field power
                        spectrum not normalized.
    Returns
    -------
    ans : three-dimensional complex ndarray
          Returns the normalized first kind estimator if assume_invariance is false. If assume_invariance
          is true, compute the normalized first kind estimator bias of that field otherwise.
    """

    Nx, Ny, Nz = field.shape[0], field.shape[1], field.shape[2]
    dx, dy, dz = box_size/Nx, box_size/Ny, box_size/Nz
    dkx, dky, dkz = (2*np.pi)/(Nx*dx), (2*np.pi)/(Ny*dy), (2*np.pi)/(Nz*dz)
    
    vol = (Nx*Ny*Nz) * (dx*dy*dz)

    ans = np.zeros((Nx, Ny, Nz), dtype='complex')

    if assume_invariance == True:

        sigma_bd_fn = sigma_bias_which_diagonal_cs
        
        ans[Nx//2, Ny//2, Nz//2] = sigma_bd_fn(field, box_size)

    else:

        input_field = fftshift((dx*dy*dz)*fftn(ifftshift(field)**2))

        ans = (1/2) * (input_field + np.conjugate(input_field))


    return ans/vol




def sigma_bias_which_diagonal_cs(pspec, box_size, diagonal=np.zeros((3))):

    """
    Compute the diagonal of the bias matrix, given the field power spectrum, using configuration space method. 
    
    Parameters
    ----------
    pspec : one-dimensiona complex ndarray
            Input one-dimensional ndarray corresponding 
            to the cosmological field Fourier transformed.
    box_size : float64
               The physical size of the box in the desired
               physical units (usually in Mpc).
    diagonal : shape (3,) ndarray, default=np.zeros((3))
               the desired diagonals of the bias matrix to be
               computed. Each entry corresponds to the different
               diagonals generated by convolving the pspec with itself.

    Returns
    -------
    ans : int
          Returns the desired diagonal of the bias matrix of that
          field in configuration space.

    """

    Nx, Ny, Nz = pspec.shape[0], pspec.shape[1], pspec.shape[2]
    dx, dy, dz = box_size/Nx, box_size/Ny, box_size/Nz
    vol = (Nx*Ny*Nz) * (dx*dy*dz)
    
    ans = 0 


    if all(d == 0 for d in diagonal):

        ans = np.sum(pspec)

    else:

        ans = 0
         
    return ans/vol