#include "ranksvmtn.h"
#include<iostream>
#include<list>
#include"../tools/matrixIO.h"

using namespace std;
using namespace Eigen;

const int maxiter = 10;
const double prec=1e-3;

int cg_solve(const MatrixXd &A, const VectorXd &b, VectorXd &x)
{
    double alpha,beta,r_1,r_2;
    int step=0;
    VectorXd q;
    VectorXd res = b - A*x; // TODO Hessian product
    VectorXd p = res;
    while (1)
    {
        // Non preconditioned version
        r_1 = res.dot(res);
        cout<<step<<":"<<r_1<<endl;
        write_stream(cout,res);
        if (r_1<1e-5) // Terminate condition
            break;
        if (step){
            beta = r_1 / r_2;
            p = res + p * beta;
        }

        q = A*p; // TODO Hessian product
        alpha = r_1/p.dot(q);
        x=x+p*alpha;
        res=res-q*alpha;
        write_stream(cout,p);
        write_stream(cout,q);
        cin.get();
        ++step;
        r_2=r_1;
    }
    return 0;
}

// Calculate objfunc gradient & support vectors
int objfunc_linear(const VectorXd &w,const MatrixXd &A,const double C,VectorXd &pred,VectorXd &grad, double &obj,MatrixXd &sv)
{
    pred = pred.cwiseMax(MatrixXd::Zero(pred.rows(),pred.cols()));
    obj = (pred.cwiseProduct(pred)*C).sum()/2 + w.dot(w)/2;
    grad = w - (((pred*C).transpose()*A)*w).transpose();
    for (int i=0;i<pred.cols();++i)
        if (pred(i)>0)
            sv(i,i)=1;
        else
            sv(i,i)=0;
    return 0;
}

// line search using newton method
int line_search(const VectorXd &w,const MatrixXd &D,const MatrixXd &A,const VectorXd &step,VectorXd &pred,double &t)
{
    double wd=w.dot(step),dd=step.dot(step);
    double g,h;
    t = 0;
    VectorXd Xd=A*(D*step);
    while (1)
    {
        pred = pred - t*Xd;
        g=wd+t*dd;
        h=dd;
        for (int i=0;i<pred.cols();++i)
            if (pred(i)>0) {
                g += pred(i)*Xd(i);
                h += Xd(i)*Xd(i);
            }
        if (g*g/h<1e-10)
            break;
    }
    return 0;
}

int train_orig(int fsize, MatrixXd &D,MatrixXd &A,VectorXd &weight){
    int iter = 0;
    double C=1;

    long n=D.rows();
    LOG(INFO) << "training with feature size:" << fsize << " Data size:" << n;
    MatrixXd sv=MatrixXd::Identity(n, n);
    VectorXd grad(fsize);
    VectorXd step(fsize);
    VectorXd pred(n);
    double obj,t;

    pred=VectorXd::Ones(n) - (A*(D*weight));

    while (true)
    {
        iter+=1;
        if (iter> maxiter)
        {
            LOG(INFO)<< "Maxiter :"<<maxiter<<" reached";
            break;
        }

        // Generate support vector matrix sv & gradient
        objfunc_linear(weight,A,C,pred,grad,obj,sv);
        step = grad*0;
        MatrixXd H = MatrixXd::Identity(grad.cols(),grad.cols());

        // Compute Hessian directly
        for (int i=0;i<n;++i)
            if (sv(i,i)>0)
            H = H + 2*C*A.row(i).transpose()*A.row(i);
        // Solve
        cg_solve(H,grad,step);
        // do line search
        line_search(weight,D,A,step,pred,t);
        weight=weight+step*t;
        // When dec is small enough
        if (-step.dot(grad) < prec * obj)
            break;
    }
    return 0;
}

int RSVMTN::train(DataList &D){
    MatrixXd Data(D.getSize(),D.getfSize()),A;
    int i=0,j=0;
    list<DataEntry*>::iterator iter,st,nx;
    for (iter= D.getData().begin();i<D.getSize();++i,++iter)
        for (j=0;j<D.getfSize();++j)
            Data(i,j)=(*iter)->feature(j);
    nx=st=iter= D.getData().begin();
    ++nx;
    int cnt=0;
    while (iter!=D.getData().end())
    {
        if ((nx == D.getData().end())||(*iter)->qid!=(*nx)->qid)
        {
            list<DataEntry*>::iterator high,low=iter;
            for (high=st;((*high)->rank)>0;++high)
                for (low=iter;((*low)->rank)<0;--low)
                    ++cnt;
            st = nx;
        }
        ++iter;
    }
    A.resize(cnt,D.getSize());
    nx=st=iter= D.getData().begin();
    ++nx;
    cnt=i=j=0;
    while (iter!=D.getData().end())
    {
        if ((nx == D.getData().end())||(*iter)->qid!=(*nx)->qid)
        {
            int v1=j,v2;
            list<DataEntry*>::iterator high,low=iter;
            for (high=st;((*high)->rank)>0;++high,++v1)
                for (low=iter,v2=i;((*low)->rank)<0;--low,--v2) {
                    A(cnt,v1) = 1;
                    A(cnt,v2) = -1;
                    ++cnt;
                }
            st = nx;
            j=i+1;
        }
        ++i;
        ++iter;
    }
    train_orig(fsize,Data,A,model.weight);
    return 0;
};

int RSVMTN::predict(DataList &D, list<double> &res){
    //TODO define A
    for (list<DataEntry*>::iterator i=D.getData().begin(), end=D.getData().end();i!=end;++i)
        res.push_back(((*i)->feature).dot(model.weight));
    return 0;
};