// ----------------------------------------------------------------------------
//
//  Copyright (C) 2003-2013 Fons Adriaensen <fons@linuxaudio.org>
//    
//  This program is free software; you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation; either version 3 of the License, or
//  (at your option) any later version.
//
//  This program 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
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program.  If not, see <http://www.gnu.org/licenses/>.
//
// ----------------------------------------------------------------------------


#define _XOPEN_SOURCE 600 // for posix_memalign()


#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "convolve.h"


#ifdef USE_SSE_AND_3DN
extern "C" void mac_sse (fftwf_complex *A, fftwf_complex *B, fftwf_complex *S, int N);
extern "C" void mac_3dn (fftwf_complex *A, fftwf_complex *B, fftwf_complex *S, int N);
#else
extern "C" void mac_sse (fftwf_complex *A, fftwf_complex *B, fftwf_complex *S, int N) {}
extern "C" void mac_3dn (fftwf_complex *A, fftwf_complex *B, fftwf_complex *S, int N) {}
#endif


Convdata::Convdata (size_t part, size_t size, unsigned int opts) :
    _refc (0), _part (part), _opts (opts), _nact (0), _buff (0), _plan (0)
{
    unsigned int i;

    if (   (part < 0x000040)
        || (part > 0x010000)
        || (part & (part - 1))
        || (size > 0x400000)
        || (size > 0x1000 * part)) abort ();

#ifndef USE_SSE_AND_3DN
    _opts = 0;
#endif

    _npar = (size + part - 1) / part; 
    _norm = 0.5f / part;
    _fftd = new fftwf_complex * [_npar]; 
    for (i = 0; i < _npar; i++) _fftd [i] = 0;
}


Convdata::~Convdata (void)
{
    unsigned int i;

    prepare_done ();
    for (i = 0; i < _npar; i++) free (_fftd [i]);
    delete[] _fftd;
}


void Convdata::prepare_part (unsigned int ipar, float gain, float *data, int step)
{
    unsigned int i;

    if (ipar >= _npar) abort ();

    if (data)
    {
        gain *= _norm;
        if (! _plan)
        { 
	    posix_memalign ((void **)(&_buff), 16, 2 * _part * sizeof (float));
            memset (_buff + _part, 0, _part * sizeof (float));  
	    _plan = fftwf_plan_dft_r2c_1d (2 * _part, _buff, _fftd [ipar], FFTW_ESTIMATE);
        }
        for (i = 0; i < _part; i++) _buff [i] = data [i * step] * gain;
        if (! _fftd [ipar]);
        {
	    posix_memalign ((void **)(_fftd + ipar), 16, (_part + 1) * sizeof (fftwf_complex));
            _nact++;
	}
        fftwf_execute_dft_r2c (_plan, _buff, _fftd [ipar]);
        if (_opts) swap (_fftd [ipar], _part);
    }
    else if (_fftd [ipar])
    {
	free (_fftd [ipar]);
        _fftd [ipar] = 0;
        _nact--;
    }
}


void Convdata::prepare_done (void)
{
    if (_plan)
    {
        fftwf_free (_buff);
	fftwf_destroy_plan (_plan);
        _buff = 0;
        _plan = 0;
    }
}


void Convdata::swap (fftwf_complex *p, size_t n)
{
    float a, b;

    while (n)
    {
	a = p [2][0];
	b = p [3][0];
        p [2][0] = p [0][1];
        p [3][0] = p [1][1];
        p [0][1] = a;
        p [1][1] = b;
	p += 4;
        n -= 4;
    }
}




Convolver::Convolver (size_t part, size_t size, unsigned int opts, unsigned int nip, unsigned int nop) :
    _part (part), _opts (opts), _nip (nip), _nop (nop)
{
    unsigned int i;

    if (   (part < 0x000040)
        || (part > 0x010000)
        || (part & (part - 1))
        || (size > 0x400000)
        || (size > 4096 * part)
        || (nip > 256)
        || (nop > 256)
        || (nip * nop > 1024)) abort ();

#ifndef USE_SSE_AND_3DN
    _opts = 0;
#endif

    _npar = (size + part - 1) / part;
    _conv = new Convdata *[_nip * _nop];
    for (i = 0; i < _nip * _nop; i++) _conv [i] = 0;
    posix_memalign ((void **)(&_ip_buff), 16, (_nip + 1) * _part * sizeof (float));
    posix_memalign ((void **)(&_oA_buff), 16, 2 * _nop * _part * sizeof (float));
    posix_memalign ((void **)(&_oB_buff), 16, 2 * _nop * _part * sizeof (float));
    posix_memalign ((void **)(&_mac_data), 16, (_part + 1) * sizeof (fftwf_complex));
    _fwd_data = new  fftwf_complex * [_nip * _npar]; 
    for (i = 0; i < _nip * _npar; i++)
    {
	posix_memalign ((void **)(_fwd_data + i), 16, (_part + 1) * sizeof (fftwf_complex));
    }
    _fwd_plan = fftwf_plan_dft_r2c_1d (2 * part, _ip_buff, _mac_data, FFTW_ESTIMATE);
    _rev_plan = fftwf_plan_dft_c2r_1d (2 * part, _mac_data, _oA_buff, FFTW_ESTIMATE);
    reset ();
}



