#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 15 19:45:33 2018

@author: anne
"""
import numpy as np
from lasp import AvPowerSpectra, Window

def test_aps1():

    nfft = 16384
    fs = 48000
    tend = 10
    t = np.linspace(0, (tend*fs)//fs, int(fs*tend), endpoint=False)
    
    w = Window.WindowType.Hann
    # w = Window.WindowType.Rectangular
    aps = AvPowerSpectra(nfft, w, 50, -1)
    
    sig = np.random.randn(int(fs*tend))
    
    cps = aps.compute(sig[:,None])
    
    pow1 = np.sum(sig**2)/sig.size
    pow2 = np.sum((cps[:,0,0]).real)

    # Check for Parseval
    assert np.isclose(pow2 - pow1,0, atol=1e-2)

def test_aps2():
    """
    Test whether we are able to estimate a simple transfer function of a gain
    of 2 between two channels
    """

    nfft = 16384
    fs = 48000
    tend = 10
    t = np.linspace(0, (tend*fs)//fs, int(fs*tend), endpoint=False)
    
    # w = Window.WindowType.Hann
    w = Window.WindowType.Rectangular
    aps = AvPowerSpectra(nfft, w, 50, -1)
    
    sig1 = np.random.randn(int(fs*tend))
    sig2 = 2*sig1

    sig = np.concatenate((sig1[None,:], sig2[None,:])).T
    
    cps = aps.compute(sig)

    H = cps[:,0,1]/cps[:,0,0]

    # Check if transfer function is obtained
    assert np.all(np.isclose(H,2))

if __name__ == '__main__':
    test_aps1()