// ========================================================= //
//                                                           //
//   File      : NT_split_ali.cxx                            //
//   Purpose   : split alignments                            //
//                                                           //
//   Coded by Ralf Westram (coder@reallysoft.de) in Jun 22   //
//   http://www.arb-home.de/                                 //
//                                                           //
// ========================================================= //

#include "NT_local_proto.h"

#include <arbdbt.h>
#include <arb_progress.h>
#include <arb_defs.h>
#include <arb_strbuf.h>
#include <pos_range.h>

#include <aw_root.hxx>
#include <aw_window.hxx>

#include <limits>
#include <vector>
#include <insdel.h>

using namespace std;

typedef vector<ExplicitRange> RangeVec;

static const int END__ = numeric_limits<int>::max();

static GB_ERROR calculate_and_check_ranges(int ali_source_len, const int *splitBefore, RangeVec& ranges) {
    GB_ERROR      error = NULp;
    ExplicitRange rest(ali_source_len);
    for (int p = 0; splitBefore[p] != END__ && !error; ++p) {
        ExplicitRange next = intersection(ExplicitRange::prior(splitBefore[p]), rest);
        if (next.is_empty()) {
            if      (rest.is_empty())             error = GBS_global_string("position %i exceeds overall length (=%i)", p ? info2bio(splitBefore[p-1]) : 0, ali_source_len);
            else if (splitBefore[p]<rest.start()) error = "invalid position order";
            else                                  error = "empty range dectected";
        }
        else {
            ranges.push_back(next);
            rest = intersection(rest, next.after());
        }
    }

    if (!error) {
        if (rest.is_empty()) error = GBS_global_string("last position exceeds overall length (=%i)", ali_source_len);
        else                 ranges.push_back(rest); //  append final range
    }

    return error;
}

CONSTEXPR_INLINE int digits(int parts) { // when using numbers from [1 .. parts]
    return parts<0 ? digits(-parts) : (parts<10 ? 1 : (parts<100 ? 2 : 3)); // works up to 999
}

inline const char *target_ali_name(const char *ali_target_prefix, int p, int digits) {
    static SmartCharPtr ali;
    ali = GBS_global_string_copy("%s_%0*i", ali_target_prefix, digits, p+1);
    return &*ali;
}

inline char *build_split_id(const RangeVec& ranges) {
    GBS_strstruct buf(100);
    for (size_t r = 0; r<ranges.size(); ++r) {
        if (r>0) buf.put('-');
        buf.putlong(ranges[r].size());
    }
    return buf.get_copy();
}

static GB_ERROR split_data_into_parts(GBDATA *gb_data, const RangeVec& ranges, const char *ali_target_prefix) {
    GB_ERROR    error      = NULp;
    char       *seq        = GB_read_as_string(gb_data);
    const char *key        = GB_read_key_pntr(gb_data);
    GB_TYPES    type       = GB_read_type(gb_data);
    int         security   = GB_read_security_write(gb_data);
    size_t      seq_len    = strlen(seq);
    const int   parts      = ranges.size();
    GBDATA     *gb_item    = GB_get_grandfather(gb_data);

    // [Note: the way elements are handled here is a bit q&d; see EditedTerminal::apply for howto properly handle all data types]

    // write splitted sequence data
    for (int p = 0; p<parts && !error; ++p) {
        const char *ali          = target_ali_name(ali_target_prefix, p, digits(parts));
        GBDATA     *gb_part_data = GBT_create_sequence_data(gb_item, ali, key, type, security);

        if (!gb_part_data) {
            UNCOVERED();
            error = GB_await_error();
        }
        else {
            char *part = ranges[p].dup_corresponding_part(seq, seq_len);
            error      = GB_write_autoconv_string(gb_part_data, part);
            free(part);
        }
    }

    free(seq);

    return error;
}

