/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * Geophysical Computational Tools & Library (GCTL)
 *
 * Copyright (c) 2022  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 "signal.h"

#ifdef GCTL_SEISMIC_MATHGL

gctl::SIG_PLOT::SIG_PLOT()
{
    sig_ptr = nullptr;
}

gctl::SIG_PLOT::~SIG_PLOT(){}

void gctl::SIG_PLOT::link_source(SIG_UNIT *sig, std::string filename)
{
    sig_ptr = sig; file = filename;
    return;
}

int gctl::SIG_PLOT::Draw(mglGraph *gr)
{
    gctl::array<double> x;
    x.sequence(sig_ptr->rt0, sig_ptr->delta, sig_ptr->val.size());

    mglData x_d, y_d;
    x_d.Set(x.get(), x.size());
    y_d.Set(sig_ptr->val.get(), sig_ptr->val.size());

    double mean = sig_ptr->val.mean();
    double mini = sig_ptr->val[0], maxi = sig_ptr->val[0];
    for (size_t i = 1; i < sig_ptr->val.size(); i++)
    {
        mini = std::min(mini, sig_ptr->val[i]);
        maxi = std::max(maxi, sig_ptr->val[i]);
    }

    gr->SetSize(800, 400);
    gr->SetRanges(x.front(), x.back(), mini - 0.5*(mean - mini), maxi + 0.5*(maxi - mean));
    gr->Title(sig_ptr->t0.time_str().c_str());
    gr->Box();
    gr->Axis("x");
    gr->Axis("y");
    gr->Grid("xy", "h:");
    gr->Label('x', "Time (s)", 0.0);
    gr->Plot(x_d, y_d, "B");

    if (file != "null")
    {
        std::string tp = file.substr(file.find_last_of('.') + 1);
        if (tp == "eps") gr->WriteEPS(file.c_str());
        if (tp == "png") gr->WritePNG(file.c_str());
        if (tp == "jpg") gr->WriteJPEG(file.c_str());
    }
    return 0;
}

#endif // GCTL_SEISMIC_MATHGL

void gctl::SIG_UNIT::init(int len, double init_val, double dt, UTC_TIME st, double rst)
{
    val.resize(len, init_val);
    delta = dt;
    t0 = st;
    te = st; te.add_duration((len-1)*delta);
    rt0 = rst;
    return;
}

void gctl::SIG_UNIT::info()
{
    std::clog << "SIG_UNIT Info\n-------------\n";
    std::clog << "T0: " << t0.time_str() << ", TE: " << te.time_str() << ", Rel-T: " << rt0 << " s\n";
    std::clog << "Delta: " << delta << " s, Duration: " << (val.size()-1)*delta << " s, Singal-Len: " << val.size() << "\n";
    std::clog << "-------------\n";
    return;
}

void gctl::SIG_UNIT::save(std::string filename)
{
    std::ofstream outfile;
    open_outfile(outfile, filename, ".txt");

    outfile << "# T0: " << t0.time_str() << "\n";
    outfile << "# Rel-T0: " << rt0 << "\n";
    outfile << "# Delta: " << delta << "\n";
    outfile << "# Length: " << val.size() << "\n";
    for (size_t i = 0; i < val.size(); i++)
    {
        outfile << rt0 + i*delta << " " << val[i] << "\n";
    }

    outfile.close();
    return;
}

int gctl::SIG_UNIT::cut_section(SIG_UNIT &sig, double cst, double cet) const
{
    if (cet <= cst || cst < 0 || cet > (double) (val.size() - 1)*delta)
    {
        return -1;
    }

    int c = 0;
    int s_id = round(cst/delta);
    int e_id = round(cet/delta);

    sig.t0 = t0;
    sig.te = t0;
    sig.t0.add_duration(s_id*delta);
    sig.te.add_duration(e_id*delta);
    
    sig.delta = delta;
    sig.val.resize(e_id - s_id + 1);
    for (int i = s_id; i <= e_id; i++)
    {
        sig.val[c] = val[i]; c++;
    }
    return 0;
}

int gctl::SIG_UNIT::cut_section(SIG_UNIT &sig, const UTC_TIME &cst, double clen) const
{
    double stt = cst.diff_sec(t0);
    double ett = stt + clen;

    if (ett <= stt || stt < 0 || (ett - (val.size() - 1)*delta) > 0.0005)
    {
        return -1;
    }

    int c = 0;
    int s_id = round(stt/delta);
    int e_id = round(ett/delta);

    sig.t0 = cst;
    sig.te = cst;
    sig.te.add_duration(clen);
    
    sig.delta = delta;
    sig.val.resize(e_id - s_id + 1);
    for (int i = s_id; i <= e_id; i++)
    {
        sig.val[c] = val[i]; c++;
    }
    return 0;
}

