/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * 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 "xcorrelation.h"

std::string gctl::seismic::seed_info::get_sac_filename(int id)
{
    return sac_t0[id].rdseed_time_str() + "." + nwk_name + "." + stn_name + ".00."  + CHANNEL_NAME[sac_chl[id]] + ".D.SAC";
}

gctl::seismic::Xcorrelation::Xcorrelation(){}

gctl::seismic::Xcorrelation::~Xcorrelation(){}

void gctl::seismic::Xcorrelation::create_stations(std::string info_list, SIG_CHANNEL working_chl)
{
    std::ifstream infile;
    gctl::open_infile(infile, info_list, ".md");

    std::string lin_s, dum_s, wrk_dir;
    std::getline(infile, lin_s);
    gctl::parse_string_to_value(lin_s, ' ', false, dum_s, dum_s, wrk_dir);

    seed_info sd;
    int tmp_len;
    double tmp_freq;
    UTC_TIME tmp_t;
    std::string t0_str, te_str, chl_str;
    while (std::getline(infile, lin_s))
    {
        if (lin_s == "") continue;

        sd.sac_t0.clear();
        sd.sac_te.clear();
        sd.sac_freq.clear();
        sd.sac_len.clear();
        sd.sac_chl.clear();

        gctl::parse_string_to_value(lin_s, ' ', false, dum_s, dum_s, sd.dir_file);
        sd.dir_file = wrk_dir + "/" + sd.dir_file;
        
        std::getline(infile, lin_s);
        gctl::parse_string_to_value(lin_s, ' ', false, dum_s, dum_s, 
            sd.stn_name, sd.nwk_name, sd.lat, sd.lon, sd.elev, 
            dum_s, dum_s, dum_s, dum_s, dum_s, 
            t0_str, te_str);

        sd.stn_t0.set_rdseed_time(t0_str);
        sd.stn_te.set_rdseed_time(te_str);

        while (std::getline(infile, lin_s))
        {
            if (lin_s == "") break;

            gctl::parse_string_to_value(lin_s, ' ', false, 
                dum_s, dum_s, dum_s, dum_s, 
                chl_str, dum_s, dum_s, dum_s, 
                t0_str, te_str, tmp_freq, tmp_len);

            tmp_t.set_rdseed_time(t0_str);
            sd.sac_t0.push_back(tmp_t);

            tmp_t.set_rdseed_time(te_str);
            sd.sac_te.push_back(tmp_t);

            sd.sac_freq.push_back(tmp_freq);
            sd.sac_len.push_back(tmp_len);

            if (chl_str == "SHN") sd.sac_chl.push_back(SHN);
            else if (chl_str == "SHE") sd.sac_chl.push_back(SHE);
            else if (chl_str == "SHZ") sd.sac_chl.push_back(SHZ);
            else if (chl_str == "BHN") sd.sac_chl.push_back(BHN);
            else if (chl_str == "BHE") sd.sac_chl.push_back(BHE);
            else if (chl_str == "BHZ") sd.sac_chl.push_back(BHZ);
            else sd.sac_chl.push_back(UNKNOWN);
        }

        seeds_.push_back(sd);
    }
    infile.close();

    station tmp_st;
    fragment_info tmp_frag;
    for (size_t i = 0; i < seeds_.size(); i++)
    {
        tmp_st.lon = seeds_[i].lon;
        tmp_st.lat = seeds_[i].lat;
        tmp_st.elev = seeds_[i].elev;
        tmp_st.name = seeds_[i].stn_name;
        tmp_st.network = seeds_[i].nwk_name;

        tmp_st.frags.clear();
        for (size_t j = 0; j < seeds_[i].sac_freq.size(); j++)
        {
            tmp_frag.seed_id = i;
            tmp_frag.t0 = seeds_[i].sac_t0[j];
            tmp_frag.te = seeds_[i].sac_te[j];
            tmp_frag.len = seeds_[i].sac_len[j];
            tmp_frag.freq= seeds_[i].sac_freq[j];
            tmp_frag.chl = seeds_[i].sac_chl[j];

            if (working_chl == SBHN && (tmp_frag.chl == SHN || tmp_frag.chl == BHN)) tmp_st.frags.push_back(tmp_frag);
            else if (working_chl == SBHE && (tmp_frag.chl == SHE || tmp_frag.chl == BHE)) tmp_st.frags.push_back(tmp_frag);
            else if (working_chl == SBHZ && (tmp_frag.chl == SHZ || tmp_frag.chl == BHZ)) tmp_st.frags.push_back(tmp_frag);
            else if (tmp_frag.chl == working_chl) tmp_st.frags.push_back(tmp_frag);
        }

        if (stations_.empty()) stations_.push_back(tmp_st);
        else
        {
            bool new_stn = true;
            for (size_t s = 0; s < stations_.size(); s++)
            {
                if (tmp_st.name == stations_[s].name && tmp_st.network == stations_[s].network)
                {
                    for (size_t j = 0; j < tmp_st.frags.size(); j++)
                    {
                        stations_[s].frags.push_back(tmp_st.frags[j]);
                    }
                    
                    new_stn = false;
                    break;
                }
            }

            if (new_stn) stations_.push_back(tmp_st);
        }
    }

    // Sort fragments according to their start times
    for (size_t i = 0; i < stations_.size(); i++)
    {
        std::sort(stations_[i].frags.begin(), stations_[i].frags.end(), 
            [](const fragment_info &a, const fragment_info &b)->bool{return (a.t0 < b.t0);});
    }
    return;
}

