/*
 *      CRFsuite C++/SWIG API wrapper.
 *
 * Copyright (c) 2007-2010, Naoaki Okazaki
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the names of the authors nor the names of its contributors
 *       may be used to endorse or promote products derived from this
 *       software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#ifndef __CRFSUITE_HPP__
#define __CRFSUITE_HPP__

#include <cmath>
#include <vector>
#include <string>
#include <stdexcept>
#include <iostream>
#include <sstream>

#include <crfsuite.h>
#include "crfsuite_api.hpp"

namespace CRFSuite
{

Trainer::Trainer()
{
    data = new crfsuite_data_t;
    if (data != NULL) {
        crfsuite_data_init(data);
    }
    tr = NULL;
}

Trainer::~Trainer()
{
    if (data != NULL) {
        clear();
        delete data;
        data = NULL;
    }
    if (tr != NULL) {
        tr->release(tr);
        tr = NULL;
    }
}

void Trainer::init()
{
    // Create an instance of attribute dictionary.
    if (data->attrs == NULL) {
        int ret = crfsuite_create_instance("dictionary", (void**)&data->attrs);
        if (!ret) {
            throw std::runtime_error("Failed to create a dictionary instance for attributes.");
        }
    }

    // Create an instance of label dictionary.
    if (data->labels == NULL) {
        int ret = crfsuite_create_instance("dictionary", (void**)&data->labels);
        if (!ret) {
            throw std::runtime_error("Failed to create a dictionary instance for labels.");
        }
    }
}

void Trainer::clear()
{
    if (data != NULL) {
        if (data->labels != NULL) {
            data->labels->release(data->labels);
            data->labels = NULL;
        }

        if (data->attrs != NULL) {
            data->attrs->release(data->attrs);
            data->attrs = NULL;
        }

        crfsuite_data_finish(data);
        crfsuite_data_init(data);
    }
}

void Trainer::append(const ItemSequence& xseq, const StringList& yseq, int group)
{
    // Create dictionary objects if necessary.
    if (data->attrs == NULL || data->labels == NULL) {
        init();
    }

    // Make sure |y| == |x|.
    if (xseq.size() != yseq.size()) {
        std::stringstream ss;
        ss << "The numbers of items and labels differ: |x| = " << xseq.size() << ", |y| = " << yseq.size();
        throw std::invalid_argument(ss.str());
    }

    // Convert instance_type to crfsuite_instance_t.
    crfsuite_instance_t _inst;
    crfsuite_instance_init_n(&_inst, xseq.size());
    for (size_t t = 0;t < xseq.size();++t) {
        const Item& item = xseq[t];
        crfsuite_item_t* _item = &_inst.items[t];

        // Set the attributes in the item.
        crfsuite_item_init_n(_item, item.size());
        for (size_t i = 0;i < item.size();++i) {
            _item->contents[i].aid = data->attrs->get(data->attrs, item[i].attr.c_str());
            _item->contents[i].value = (floatval_t)item[i].value;
        }

        // Set the label of the item.
        _inst.labels[t] = data->labels->get(data->labels, yseq[t].c_str());
    }
    _inst.group = group;

    // Append the instance to the training set.
    crfsuite_data_append(data, &_inst);

    // Finish the instance.
    crfsuite_instance_finish(&_inst);
}

bool Trainer::select(const std::string& algorithm, const std::string& type)
{
    int ret;

    // Release the trainer if it is already initialized.
    if (tr != NULL) {
        tr->release(tr);
        tr = NULL;
    }

    // Build the trainer string ID.
    std::string tid = "train/";
    tid += type;
    tid += '/';
    tid += algorithm;

    // Create an instance of a trainer.
    ret = crfsuite_create_instance(tid.c_str(), (void**)&tr);
    if (!ret) {
        return false;
    }

    // Set the callback function for receiving messages.
    tr->set_message_callback(tr, this, __logging_callback);

    return true;
}

int Trainer::train(const std::string& model, int holdout)
{
    int ret;

    if (tr == NULL) {
        std::stringstream ss;
        ss << "The trainer is not initialized. Call Trainer::select before Trainer::train.";
        throw std::invalid_argument(ss.str());
    }

    if (data->attrs == NULL || data->labels == NULL) {
        std::stringstream ss;
        ss << "The data is empty. Call Trainer::append before Trainer::train.";
        throw std::invalid_argument(ss.str());
    }

    // Run the training algorithm.
    ret = tr->train(tr, data, model.c_str(), holdout);

    return ret;
}

StringList Trainer::params()
{
    StringList pars;
    crfsuite_params_t* params = tr->params(tr);
    int n = params->num(params);
    for (int i = 0;i < n;++i) {
        char *name = NULL;
        params->name(params, i, &name);
        pars.push_back(name);
        params->free(params, name);
    }
    return pars;
}

void Trainer::set(const std::string& name, const std::string& value)
{
    crfsuite_params_t* params = tr->params(tr);
    if (params->set(params, name.c_str(), value.c_str()) != 0) {
        std::stringstream ss;
        ss << "Parameter not found: " << name << " = " << value;
        params->release(params);
        throw std::invalid_argument(ss.str());
    }
    params->release(params);
}

std::string Trainer::get(const std::string& name)
{
    std::string value;
    char *_value = NULL;
    crfsuite_params_t* params = tr->params(tr);
    if (params->get(params, name.c_str(), &_value) != 0) {
        std::stringstream ss;
        ss << "Parameter not found: " << name << " = " << value;
        params->release(params);
        throw std::invalid_argument(ss.str());
    }
    value = _value;
    params->free(params, _value);
    params->release(params);
    return value;
}

std::string Trainer::help(const std::string& name)
{
    std::string str;
    crfsuite_params_t* params = tr->params(tr);
    char *_str = NULL;
    if (params->help(params, name.c_str(), NULL, &_str) != 0) {
        std::stringstream ss;
        ss << "Parameter not found: " << name;
        params->release(params);
        throw std::invalid_argument(ss.str());
    }
    str = _str;
    params->free(params, _str);
    params->release(params);
    return str;
}

void Trainer::message(const std::string& msg)
{
}

int Trainer::__logging_callback(void *instance, const char *format, va_list args)
{
    char buffer[65536];
    vsnprintf(buffer, sizeof(buffer)-1, format, args);
    reinterpret_cast<Trainer*>(instance)->message(buffer);
    return 0;
}



Tagger::Tagger()
{
    model = NULL;
    tagger = NULL;
}

Tagger::~Tagger()
{
    this->close();
}

bool Tagger::open(const std::string& name)
{
    int ret;

    // Close the model if it is already opened.
    this->close();

    // Open the model file.
    if ((ret = crfsuite_create_instance_from_file(name.c_str(), (void**)&model))) {
        return false;
    }

    // Obtain the tagger interface.
    if ((ret = model->get_tagger(model, &tagger))) {
        throw std::runtime_error("Failed to obtain the tagger interface");
    }

    return true;
}

bool Tagger::open(const void* data, std::size_t size)
{
    int ret;

    // Close the model if it is already opened.
    this->close();

    // Open the model.
    if ((ret = crfsuite_create_instance_from_memory(data, size, (void**)&model))) {
        return false;
    }

    // Obtain the tagger interface.
    if ((ret = model->get_tagger(model, &tagger))) {
        throw std::runtime_error("Failed to obtain the tagger interface");
    }

    return true;
}

void Tagger::close()
{
    if (tagger != NULL) {
        tagger->release(tagger);
        tagger = NULL;
    }
    if (model != NULL) {
        model->release(model);
        model = NULL;
    }
}

StringList Tagger::labels()
{
    int ret;
    StringList lseq;
    crfsuite_dictionary_t *labels = NULL;

    if (model == NULL) {
        throw std::invalid_argument("The tagger is not opened");
    }

    // Obtain the dictionary interface representing the labels in the model.
    if ((ret = model->get_labels(model, &labels))) {
        throw std::runtime_error("Failed to obtain the dictionary interface for labels");
    }

    // Collect all label strings to lseq.
    for (int i = 0;i < labels->num(labels);++i) {
        const char *label = NULL;
        if (labels->to_string(labels, i, &label) != 0) {
            labels->release(labels);
            throw std::runtime_error("Failed to convert a label identifier to string.");
        }
        lseq.push_back(label);
        labels->free(labels, label);
    }

    labels->release(labels);
    return lseq;
}

StringList Tagger::tag(const ItemSequence& xseq)
{
    set(xseq);
    return viterbi();
}

void Tagger::set(const ItemSequence& xseq)
{
    int ret;
    StringList yseq;
    crfsuite_instance_t _inst;
    crfsuite_dictionary_t *attrs = NULL;

    if (model == NULL || tagger == NULL) {
        throw std::invalid_argument("The tagger is not opened");
    }

    // Obtain the dictionary interface representing the attributes in the model.
    if ((ret = model->get_attrs(model, &attrs))) {
        throw std::runtime_error("Failed to obtain the dictionary interface for attributes");
    }

    // Build an instance.
    crfsuite_instance_init_n(&_inst, xseq.size());
    for (size_t t = 0;t < xseq.size();++t) {
        const Item& item = xseq[t];
        crfsuite_item_t* _item = &_inst.items[t];

        // Set the attributes in the item.
        crfsuite_item_init(_item);
        for (size_t i = 0;i < item.size();++i) {
            int aid = attrs->to_id(attrs, item[i].attr.c_str());
            if (0 <= aid) {
                crfsuite_attribute_t cont;
                crfsuite_attribute_set(&cont, aid, item[i].value);
                crfsuite_item_append_attribute(_item, &cont);
            }
        }
    }

    // Set the instance to the tagger.
    if ((ret = tagger->set(tagger, &_inst))) {
        crfsuite_instance_finish(&_inst);
        attrs->release(attrs);
        throw std::runtime_error("Failed to set the instance to the tagger.");
    }

    crfsuite_instance_finish(&_inst);
    attrs->release(attrs);
}

StringList Tagger::viterbi()
{
    int ret;
    StringList yseq;
    crfsuite_dictionary_t *labels = NULL;

    if (model == NULL || tagger == NULL) {
        throw std::invalid_argument("The tagger is not opened");
    }

    // Make sure that the current instance is not empty.
    const size_t T = (size_t)tagger->length(tagger);
    if (T <= 0) {
        return yseq;
    }

    // Obtain the dictionary interface representing the labels in the model.
    if ((ret = model->get_labels(model, &labels))) {
        throw std::runtime_error("Failed to obtain the dictionary interface for labels");
    }

    // Run the Viterbi algorithm.
    floatval_t score;
    int *path = new int[T];
    if ((ret = tagger->viterbi(tagger, path, &score))) {
        delete[] path;
        labels->release(labels);
        throw std::runtime_error("Failed to find the Viterbi path.");
    }

    // Convert the Viterbi path to a label sequence.
    yseq.resize(T);
    for (size_t t = 0;t < T;++t) {
        const char *label = NULL;
        if (labels->to_string(labels, path[t], &label) != 0) {
            delete[] path;
            labels->release(labels);
            throw std::runtime_error("Failed to convert a label identifier to string.");
        }
        yseq[t] = label;
        labels->free(labels, label);
    }

    labels->release(labels);
    delete[] path;
    return yseq;
}

double Tagger::probability(const StringList& yseq)
{
    int ret;
    size_t T;
    int *path = NULL;
    std::stringstream msg;
    floatval_t score, lognorm;
    crfsuite_dictionary_t *labels = NULL;

    if (model == NULL || tagger == NULL) {
        msg << "The tagger is not opened";
        throw std::invalid_argument(msg.str());
    }

    // Make sure that the current instance is not empty.
    T = (size_t)tagger->length(tagger);
    if (T <= 0) {
        return 0.;
    }

    // Make sure that |y| == |x|.
    if (yseq.size() != T) {
        msg << "The numbers of items and labels differ: |x| = " << T << ", |y| = " << yseq.size();
        throw std::invalid_argument(msg.str());
    }

    // Obtain the dictionary interface representing the labels in the model.
    if ((ret = model->get_labels(model, &labels))) {
        msg << "Failed to obtain the dictionary interface for labels";
        goto error_exit;
    }

    // Convert string labels into label IDs.
    path = new int[T];
    for (size_t t = 0;t < T;++t) {
        int l = labels->to_id(labels, yseq[t].c_str());
        if (l < 0) {
            msg << "Failed to convert into label identifier: " << yseq[t];
            goto error_exit;
        }
        path[t] = l;
    }

    // Compute the score of the path.
    if ((ret = tagger->score(tagger, path, &score))) {
        msg << "Failed to score the label sequence";
        goto error_exit;
    }

    // Compute the partition factor.
    if ((ret = tagger->lognorm(tagger, &lognorm))) {
        msg << "Failed to compute the partition factor";
        goto error_exit;
    }

    labels->release(labels);
    delete[] path;
    return std::exp((double)(score - lognorm));

error_exit:
    if (labels != NULL) {
        labels->release(labels);
        labels = NULL;
    }
    delete[] path;
    throw std::runtime_error(msg.str());
}

double Tagger::marginal(const std::string& y, const int t)
{
    int l, ret, T;
    floatval_t prob;
    std::stringstream msg;
    crfsuite_dictionary_t *labels = NULL;

    if (model == NULL || tagger == NULL) {
        msg << "The tagger is not opened";
        throw std::invalid_argument(msg.str());
    }

    // Make sure that the current instance is not empty.
    T = tagger->length(tagger);
    if (T <= 0) {
        return 0.;
    }

    // Make sure that 0 <= t < |x|.
    if (t < 0 || T <= t) {
        msg << "The position, " << t << "is out of range of " << T;
        throw std::invalid_argument(msg.str());
    }

    // Obtain the dictionary interface representing the labels in the model.
    if ((ret = model->get_labels(model, &labels))) {
        msg << "Failed to obtain the dictionary interface for labels";
        goto error_exit;
    }

    // Convert string labels into label IDs.
    l = labels->to_id(labels, y.c_str());
    if (l < 0) {
        msg << "Failed to convert into label identifier: " << y;
        goto error_exit;
    }

    // Compute the score of the path.
    if ((ret = tagger->marginal_point(tagger, l, t, &prob))) {
        msg << "Failed to compute the marginal probability of '" << y << "' at " << t;
        goto error_exit;
    }

    labels->release(labels);
    return prob;

error_exit:
    if (labels != NULL) {
        labels->release(labels);
        labels = NULL;
    }
    throw std::runtime_error(msg.str());
}


std::string version()
{
    return CRFSUITE_VERSION;
}

}

#endif/*__CRFSUITE_HPP__*/