static GB_ERROR copy_field_into_partial_alignments(GBDATA *gb_data, const int parts, const char *ali_target_prefix) {
    GB_ERROR    error   = NULp;
    GB_TYPES    type    = GB_read_type(gb_data);
    GBDATA     *gb_item = GB_get_grandfather(gb_data);
    const char *key     = GB_read_key_pntr(gb_data);

    char *target_val;
    if (type == GB_STRING) target_val = GBS_global_string_copy("[splitted] %s", GB_read_char_pntr(gb_data));
    else {
        target_val = GBS_global_string_copy("[split skipped copy of field '%s' (type: %i)]", key, type);
        type       = GB_STRING;
    }

    const int security = GB_read_security_write(gb_data);
    for (int p = 0; p<parts && !error; ++p) {
        const char *ali          = target_ali_name(ali_target_prefix, p, digits(parts));
        GBDATA     *gb_part_data = GBT_create_sequence_data(gb_item, ali, key, type, security);

        error = GB_write_string(gb_part_data, target_val);
    }
    free(target_val);

    return error;
}

static GB_ERROR split_alignment(GBDATA *gb_main, const char *ali_source, const char *ali_target_prefix, const int *splitBefore, bool markedOnly, bool overwriteAlignments) {
    /*! splits one alignment into multiple alignments.
     * @param ali_source name of source alignment
     * @param ali_target_prefix prefix of target alignment. will be extended by '_1', '_2', ...
     * @param splitBefore list of positions [0..N-1] where 2nd, 3rd, ... part shall begin. Terminate with 'END__'.
     * @param markedOnly true->operate only on marked species, false->operate on all species
     * @param overwriteAlignments true->delete+recreate existing alignments, false->bail-out with error
     * @return error if sth went wrong.
     */

    GB_ERROR       error = NULp;
    GB_transaction ta(gb_main);

    arb_progress progress_main("Splitting sequences");

    // check source alignment exists:
    long ali_source_len = GBT_get_alignment_len(gb_main, ali_source);
    if (ali_source_len<=0) {
        error = GB_await_error();
    }
    else {
        RangeVec ranges;
        error = calculate_and_check_ranges(ali_source_len, splitBefore, ranges);

        const int parts = ranges.size();
        if (parts<2 && !error) {
            error = GBS_global_string("Not enough parts (%i) specified to split alignment", parts);
        }

        // @@@ read the following values from source alignment:
        bool        is_aligned = true; // @@@ if source alignment is not formatted -> abort with error (for safety and to avoid special cases here)
        int         security   = 3;
        const char *ali_type   = "rna";

        char *split_id = build_split_id(ranges);

        // create alignments:
        for (int p = 0; p<parts && !error; ++p) {
            const char *ali = target_ali_name(ali_target_prefix, p, digits(parts));
            {
                GBDATA *gb_ali = GBT_get_alignment(gb_main, ali);
                if (!gb_ali) {
                    GB_clear_error(); // from GBT_get_alignment
                }
                else {
                    if (overwriteAlignments) {
                        error = GBT_delete_alignment(gb_main, ali); // delete existing alignment
                    }
                    else {
                        error = GBS_global_string("alignment '%s' already exists", ali);
                    }
                }
            }
            if (!error) {
                // create target alignment:
                const ExplicitRange&  range       = ranges[p];
                char                 *why_created = GBS_global_string_copy("while splitting %s into %s [%i-%i]", ali_source, split_id, info2bio(range.start()), info2bio(range.end()));

                GBDATA *gb_ali = GBT_create_new_alignment(gb_main, ali, ranges[p].size(), is_aligned, security, ali_type, why_created);
                if (!gb_ali) {
                    error = GB_await_error();
                }
                else {
                    error = GBT_add_alignment_changekeys(gb_main, ali);
                }

                free(why_created);
            }
        }

        if (!error) {
            size_t data_count         = 0; // count species with sequence data in alignment (during pass 1)
            size_t sai_data_count     = 0; // count sequence associated data in alignment containers of all SAIs (during pass 1)
            size_t sequences_splitted = 0; // count splitted sequences (species)
            size_t sai_data_splitted  = 0; // count splitted data entries (SAI; may be multiple per SAI)

            SmartPtr<arb_progress> progress;

            for (int pass = 1; pass<=2 && !error; ++pass) {
                if (pass == 2) {
                    progress = new arb_progress(data_count+sai_data_count);
                    progress->auto_subtitles("seq");
                }

                // generate sequence data in splitted alignments:
                for (GBDATA *gb_species = markedOnly ? GBT_first_marked_species(gb_main) : GBT_first_species(gb_main);
                     gb_species && !error;
                     gb_species = markedOnly ? GBT_next_marked_species(gb_species) : GBT_next_species(gb_species))
                {
                    GBDATA *gb_seq = GBT_find_sequence(gb_species, ali_source);
                    if (gb_seq) {
                        if (pass == 1) ++data_count;
                        else           {
                            error = split_data_into_parts(gb_seq, ranges, ali_target_prefix);
                            if (!error) ++sequences_splitted; // count splits
                            progress->inc_and_check_user_abort(error);
                        }
                    }
                }

                for (GBDATA *gb_sai = GBT_first_SAI(gb_main);  gb_sai && !error; gb_sai = GBT_next_SAI(gb_sai)) {
                    GBDATA *gb_ali = GB_entry(gb_sai, ali_source);
                    if (gb_ali) {
                        for (GBDATA *gb_data = GB_child(gb_ali); gb_data && !error; gb_data = GB_nextChild(gb_data)) {
                            if (ARB_is_alignment_relative_data(gb_data)) {
                                if (pass == 1) {
                                    ++sai_data_count;
                                }
                                else {
                                    error = split_data_into_parts(gb_data, ranges, ali_target_prefix);
                                    if (!error) ++sai_data_splitted; // count splits
                                    progress->inc_and_check_user_abort(error);
                                }
                            }
                            else if (pass == 2) {
                                error = copy_field_into_partial_alignments(gb_data, parts, ali_target_prefix);
                            }
                        }
                    }
                }
            }

            if (!error) {
                size_t overall_data_splitted = sequences_splitted + sai_data_splitted;
                if (overall_data_splitted) {
                    size_t split_count = overall_data_splitted*parts;
                    fprintf(stdout, "Summary: splitted %zu element%s into %zu parts.\n", overall_data_splitted, plural(overall_data_splitted), split_count);
                }
                else {
                    UNCOVERED();
                    error = "No sequences were splitted (none of the handled species contained data in alignment; check settings)";
                }
            }
        }

        free(split_id);
    }
    error = ta.close(error);

    return error;
}