Convolver::~Convolver (void)
{
    unsigned int i;

    fftwf_destroy_plan (_fwd_plan);
    fftwf_destroy_plan (_rev_plan);
    for (i = 0; i < _nip * _npar; i++) free (_fwd_data [i]);
    delete[] _fwd_data;
    free (_ip_buff);
    free (_oA_buff);
    free (_oB_buff);
    free (_mac_data);
    for (i = 0; i < _nip * _nop; i++)
    {
	if (_conv [i] && (_conv [i]->dec_refc () == 0)) delete _conv [i];
    }
    delete[] _conv;
}


int Convolver::set_conv (unsigned int ip, unsigned int op, Convdata *C)
{
    unsigned int i;

    if ((ip >= _nip) || (op >= _nop)) return -1;
    i = ip + op * _nip;

    if (_conv [i] && (_conv [i]->dec_refc () == 0))
    {
        delete _conv [i];
        _conv [i] = 0;
    }
    if (C)
    {
        if (C->_part != _part) return -1; 
	C->inc_refc ();
        _conv [i] = C;
    }
    return 0;
}


void Convolver::reset ()
{
    unsigned int i;

    memset (_ip_buff, 0, (_nip + 1) * _part * sizeof (float));
    memset (_oA_buff, 0,   2 * _nop * _part * sizeof (float));
    memset (_oB_buff, 0,   2 * _nop * _part * sizeof (float));
    for (i = 0; i < _nip * _npar; i++)
    {
	memset (_fwd_data [i], 0, (_part + 1) * sizeof (fftwf_complex));
    }
    _iter = 0;
    _offs = 0; 
    _op_buff = _oB_buff;
}


void Convolver::process ()
{
    unsigned int    g, h, i, j;
    float           *p1, *p2;
    fftwf_complex   *q1, *q2; 
    Convdata        *C;

    if (_offs == 0) _offs = _npar;
    _offs--;

    p1 = _ip_buff + _nip * _part;
    for (i = 0; i < _nip;  i++)
    {
	memset (p1, 0, _part * sizeof (float));
	p1 -= _part;
        q1 = _fwd_data [_offs + (_nip - 1 - i) * _npar]; 
        fftwf_execute_dft_r2c (_fwd_plan, p1, q1);
        if (_opts) Convdata::swap (q1, _part);
    }

    if (++_iter & 1)
    {  
	p1 = _oA_buff;
	p2 = _oB_buff + _part;
    }
    else 
    {
	p1 = _oB_buff;
	p2 = _oA_buff + _part;
    }
    _op_buff = p1;
    for (j = 0; j < _nop; j++)
    {
        memset (_mac_data, 0, (_part + 1) * sizeof (fftwf_complex));
	for (i = 0; i < _nip; i++)
	{
            C = _conv [i + j * _nip];
            if (C && C->_nact)
            {
                for (h = 0; h < _npar; h++)
		{
                    q2 = C->_fftd [h];
		    if (q2)
		    {
                        g = h + _offs;
                        if (g >= _npar) g -= _npar;            
                        q1 = _fwd_data [g + i * _npar];
                        if      (_opts == CONV_OPT_SSE) mac_sse (q1, q2, _mac_data, _part >> 2); 
                        else if (_opts == CONV_OPT_3DN) mac_3dn (q1, q2, _mac_data, _part >> 2);
                        else for (g = 0; g <= _part; g++)
                        {
			    _mac_data [g][0] += q1 [g][0] * q2 [g][0] - q1 [g][1] * q2 [g][1];
			    _mac_data [g][1] += q1 [g][0] * q2 [g][1] + q1 [g][1] * q2 [g][0];
			}
		    }
		}  
	    }
	}

        if (_opts) Convdata::swap (_mac_data, _part);
        fftwf_execute_dft_c2r (_rev_plan, _mac_data, p1);
        for (g = 0; g < _part; g++) *p1++ += *p2++;
        p1 += _part;
        p2 += _part;
    }
}
