/*  FreeContact - program to predict protein residue contacts from a sufficiently large multiple alignment
*   Copyright (C) 2013 by Laszlo Kajan, Technical University of Munich, Germany
*
*   This program is free software: you can redistribute it and/or modify
*   it under the terms of the GNU General Public License as published by
*   the Free Software Foundation, either version 3 of the License, or
*   (at your option) any later version.
*
*   This program is distributed in the hope that it will be useful,
*   but WITHOUT ANY WARRANTY; without even the implied warranty of
*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*   GNU General Public License for more details.
*
*   You should have received a copy of the GNU General Public License
*   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
#include <boost/python.hpp>
#include <boost/python/raw_function.hpp>
#include <freecontact.h>

namespace bp = boost::python;
namespace fc = freecontact;

typedef freecontact::predictor::cont_res_t	cont_res_t;
typedef freecontact::predictor::time_res_t	time_res_t;
typedef freecontact::predictor::freq_vec_t	freq_vec_t;

char const* greet(int hel, int word) { return "hello, world"; }

bp::dict            _get_ps(const fc::parset_t& __ps)
{
    bp::dict d;
    d["clustpc"] = __ps.clustpc;
    d["density"] = __ps.density;
    d["gapth"] = __ps.gapth;
    d["mincontsep"] = __ps.mincontsep;

    d["pseudocnt"] = __ps.pseudocnt;
    d["pscnt_weight"] = __ps.pscnt_weight;
    d["estimate_ivcov"] = __ps.estimate_ivcov;
    d["shrink_lambda"] = __ps.shrink_lambda;

    d["cov20"] = __ps.cov20;
    d["apply_gapth"] = __ps.apply_gapth;
    d["rho"] = __ps.rho;
    return d;
}

bp::dict            get_ps_evfold() { return _get_ps(fc::ps_evfold); }
bp::dict            get_ps_psicov() { return _get_ps(fc::ps_psicov); }
bp::dict            get_ps_psicov_sd() { return _get_ps(fc::ps_psicov_sd); }

class __attribute__ ((visibility ("hidden"))) py_predictor : public fc::predictor {
    public:
                    py_predictor(bool __dbg = false) : fc::predictor(__dbg){}

    bp::dict        py_get_seq_weights(const bp::list& __ali, double __clustpc,
                        bool __veczw, int __num_threads);

/*    bp::dict        py_run_with_seq_weights(const bp::list& __ali, const bp::list& __aliw, double __wtot,
                        bp::dict& __timing,
                        double __density, double __gapth, uint16_t __mincontsep,
                        double __pseudocnt, double __pscnt_weight, bool __estimate_ivcov, double __shrink_lambda,
                        bool __cov20, bool __apply_gapth, double __rho,
                        int __num_threads, time_t __icme_timeout)*/
};

fc::ali_t           input_t_ali(const bp::list& __ali)
{
    size_t seqcnt = bp::len(__ali);
    if(!seqcnt) return fc::ali_t();

    //            alilen
    fc::ali_t ali(bp::len(__ali[0])); for(size_t i = 0, i_e = seqcnt; i < i_e; ++i) ali.push(bp::extract<std::string>(__ali[i]));
    return ali;
}

bp::dict            py_predictor::py_get_seq_weights(const bp::list& __ali, double __clustpc,
                        bool __veczw, int __num_threads)
{
    freq_vec_t aliw; double wtot = 0;
    bp::list aliw_list;
    size_t seqcnt = bp::len(__ali);
    if(seqcnt)
    {
        fc::ali_t ali = input_t_ali(__ali);
        get_seq_weights(aliw, wtot, ali, __clustpc, __veczw, __num_threads);

        for(size_t i = 0, i_e = aliw.size(); i < i_e; ++i) aliw_list.append(aliw[i]);
    }
    bp::dict aliw_wtot;
    aliw_wtot["aliw"] = aliw_list;
    aliw_wtot["wtot"] = wtot;
    return aliw_wtot;
}

void                translate_runtime_error(const std::runtime_error& __e)
{
    PyErr_SetString(PyExc_RuntimeError, __e.what());
}