// --------------------------------------------------------------------------------

#ifdef UNIT_TESTS
#include <arb_diff.h>
#include <arb_file.h>
#ifndef TEST_UNIT_H
#include <test_unit.h>
#endif

// #define TEST_AUTO_UPDATE // uncomment to auto-update expected result DB

void TEST_split_alignment() {
    GB_shell  shell;
    GBDATA   *gb_main = GB_open("TEST_prot_tiny.arb", "rw"); // source: ../UNIT_TESTER/run/TEST_prot_tiny.arb
                                                             // target: ../UNIT_TESTER/run/TEST_split_expected.arb

    // alignments existing in DB:
    // --------------------------
    // alignment_name   "ali_dna"            alignment_len    %i 337
    // alignment_name   "ali_prot"           alignment_len    %i 112
    // alignment_name   "ali_dna_incomplete" alignment_len    %i 337

    {
        GB_securityLevel raised(gb_main, 3);

        // split some alignment
        {
            GB_transaction ta(gb_main);
            TEST_EXPECT_NO_ERROR(GBT_restore_marked_species(gb_main, "CytLyti6;BctFra12;StrCoel9;MucRace2"));

            int split[] = { 41, 175, END__ };
            TEST_EXPECT_NO_ERROR(split_alignment(gb_main, "ali_dna", "ali_dna_split1", split, true, false));
        }
        {
            GB_transaction ta(gb_main);
            TEST_EXPECT_NO_ERROR(GBT_restore_marked_species(gb_main, "CytLyti6;TaxOcell;BctFra12;StrRamo3"));

            int split[] = { 168, END__ }; // approx. split in the middle of the sequence
            TEST_EXPECT_NO_ERROR(split_alignment(gb_main, "ali_dna", "ali_dna_split2", split, true, false));
            TEST_EXPECT_NO_ERROR(split_alignment(gb_main, "ali_dna", "ali_overwrite", split, true, false)); // created to be overwritten below
        }
        {
            GB_transaction ta(gb_main);
            TEST_EXPECT_NO_ERROR(GBT_restore_marked_species(gb_main, "TaxOcell"));

            int split[] = { 1, 160, 161, 336, END__ }; // test short parts (at 5'-end, in the middle and at 3'-end)
            TEST_EXPECT_NO_ERROR(split_alignment(gb_main, "ali_dna", "ali_dna_split3", split, true, false));
        }
        {
            GB_transaction ta(gb_main);

            int split[] = { 25, 65, END__ };
            TEST_EXPECT_NO_ERROR(split_alignment(gb_main, "ali_prot", "ali_prot_split", split, false, false));
        }
        { // split alignment not occurring in all species (applied to all species)
            GB_transaction ta(gb_main);

            int split[] = { 150, END__ };
            TEST_EXPECT_NO_ERROR(split_alignment(gb_main, "ali_dna_incomplete", "ali_dna_incomplete_split", split, false, false));
        }


        // provoke errors (invalid ali positions):
        {
            int split[] = { END__ };
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_any", split, true, false),
                                       "Not enough parts (1) specified to split alignment");
        }
        {
            int split[] = { 41, 470, END__ }; // requests invalid split (beyond 3-end) -> error wanted
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_any", split, true, false),
                                       "last position exceeds overall length (=337)");
        }

        {
            int split[] = { 336, 4100, END__ }; // tests that offset 336 is a legal split position
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_any", split, true, false),
                                       "last position exceeds overall length (=337)"); // 2nd offset causes this failure
        }
        {
            int split[] = { 337, 4100, END__ }; // test that offset 337 is NO legal split position (no data behind)
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_any", split, true, false),
                                       "position 338 exceeds overall length (=337)");
        }
        {
            int split[] = { 338, 4100, END__ }; // requests invalid split (beyond 3-end) with first offset -> error wanted
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_any", split, true, false),
                                       "position 339 exceeds overall length (=337)");
        }
        {
            int split[] = { 0, 4100, END__ }; // requests invalid split (at 5-end) with first offset -> error wanted
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_any", split, true, false),
                                       "empty range dectected");
        }

        {
            int split[] = { 41, 30, 175, END__ }; // 2nd offset < 1st offset -> error wanted
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_any", split, true, false),
                                       "invalid position order");
        }
        {
            int split[] = { 41, 41, 175, END__ }; // empty range -> error wanted
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_any", split, true, false),
                                       "empty range dectected");
        }

        {
            // provoke error (no such alignment)
            int split[] = { 41, 175, END__ };
            TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_nosuch", "ali_any", split, true, false),
                                       "alignment 'ali_nosuch' not found");
        }

        {
            // provoke error (target alignment exists)

            const char *ali_conflict    = "ali_existing_3";
            GBDATA     *gb_ali_conflict = NULp;
            {
                GB_transaction ta(gb_main);
                gb_ali_conflict = GBT_create_new_alignment(gb_main, ali_conflict, 123, 0, 0, "dna", "by TEST_split_alignment");
                TEST_REJECT_NULL(gb_ali_conflict);
            }

            {
                int split[] = { 41, 175, END__ };
                TEST_EXPECT_ERROR_CONTAINS(split_alignment(gb_main, "ali_dna", "ali_existing", split, true, false),
                                           "alignment 'ali_existing_3' already exists");
            }

            // overwrite existing alignment
            {
                GB_transaction ta(gb_main);
                TEST_EXPECT_NO_ERROR(GBT_restore_marked_species(gb_main, "CytLyti6;BctFra12;MucRace2"));
            }
            {
                int split[] = { 170, END__ }; // approx. split in the middle of the sequence (but different than in 1st split to ali_overwrite..)

                GB_topSecurityLevel raise(gb_main);
                TEST_EXPECT_NO_ERROR(split_alignment(gb_main, "ali_dna", "ali_overwrite", split, true, true));
            }
        }
    }

    GBT_mark_all(gb_main, 0); // clear marks before saving database

    {
        // save as ascii with different name + compare result with expected file

        char *savename = ARB_strdup(GB_path_in_ARBHOME("UNIT_TESTER/run/TEST_split.arb"));
        char *expected = ARB_strdup(GB_path_in_ARBHOME("UNIT_TESTER/run/TEST_split_expected.arb"));

        TEST_EXPECT_NO_ERROR(GB_save_as(gb_main, savename, "a"));

#if defined(TEST_AUTO_UPDATE)
        TEST_COPY_FILE(savename, expected);
#else // !defined(TEST_AUTO_UPDATE)
        TEST_EXPECT_TEXTFILES_EQUAL(savename, expected);
#endif
        TEST_EXPECT_ZERO_OR_SHOW_ERRNO(GB_unlink(savename));

        free(expected);
        free(savename);
    }

    GB_close(gb_main);
}