int gctl::SIG_UNIT::cut_multiple_sections(time_unit roundup, double clen, array<SIG_UNIT> &sigs)
{
    UTC_TIME tmp_t, cst = t0;
    cst.round_up_to(roundup);
    tmp_t = cst;

    UTC_TIME et = t0;
    et.add_duration(delta*(val.size() - 1));

    int ss = 0;
    do
    {
        tmp_t.add_duration(clen);
        if (et.diff_sec(tmp_t) >= 0) ss++;
        else break; 
    } while (1);

    sigs.resize(ss);
    for (size_t i = 0; i < ss; i++)
    {
        if (cut_section(sigs[i], cst, clen)) return -1;
        cst.add_duration(clen);
    }
    return 0;
}

int gctl::SIG_UNIT::down_sampling(int factor)
{
    if (factor < 1)
    {
        throw std::invalid_argument("[GCTL_SEISMIC] Invalid parameter for gctl::SIG_UNIT::down_sampling(...)");
        return -1;
    }

    int s = 1 + (val.size() - 1)/factor;
    if (s <= 1) return -1;

    gctl::array<double> ret(s);
    for (size_t i = 0; i < ret.size(); i++)
    {
        ret[i] = val[factor*i];
    }

    delta *= factor;
    val = ret;

    double dur = delta*(s - 1);
    te = t0;
    te.add_duration(dur);
    return 0;
}

int gctl::SIG_UNIT::resampling(double dt, time_unit roundup)
{
    if (fabs(dt - delta) < 0.001) return 0;

    if (dt < delta)
    {
        throw std::invalid_argument("[GCTL_SEISMIC] Invalid parameter for gctl::SIG_UNIT::resampling(...)");
        return -1;
    }

    UTC_TIME st = t0;
    st.round_up_to(roundup);

    double stt = st.diff_sec(t0);
    int sid = round(stt/delta);
    int dn = round(dt/delta);

    int s = 0;
    int id = sid;
    int vnum = val.size();
    while (id < vnum)
    {
        id += dn;
        s++;
    }

    if (s < 2) return -1;

    gctl::array<double> ret(s);
    for (size_t i = 0; i < ret.size(); i++)
    {
        ret[i] = val[sid + i*dn];
    }

    val = ret;
    delta = dt;
    t0 = st;

    double dur = delta*(s - 1);
    te = t0;
    te.add_duration(dur);
    return 0;
}

void gctl::SIG_UNIT::remove_mean(int win_size)
{
    double mean;
    if (win_size <= 1)
    {
        mean = val.mean();
        for (size_t i = 0; i < val.size(); i++)
        {
            val[i] -= mean;
        }
    }
    else
    {
        double sum;
        int val_s = val.size();
        int ht = win_size/2;
        int id, cnt;
        
        for (size_t i = 0; i < val_s; i++)
        {
            sum = 0.0;
            cnt = 0;
            for (size_t j = 0; j < win_size; j++)
            {
                // 一种滑动窗口算法 窗口会延时启动（小于半个窗口时不开始滑动）并提前终止（小于半个窗口时不再滑动）
                //id = i + j - ht;
                //if (i <= ht) id = j;
                //if (val_s - i <= ht) id = val_s  - win_size + j;
                //sum += val[id];

                id = i + j - ht; 
                if (id >= 0 && id < val_s) {sum += val[id]; cnt++;}
            }

            //val[i] -= sum/win_size;
            val[i] -= sum/cnt;
        }
    }
    return;
}

void gctl::SIG_UNIT::remove_outliers(double factor)
{
    double std = val.std();
    double mean = val.mean();

    int vs = val.size();
    for (size_t i = 0; i < vs; i++)
    {
        if (fabs(val[i] - mean) > factor*std)
        {
            val[i] = mean;
        }        
    }
    return;
}

void gctl::SIG_UNIT::normalize(double norm)
{
    norm = fabs(norm);

    double maxi = val[0];
    double mini = val[0];
    for (size_t i = 1; i < val.size(); i++)
    {
        maxi = std::max(maxi, val[i]);
        mini = std::min(mini, val[i]);
    }

    for (size_t i = 0; i < val.size(); i++)
    {
        val[i] = 2.0*norm*(val[i] - mini)/(maxi - mini) - norm;
    }
    return;
}