freq_vec_t          input_freq_vec_t(const bp::list& __aliw)
{
    freq_vec_t aliw(len(__aliw));
    for(size_t i = 0, i_e = len(__aliw); i < i_e; ++i) aliw[i] = bp::extract<freq_vec_t::value_type>(__aliw[i]);
    return aliw;
}

bp::dict            output_cont_res_t(const cont_res_t& __cont_res)
{
    if(__cont_res.empty()) return bp::dict();
    else
    {
   		// return { 'l1norm' => [ [ 0, 1, 45.786 ], [ 0, 2, 7.993 ], ... ], ... }
		bp::dict res;
		for(cont_res_t::const_iterator cr_i = __cont_res.begin(), cr_e = __cont_res.end(); cr_i != cr_e; ++cr_i)
		{
			bp::list contacts;
			for(std::vector<fc::contact_t>::const_iterator ct_i = cr_i->second.begin(), ct_e = cr_i->second.end(); ct_i != ct_e; ++ct_i)
			{
				bp::list ijscore;

				ijscore.append(ct_i->i);
				ijscore.append(ct_i->j);
				ijscore.append(ct_i->score);
				//
				contacts.append(ijscore);
			}
			res[cr_i->first] = contacts;
		}
        return res;
    }
}

void                output_time_res_t(const time_res_t* __p_timing, bp::dict __py_timing)
{
    if(__p_timing)
        for(time_res_t::const_iterator tr_i = __p_timing->begin(), tr_e = __p_timing->end(); tr_i != tr_e; ++tr_i)
            __py_timing[tr_i->first] = tr_i->second;
}

void                input_parset_t(bp::dict& __kw, fc::parset_t& __ps)
{
    if(__kw.has_key("clustpc")) __ps.clustpc = bp::extract<double>(__kw["clustpc"]);

    if(__kw.has_key("density")) __ps.density = bp::extract<double>(__kw["density"]);
    if(__kw.has_key("gapth")) __ps.gapth = bp::extract<double>(__kw["gapth"]);
    if(__kw.has_key("mincontsep")) __ps.mincontsep = bp::extract<uint16_t>(__kw["mincontsep"]);

    if(__kw.has_key("pseudocnt")) __ps.pseudocnt = bp::extract<double>(__kw["pseudocnt"]);
    if(__kw.has_key("pscnt_weight")) __ps.pscnt_weight = bp::extract<double>(__kw["pscnt_weight"]);
    if(__kw.has_key("estimate_ivcov")) __ps.estimate_ivcov = bp::extract<bool>(__kw["estimate_ivcov"]);
    if(__kw.has_key("shrink_lambda")) __ps.shrink_lambda = bp::extract<double>(__kw["shrink_lambda"]);

    if(__kw.has_key("cov20")) __ps.cov20 = bp::extract<bool>(__kw["cov20"]);
    if(__kw.has_key("apply_gapth")) __ps.apply_gapth = bp::extract<bool>(__kw["apply_gapth"]);
    if(__kw.has_key("rho")) __ps.rho = bp::extract<double>(__kw["rho"]);
}

bp::dict            py_run(bp::tuple __args, bp::dict __kw)
{
    py_predictor& self = bp::extract<py_predictor&>(__args[0])();

    if(!__kw.has_key("ali")) throw std::runtime_error("required argument 'ali' is missing");
    if(!bp::extract<bp::list>(__kw["ali"]).check()) throw std::runtime_error("argument 'ali' is not a list");

    const bp::list& py_ali = bp::extract<bp::list>(__kw["ali"]);
    size_t seqcnt = bp::len(py_ali);
    if(seqcnt)
    {
        time_res_t timing, *p_timing = NULL;
        fc::parset_t ps = fc::ps_evfold;

        fc::ali_t ali = input_t_ali(py_ali);
        input_parset_t(__kw, ps);
        if(__kw.has_key("timing") && bp::extract<bp::dict>(__kw["timing"]).check()) p_timing = &timing;

        cont_res_t cont_res = self.run(
                ali,
                ps.clustpc,
                ps.density, ps.gapth, ps.mincontsep,
                ps.pseudocnt, ps.pscnt_weight, ps.estimate_ivcov, ps.shrink_lambda,
                ps.cov20, ps.apply_gapth, ps.rho,
                __kw.has_key("veczw") ? bp::extract<bool>(__kw["veczw"])() : true, __kw.has_key("num_threads") ? bp::extract<int>(__kw["num_threads"])() : 0, __kw.has_key("icme_timeout") ? bp::extract<time_t>(__kw["icme_timeout"])() : 1800, p_timing
        );

        if(p_timing) output_time_res_t(p_timing, bp::extract<bp::dict>(__kw["timing"]));

        return output_cont_res_t(cont_res);
    }
    else return bp::dict();
}