#endif // UNIT_TESTS

// --------------------------------------------------------------------------------

#include "NT_local.h"

#include <arb_strarray.h>

#include <aw_awar.hxx>
#include <aw_awar_defs.hxx>
#include <aw_msg.hxx>


#define AWAR_SPLIT_ALI_BASE "split_ali/"
#define AWAR_SPLIT_ALI_BTMP "tmp/" AWAR_SPLIT_ALI_BASE

#define AWAR_SPLIT_ALI_TARGET    AWAR_SPLIT_ALI_BTMP "target"
#define AWAR_SPLIT_ALI_POSITIONS AWAR_SPLIT_ALI_BASE "pos"
#define AWAR_SPLIT_ALI_MARKED    AWAR_SPLIT_ALI_BASE "marked"
#define AWAR_SPLIT_ALI_OVERWRITE AWAR_SPLIT_ALI_BASE "overwrite"

static void create_splitAlignment_awars(AW_root *root, AW_default properties, GBDATA *gb_main) {
    // use current alignment to init default for target alignment:
    const char *ali_part = "ali_example_part";
    {
        char *ali_source = GBT_get_default_alignment(gb_main);
        if (!ali_source) GB_clear_error();
        else             ali_part = GBS_global_string("%s_part", ali_source);
        free(ali_source);
    }

    root->awar_string(AWAR_SPLIT_ALI_TARGET,    ali_part, properties);
    root->awar_string(AWAR_SPLIT_ALI_POSITIONS, "",       properties);
    root->awar_int   (AWAR_SPLIT_ALI_MARKED,    0,        properties);
    root->awar_int   (AWAR_SPLIT_ALI_OVERWRITE, 0,        properties);

    root->awar_string(AWAR_RANGE, "", gb_main);
}