void gctl::seismic::Xcorrelation::save_stations(std::string stn_file)
{
    std::ofstream ofile;
    open_outfile(ofile, stn_file, ".md");

    ofile << "##### seed list\n";
    ofile << "seed number: " << seeds_.size() << "\n\n";
    for (size_t i = 0; i < seeds_.size(); i++)
    {
        ofile << i+1 << ". " << seeds_[i].dir_file << "\n";
    }
    ofile << std::endl;

    int l;
    ofile << "##### station list\n";
    ofile << "station number: " << stations_.size() << "\n\n";

    for (size_t i = 0; i < stations_.size(); i++)
    {
        ofile << i+1 << " " << stations_[i].network << " " << stations_[i].name << "  \n";
        ofile << stations_[i].lon << " " << stations_[i].lat << " " << stations_[i].elev << "  \n";

        ofile << "fragment number: " << stations_[i].frags.size() << "\n\n";
        l = 0;
        for (size_t j = 0; j < stations_[i].frags.size(); j++)
        {
            ofile << l+1 << ". " << stations_[i].frags[j].seed_id + 1 << " "
                << CHANNEL_NAME[stations_[i].frags[j].chl] << " " 
                << stations_[i].frags[j].t0.time_str() << " " 
                << stations_[i].frags[j].te.time_str() << " " 
                << stations_[i].frags[j].freq << " " 
                << stations_[i].frags[j].len << "\n";
            l++;
        }
        ofile << std::endl;
    }

    ofile.close();
    return;
}

void gctl::seismic::Xcorrelation::create_station_sac(std::string working_dir, std::string out_dir, 
    sig_unit_process_ptr pre_merge_func, sig_unit_process_ptr after_merge_func)
{
    int seed_id;
    std::string sac_name, sys_cmd;
    SAC tmp_sac, stn_sac;
    std::vector<SIG_UNIT> sigs;

    size_t i, j, s;
//#pragma omp parallel for private (i, j, s, seed_id, sac_name, sys_cmd, tmp_sac, stn_sac, sigs) schedule(guided)
    for (i = 0; i < stations_.size(); i++)
    {
        seed_id = -1;
        sigs.clear();
        sigs.reserve(stations_[i].frags.size());

        for (j = 0; j < stations_[i].frags.size(); j++)
        {
            sac_name = working_dir + "/" + stations_[i].frags[j].get_sac_filename(stations_[i].name, stations_[i].network);
                    
            if (!access(sac_name.c_str(), R_OK))
            {
                tmp_sac.read(sac_name);
            }
            else
            {
                // delete all sac files opened by the last recorded seed file
                if (seed_id != -1)
                {
                    for (s = 0; s < seeds_[seed_id].sac_t0.size(); s++)
                    {
                        sys_cmd = "rm " + working_dir + "/" + seeds_[seed_id].get_sac_filename(s);
                        system(sys_cmd.c_str());
                    }
                }

                // Call rdseed to read seed file and output SAC files into the working directory
                seed_id = stations_[i].frags[j].seed_id;
                sys_cmd = "rdseed -f " + seeds_[seed_id].dir_file + " -dq " + working_dir;
                system(sys_cmd.c_str());
                tmp_sac.read(sac_name);
            }

            if (pre_merge_func(tmp_sac.signal)) continue;
            if (vecvalid(tmp_sac.signal.val)) sigs.push_back(tmp_sac.signal);
        }
        
        stn_sac.head = tmp_sac.head;
        stn_sac.signal.merge_sig_units(sigs);
        if (after_merge_func(stn_sac.signal))
        {
            throw std::runtime_error("[GCTL_SEIMIC] gctl::seismic::Xcorrelation::create_station_sac(...) stopped by user function.");
        }
        stn_sac.sync();

        sac_name = out_dir + "/MERGED." + stations_[i].network + "." + stations_[i].name + "." + stn_sac.signal.t0.rdseed_time_str() + ".SAC";
        stn_sacs_[stations_[i].name] = sac_name;
        stn_sac.save(sac_name);

        // delete all sac files opened by the last recorded seed file
        if (seed_id != -1)
        {
            for (s = 0; s < seeds_[seed_id].sac_t0.size(); s++)
            {
                sys_cmd = "rm " + working_dir + "/" + seeds_[seed_id].get_sac_filename(s);
                system(sys_cmd.c_str());
            }
        }
    }

    std::ofstream info_out;
    open_outfile(info_out, out_dir + "/station_sac_list.txt");
    info_out << stn_sacs_.size() << std::endl;
    
    std::map<std::string, std::string>::iterator iter;
    for (iter = stn_sacs_.begin(); iter != stn_sacs_.end(); ++iter)
    {
        info_out << iter->first << " " << iter->second << std::endl;
    }
    info_out.close();
    return;
}

