""" **Description**

        Test unit models.

    **Example**

        ::

            pytest --capture=no --verbose

    **License**

        `BSD-3C. <https://github.com/larryturner/diamondback/blob/master/license>`_

        © 2018 - 2021 Larry Turner, Schneider Electric Industries SAS. All rights reserved.

    **Author**

        Larry Turner, Schneider Electric, Analytics & AI, 2018-04-03.

    **Definition**

"""

from diamondback import ComplexExponentialFilter, WindowFilter
from diamondback import ComplexTransform, FourierTransform, PowerSpectrumTransform, WaveletTransform
from diamondback import ZTransform
import math
import numpy


class Test( object ) :

    """ Test.
    """

    def test_ComplexTransform( self ) :

        """ Test ComplexTransform.
        """

        count = 32

        frequency = 0.1

        x = ComplexExponentialFilter( 0.5 ).filter( numpy.linspace( -1.0e-4, 1.0e-4, count ) + frequency )

        for ii in range( 0, 2 ) :

            y = ComplexTransform.transform( x, ii == 0 )

            assert y.shape == ( 3, count )

            assert not isinstance( y[ 0 ], complex )

            z = ComplexTransform.transform( y, ii == 0 )

            assert z.shape == ( count, )

            assert isinstance( z[ 0 ], complex )

            assert numpy.allclose( z, x )


    def test_FourierTransform( self ) :

        """ Test FourierTransform.
        """

        count = 128

        frequency = ( 0.12, 0.23 )

        b = numpy.ones( count )

        for ii in range( 0, 4 ) :

            x = numpy.random.rand( ) * ComplexExponentialFilter( numpy.random.rand( ) ).filter( numpy.linspace( frequency[ 0 ], frequency[ 1 ], count ) )

            if ( ii % 2 ) :

                x = x.real

            if ( ii >= 2 ) :

                b = WindowFilter.Factory.instance( WindowFilter, 'Hann', count - 1 ).b

            y, f = FourierTransform.transform( x, b )

            assert ( ( len( y ) == count ) and ( len( f ) == count ) )

            assert isinstance( y[ 0 ], complex )

            assert numpy.isclose( f[ 0 ], -1.0 ) and numpy.isclose( f[ -1 ], 1.0 - 2.0 / count )

            h = numpy.argmax( abs( y[ count // 2 : ] ) )

            assert ( h >= numpy.argmin( abs( f[ count // 2 : ] - frequency[ 0 ] ) ) ) and ( h <= numpy.argmin( abs( f[ count // 2 : ] - frequency[ 1 ] ) ) )

            z = FourierTransform.transform( y, b, True )[ 0 ]

            if ( ii % 2 ) :

                z = z.real

            assert numpy.allclose( z[ 1 : -1 ], x[ 1 : -1 ] )


    def test_PowerSpectrumTransform( self ) :

        """ Test PowerSpectrumTransform.
        """

        count = 128

        frequency, index = ( 0.12, 0.23 ), 64

        b = WindowFilter.Factory.instance( WindowFilter, 'Hann', count - 1 ).b

        for ii in range( 0, 2 ) :

            x = numpy.random.rand( ) * ComplexExponentialFilter( numpy.random.rand( ) ).filter( numpy.linspace( frequency[ 0 ], frequency[ 1 ], count * 8 ) )

            if ( ii > 0 ) :

                x = x.real

            y, f = PowerSpectrumTransform.transform( x, b, index )

            assert ( ( len( y ) == count ) and ( len( f ) == count ) )

            assert not isinstance( y[ 0 ], complex )

            assert numpy.isclose( f[ 0 ], -1.0 ) and numpy.isclose( f[ -1 ], 1.0 - 2.0 / count )

            h = numpy.argmax( y[ len( b ) // 2 : ] )

            assert ( h >= numpy.argmin( abs( f[ count // 2 : ] - frequency[ 0 ] ) ) ) and ( h <= numpy.argmin( abs( f[ count // 2 : ] - frequency[ 1 ] ) ) )


    def test_WaveletTransform( self ) :

        """ Test WaveletTransform.
        """

        count = 3

        h = numpy.array( [  0.235233603892082,  0.570558457915722, 0.325182500263116, -0.095467207784164,
                           -0.060416104155198, 0.024908749868442 ] )

        g = numpy.array( [ -0.024908749868442, -0.060416104155198, 0.095467207784164,  0.325182500263116,
                           -0.570558457915722, 0.235233603892082 ] )

        obj = WaveletTransform.Factory.instance( WaveletTransform, 'Daubechies', len( h ) - 1 )

        assert numpy.allclose( obj.b[ 0 ][ 0 ], h )

        assert numpy.allclose( obj.b[ 0 ][ 1 ], g )

        h = numpy.array( [ 0.024908749868442, -0.060416104155198, -0.095467207784164, 0.325182500263116,
                           0.570558457915722,  0.235233603892082 ] )

        g = numpy.array( [ 0.235233603892082, -0.570558457915722,  0.325182500263116, 0.095467207784164,
                           -0.060416104155198, -0.024908749868442 ] )

        assert numpy.allclose( obj.b[ 1 ][ 0 ], h )

        assert numpy.allclose( obj.b[ 1 ][ 1 ], g )

        h = numpy.array( [ 0.5, 0.5 ] )

        g = numpy.array( [ 0.5, -0.5 ] )

        obj = WaveletTransform.Factory.instance( WaveletTransform, 'Haar', len( h ) - 1 )

        assert numpy.allclose( obj.b[ 0 ][ 0 ], h )

        assert numpy.allclose( obj.b[ 0 ][ 1 ], g )

        h = numpy.array( [ 0.5, 0.5 ] )

        g = numpy.array( [ -0.5, 0.5 ] )

        assert numpy.allclose( obj.b[ 1 ][ 0 ], h )

        assert numpy.allclose( obj.b[ 1 ][ 1 ], g )

        u = numpy.array( [ -0.677118658500953,  0.345009188833763,  0.118882064536892, -0.537364409907213,
                           -0.250663848240361,  0.491955501906881, -0.545336874797487,  0.390418096834100,
                           -0.271946313073118,  0.047745751406263,  0.301454810260592,  0.138563567406933,
                           -0.215817816000671, -0.271946313073118,  0.047745751406264,  0.301454810260592,
                           -0.154508497187474, -0.110615871041237, -0.024471741852423,  0.071019760960103,
                            0.139384128958763,  0.154508497187474,  0.110615871041237,  0.024471741852423,
                           -0.071019760960103, -0.139384128958763, -0.154508497187474, -0.110615871041237,
                           -0.024471741852423,  0.071019760960104,  0.139384128958763,  0.154508497187474 ] )

        x = ComplexExponentialFilter( 0.5 ).filter( numpy.ones( len( u ) ) * 0.1 ).real

        y = obj.transform( x, count, False )

        assert numpy.allclose( y, u )

        z = obj.transform( y, count, True )

        assert numpy.allclose( z, x )

        x = numpy.zeros( ( len( u ), len( u ) ) )

        for ii in range( 0, x.shape[ 0 ] ) :

            x[ ii, : ] = numpy.random.rand( ) * ComplexExponentialFilter( numpy.random.rand( ) * 2.0 - 1.0 ).filter( numpy.ones( len( u ) ) * numpy.random.rand( ) * 0.5 ).real

        z = obj.transform( obj.transform( x, count, False ), count, True )

        assert numpy.allclose( z, x )


    def test_ZTransform( self ) :

        """ Test ZTransform.
        """

        frequency, ripple = 0.05, 0.125

        u = numpy.array( [  0.000000000000000, 1.868522837482727, -1.863713887857194, 0.948381518868702,
                           -0.212543935987133 ] )

        v = numpy.array( [ 0.016209591718306, 0.064838366873225, 0.097257550309837, 0.064838366873225,
                           0.01620959171830 ] )

        s = numpy.array( [ numpy.exp( 1j * math.pi * x / ( ( len( u ) - 1 ) * 2 ) ) for x in range( 1, ( len( u ) - 1 ) * 2, 2 ) ] )

        t = math.asinh( 1.0 / ( ( 10.0 ** ( 0.1 * ripple ) - 1.0 ) ** 0.5 ) ) / ( len( u ) - 1 )

        a = ( numpy.poly( ( -math.sinh( t ) * s.imag + 1j * math.cosh( t ) * s.real ) * 2.0 * math.pi ) ).real

        a /= a[ -1 ]

        b = [ 1.0 ]

        a, b = ZTransform.transform( a, b, frequency, True )

        b = numpy.poly( -numpy.ones( ( len( u ) - 1 ) ) )

        b *= ( 1.0 - sum( a ) ) / sum( b )

        assert numpy.allclose( a, u )

        assert numpy.allclose( b, v )

        frequency = 0.025

        u = numpy.array( [ 0.000000000000000, 4.913032434807922, -10.502699472790392, 12.647757190438988,
                           -9.256183138855487, 4.112750158020544, -1.0263940516247210, 0.110901278364195 ] )

        v = numpy.array( [ 0.000000000160721, 0.002515288055351, 0.103543670353901, 0.397913857193261,
                           0.290734727994417, 0.040371147586715, 0.000522947605366, 0.000000000000000 ] ) * 1.0e-3

        s, a = numpy.ones( 1 ), numpy.ones( 2 )

        for ii in range( 2, len( u ) ) :

            x = numpy.concatenate( ( s, numpy.zeros( 2 ) ) ) + numpy.concatenate( ( [ 0.0 ], ( ( 2.0 * ii ) - 1.0 ) * a ) )

            s, a = a, x

        a /= a[ -1 ]

        b = [ 1.0 ]

        a, b = ZTransform.transform( a, b, frequency, False )

        assert numpy.allclose( a, u )

        assert numpy.allclose( b, v )