static void split_ali_cb(AW_window *aww) {
    AW_root  *awr     = aww->get_root();
    GBDATA   *gb_main = GLOBAL.gb_main;
    GB_ERROR  error   = NULp;

    char *ali_source = GBT_get_default_alignment(gb_main);
    if (!ali_source) {
        error = GB_await_error();
    }
    else {
        char *ali_target_prefix = awr->awar(AWAR_SPLIT_ALI_TARGET)->read_string();

        ConstStrArray numbers;
        {
            char *split_pos_list    = awr->awar(AWAR_SPLIT_ALI_POSITIONS)->read_string();
            GBT_splitNdestroy_string(numbers, split_pos_list, ", ", SPLIT_DROPEMPTY);
        }

        int split_positions[numbers.size()+1];
        for (size_t n = 0; n<numbers.size() && !error; ++n) {
            const char *num    = numbers[n];
            split_positions[n] = bio2info(atoi(num));
            if (split_positions[n]<1) {
                error = GBS_global_string("invalid position '%s'", num); // shows biological position
            }
        }
        split_positions[numbers.size()] = END__;

        if (!error) {
            bool markedOnly = awr->awar(AWAR_SPLIT_ALI_MARKED)->read_int();
            bool overwrite  = awr->awar(AWAR_SPLIT_ALI_OVERWRITE)->read_int();

            error = split_alignment(gb_main, ali_source, ali_target_prefix, split_positions, markedOnly, overwrite);
        }

        free(ali_target_prefix);
        free(ali_source);
    }

    aw_message_if(error);
}