bp::dict            py_run_with_seq_weights(bp::tuple __args, bp::dict __kw)
{
    py_predictor& self = bp::extract<py_predictor&>(__args[0])();

    if(!__kw.has_key("ali")) throw std::runtime_error("required argument 'ali' is missing");
    if(!bp::extract<bp::list>(__kw["ali"]).check()) throw std::runtime_error("argument 'ali' is not a list");

    const bp::list& py_ali = bp::extract<bp::list>(__kw["ali"]);
    size_t seqcnt = bp::len(py_ali);
    if(seqcnt)
    {
        time_res_t timing, *p_timing = NULL;
        fc::parset_t ps = fc::ps_evfold;

        fc::ali_t ali = input_t_ali(py_ali);
        input_parset_t(__kw, ps);
        if(__kw.has_key("timing") && bp::extract<bp::dict>(__kw["timing"]).check()) p_timing = &timing;

        if(!__kw.has_key("aliw")) throw std::runtime_error("required argument 'aliw' is missing");
        if(!bp::extract<bp::list>(__kw["aliw"]).check()) throw std::runtime_error("argument 'aliw' is not a list");
        if(!__kw.has_key("wtot")) throw std::runtime_error("required argument 'wtot' is missing");

        freq_vec_t aliw = input_freq_vec_t(bp::extract<bp::list>(__kw["aliw"]));

        cont_res_t cont_res = self.run(
                ali,
                aliw, bp::extract<double>(__kw["wtot"]),
                ps.density, ps.gapth, ps.mincontsep,
                ps.pseudocnt, ps.pscnt_weight, ps.estimate_ivcov, ps.shrink_lambda,
                ps.cov20, ps.apply_gapth, ps.rho,
                __kw.has_key("num_threads") ? bp::extract<int>(__kw["num_threads"])() : 0, __kw.has_key("icme_timeout") ? bp::extract<time_t>(__kw["icme_timeout"])() : 1800, p_timing
        );

        if(p_timing) output_time_res_t(p_timing, bp::extract<bp::dict>(__kw["timing"]));

        return output_cont_res_t(cont_res);
    }
    else return bp::dict();
}