void gctl::seismic::Xcorrelation::Xcorrelation_two_stations(std::string stn_sac_list, std::string stn1_name, 
    std::string stn2_name, SIG_UNIT &out_xc, time_unit roundup, double frag_sec, double over_lap, double sigma, 
    sig_unit_process_ptr sig1_func, sig_unit_process_ptr sig2_func)
{
    std::ifstream list_in;
    open_infile(list_in, stn_sac_list);

    int stn_num;
    std::vector<std::string> stn_names, stn_files;
    
    list_in >> stn_num;
    stn_names.resize(stn_num);
    stn_files.resize(stn_num);
    for (size_t i = 0; i < stn_num; i++)
    {
        list_in >> stn_names[i] >> stn_files[i];
    }
    list_in.close();

    int f = 0;
    std::string stn1_sac, stn2_sac;
    for (size_t i = 0; i < stn_num; i++)
    {
        if (stn1_name == stn_names[i]) {stn1_sac = stn_files[i]; f++;}
        if (stn2_name == stn_names[i]) {stn2_sac = stn_files[i]; f++;}
        if (f > 1) break;
    }

    if (f < 2)
    {
        throw std::runtime_error("Inquired station(s) not found. From gctl::seismic::Xcorrelation::Xcorrelation_two_stations(...)");
    }

    SAC sac1, sac2;
    SIG_UNIT sig1, sig2, tmp_xc;
    UTC_TIME stn1_t0, stn2_t0, cut_t0; // Start time of the cutting section
    std::vector<SIG_UNIT> sections_1, sections_2, stored_xcs;

    // read station sacs
    sac1.read(stn1_sac);
    sac2.read(stn2_sac);
    stn1_t0 = sac1.signal.t0;
    stn2_t0 = sac2.signal.t0;

    cut_t0 = stn1_t0<stn2_t0?stn2_t0:stn1_t0; // Set start time
    cut_t0.round_up_to(roundup); // Set start time to round up

    double sig_std1, sig_std2, day_std1, day_std2;
    
    SIG_UNIT day_sig1, day_sig2;
    UTC_TIME cut_t00;
    while ((!sac1.signal.cut_section(day_sig1, cut_t0, 24*3600)) && 
        (!sac2.signal.cut_section(day_sig2, cut_t0, 24*3600)))
    {
        day_std1 = day_sig1.val.std();
        day_std2 = day_sig2.val.std();

        cut_t00 = cut_t0;
        while((!day_sig1.cut_section(sig1, cut_t00, frag_sec)) && 
        (!day_sig2.cut_section(sig2, cut_t00, frag_sec))) // Both successed
        {
            // Skiped by threshold selection
            sig_std1 = sig1.val.std();
            sig_std2 = sig2.val.std();
            if (sig_std1 > sigma*day_std1 || sig_std2 > sigma*day_std2)
            {
                cut_t00.add_duration(frag_sec*over_lap); continue;
            }

            // Skiped by user
            if (sig1_func(sig1)) {cut_t00.add_duration(frag_sec*over_lap); continue;}
            if (sig2_func(sig2)) {cut_t00.add_duration(frag_sec*over_lap); continue;}

            // Stored them to sections
            if (vecvalid(sig1.val) && vecvalid(sig2.val))
            {
                sections_1.push_back(sig1);
                sections_2.push_back(sig2);
            }
            // move cut time
            cut_t00.add_duration(frag_sec*over_lap);
        }

        // move cut time
        cut_t0.add_duration(24*3600);
    }

    for (size_t i = 0; i < sections_1.size(); i++)
    {
        std::clog << "\r" << i+1 << "/" << sections_1.size() << "\n";
        tmp_xc.linear_cross_correlation(sections_1[i], sections_2[i]);
        //tmp_xc.linear_cross_correlation_freq(sections_1[i], sections_2[i], 20);
        stored_xcs.push_back(tmp_xc);
    }

    out_xc.stack_sig_units(stored_xcs);
    return;
}

