# -*- coding: utf-8 -*-
# _simulateSNR.py
# Module providing the simulateSNR function
# Copyright 2013 Giuseppe Venturini
# This file is part of python-deltasigma.
#
# python-deltasigma is a 1:1 Python replacement of Richard Schreier's
# MATLAB delta sigma toolbox (aka "delsigma"), upon which it is heavily based.
# The delta sigma toolbox is (c) 2009, Richard Schreier.
#
# python-deltasigma is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# LICENSE file for the licensing terms.

"""Module providing the simulateSNR() function
"""

from __future__ import division, print_function

import collections
from warnings import warn

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

from ._calculateSNR import calculateSNR
from ._mapQtoR import mapQtoR
from ._simulateDSM import simulateDSM
from ._simulateQDSM import simulateQDSM
from ._utils import _get_zpk


def simulateSNR(arg1, osr, amp=None, f0=0, nlev=2, f=None, k=13,
                quadrature=False):
    """Determine the SNR for a delta-sigma modulator by using simulations.

    Simulate a delta-sigma modulator with sine wave inputs of various
    amplitudes and calculate the signal-to-noise ratio (SNR) in dB for each 
    input.

    Three alternative descriptions of the modulator can be used:

     * The modulator is described by a noise transfer function (NTF), provided
       as ``arg1`` and the number of quantizer levels (``nlev``).

     * Alternatively, the first argument to simulateSNR may be an ABCD matrix. 

     * Lastly, ``arg1`` may be a function taking the input signal as its
       sole argument.

    The band of interest is defined by the oversampling ratio (``osr``)
    and the center frequency (``f0``).

    The input signal is characterized by the ``amp`` vector and the ``f`` variable.
    A default value for ``amp`` is used if not supplied.

    ``f`` is the input frequency, normalized such that 1 -> fs;
    ``f`` is rounded to an FFT bin.

    Using sine waves located in FFT bins, the SNR is calculated as the ratio
    of the sine wave power to the power in all in-band bins other than those
    associated with the input tone. Due to spectral smearing, the input tone
    is not allowed to lie in bins 0 or 1. The length of the FFT is :math:`2^k`.

    If the NTF is complex, :func:`simulateQDSM` (which is slow, also available
    in a future release) is called.

    If ABCD is complex, :func:`simulateDSM` is used with the real equivalent
    of ABCD in order to speed up simulations.

    Future versions may accommodate STFs.

    **Parameters:**

    arg1 : scipy 'lti' object, or ndarray
        The first argument may be one of the various supported representations
        for an NTF or an ABCD matrix. Quadrature modulators are supported both
        with quadrature TFs or with quadrature ABCD matrices, the latter being
        the recommended description.

    osr : int
        The over-sampling ratio.

    amp : sequence, optional
        The amplitudes in dB, referred to the FS, for which the SNR is to be 
        evaluated. ``amp`` defaults to [-120 -110...-20 -15 -10 -9 -8 ... 0]dB,
        where 0 dB means a full-scale (peak value = nlev-1) sine wave.

    f0 : float, optional
        The center frequency. Normalized. Defaults to 0.

    nlev : int, optional
        Number of quantizer levels, defaults to 2.

    f : float, optional
        Test signal input frequency. Normalized. Rounded to an FFT bin.
        Defaults to:

    .. math::

        f = \\frac{1}{4\ \mathrm{OSR}}

    k : int, optional
        The number of samples used to compute the FFT is set by the integer `k`
        - default value 13 - through:

    .. math::

        N_{samples} = 2^k

    quadrature : boolean, optional
        Whether the delta sigma modulator is a quadrature modulator or not.
        Defaults to ``False``.

    **Returns:**

    snr : ndarray
        The SNR, from simulation.

    amp : ndarray
        The amplitudes corresponding to the SNR values.

    .. rubric:: Example:

    Compare the SNR vs input amplitude curve for a fifth-order modulator, as 
    determined by the describing function method (:func:`predictSNR`) with
    that determined by simulation (:func:`simulateSNR`).::

        import pylab as plt
        from deltasigma import *
        OSR = 32
        H = synthesizeNTF(5, OSR, 1)
        snr_pred, amp, _, _, _ = predictSNR(H,OSR)
        snr, amp = simulateSNR(H, OSR)
        plt.plot(amp, snr_pred, 'b', label='Predicted')
        plt.hold(True)
        plt.plot(amp, snr, 'go', label='Simulated')
        plt.grid(True)
        figureMagic([-100, 0], 10, None,
                    [0, 100], 10, None)
        plt.xlabel('Input Level, dB')
        plt.ylabel('SNR, dB')
        s = 'peak SNR = %4.1fdB\\n' % max(snr)
        plt.text(-65, 15, s, horizontalalignment='left')
        plt.legend(loc='best')

    .. plot::

        import pylab as plt
        from deltasigma import *
        OSR = 32
        H = synthesizeNTF(5, OSR, 1)
        snr_pred, amp, _, _, _ = predictSNR(H,OSR)
        snr, amp = simulateSNR(H, OSR)
        plt.plot(amp, snr_pred, 'b', label='Predicted')
        plt.hold(True)
        plt.plot(amp, snr, 'go', label='Simulated')
        plt.grid(True)
        figureMagic([-100, 0], 10, None,
                    [0, 100], 10, None)
        plt.xlabel('Input Level, dB')
        plt.ylabel('SNR, dB')
        s = 'peak SNR = %4.1fdB\\n' % max(snr)
        plt.text(-65, 15, s, horizontalalignment='left')
        plt.legend(loc='best')

    """
    # Look at arg1 and decide if the system is quadrature
    quadrature_ntf = False
    if callable(arg1):
        raise ValueError('There is no support for NTFs described through ' +
                         'a function.')
    elif not isinstance(arg1, np.ndarray):
        # NTF of LTI type or zpk tuple, or (num,den) etc...
        # in this case, if the modulator is in quadrature, 
        # we will have to call simulateQDSM
        for roots in _get_zpk(arg1)[:2]:
            if np.any(np.abs(np.imag(np.poly(roots))) > 0.0001):
                quadrature = True
                quadrature_ntf = True
    else: # ABCD matrix, no matter what, here we will get through the
        # simulation with simulateDSM()
        arg1 = np.real_if_close(arg1) # tiny imaginary parts due to rounding
        if not np.all(np.imag(arg1) == 0):
            quadrature = True

    if amp is None:
        amp = np.concatenate((np.arange(- 120, -20 + 1, 10),
                              np.array((-15,)),
                              np.arange(-10, 1)))
    elif not isinstance(amp, collections.Iterable):
        amp = np.array((amp, ))
    else:
        amp = np.asarray(amp)
    osr_mult = 2
    if f0 != 0 and not quadrature:
        osr_mult = 2*osr_mult
    if f is None or np.isnan(f):
        f = f0 + 0.5/(osr*osr_mult) # Halfway across the band
    M = nlev - 1
    if quadrature and not quadrature_ntf:
        # Modify arg1 (ABCD) and nlev so that simulateDSM can be used
        nlev = np.array([nlev, nlev])
        arg1 = mapQtoR(arg1)
    if abs(f - f0) > 1./(osr*osr_mult):
        warn('The input tone is out-of-band.')
    N = 2**k
    if N < 8*2*osr:
        warn('Increasing k to accommodate a large oversampling ratio.')
        k = np.array(np.ceil(np.log2(8*2*osr)), dtype=np.int64)
        N = 2**k
    F = np.round(f*N)
    if np.abs(F) <= 1:
        warn('Increasing k to accommodate a low input frequency.')
        # Want f*N >= 1
        k = np.ceil(np.log2(1./f))
        N = 2**k
        F = 2
    Ntransient = 100
    soft_start = 0.5*(1 - np.cos(2*np.pi/Ntransient * \
                                 np.arange(Ntransient/2)))
    if not quadrature:
        tone = M*np.sin(2*np.pi*F/N*np.arange(N + Ntransient))
        tone[:Ntransient/2] = tone[:Ntransient/2] * soft_start
    else:
        tone = M*np.exp(2j*np.pi*F/N * np.arange(N + Ntransient))
        tone[:Ntransient/2] = tone[:Ntransient/2] * soft_start
        if not quadrature_ntf:
            tone = tone.reshape((1, -1))
            tone = np.vstack((np.real(tone), np.imag(tone)))
    # create a Hann window
    window = 0.5*(1 - np.cos(2*np.pi*np.arange(N)/N))
    if f0 == 0:
        # Exclude DC and its adjacent bin
        inBandBins = int(N/2) + np.arange(3,
                                     np.round(N/osr_mult/osr) + 1,
                                     dtype=np.int32)
        F = F - 2
    else:
        f1 = np.round(N*(f0 - 1./osr_mult/osr))
        # Should exclude DC
        inBandBins = int(N/2) + np.arange(f1,
                                     np.round(N*(f0 + 1./osr_mult/osr)) + 1,
                                     dtype=np.int32)
        F = F - f1 + 1
    snr = np.zeros(amp.shape)
    for i, A in enumerate(np.power(10.0, amp/20)):
        if quadrature_ntf:
            v, _, _, _ = simulateQDSM(A*tone, arg1, nlev)
        else:
            v, _, _, _ = simulateDSM(A*tone, arg1, nlev)
            if quadrature:
                v = v[0, :] + 1j*v[1, :]
        hwfft = fftshift(fft(window*v[Ntransient:N + Ntransient]))
        snr[i] = calculateSNR(hwfft[inBandBins - 1], F)
    return snr, amp

