#include <shogun/latent/LatentSOModel.h>

using namespace shogun;

CLatentSOModel::CLatentSOModel()
    : CLatentModel(),
    m_sum_argmax_h(0.0)
{
    register_parameters();
}

CLatentSOModel::CLatentSOModel(CLatentFeatures* feats, CLatentLabels* labels)
    : CLatentModel(feats, labels),
    m_sum_argmax_h(0.0)
{
    register_parameters();
}

CLatentSOModel::~CLatentSOModel()
{

}

void CLatentSOModel::register_parameters()
{

}

float64_t CLatentSOModel::get_sum_argmax_h() const
{
    return m_sum_argmax_h;
}

void CLatentSOModel::argmax_h(const SGVector<float64_t>& w)
{
    int32_t num = get_num_vectors();
    ASSERT(num > 0);

    for (int32_t i = 0; i < num; ++i)
    {
        // infer h and set it for the argmax_h <w,psi(x,h)>
        CData* latent_data = infer_latent_variable(w, i);
        m_labels->set_latent_label(i, latent_data);
    }
}

CLSOModel::CLSOModel()
    : CStructuredModel()
{

}

CLSOModel::CLSOModel(CLatentSOModel* latent_model)
    : CStructuredModel(),
    m_latent_model(latent_model)
{
    SG_REF(m_latent_model);
}

CLSOModel::~CLSOModel()
{
    SG_UNREF(m_latent_model);
}

int32_t CLSOModel::get_dim() const
{
    m_latent_model->get_dim();
}

CResultSet* CLSOModel::argmax(SGVector< float64_t > w, int32_t feat_idx, bool const training)
{
    CStructuredData* ybar, y;
    CData* hbar;

    m_latent_model->get_most_violated_constraint(w, feat_idx, ybar, hbar);

    CResultSet* ret = new CResultSet();
    SG_REF(ret);
    ret->score = ;
    ret->delta = m_latent_model->delta(y, ybar, hbar);

    return ret;
}

float64_t CLSOModel::risk(float64_t* subgrad, float64_t* W, TMultipleCPinfo* info)
{
    int32_t w_dim = get_dim();
    float64_t R = 0.0;
    for (int32_t i=0; i < w_dim; i++)
        subgrad[i] = 0;

    for (index_t i = 0; i < m_latent_model->get_num_vectors(); ++i)
    {
        CResultSet* ret = argmax(SGVector<float64_t>(W, w_dim, false), i);
        R += (ret->score - ret->delta);
        SG_UNREF(ret);
    }

    /* calculate \f$ \sum_{i=0}^N w\cdot\Psi(x,y,h*) \f$ */
    CDotFeatures* feats = m_latent_model->get_psi_feature_vectors();
    SGVector<float64_t> v(m_latent_model->get_num_vectors());
    feats->dense_dot_range(v.vector, 0, m_latent_model->get_num_vectors(), NULL, W, get_dim(), 0.0);
    R -= SGVector<float64_t>::sum(v.vector, v.vlen);

    return R;
}