void gctl::SIG_UNIT::normalize_1bit()
{
    for (size_t i = 0; i < val.size(); i++)
    {
        val[i] = val[i]>0?1:-1;
    }
    return;
}

void gctl::SIG_UNIT::spectral_whitening(int avg_num)
{
    _1d_array V;
    cosine_extrapolate_1d(val, V);

    int vnum = val.size();
    int Vnum = V.size();
    int l_num = (Vnum - vnum)/2;

    _1cd_array spec;
    gctl::dft_r2c_1d(V, spec);

    int f_size = spec.size();
    int ht = avg_num/2;
    //int id;
    int id, cnt;
    double sum;
    for (size_t i = 0; i < f_size; i++)
    {
        sum = 1e-8;
        cnt = 0;
        for (size_t j = 0; j < avg_num; j++)
        {
            // 一种滑动窗口算法 窗口会延时启动（小于半个窗口时不开始滑动）并提前终止（小于半个窗口时不再滑动）
            id = i + j - ht;
            if (i <= ht) id = j;
            if (f_size - i <= ht) id = f_size  - avg_num + j;
            sum += sqrt(power2(spec[id].real()) + power2(spec[id].imag()));

            //id = i + j - ht; 
            //if (id >= 0 && id < f_size) {sum += sqrt(power2(spec[id].real()) + power2(spec[id].imag())); cnt++;}
        }
        sum /= avg_num;
        //sum /= cnt;

        spec[i].real(spec[i].real()/sum);
        spec[i].imag(spec[i].imag()/sum);
    } 

    gctl::dft_c2r_1d(spec, V);
    for (size_t i = 0; i < vnum; i++)
    {
        val[i] = V[l_num + i];
    }
    return;
}

void gctl::SIG_UNIT::filter(gctl::filting_type_e fi_type, gctl::filter_type_e fr_type, int order, double f1, double f2)
{
    fft_filter ft(fi_type, fr_type, order, 1.0/delta, f1, f2);
    
    array<double> ret;
    ft.filter(val, ret);

    val = ret;
    return;
}

void gctl::SIG_UNIT::stack_sig_units(const gctl::array<SIG_UNIT*> &sigs)
{
    int ls = sigs.size();
    if (sigs.empty())
    {
        throw std::runtime_error("[GCTL_SEISMIC] Invalid stacking list for gctl::SIG_UNIT::stack_sig_units(...)");
    }
    
    for (size_t i = 0; i < ls - 1; i++)
    {
        if ((sigs[i]->t0 != sigs[i+1]->t0) || (sigs[i]->delta != sigs[i+1]->delta) || (sigs[i]->val.size() != sigs[i+1]->val.size()))
        {
            throw std::runtime_error("[GCTL_SEISMIC] Stacking singals do not match for gctl::SIG_UNIT::stack_sig_units(...)");
        }
    }

    rt0 = sigs[0]->rt0;
    t0 = sigs[0]->t0;
    te = sigs[0]->te;
    delta = sigs[0]->delta;
    val.resize(sigs[0]->val.size());
    for (size_t i = 0; i < val.size(); i++)
    {
        val[i] = 0.0;
        for (size_t j = 0; j < ls; j++)
        {
            val[i] += sigs[j]->val[i];
        }
        val[i] /= ls;
    }
    return;
}

void gctl::SIG_UNIT::stack_sig_units(const array<SIG_UNIT> &sigs)
{
    int ls = sigs.size();
    if (sigs.empty())
    {
        throw std::runtime_error("[GCTL_SEISMIC] Invalid stacking list for gctl::SIG_UNIT::stack_sig_units(...)");
    }
    
    for (size_t i = 0; i < ls - 1; i++)
    {
        if ((sigs[i].t0 != sigs[i+1].t0) || (sigs[i].delta != sigs[i+1].delta) || (sigs[i].val.size() != sigs[i+1].val.size()))
        {
            throw std::runtime_error("[GCTL_SEISMIC] Stacking singals do not match for gctl::SIG_UNIT::stack_sig_units(...)");
        }
    }

    rt0 = sigs[0].rt0;
    t0 = sigs[0].t0;
    te = sigs[0].te;
    delta = sigs[0].delta;
    val.resize(sigs[0].val.size());
    for (size_t i = 0; i < val.size(); i++)
    {
        val[i] = 0.0;
        for (size_t j = 0; j < ls; j++)
        {
            val[i] += sigs[j].val[i];
        }
        val[i] /= ls;
    }
    return;
}