void gctl::seismic::Xcorrelation::Xcorrelation_all_stations(std::string working_dir, 
    std::string out_dir, time_unit roundup, double frag_sec)
{
    UTC_TIME stn1_t0, stn2_t0, cut_t0, cut_te; // Start time of the cutting section
    std::vector<SIG_UNIT> sections_1, sections_2, stored_xcs;
    SIG_UNIT tmp_xc, station_xc;

    // Loop over pairs between any two stations
    int seed1_id, seed2_id;
    int sac1_id, sac2_id;
    int sac1_num, sac2_num;
    bool new_sac1, new_sac2, end_of_station;
    SAC sac1, sac2;
    SIG_UNIT sac1_sig, sac2_sig; 
    std::string sac_name, sys_cmd;

    sys_cmd = "rm " + working_dir + "/*.SAC";
    system(sys_cmd.c_str());

    for (size_t s = 0; s < stations_.size()-1; s++)
    {
        //for (size_t n = s+1; n < stations_.size(); n++)
        for (size_t n = s+1; n < 2; n++)
        {
            //std::clog << "Calculating Xcorrleation between stations " << stations_[s].name << " and "  << stations_[n].name << "...\n";

            seed1_id = seed2_id = -1;
            new_sac1 = new_sac2 = true;
            end_of_station = false;
            sac1_num = stations_[s].frags.size();
            sac2_num = stations_[n].frags.size();
            stored_xcs.clear();
            stored_xcs.reserve(1000*(sac1_num>sac2_num?sac1_num:sac2_num));

            // Continue if no frags are found
            if (sac1_num == 0 || sac2_num == 0) continue;

            // Find out the start time of the two comparing stations
            stn1_t0 = stations_[s].frags[0].t0;
            stn2_t0 = stations_[n].frags[0].t0;

            cut_t0 = stn1_t0<stn2_t0?stn2_t0:stn1_t0; // Set start time
            cut_t0.round_up_to(roundup); // Set start time to round up

            // Determine the inital sac ids;
            for (size_t i = 0; i < sac1_num; i++)
            {
                if (stations_[s].frags[i].t0 <= cut_t0 && cut_t0 < stations_[s].frags[i].te)
                {
                    sac1_id = i; break;
                }
            }

            for (size_t i = 0; i < sac2_num; i++)
            {
                if (stations_[n].frags[i].t0 <= cut_t0 && cut_t0 < stations_[n].frags[i].te)
                {
                    sac2_id = i; break;
                }
            }

            while (!end_of_station)
            {
                //std::cout << sac1_id << " " << sac2_id << "\n";
                // Get the current SAC file for the first track
                if (new_sac1)
                {
                    sac_name = working_dir + "/" + stations_[s].frags[sac1_id].get_sac_filename(stations_[s].name, stations_[s].network);
                    
                    if (!access(sac_name.c_str(), R_OK))
                    {
                        sac1.read(sac_name);
                    }
                    else
                    {
                        // delete all sac files opened by the last recorded seed file
                        if (seed1_id != -1)
                        {
                            for (size_t i = 0; i < seeds_[seed1_id].sac_t0.size(); i++)
                            {
                                sys_cmd = "rm " + working_dir + "/" + seeds_[seed1_id].get_sac_filename(i);
                                system(sys_cmd.c_str());
                            }
                        }
                        
                        // Call rdseed to read seed file and output SAC files into the working directory
                        seed1_id = stations_[s].frags[sac1_id].seed_id;
                        sys_cmd = "rdseed -f " + seeds_[seed1_id].dir_file + " -dq " + working_dir;
                        system(sys_cmd.c_str());
                        sac1.read(sac_name);
                    }
                    new_sac1 = false;
                }
                
                if (new_sac2)
                {
                    // Get the current SAC file for the second track
                    sac_name = working_dir + "/" + stations_[n].frags[sac2_id].get_sac_filename(stations_[n].name, stations_[n].network);
                    
                    if (!access(sac_name.c_str(), R_OK))
                    {
                        sac2.read(sac_name);
                    }
                    else
                    {
                        // delete all sac files opened by the last recorded seed file
                        if (seed2_id != -1)
                        {
                            for (size_t i = 0; i < seeds_[seed2_id].sac_t0.size(); i++)
                            {
                                sys_cmd = "rm " + working_dir + "/" + seeds_[seed2_id].get_sac_filename(i);
                                system(sys_cmd.c_str());
                            }
                        }

                        // Call rdseed to read seed file and output SAC files into the working directory
                        seed2_id = stations_[n].frags[sac2_id].seed_id;
                        sys_cmd = "rdseed -f " + seeds_[seed2_id].dir_file + " -dq " + working_dir;
                        system(sys_cmd.c_str());
                        sac2.read(sac_name);
                    }
                    new_sac2 = false;
                }

                // Find out the start time of the two comparing stations
                stn1_t0 = sac1.signal.t0;
                stn2_t0 = sac2.signal.t0;

                cut_t0 = stn1_t0<stn2_t0?stn2_t0:stn1_t0; // Set start time
                cut_t0.round_up_to(roundup); // Set start time to round up

                //sac1.signal.down_sampling(100);
                //sac2.signal.down_sampling(100);
                //sac1.signal.remove_mean();
                //sac2.signal.remove_mean();
                //sac1.signal.normalize();
                //sac2.signal.normalize();

                sections_1.clear();
                sections_2.clear();
                sections_1.reserve(100);
                sections_2.reserve(100);

                while((!sac1.signal.cut_section(sac1_sig, cut_t0, frag_sec)) && 
                    (!sac2.signal.cut_section(sac2_sig, cut_t0, frag_sec))) // Both successed
                {
                    sac1_sig.down_sampling(100);
                    sac2_sig.down_sampling(100);
                    sac1_sig.remove_outliers(2.0);
                    sac2_sig.remove_outliers(2.0);
                    sac1_sig.remove_mean();
                    sac2_sig.remove_mean();
                    sac1_sig.normalize();
                    sac2_sig.normalize();

                    // Stored them to sections
                    sections_1.push_back(sac1_sig);
                    sections_2.push_back(sac2_sig);
                    // move cut time
                    cut_t0.add_duration(frag_sec);
                }

                for (size_t i = 0; i < sections_1.size(); i++)
                {
                    tmp_xc.linear_cross_correlation(sections_1[i], sections_2[i]);
                    stored_xcs.push_back(tmp_xc);
                }

                // determine if we need to switch to the next sac file
                cut_te = cut_t0;
                cut_te.add_duration(frag_sec);

                //std::cout << sac1.signal.te.time_str() << " " << cut_te.time_str() << "\n";
                //std::cout << sac2.signal.te.time_str() << " " << cut_te.time_str() << "\n";

                if (sac1.signal.te < cut_te) // we need to switch
                {
                    new_sac1 = true; sac1_id++;
                    if (sac1_id > sac1_num-1) end_of_station = true;
                }

                if (sac2.signal.te < cut_te)
                {
                    new_sac2 = true; sac2_id++;
                    if (sac2_id > sac2_num-1) end_of_station = true;
                }
            }

            station_xc.stack_sig_units(stored_xcs);
            station_xc.save(out_dir + "/xc_" + stations_[s].name + "_" + stations_[n].name + "_" + stations_[s].network);

            // delete all sac files opened by the last recorded seed file
            if (seed1_id != -1)
            {
                for (size_t i = 0; i < seeds_[seed1_id].sac_t0.size(); i++)
                {
                    sys_cmd = "rm " + working_dir + "/" + seeds_[seed1_id].get_sac_filename(i);
                    system(sys_cmd.c_str());
                }
            }

            // delete all sac files opened by the last recorded seed file
            if (seed2_id != -1)
            {
                for (size_t i = 0; i < seeds_[seed2_id].sac_t0.size(); i++)
                {
                    sys_cmd = "rm " + working_dir + "/" + seeds_[seed2_id].get_sac_filename(i);
                    system(sys_cmd.c_str());
                }
            }
        }
    }
    return;
}