# Author: Eric Bezzam
# Date: July 15, 2016

import numpy as np
from .music import MUSIC
from scipy.linalg import svdvals

from scipy import linalg


class TOPS(MUSIC):
    """
    Class to apply Test of Orthogonality of Projected Subspaces [TOPS]_ for 
    Direction of Arrival (DoA) estimation.

    .. note:: Run locate_source() to apply the TOPS algorithm.

    Parameters
    ----------
    L: numpy array
        Microphone array positions. Each column should correspond to the 
        cartesian coordinates of a single microphone.
    fs: float
        Sampling frequency.
    nfft: int
        FFT length.
    c: float
        Speed of sound. Default: 343 m/s
    num_src: int
        Number of sources to detect. Default: 1
    mode: str
        'far' or 'near' for far-field or near-field detection 
        respectively. Default: 'far'
    r: numpy array
        Candidate distances from the origin. Default: np.ones(1)
    azimuth: numpy array
        Candidate azimuth angles (in radians) with respect to x-axis.
        Default: np.linspace(-180.,180.,30)*np.pi/180
    colatitude: numpy array
        Candidate elevation angles (in radians) with respect to z-axis.
        Default is x-y plane search: np.pi/2*np.ones(1)

    References
    ----------

    .. [TOPS] Y. Yeo-Sun, L. M. Kaplan, J. H. McClellan, *TOPS: New DOA estimator for wideband 
        signals*, IEEE Trans. Signal Process., Vol. 54, Num 6., pp 1977--1989, 2006

    """

    def __init__(
        self,
        L,
        fs,
        nfft,
        c=343.0,
        num_src=1,
        mode="far",
        r=None,
        azimuth=None,
        colatitude=None,
        **kwargs
    ):

        MUSIC.__init__(
            self,
            L=L,
            fs=fs,
            nfft=nfft,
            c=c,
            num_src=num_src,
            mode=mode,
            r=r,
            azimuth=azimuth,
            colatitude=colatitude,
            **kwargs
        )

    def _process(self, X):
        """
        Perform TOPS for a given frame in order to estimate steered response 
        spectrum.
        """

        # need more than 1 frequency band
        if self.num_freq < 2:
            raise ValueError("Need more than one frequency band!")

        # select reference frequency (largest power)
        max_bin = np.argmax(
            np.sum(np.sum(abs(X[:, self.freq_bins, :]), axis=0), axis=1)
        )
        f0 = self.freq_bins[max_bin]
        freq = list(self.freq_bins)
        freq.remove(f0)

        # compute empirical cross correlation matrices
        C_hat = self._compute_correlation_matrices(X)

        # compute signal and noise subspace for each frequency band
        F = np.zeros((self.num_freq, self.M, self.num_src), dtype="complex64")
        W = np.zeros((self.num_freq, self.M, self.M - self.num_src), dtype="complex64")
        for k in range(self.num_freq):
            # subspace decomposition
            F[k, :, :], W[k, :, :], ws, wn = self._subspace_decomposition(
                C_hat[k, :, :]
            )

        # create transformation matrix
        f = (
            1.0
            / self.nfft
            / self.c
            * 1j
            * 2
            * np.pi
            * self.fs
            * (np.linspace(0, self.nfft // 2, self.nfft // 2 + 1) - f0)
        )

        Phi = np.zeros((len(f), self.M, self.grid.n_points), dtype="complex64")

        for n in range(self.grid.n_points):
            p_s = self.grid.cartesian[: self.grid.dim, n]
            proj = np.dot(p_s, self.L[: self.grid.dim, :])
            for m in range(self.M):
                Phi[:, m, n] = np.exp(f * proj[m])

        # determine direction using rank test
        for n in range(self.grid.n_points):
            # form D matrix
            D = np.zeros(
                (self.num_src, (self.M - self.num_src) * (self.num_freq - 1)),
                dtype="complex64",
            )
            for k in range(self.num_freq - 1):
                Uk = np.conjugate(np.dot(np.diag(Phi[k, :, n]), F[max_bin, :, :])).T
                # F[max_bin,:,:])).T
                idx = range(
                    k * (self.M - self.num_src), (k + 1) * (self.M - self.num_src)
                )
                D[:, idx] = np.dot(Uk, W[k, :, :])
            s = svdvals(D)
            self.grid.values[n] = 1.0 / s[-1]
