import xarray as xr
import numpy as np

from datetime import datetime

import xrft
import finufft

    
def fft(dataArray, **kwargs):
    return xrft.fft(dataArray, **kwargs)


def ifft(dataArray, **kwargs):
    return xrft.ifft(dataArray, **kwargs)


def fft_nutou(dataArray, modeNum, **kwargs):
    
    data = dataArray.to_numpy()
    data = data.astype('complex128')
    
    time = dataArray[dataArray.dims[0]].to_numpy()
    
    if isinstance(time[0], type(np.datetime64(500,'s'))):
        time = time.astype(float)
        time = time - time[0]
        freqUpLim = 1 / np.min(np.abs(time - np.roll(time, 1))) * 1e9
    else:
        time = time.astype(float)
        time = time - time[0]
        freqUpLim = 1 / np.min(np.abs(time - np.roll(time, 1)))
        
    print(freqUpLim)
    
    time = time / time.max() * 2 * np.pi

    # calculate the transform
    res = xr.DataArray(
        data=finufft.nufft1d1(time, data, modeNum, **kwargs),
        dims=['freq'],
        coords={
            "freq":np.linspace(-freqUpLim/2, freqUpLim/2, modeNum)
        }
    )

    return res


def ifft_nutou(dataArray, modeNum):
    pass