source: trunk/TOOLS/arb_export_seq_filtered.cxx

// ============================================================ //
//                                                              //
//   File      : arb_export_seq_filtered.cxx                    //
//   Purpose   : export sequences filtered by SAI               //
//                                                              //
//   Coded by Ralf Westram (coder@reallysoft.de) in June 2017   //
//   http://www.arb-home.de/                                    //
//                                                              //
// ============================================================ //

#include <FilteredExport.h>
#include <arbdbt.h>
#include <arb_file.h>
#include <string>
#include <list>
#include <cstdlib>

using namespace std;

class CLI : virtual Noncopyable {
    bool     helpWanted;
    GB_ERROR error;

    string database;   // name of input database
    string fasta;      // name of FASTA output file
    string aliname;    // name of alignment to use
    string head_aci;   // ACI used to create content of fasta header
    string seq_aci;    // ACI used to post-process sequence data

    bool accept_missing_data; // true -> skip sequences with less bases in data

    int    min_bases; // skip sequences with less than 'min_bases' base-characters in data
    string bases;     // defines what counts as base-character

    string filterby;            // name of last SAI specified with "--filterby"

    typedef SmartPtr<FilterDefinition>     FilterDefinitionPtr;
    typedef std::list<FilterDefinitionPtr> FilterDefinitionList;

    FilterDefinitionList def_filters;

    static inline const char *getarg(int& argc, const char**& argv) {
        return argc>0 ? (--argc,*argv++) : NULp;
    }
    inline const char *expect_arg(int& argc, const char**& argv) {
        const char *arg = getarg(argc, argv);
        if (!arg) {
            error = "expected argument missing";
            arg   = "";
        }
        return arg;
    }
    void show_help() const {
        fputs("\n"
              "arb_export_seq_filtered -- export sequences filtered by SAI\n"
              "Usage: arb_export_seq_filtered [switches]\n"
              "\n"
              "general:\n"
              " <chars> defines a set of characters (for blocking/passthrough).\n"
              "         It may contain alphanumeric ranges, e.g. '0-5' which is equiv. to '012345'.\n"
              "         To specify a literal '-' put it at start or end of <chars>.\n"
              "\n"
              "switches:\n"
              "--db <dbname>              ARB database to export from\n"
              "--fasta <outname>          name of generated fasta file\n"
              "\n"
              "--id <aci>                 specify content for fasta header using ACI\n"
              "                           (default: \"readdb(name)\"; see http://help.arb-home.de/aci.html)\n"
              "\n"
              "--ali <aliname>            name of alignment to use (default: use default alignment)\n"
              "--accept-missing-data      silently skip species w/o data (default: abort with error)\n"
              "--seqpp <aci>              specify ACI to postprocess sequence data (default: none; example dot->hyphen: \":.=-\")\n"
              "\n"
              "--filterby <SAI>           filter sequences using <SAI> (SAI has to be stored in database)\n"
              "--block [allbut] <chars>   exclude columns where <SAI> contains any of <chars>\n"
              "--pass [allbut] <chars>    include columns where <SAI> contains any of <chars>\n"
              "                           (specifying 'allbut' will invert the set of <chars>)\n"
              "\n"
              "--min-bases <count>        min. base count of exported data (after filtering).\n"
              "                           Export of species is skipped if not reached and\n"
              "                           a diagnostic message gets written to stderr.\n"
              "--count-bases <chars>      defines base-characters to count (default: \"a-zA-Z\")\n"
              "\n"
              ,stderr);
    }

    const char *missing_passBlock() { return "--filterby has to be followed by --pass or --block"; }
    const char *missing_filterby() { return "--pass and --block have to be preceeded by --filterby"; }

    void store_filterdef(FilterDefinition *def) {
        if (def) def_filters.push_back(def);
    }
    FilterDefinition *expect_filterdef_args(FilterDefType type, int& argc, const char**& argv) {
        FilterDefinition *result = NULp;
        string            arg    = expect_arg(argc, argv);
        if (!error) {
            bool filter_chars = true;

            if (arg == "allbut") {
                filter_chars = false;
                arg          = expect_arg(argc, argv);
            }

            if (!error) {
                if (filterby.empty()) {
                    error = missing_filterby();
                }
                else {
                    result = new FilterDefinition(filterby.c_str(), type, filter_chars, arg.c_str());
                }
            }
        }
        filterby = ""; // forget last filterby (even in error-case -> avoids some weird error-messages)
        return result;
    }