static void use_editor_range_cb(AW_window *aww) {
    AW_root    *awr          = aww->get_root();
    GB_ERROR    error        = NULp;
    const char *editor_range = awr->awar(AWAR_RANGE)->read_char_pntr();
    const char *hyphen       = strchr(editor_range, '-');

    if (!hyphen) {
        if (editor_range[0]) {
            error = "expected range to contain '-'.";
        }
        else {
            error = "No block has been highlighted in sequence editor.";
        }
    }
    else {
        long i1 = atol(editor_range);
        long i2 = atol(hyphen+1);

        const char *pos_list = GBS_global_string("%li,%li", i1, i2+1); // split BEFORE i1 + AFTER i2
        awr->awar(AWAR_SPLIT_ALI_POSITIONS)->write_string(pos_list);
    }

    aw_message_if(error);
}

AW_window *NT_create_splitAlignment_window(AW_root *awr) {
    create_splitAlignment_awars(awr, AW_ROOT_DEFAULT, GLOBAL.gb_main);

    AW_window_simple *aws = new AW_window_simple;
    aws->init(awr, "SPLIT_ALIGNMENT", "Split alignment");

    const int PAD = 10;

    const int BUTTON_LENGTH = 10;
    const int FIELD_LENGTH  = 40;
    const int LABEL_LENGTH  = 30;

    aws->at(PAD, PAD);
    aws->auto_space(PAD/2, PAD);

    aws->button_length(BUTTON_LENGTH);

    aws->callback(AW_POPDOWN);
    aws->create_button("CLOSE", "Close", "C");

    aws->callback(makeHelpCallback("split.hlp"));
    aws->create_button("HELP", "Help", "H");

    aws->at_newline();

    aws->label_length(LABEL_LENGTH);

    aws->label("Target alignment prefix:");
    aws->create_input_field(AWAR_SPLIT_ALI_TARGET, FIELD_LENGTH);
    aws->at_newline();

    aws->label("Position(s) to split before:\n"
               "(comma-separated)");
    aws->create_input_field(AWAR_SPLIT_ALI_POSITIONS, FIELD_LENGTH);
    aws->at_newline();

    aws->button_length(LABEL_LENGTH-1);
    aws->callback(use_editor_range_cb);
    aws->create_button("USE_EDITOR_RANGE", "Use editor range:");
    aws->button_length(FIELD_LENGTH);
    aws->create_button(NULp, AWAR_RANGE, "", "-");
    aws->button_length(BUTTON_LENGTH);
    aws->at_newline();

    aws->label("Split marked only:");
    aws->create_toggle(AWAR_SPLIT_ALI_MARKED);
    aws->at_newline();

    aws->label("Overwrite existing alignments:\n"
               "(deletes ALL data)");
    aws->create_toggle(AWAR_SPLIT_ALI_OVERWRITE);
    aws->at_newline();

    aws->callback(split_ali_cb);
    aws->create_button("GO", "GO", "G");

    return aws;
}