void gctl::SIG_UNIT::stack_sig_units(const std::vector<SIG_UNIT> &sigs)
{
    int ls = sigs.size();
    if (sigs.empty())
    {
        throw std::runtime_error("[GCTL_SEISMIC] Invalid stacking list for gctl::SIG_UNIT::stack_sig_units(...)");
    }
    
    for (size_t i = 0; i < ls - 1; i++)
    {
        if ((sigs[i].t0 != sigs[i+1].t0) || (sigs[i].delta != sigs[i+1].delta) || (sigs[i].val.size() != sigs[i+1].val.size()))
        {
            throw std::runtime_error("[GCTL_SEISMIC] Stacking singals do not match for gctl::SIG_UNIT::stack_sig_units(...)");
        }
    }

    rt0 = sigs[0].rt0;
    t0 = sigs[0].t0;
    te = sigs[0].te;
    delta = sigs[0].delta;
    val.resize(sigs[0].val.size());
    for (size_t i = 0; i < val.size(); i++)
    {
        val[i] = 0.0;
        for (size_t j = 0; j < ls; j++)
        {
            val[i] += sigs[j].val[i];
        }
        val[i] /= ls;
    }
    return;
}

void gctl::SIG_UNIT::merge_sig_units(const array<SIG_UNIT> &sigs)
{
    if (sigs.empty()) return;

    delta = sigs[0].delta;
    for (size_t i = 1; i < sigs.size(); i++)
    {
        if (sigs[i].delta != delta)
        {
            throw std::runtime_error("[GCTL_SEISMIC] Can't merge signal units with different sampling rates. From gctl::SIG_UNIT::merge_sig_units(...)");
        }
    }

    t0 = sigs[0].t0;
    te = sigs[0].te;
    for (size_t i = 1; i < sigs.size(); i++)
    {
        if (sigs[i].t0 < t0) t0 = sigs[i].t0;
        if (te < sigs[i].te) te = sigs[i].te;
    }

    int sig_n = round(te.diff_sec(t0)/delta) + 1;
    array<int> cnt(sig_n, 0);
    
    int s_id;
    val.resize(sig_n, 0.0);
    for (size_t i = 0; i < sigs.size(); i++)
    {
        s_id = round(sigs[i].t0.diff_sec(t0)/delta);
        for (size_t n = 0; n < sigs[i].val.size(); n++)
        {
            val[s_id+n] = (val[s_id+n]*cnt[s_id+n] + sigs[i].val[n])/(cnt[s_id+n] + 1);
            cnt[s_id+n] += 1;
        }
    }
    return;
}

void gctl::SIG_UNIT::merge_sig_units(const std::vector<SIG_UNIT> &sigs)
{
    if (sigs.empty()) return;

    delta = sigs[0].delta;
    for (size_t i = 1; i < sigs.size(); i++)
    {
        if (sigs[i].delta != delta)
        {
            throw std::runtime_error("[GCTL_SEISMIC] Can't merge signal units with different sampling rates. From gctl::SIG_UNIT::merge_sig_units(...)");
        }
    }

    t0 = sigs[0].t0;
    te = sigs[0].te;
    for (size_t i = 1; i < sigs.size(); i++)
    {
        if (sigs[i].t0 < t0) t0 = sigs[i].t0;
        if (te < sigs[i].te) te = sigs[i].te;
    }

    int sig_n = round(te.diff_sec(t0)/delta) + 1;
    array<int> cnt(sig_n, 0);

    int s_id;
    val.resize(sig_n, 0.0);
    for (size_t i = 0; i < sigs.size(); i++)
    {
        s_id = round(sigs[i].t0.diff_sec(t0)/delta);
        for (size_t n = 0; n < sigs[i].val.size(); n++)
        {
            val[s_id+n] = (val[s_id+n]*cnt[s_id+n] + sigs[i].val[n])/(cnt[s_id+n] + 1);
            cnt[s_id+n] += 1;
        }
    }
    return;
}