    void parse(int& argc, const char**& argv) {
        const char *arg = getarg(argc, argv);
        if (arg) {
            if      (strcmp(arg, "--db")          == 0) database = expect_arg(argc, argv);
            else if (strcmp(arg, "--fasta")       == 0) fasta    = expect_arg(argc, argv);
            else if (strcmp(arg, "--id")          == 0) head_aci = expect_arg(argc, argv);
            else if (strcmp(arg, "--seqpp")       == 0) seq_aci  = expect_arg(argc, argv);
            else if (strcmp(arg, "--ali")         == 0) aliname  = expect_arg(argc, argv);
            else if (strcmp(arg, "--count-bases") == 0) bases    = expect_arg(argc, argv);

            else if (strcmp(arg, "--filterby") == 0) {
                if (!filterby.empty()) error = missing_passBlock();
                else filterby                = expect_arg(argc, argv);
            }

            else if (strcmp(arg, "--block") == 0) store_filterdef(expect_filterdef_args(BLOCK, argc, argv));
            else if (strcmp(arg, "--pass")  == 0) store_filterdef(expect_filterdef_args(PASS,  argc, argv));

            else if (strcmp(arg, "--accept-missing-data") == 0) accept_missing_data = true;

            else if (strcmp(arg, "--min-bases") == 0) min_bases = atoi(expect_arg(argc, argv));

            else if (strcmp(arg, "--help") == 0 || strcmp(arg, "-h") == 0) helpWanted = true;

            else {
                error = GBS_global_string("unexpected argument '%s'", arg);
            }
        }
    }
    void check_required_arguments() {
        if (database.empty()) error = "no input database specified";
        else if (fasta.empty()) error = "no output file specified";
    }

public:
    CLI(int argc, const char **argv) :
        helpWanted(false),
        error(NULp),
        accept_missing_data(false),
        min_bases(0),
        bases("a-zA-Z")
    {
        --argc; ++argv;
        while (!error && argc>0 && !helpWanted) {
            parse(argc, argv);
        }

        if (!helpWanted) { // do not check extended conditions, if '--help' seen
            if (!error && !filterby.empty()) {
                error = missing_passBlock();
            }
            if (!error) {
                check_required_arguments();
                if (error) helpWanted = true;
            }
        }
    }

    bool help_wanted() const { return helpWanted; }
    void show_help_if_useful() const { if (helpWanted) show_help(); }
    GB_ERROR get_error() const { return error; }

    const char *get_database()  const { return database.c_str(); }
    const char *get_fasta()     const { return fasta.c_str(); }
    const char *get_aliname()   const { return aliname.c_str(); }
    const char *get_head_aci()  const { return head_aci.c_str(); }
    const char *get_seq_aci()   const { return seq_aci.c_str(); }
    const char *get_basechars() const { return bases.c_str(); }

    bool shall_accept_missing_data() const { return accept_missing_data; }

    int get_min_bases_count() const { return min_bases; }

    bool has_filters() const { return !def_filters.empty(); }
    GB_ERROR add_filters_to(FilteredExport& exporter) const {
        arb_assert(!error);
        GB_ERROR err = NULp;
        for (FilterDefinitionList::const_iterator d = def_filters.begin(); !err && d != def_filters.end(); ++d) {
            err = exporter.add_SAI_filter(**d);
        }
        return err;
    }
};

static GB_ERROR export_seq_data(const CLI& args) {
    GB_ERROR error = NULp;

    const char *dbname  = args.get_database();
    GB_shell    shell;
    GBDATA     *gb_main = GB_open(dbname, "r");

    if (!gb_main) {
        error = GB_await_error();
    }
    else {
        {
            GB_transaction ta(gb_main);

            char *aliname = strdup(args.get_aliname());
            if (aliname[0]) {
                if (!GBT_get_alignment(gb_main, aliname)) {
                    error = GB_await_error();
                }
            }
            else {
                freeset(aliname, GBT_get_default_alignment(gb_main));
                if (!aliname || !aliname[0]) {
                    GB_clear_error();
                    error = "no default alignment defined (use --ali to specify alignment name)";
                }
                else {
                    fprintf(stderr, "Using default alignment '%s'\n", aliname);
                }
            }

            if (!error) {
                arb_assert(aliname && aliname[0]);

                long alisize = GBT_get_alignment_len(gb_main, aliname);
                if (alisize<=0) {
                    error = GB_await_error();
                }
                else {
                    const char *outname = args.get_fasta();
                    FILE       *out     = fopen(outname, "wt");

                    if (!out) {
                        error = GB_IO_error("writing", outname);
                    }
                    else {
                        FilteredExport exporter(gb_main, aliname, alisize);

                        if (args.get_head_aci()[0]) {
                            exporter.set_header_ACI(args.get_head_aci());
                        }
                        if (args.shall_accept_missing_data()) {
                            exporter.do_accept_missing_data();
                        }
                        if (args.get_min_bases_count()>0) {
                            exporter.set_required_baseCount(args.get_basechars(), args.get_min_bases_count());
                        }
                        if (args.get_seq_aci()[0]) {
                            exporter.set_sequence_ACI(args.get_seq_aci());
                        }
                        if (args.has_filters()) {
                            arb_assert(!error);
                            error = args.add_filters_to(exporter);
                        }

                        if (!error) {
                            error = exporter.write_fasta(out);
                        }
                        fclose(out);

                        if (error) GB_unlink_or_warn(outname, NULp);
                    }
                }
            }
            free(aliname);
        }
        GB_close(gb_main);
    }

    return error;
}

int main(int argc, char **argv) {
    GB_ERROR error = NULp;
    CLI      args(argc, const_cast<const char**>(argv));

    if (!error) error = args.get_error();

    if (!error && args.help_wanted()) {
        args.show_help_if_useful();
        return EXIT_FAILURE;
    }

    if (!error) error = export_seq_data(args);

    if (error) {
        args.show_help_if_useful();
        fprintf(stderr, "Error in arb_export_seq_filtered: %s\n", error);
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}