/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * Geophysical Computational Tools & Library (GCTL)
 *
 * Copyright (c) 2023  Yi Zhang (yizhang-geo@zju.edu.cn)
 *
 * GCTL is distributed under a dual licensing scheme. You can redistribute 
 * it and/or modify it under the terms of the GNU Lesser General Public 
 * License as published by the Free Software Foundation, either version 2 
 * of the License, or (at your option) any later version. You should have 
 * received a copy of the GNU Lesser General Public License along with this 
 * program. If not, see <http://www.gnu.org/licenses/>.
 * 
 * If the terms and conditions of the LGPL v.2. would prevent you from using 
 * the GCTL, please consider the option to obtain a commercial license for a 
 * fee. These licenses are offered by the GCTL's original author. As a rule, 
 * licenses are provided "as-is", unlimited in time for a one time fee. Please 
 * send corresponding requests to: yizhang-geo@zju.edu.cn. Please do not forget 
 * to include some description of your company and the realm of its activities. 
 * Also add information on how to contact you by electronic and paper mail.
 ******************************************************/

#include "fft_filter.h"

#ifdef GCTL_FFTW3

gctl::fft_filter::fft_filter(){}

gctl::fft_filter::~fft_filter(){}

gctl::fft_filter::fft_filter(filting_type_e fi_type, filter_type_e fr_type, int order, double fs, double f1, double f2)
{
    init(fi_type, fr_type, order, fs, f1, f2);
}

void gctl::fft_filter::init(filting_type_e fi_type, filter_type_e fr_type, int order, double fs, double f1, double f2)
{
    if (fs <= 0 || f1 <= 0 || f2 < 0 || f1/fs >= 1)
    {
        throw std::invalid_argument("[GCTL] Invalid filtering frequency settings for gctl::fft_filter");
    }

    if (f2 > 0 && (f2/fs >= 1 || f1/fs > f2/fs))
    {
        throw std::invalid_argument("[GCTL] Invalid filtering frequency settings for gctl::fft_filter");
    }

    fi_type_ = fi_type;
    fr_type_ = fr_type;
    odr_ = order;
    fs_ = fs;

    if (fi_type_ == LowPass) fl_ = f1;
    else if (fi_type_ == HighPass) fh_ = f1;
    else if (fi_type_ == BandPass) {fl_ = f1, fh_ = f2;}
    else if (fi_type_ == BandStop) {fl_ = f1, fh_ = f2;}
    else throw std::runtime_error("[GCTL] Unkown filting type for gctl::fft_filter");

    return;
}

void gctl::fft_filter::filter(const array<double> &in, array<double> &out)
{
    cosine_extrapolate_1d(in, din_);

    size_t is = in.size();
    size_t dis = din_.size();

    int l_num = floor((dis - is)/2.0);

    dft_r2c_1d(din_, freqs_, fs_, &fid_);

    double re, im, f;
    if (fi_type_ == LowPass)
    {
        for (size_t i = 0; i < freqs_.size(); i++)
        {
            if (fr_type_ == Butterworth) f = 1.0/(1.0 + pow(fid_[i]/fl_, 2*odr_));
            else if (fr_type_ == Gaussian) f = exp(-0.5*fid_[i]*fid_[i]/(fl_*fl_));

            freqs_[i].real(f*freqs_[i].real());
            freqs_[i].imag(f*freqs_[i].imag());
        }
    }
    else if (fi_type_ == HighPass)
    {
        for (size_t i = 0; i < freqs_.size(); i++)
        {
            if (fr_type_ == Butterworth) f = 1.0/(1.0 + pow(fh_/fid_[i], 2*odr_));
            else if (fr_type_ == Gaussian) f = 1.0 - exp(-0.5*fid_[i]*fid_[i]/(fh_*fh_));

            freqs_[i].real(f*freqs_[i].real());
            freqs_[i].imag(f*freqs_[i].imag());
        }
    }
    else if (fi_type_ == BandPass)
    {
        for (size_t i = 0; i < freqs_.size(); i++)
        {
            if (fr_type_ == Butterworth) f = std::max(1.0/(1.0 + pow(fid_[i]/fl_, 2*odr_)), 1.0/(1.0 + pow(fh_/fid_[i], 2*odr_)));
            else if (fr_type_ == Gaussian) f = std::max(exp(-0.5*fid_[i]*fid_[i]/(fl_*fl_)), 1.0 - exp(-0.5*fid_[i]*fid_[i]/(fh_*fh_)));

            f = 1.0 - f;

            freqs_[i].real(f*freqs_[i].real());
            freqs_[i].imag(f*freqs_[i].imag());
        }
    }
    else if (fi_type_ == BandStop)
    {
        for (size_t i = 0; i < freqs_.size(); i++)
        {
            if (fr_type_ == Butterworth) f = std::max(1.0/(1.0 + pow(fid_[i]/fl_, 2*odr_)), 1.0/(1.0 + pow(fh_/fid_[i], 2*odr_)));
            else if (fr_type_ == Gaussian) f = std::max(exp(-0.5*fid_[i]*fid_[i]/(fl_*fl_)), 1.0 - exp(-0.5*fid_[i]*fid_[i]/(fh_*fh_)));

            freqs_[i].real(f*freqs_[i].real());
            freqs_[i].imag(f*freqs_[i].imag());
        }
    }
    else throw std::runtime_error("[GCTL] Unkown filting type for gctl::fft_filter");

    dft_c2r_1d(freqs_, dout_);

    out.resize(is);
    for (size_t i = 0; i < is; i++)
    {
        out[i] = dout_[l_num + i];
    }
    return;
}

#endif // GCTL_FFTW3