void gctl::SIG_UNIT::linear_cross_correlation(const SIG_UNIT &a, const SIG_UNIT &b)
{
    if (a.t0 != b.t0 || a.val.size() != b.val.size() || fabs(a.delta - b.delta) > 1e-6 || fabs(a.rt0 - b.rt0) > 1e-6)
    {
        throw std::runtime_error("[GCTL_SEISMIC] Singals do not match for gctl::SIG_UNIT::linear_cross_correlation(...)");
    }

    int hs = a.val.size();

    // There is no excat start time for a result from the linear cross correlation. So it will be set to the start of year 1900.
    t0 = UTC_START;
    te = UTC_START;
    delta = a.delta;
    rt0 = -1*(hs - 1)*delta;
    val.resize(2*hs - 1);

    int s, as, bs;
    for (size_t i = 0; i < val.size(); i++)
    {
        val[i] = 0;
        s = i + 1 - 2*((i+1)%hs)*(i/hs); // dot size
        as = (i/hs)*((i+1)%hs); // start index for 'a' which is not moving
        bs = ((2*hs-2-i)/hs)*((2*hs-1-i)%hs); // start index for 'b' which is moving
        for (size_t n = 0; n < s; n++)
        {
            val[i] += a.val[as+n]*b.val[bs+n];
        }
    }
    return;
}

void gctl::SIG_UNIT::linear_cross_correlation_freq(const SIG_UNIT &a, const SIG_UNIT &b, int avg_num)
{
    _1d_array A, B;
    cosine_extrapolate_1d(a.val, A);
    cosine_extrapolate_1d(b.val, B);

    int anum = a.val.size();
    int Anum = A.size();
    int l_num = (Anum - anum)/2;

    _1cd_array a_spec, b_spec, cohere;
    gctl::dft_r2c_1d(A, a_spec);
    gctl::dft_r2c_1d(B, b_spec);

    int f_size = a_spec.size();
    int ht = avg_num/2;
    int id;
    double a_sum, b_sum;

    cohere.resize(f_size);
    for (size_t i = 0; i < f_size; i++)
    {
        a_sum = b_sum = 0.0;
        for (size_t j = 0; j < avg_num; j++)
        {
            // 一种滑动窗口算法 窗口会延时启动（小于半个窗口时不开始滑动）并提前终止（小于半个窗口时不再滑动）
            id = i + j - ht;
            if (i <= ht) id = j;
            if (f_size - i <= ht) id = f_size  - avg_num + j;

            a_sum += sqrt(power2(a_spec[id].real()) + power2(a_spec[id].imag()));
            b_sum += sqrt(power2(b_spec[id].real()) + power2(b_spec[id].imag()));
        }

        a_sum /= avg_num;
        b_sum /= avg_num;
        
        cohere[i].real(0.5*Anum*(a_spec[i].real()*b_spec[i].real() + a_spec[i].imag()*b_spec[i].imag())/(a_sum*b_sum + 1e-8));
        cohere[i].imag(0.5*Anum*(a_spec[i].imag()*b_spec[i].real() - a_spec[i].real()*b_spec[i].imag())/(a_sum*b_sum + 1e-8));
    }

    array<double> V;
    gctl::dft_c2r_1d(cohere, V);
    
    val.resize(anum);
    for (size_t i = 0; i < anum; i++)
    {
        val[i] = V[l_num + i];
    }
    return;
}

void gctl::SIG_UNIT::circular_cross_correlation(const SIG_UNIT &a, const SIG_UNIT &b)
{
    if (a.t0 != b.t0 || a.val.size() != b.val.size() || fabs(a.delta - b.delta) > 1e-6 || fabs(a.rt0 - b.rt0) > 1e-6)
    {
        throw std::runtime_error("[GCTL_SEISMIC] Singals do not match for gctl::SIG_UNIT::circular_cross_correlation(...)");
    }

    int hs = a.val.size();

    t0 = a.t0;
    te = a.te;
    rt0 = a.rt0;
    delta = a.delta;
    val.resize(hs);

    int bs;
    for (size_t i = 0; i < hs; i++)
    {
        val[i] = 0;
        bs = (hs - i)%hs;
        for (size_t n = 0; n < hs; n++)
        {
            val[i] += a.val[n]*b.val[(bs+n)%hs];
        }
    }
    return;
}

#ifdef GCTL_SEISMIC_MATHGL

int gctl::SIG_UNIT::plot(std::string filename)
{
    plt.link_source(this, filename);
    mglFLTK gr(&plt, "SIG_UNIT plot");
    return gr.Run();
}

#endif // GCTL_SEISMIC_MATHGL