BOOST_PYTHON_MODULE(freecontact)
{
    using namespace boost::python;
    register_exception_translator<std::runtime_error>(&translate_runtime_error);

    // http://www.boost.org/doc/libs/1_49_0/libs/python/doc
    // http://www.boost.org/doc/libs/1_49_0/libs/python/doc/tutorial
    def("get_ps_evfold", get_ps_evfold,
"Get parameters for EVfold-mfDCA operating mode."
    );
    def("get_ps_psicov", get_ps_psicov,
"Get parameters for PSICOV 'improved results' operating mode."
    );
    def("get_ps_psicov_sd", get_ps_psicov_sd,
"Get parameters for PSICOV 'sensible default' operating mode. This is much faster than 'improved results' for a "
"slight loss of precision.\n"
"\n"
"These get_ps_() functions return a dictionary of arguments {'clustpc': num,...,'rho': num} that can be used "
"with get_seq_weights(), run() or run_with_seq_weights(). The arguments correspond to the published parametrization of the "
"respective method."
    );

    class_<py_predictor>("Predictor", init<optional<bool> >(arg("dbg")))
        .def_readwrite("dbg", &py_predictor::dbg)
        .def("get_seq_weights", &py_predictor::py_get_seq_weights,
                (arg("ali"), arg("clustpc") = fc::ps_evfold.clustpc,
                 arg("veczw") = true, arg("num_threads") = 0),
"Defaults for the arguments are obtained with get_ps_evfold()."
        )
        .def("run", raw_function(py_run, 1),
"run(ali = [], clustpc = dbl, "
"density = dbl, gapth = dbl, mincontsep = uint, "
"pseudocnt = dbl, pscnt_weight = dbl, estimate_ivcov = bool, shrink_lambda = dbl, "
"cov20 = bool, apply_gapth = bool, rho = dbl, "
"[veczw = bool], [num_threads = int], [icme_timeout = int], [timing = {}])\n"
"\n"
"Defaults for the arguments are obtained with get_ps_evfold().\n"
"\n"
"ali\n"
"    List holding alignment rows as strings. The first "
"row must hold the query, with no gaps.\n"
"\n"
"clustpc\n"
"    BLOSUM-style clustering similarity threshold [0-1].\n"
"\n"
"icme_timeout\n"
"    Inverse covariance matrix estimation timeout in seconds. Default: 1800.\n"
"\n"
"    The estimation sometimes gets stuck. If the timeout is reached, the run() "
"method raises a RuntimeError exception with \"inverse covariance matrix estimation exceeded time limit ...\". You can catch "
"this exception and handle it as needed, e.g. by setting a higher rho value.\n"
"\n"
"num_threads\n"
"\n"
"    Number of OpenMP threads to use. If unset, all CPUs are used.\n"
"\n"
"timing\n"
"\n"
"If given, this dictionary is filled with data containing wall clock "
"timing results in seconds:\n"
"\n"
"      {\n"
"        'num_threads':  NUM,\n"
"        'seqw':         NUM,\n"
"        'pairfreq':     NUM,\n"
"        'shrink':       NUM,\n"
"        'inv':          NUM,\n"
"        'all':          NUM\n"
"      }\n"
"\n"
"run() returns a dictionary of contact prediction results:\n"
"\n"
"      {\n"
"        'fro': [  # identifier of scoring scheme\n"
"          [\n"
"            I,    # 0-based index of amino acid i\n"
"            J,    # 0-based index of amino acid j\n"
"            SCORE # contact score\n"
"          ], ...\n"
"        ],\n"
"        'MI': ...,\n"
"        'l1norm': ...\n"
"      }\n"
"    \n"
"Use 'fro' scores with EVfold."
        )
        .def("run_with_seq_weights", raw_function(py_run_with_seq_weights, 1))
    ;

    // set the docstring of the current module scope http://wiki.python.org/moin/boost.python/module
    scope().attr("__doc__") = "fast protein contact predictor\n"
"\n"
"FreeContact is a protein residue contact predictor optimized for speed. "
"Its input is a multiple sequence alignment. FreeContact can function as an "
"accelerated drop-in for the published contact predictors "
"EVfold-mfDCA of DS. Marks (2011) and "
"PSICOV of D. Jones (2011). "
"FreeContact is accelerated by a combination of vector instructions, multiple "
"threads, and faster implementation of key parts. "
"Depending on the alignment, 10-fold or higher speedups are possible.\n"
"\n"
"A sufficiently large alignment is required for meaningful results. "
"As a minimum, an alignment with an effective (after-weighting) sequence count "
"bigger than the length of the query sequence should be used. Alignments with "
"tens of thousands of (effective) sequences are considered good input.\n"
"\n"
"jackhmmer(1) from the hmmer package, or hhblits(1) from hhsuite "
"can be used to generate the alignments, for example.\n"
"\n"
"SYNOPSIS\n"
"========\n"
"import freecontact\n"
"\n"
"EXAMPLE = open('/usr/share/doc/python-freecontact/examples/demo_1000.aln', 'r')\n"
"aln = EXAMPLE.readlines(); aln = map(lambda s: s.rstrip(), aln)\n"
"EXAMPLE.close()\n"
"\n"
"contacts = freecontact.Predictor().run(ali = aln)\n"
"\n"
"predictor = freecontact.Predictor()\n"
"parset = freecontact.get_ps_evfold()\n"
"parset.update({'ali': aln, 'num_threads': 1})\n"
"contacts = predictor.run(**parset)\n"
"\n"
"predictor = freecontact.Predictor()\n"
"aliw_wtot = predictor.get_seq_weights(ali = aln, num_threads = 1)\n"
"contacts = predictor.run_with_seq_weights(ali = aln, aliw = aliw_wtot['aliw'], wtot = aliw_wtot['wtot'], num_threads = 1)"
;
}

// vim:et:ts=4:ai:
