# Path: thrifty/carrier_sync.py
# class Synchronizer(object):
# class DefaultSynchronizer(Synchronizer):
#     def __init__(self, detector, interpolator, shifter):
#     def sync(self, signal):
#     def __call__(self, signal):
#     def __init__(self, thresh_coeffs, window, block_len, carrier_len):
#     def detect(self, fft_mag):
# def dirichlet_kernel(xdata, block_len, carrier_len):
# def dirichlet_weights(filter_len, block_len, carrier_len):
# def make_dirichlet_interpolator(block_len, carrier_len,
#                                 width=6, return_amplitude=False):
#     def _fit_model(xdata, amplitude, time_offset):
#     def _interpolator(fft_mag, peak_idx):
# def parabolic_interpolator(fft_mag, peak_idx):
# def make_polyfit_interpolator(width):
#     def _interpolator(fft_mag, peak_idx):
# def freq_shift(signal, shift):
# def freq_shift_integer(signal, shift):
#     N, W = block_len, carrier_len
# 
# Path: thrifty/signal_utils.py
# class Signal(np.ndarray):
#     """Representation of signal with caching and on-demand calculations."""
#     # Based on https://docs.scipy.org/doc/numpy/user/basics.subclassing.html
# 
#     def __new__(cls, input_array, fft=None, ifft=None):
#         obj = np.asarray(input_array).view(cls)
# 
#         if fft is not None:
#             obj._fft = fft
#             fft._ifft = obj
# 
#         if ifft is not None:
#             obj._ifft = ifft
#             ifft._fft = obj
# 
#         return obj
# 
#     def __array_finalize__(self, obj):
#         if obj is None:
#             return
#         self._reset()
# 
#     def __array_wrap__(self, obj, context=None):
#         # http://stackoverflow.com/questions/19223926/
# 
#         if obj.shape == ():
#             return obj[()]  # if ufunc output is scalar, return it
#         else:
#             ret = np.ndarray.__array_wrap__(self, obj, context)
#             if isinstance(ret, Signal):
#                 ret._reset()
#             return ret
# 
#     def _reset(self):
#         self._fft = None
#         self._ifft = None
#         self._rms = None
#         self._mag = None
#         self._power = None
#         self._conj = None
# 
#     @property
#     def fft(self):
#         """FFT of signal."""
#         if self._fft is None:
#             self._fft = Signal(compute_fft(self), ifft=self)
#         return self._fft
# 
#     @property
#     def ifft(self):
#         """IFFT of signal."""
#         if self._ifft is None:
#             self._ifft = Signal(compute_ifft(self), fft=self)
#         return self._ifft
# 
#     @property
#     def mag(self):
#         """Calculate the magnitude of the samples."""
#         if self._mag is None:
#             if self._power is not None:
#                 self._mag = np.sqrt(self._power)
#             else:
#                 self._mag = np.abs(self)  # TODO: Convert to Signal?
#         return self._mag
# 
#     @property
#     def power(self):
#         """Calculate the power (magnitude^2) of the samples."""
#         if self._power is None:
#             # Alternative: self._power = (samples * np.conj(samples)).real
#             # or: self._power = np.abs(self)**2
#             self._power = self.mag**2
#         return self._power
# 
#     @property
#     def rms(self):
#         """Calculate signal's root mean square value."""
#         if self._rms is None:
#             if self._fft is not None and self._fft._rms is not None:
#                 self._rms = self._fft._rms / np.sqrt(len(self))
#             elif self._ifft is not None and self._ifft._rms is not None:
#                 self._rms = self._ifft._rms * np.sqrt(len(self))
#             else:
#                 self._rms = np.sqrt(np.mean(self.power))
#         return self._rms
# 
#     @property
#     def conj(self):
#         """Calculate and cache the complex conjugate."""
#         if self._conj is None:
#             self._conj = np.ndarray.conj(self)
#             self._conj._conj = self
#         return self._conj


import pytest
import numpy as np
from thrifty import carrier_sync
from thrifty.signal_utils import Signal
    # plt.plot(np.abs(signal_fft), label="Signal")
    # plt.plot(np.abs(expected_fft), label="Expected")
    # plt.plot(np.abs(got), label="Got")
    # plt.legend()
    # plt.show()

    np.testing.assert_allclose(np.abs(got), np.abs(expected_fft),
                               atol=1e-6, rtol=1e-6)


def test_dirichlet_kernel():
    """Test dirichlet_kernel with specific parameters."""
    expected = np.array([-0.1711, 0.0164, 0.3164, 0.6468, 0.9034, 1.,
                         0.9034, 0.6468, 0.3164, 0.0164, -0.1711])
    got = carrier_sync.dirichlet_kernel(np.arange(-5, 6), 8192, 2015)
    np.testing.assert_allclose(got, expected, rtol=2e-3)


INTERPOLATOR_OFFSETS = [-0.51, -0.5, -0.25, -0.1263, -0.1, 0.,
                        0.001, 0.2, 0.4995, 0.56]


@pytest.mark.parametrize("offset", INTERPOLATOR_OFFSETS)
def test_dirichlet_interpolator(offset):
    """Test Dirichlet interpolator with signal with different freq shifts."""
    peak_idx, width, block_len, carrier_len = 10, 6, 8192, 2024
    freq = (1.*offset+peak_idx)*carrier_len/block_len
    carrier = np.exp(2j*np.pi*np.arange(carrier_len)/carrier_len*freq)
    signal = np.concatenate([carrier, np.zeros(block_len-carrier_len)])
    signal_fft = np.abs(np.fft.fft(signal))

# Path: thrifty/carrier_sync.py
# class Synchronizer(object):
# class DefaultSynchronizer(Synchronizer):
#     def __init__(self, detector, interpolator, shifter):
#     def sync(self, signal):
#     def __call__(self, signal):
#     def __init__(self, thresh_coeffs, window, block_len, carrier_len):
#     def detect(self, fft_mag):
# def dirichlet_kernel(xdata, block_len, carrier_len):
# def dirichlet_weights(filter_len, block_len, carrier_len):
# def make_dirichlet_interpolator(block_len, carrier_len,
#                                 width=6, return_amplitude=False):
#     def _fit_model(xdata, amplitude, time_offset):
#     def _interpolator(fft_mag, peak_idx):
# def parabolic_interpolator(fft_mag, peak_idx):
# def make_polyfit_interpolator(width):
#     def _interpolator(fft_mag, peak_idx):
# def freq_shift(signal, shift):
# def freq_shift_integer(signal, shift):
#     N, W = block_len, carrier_len
# 
# Path: thrifty/signal_utils.py
# class Signal(np.ndarray):
#     """Representation of signal with caching and on-demand calculations."""
#     # Based on https://docs.scipy.org/doc/numpy/user/basics.subclassing.html
# 
#     def __new__(cls, input_array, fft=None, ifft=None):
#         obj = np.asarray(input_array).view(cls)
# 
#         if fft is not None:
#             obj._fft = fft
#             fft._ifft = obj
# 
#         if ifft is not None:
#             obj._ifft = ifft
#             ifft._fft = obj
# 
#         return obj
# 
#     def __array_finalize__(self, obj):
#         if obj is None:
#             return
#         self._reset()
# 
#     def __array_wrap__(self, obj, context=None):
#         # http://stackoverflow.com/questions/19223926/
# 
#         if obj.shape == ():
#             return obj[()]  # if ufunc output is scalar, return it
#         else:
#             ret = np.ndarray.__array_wrap__(self, obj, context)
#             if isinstance(ret, Signal):
#                 ret._reset()
#             return ret
# 
#     def _reset(self):
#         self._fft = None
#         self._ifft = None
#         self._rms = None
#         self._mag = None
#         self._power = None
#         self._conj = None
# 
#     @property
#     def fft(self):
#         """FFT of signal."""
#         if self._fft is None:
#             self._fft = Signal(compute_fft(self), ifft=self)
#         return self._fft
# 
#     @property
#     def ifft(self):
#         """IFFT of signal."""
#         if self._ifft is None:
#             self._ifft = Signal(compute_ifft(self), fft=self)
#         return self._ifft
# 
#     @property
#     def mag(self):
#         """Calculate the magnitude of the samples."""
#         if self._mag is None:
#             if self._power is not None:
#                 self._mag = np.sqrt(self._power)
#             else:
#                 self._mag = np.abs(self)  # TODO: Convert to Signal?
#         return self._mag
# 
#     @property
#     def power(self):
#         """Calculate the power (magnitude^2) of the samples."""
#         if self._power is None:
#             # Alternative: self._power = (samples * np.conj(samples)).real
#             # or: self._power = np.abs(self)**2
#             self._power = self.mag**2
#         return self._power
# 
#     @property
#     def rms(self):
#         """Calculate signal's root mean square value."""
#         if self._rms is None:
#             if self._fft is not None and self._fft._rms is not None:
#                 self._rms = self._fft._rms / np.sqrt(len(self))
#             elif self._ifft is not None and self._ifft._rms is not None:
#                 self._rms = self._ifft._rms * np.sqrt(len(self))
#             else:
#                 self._rms = np.sqrt(np.mean(self.power))
#         return self._rms
# 
#     @property
#     def conj(self):
#         """Calculate and cache the complex conjugate."""
#         if self._conj is None:
#             self._conj = np.ndarray.conj(self)
#             self._conj._conj = self
#         return self._conj


# Path: tests/test_carrier_sync.py
import pytest
import numpy as np
from thrifty import carrier_sync
from thrifty.signal_utils import Signal

    # plt.plot(np.abs(signal_fft), label="Signal")
    # plt.plot(np.abs(expected_fft), label="Expected")
    # plt.plot(np.abs(got), label="Got")
    # plt.legend()
    # plt.show()

    np.testing.assert_allclose(np.abs(got), np.abs(expected_fft),
                               atol=1e-6, rtol=1e-6)


def test_dirichlet_kernel():
    """Test dirichlet_kernel with specific parameters."""
    expected = np.array([-0.1711, 0.0164, 0.3164, 0.6468, 0.9034, 1.,
                         0.9034, 0.6468, 0.3164, 0.0164, -0.1711])
    got = carrier_sync.dirichlet_kernel(np.arange(-5, 6), 8192, 2015)
    np.testing.assert_allclose(got, expected, rtol=2e-3)


INTERPOLATOR_OFFSETS = [-0.51, -0.5, -0.25, -0.1263, -0.1, 0.,
                        0.001, 0.2, 0.4995, 0.56]


@pytest.mark.parametrize("offset", INTERPOLATOR_OFFSETS)
def test_dirichlet_interpolator(offset):
    """Test Dirichlet interpolator with signal with different freq shifts."""
    peak_idx, width, block_len, carrier_len = 10, 6, 8192, 2024
    freq = (1.*offset+peak_idx)*carrier_len/block_len
    carrier = np.exp(2j*np.pi*np.arange(carrier_len)/carrier_len*freq)
    signal = np.concatenate([carrier, np.zeros(block_len-carrier_len)])
    signal_fft = np.abs(np.fft.fft(signal))

    interpolator = carrier_sync.make_dirichlet_interpolator(