/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * 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 "../lib/core.h"
#include "../lib/io.h"
#include "../lib/algorithm.h"

using namespace gctl;

double module(std::complex<double> a)
{
    return sqrt(a.real()*a.real() + a.imag()*a.imag());
}

double sin_wave(double t, double f, double phi = 0, double A = 1.0)
{
    return A*sin(2*GCTL_Pi*f*t + phi);
}

void average_radian_spec(const array<double> &in_freq, const array<double> &in_power, 
    const array<double> &out_freq, array<double> &out_power, array<double> &out_std)
{
    int in_num = in_freq.size();
    int out_num = out_freq.size();
    out_power.resize(out_num, 0);
    out_std.resize(out_num, 0);

    std::vector<double> bin;
    _1d_array bin_arr;
    for (size_t i = 0; i < out_num; i++)
    {
        bin.clear();
        for (size_t j = 0; j < in_num; j++)
        {
            if (in_freq[j] > out_freq[i] - 0.5 && in_freq[j] <= out_freq[i] + 0.5)
            {
                bin.push_back(in_power[j]);
            }
        }

        if (!bin.empty())
        {
            bin_arr.import_vector(bin);
            out_power[i] = bin_arr.mean();
            out_std[i] = bin_arr.std();
        }
    }
    return;
}

int main(int argc, char const *argv[]) try
{
    double T = 2.0; // 采样时间
    double dT = 0.01; // 采样间隔
    double freq = 1.0/dT; // 采样频率

    int M = round(T/dT) + 1, N = round(T/dT) + 1;

    _2d_matrix sig(M, N);
    _1d_array X(N), Y(M);
    _1d_array U(N), V(M);
    X.sequence(-1.0, dT);
    Y.sequence(-1.0, dT);
    U.sequence(-1.0*freq, 2.0*freq/(N - 1));
    V.sequence(-1.0*freq, 2.0*freq/(M - 1));

    double t;
    double f1 = 40.0;
    double f2 = 20.0;
    double f3 = 10.0;
    for (size_t i = 0; i < M; i++)
    {
        for (size_t j = 0; j < N; j++)
        {
            t = sqrt(X[j]*X[j] + Y[i]*Y[i]);
            sig[i][j] = pow(-1, i+j)*(sin_wave(t, f1) + sin_wave(t, f2) + sin_wave(t, f3));
        }
    }

    _1cd_array spectrum;
    dft2d(sig, spectrum);

    _1d_array spec_real(M*N), spec_imag(M*N), spec_power(M*N);
    for (size_t i = 0; i < M*N; i++)
    {
        spec_real[i] = spectrum[i].real();
        spec_imag[i] = spectrum[i].imag();
        spec_power[i] = module(spectrum[i]);
    }

    save_netcdf_grid("fft2d_sig", sig, U, V, "u", "v");
    append_netcdf_grid("fft2d_sig", spec_real, "u", "v", "real");
    append_netcdf_grid("fft2d_sig", spec_imag, "u", "v", "imag");
    append_netcdf_grid("fft2d_sig", spec_power, "u", "v", "power");

    _1d_array pro_freq(M*N), pro_spec(M*N);
    for (size_t i = 0; i < M; i++)
    {
        for (size_t j = 0; j < N; j++)
        {
            pro_freq[j + i*N] = sqrt(U[j]*U[j] + V[i]*V[i]);
            pro_spec[j + i*N] = spec_power[j + i*N];
        }
    }

    _1d_array out_freq, out_power, out_std;
    out_freq.resize(150, 1.0, 1.0);
    average_radian_spec(pro_freq, pro_spec, out_freq, out_power, out_std);

    std::ofstream ofile;
    open_outfile(ofile, "fft2d_radian_spec", ".txt");
    for (size_t i = 0; i < out_freq.size(); i++)
    {
        ofile << out_freq[i] << " " <<  out_power[i] << " " << out_std[i] << "\n";
    }
    return 0;
}
catch(std::exception &e)
{
    GCTL_ShowWhatError(e.what(), GCTL_ERROR_ERROR, 0, 0, 0);
}