codebase / stimlib

#ifndef STIM_ARGUMENTS
#define STIM_ARGUMENTS

#include <string>
#include <vector>
#include <iostream>
#include <iomanip>
#include <sstream>
#include <iterator>
#include <algorithm>

#ifdef _WIN32
#include <Windows.h>
#endif

namespace stim{

	class cmd_option
	{
	private:
		bool ansi;

		//argument name
		std::string name;

		//description of the argument
		std::vector<std::string> desc;

		//argument values
		std::vector<std::string> vals;

		//range or example
		std::string range;

		//flag is true when the argument is user-specified
		bool flag;

		void parse_val(const std::string &s){

			vals.clear();

			std::stringstream ss(s);
			std::string item;
			while (std::getline(ss, item, ' ')) {
				vals.push_back(item);
			}
		}

		void parse_desc(const std::string &s){

			desc.clear();

			std::stringstream ss(s);
			std::string item;
			while (std::getline(ss, item, '\n')) {
				desc.push_back(item);
			}
		}

	public:
		void set_ansi(bool b){ ansi = b; }
        //create an option with a given name, description, and default value
		cmd_option(std::string _name, std::string _desc, std::string _default = "", std::string _range = "")
		{
			name = _name;
			parse_desc(_desc);
			parse_val(_default);

			//if a default value is provided, set the flag
			if(_default != "")
                flag = true;
            else flag = false;

			range = _range;


		}

		size_t nargs()
		{
            return vals.size();
        }

		//return the value of a text option
		std::string as_string(size_t n = 0)
		{
            if(!flag)
            {
                std::cout<<"ERROR - Option requested without being set: "<<name<<std::endl;
                exit(1);
            }

            if(vals.size() > n)
                return vals[n];

            else return "";
		}

        //return the value of a floating point option
		double as_float(size_t n = 0)
		{
            if(!flag)
            {
                std::cout<<"ERROR - option requested without being set: "<<name<<std::endl;
                exit(1);
            }

            if(vals.size() > n)
            {
                float r;
                if ( ! (std::istringstream(vals[n]) >> r) ) r = 0;
                return r;
            }

            else return 0;
		}

		//return the value of an integer option
		int as_int(size_t n = 0)
		{
            if(!flag)
            {
                std::cout<<"ERROR - option requested without being set: "<<name<<std::endl;
                exit(1);
            }

            if(vals.size() > n)
            {
                int r;
                if ( ! (std::istringstream(vals[n]) >> r) ) r = 0;
                return r;
            }

            else return 0;
		}

		//get the width of the left column
		size_t col_width()
		{
            size_t n = 3;
            //add the length of the option name
            n += name.size();

            //if there are any default parameters
            if(vals.size() > 0)
            {
                //padding (parenthesis, =, etc.)
                n += 6;

                //for each default option value
                for(size_t v=0; v<vals.size(); v++)
                    n += vals[v].size() + 1;
            }

            //add a buffer of 4 characters
            n += 4;

            return n;
		}


		//string output
		std::string toStr(size_t width = 0)
		{
			std::stringstream ss;

			int color_size = 0;


			//create the left column
			std::string left_part = std::string(" --") + name;
			if(vals.size() != 0)
			{
				if(ansi)
					left_part += "\033[1;32m";
				left_part += " ( = ";
				for(size_t v=0; v<vals.size(); v++)
					left_part += vals[v] + std::string(" ");
				left_part += ")";
				if(ansi)
					left_part += "\033[0m";
				if(ansi)
					color_size = 11;
			}
			else
                color_size = 0;

			//if no width is passed, put 4 spaces between left and right columns
			if(width == 0) width = col_width();

			ss<<std::left<<std::setw(width + color_size)<<left_part;

			//output right column
			for(size_t d=0; d<desc.size(); d++)
			{
				if(d == 0)
					ss<<desc[0];
				else
					ss<<std::endl<<std::setfill(' ')<<std::setw(width)<<" "<<desc[d];

			}

			//output the range in the right column
			if(range != "" && ansi)
                ss<<std::endl<<std::setfill(' ')<<std::setw(width)<<" "<<"    "<<std::string("\033[1;36m") + range + "\033[0m";
			else if(range != "")
				ss<<std::endl<<std::setfill(' ')<<std::setw(width)<<" "<<"    "<<range;

			return ss.str();
		}

		//compare the name of the option to a string
		bool operator==(std::string rhs)
		{
            return (name == rhs);
		}

		//set the option to a given value
		void set(std::string _value)
		{
            parse_val(_value);

            //set the flag
            flag = true;
		}

		bool is_set()
		{
            return flag;
        }

	};

	struct argsection
	{
		std::string name;
		size_t index;
	};

    /**The arglist class implements command line arguments.
        Example:

        1) Create an arglist instance:

            stim::arglist args;

        2) Add arguments:

            args.add("help", "prints this help");
            args.add("foo", "foo takes a single integer value", "", "[intval]");
            args.add("bar", "bar takes two floating point values", "", "[value1], [value2]");

        3) Parse the command line:

            args.parse(argc, argv);

        4) You generally want to immediately test for help and output available arguments:

            if(args["help"].is_set())
                std::cout<<args.str();



        5)  Retrieve values:

            int foo;
            float bar1, bar2;
            if(args["foo"])
                foo = args["foo"].as_int();
            if(args["bar"]){
                bar1 = args["bar"].as_float(0);
                bar2 = args["bar"].as_float(1);
            }


    **/

	class arglist
	{
    private:
		bool ansi;

		//vector of options
        std::vector<cmd_option> opts;
        std::vector<std::string> args;

		//column width of the longest option
        size_t col_width;

		//list of sections
		std::vector<argsection> sections;

    public:

        arglist(){
        	col_width = 0;

			ansi = true;

			//set ansi to false by default if this is a Windows system
			//		(console doesn't support ANSI colors)
			#ifdef _WIN32
				ansi=false;
			#endif

        }

        ///Sets ANSI support (default is on), which allows colored values in the help output.

        /// @param b is a boolean value specifying ANSI support (true is on, false is off)
		void set_ansi(bool b)
		{
			ansi = b;
			for(unsigned int i=0; i<opts.size(); i++)
				opts[i].set_ansi(ansi);
		}

        ///Add a supported command line argument.

        /// @param _name is the name of the command line argument (supplied as --name)
        /// @param _desc is a text description of the argument
        /// @param _default is the default value of the argument
        /// @param _range is the supported range of values, list of values, etc. It will be displayed to the user.
        void add(std::string _name, std::string _desc, std::string _default = "", std::string _range = "")
        {
            cmd_option opt(_name, _desc, _default, _range);
			opt.set_ansi(ansi);
            opts.push_back(opt);

            col_width = std::max<size_t>(col_width, opt.col_width());
        }

        ///Specifies a section name that can be used to organize parameters in the output.

        /// @param _name is the name of the section, which will be displayed to the user
		void section(std::string _name)
		{
			argsection s;
			s.name = _name;
			s.index = opts.size();
			sections.push_back(s);
		}

        ///Outputs supported arguments. If ANSI support is available, they will be color-coded. Generally this is called in response to "--help"
        std::string str()
        {
            std::stringstream ss;

            int si = -1;

            if(sections.size() > 0)
                si = 0;

            //for each option
            for(unsigned int a=0; a<opts.size(); a++)
            {
                if(si != -1 && a == sections[si].index)
                {
					if(ansi)
						ss<<std::endl<<std::left<<std::setw(col_width)<<std::string("\033[1;31m") + sections[si].name<<"\033[0m"<<std::endl;
					else
						ss<<std::endl<<std::left<<std::setw(col_width)<<sections[si].name<<std::endl;
                    si++;
                    if(si == (int)sections.size()) si = -1;
                }

                ss<<opts[a].toStr(col_width)<<std::endl;
            }

            return ss.str();
        }

        ///Retrieves the index for a supported argument, given that argument's name.

        /// @param _name is the name of the requested argument
        size_t index(std::string _name)
        {
        	size_t i = find(opts.begin(), opts.end(), _name) - opts.begin();

            if(i >= opts.size()){
				std::cout<<"ERROR stim::arglist: option name '"<<_name<<"' not found"<<std::endl; 
                exit(1);
			}

            return i;
        }

        ///Sets an argument to a specified value

        /// @param _name is the name of the argument to be set
        /// @param _value is the value that it is given
        void set(std::string _name, std::string _value){
            size_t i = index(_name);

			opts[i].set(_value);
			//adjust the column width if necessary
			col_width = (std::max)(col_width, opts[i].col_width());
        }

        ///Parses the command line

        /// @param argc is the number of command line arguments (provided by the OS)
        /// @param argv[] is the list of command line arguments (provided by the OS)
        void parse(int argc, char* argv[])
        {
            //if the number of options is 1, we're done
            if(argc <= 1) return;

            std::string name;
            std::string params;

            bool args_done = false;		//create a flag that turns true when the first option is encountered

            for(int i=1; i<argc; i++)
            {
                //if the argument is an option
                if(argv[i][0] == '-' && argv[i][1] == '-')
                {
                	args_done = true;				//arguments for the executable are done, all options now
                    //add any previous options
                    if(name != "")
                        set(name, params);
                    //set the current option to this name
                    name = argv[i]+2;
                    //clear the parameters list
                    params = "";
                }
                else if(!args_done){
                	args.push_back(argv[i]);
                }
                else{	//everything else is an arg for the most recent option
					if(params != "")
						params += " ";
                    params += argv[i];
                }
            }

            //set the last option
            if(name != "")
            	set(name, params);
        }

        ///Determines of a parameter has been set and returns true if it has

        /// @param _name is the name of the argument
        bool operator()(std::string _name)
        {
            size_t i = find(opts.begin(), opts.end(), _name) - opts.begin();

            if(i < 0){
                std::cout<<"ERROR - Unspecified parameter name: "<<_name<<std::endl;
                exit(1);
            }

            return opts[i].is_set();
        }

        ///Returns the number of parameters for a specified argument

        /// @param _name is the name of the argument whose parameter number will be returned
        size_t nargs(std::string _name)
        {
            size_t i = find(opts.begin(), opts.end(), _name) - opts.begin();

            if(i < 0){
                std::cout<<"ERROR - Unspecified parameter name: "<<_name<<std::endl;
                exit(1);
            }

            return opts[i].nargs();
        }

        ///Returns the number of arguments that have been set
        size_t nargs(){
        	return args.size();
        }

        /// Returns the number of options that are set
        size_t nopts(){
        	size_t n = 0;									//initialize the counter for the number of options
        	for(size_t i = 0; i < opts.size(); i++)		//go through each option
        		if(opts[i].is_set()) n++;						//if a value is specified, increment the counter
        	return n;
        }

        ///Returns the name of an argument, given its index

        /// @param a is the index of the requested argument
        std::string arg(size_t a){
        	return args[a];
        }

        ///Returns an object describing the argument

        /// @param _name is the name of the requested argument
        cmd_option operator[](std::string _name)
        {
            size_t i = find(opts.begin(), opts.end(), _name) - opts.begin();

            if(i < 0 || i >= opts.size())
            {
                std::cout<<"ERROR - Unspecified parameter name: "<<_name<<std::endl;
                exit(1);
            }

            return opts[i];
        }


	};



}	//end namespace stim

#endif