#include <cstring>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <string>
#include <iomanip>
#include <stdexcept>
#include <vector>

#include "mf.h"

using namespace std;
using namespace mf;

struct Option
{
    Option() : eval(RMSE) {}
    string test_path, model_path, output_path;
    mf_int eval;
};

string predict_help()
{
    return string(
"usage: mf-predict [options] test_file model_file [output_file]\n"
"\n"
"options:\n"
"-e <eval>: specify the evaluation criterion (default 0)\n"
"\t 0 -- root mean square error\n"
"\t 1 -- mean absolute error\n"
"\t 2 -- generalized KL-divergence\n"
"\t 5 -- logarithmic error\n"
"\t 6 -- accuracy\n"
"\t10 -- row-wise mean percentile rank\n"
"\t11 -- column-wise mean percentile rank\n"
"\t12 -- row-wise area under the curve\n"
"\t13 -- column-wise area under the curve\n");
}

Option parse_option(int argc, char **argv)
{
    vector<string> args;
    for(int i = 0; i < argc; i++)
        args.push_back(string(argv[i]));

    if(argc == 1)
        throw invalid_argument(predict_help());

    Option option;

    mf_int i;
    for(i = 1; i < argc; i++)
    {
        if(args[i].compare("-e") == 0)
        {
            if((i+1) >= argc)
                throw invalid_argument("need to specify evaluation criterion after -e");
            i++;
            option.eval = atoi(argv[i]);
            if(option.eval != RMSE &&
               option.eval != MAE &&
               option.eval != GKL &&
               option.eval != LOGLOSS &&
               option.eval != ACC &&
               option.eval != ROW_AUC &&
               option.eval != COL_AUC &&
               option.eval != ROW_MPR &&
               option.eval != COL_MPR)
                throw invalid_argument("unknown evaluation criterion");
        }
        else
            break;
    }
    if(i >= argc-1)
        throw invalid_argument("testing data and model file not specified");
    option.test_path = string(args[i++]);
    option.model_path = string(args[i++]);

    if(i < argc)
    {
        option.output_path = string(args[i]);
    }
    else if(i == argc)
    {
        const char *ptr = strrchr(&*option.test_path.begin(), '/');
        if(!ptr)
            ptr = option.test_path.c_str();
        else
            ++ptr;
        option.output_path = string(ptr) + ".out";
    }
    else
    {
        throw invalid_argument("invalid argument");
    }

    return option;
}

void predict(string test_path, string model_path, string output_path, mf_int eval)
{
    mf_problem prob = read_problem(test_path);

    ofstream f_out(output_path);
    if(!f_out.is_open())
        throw runtime_error("cannot open " + output_path);

    mf_model *model = mf_load_model(model_path.c_str());
    if(model == nullptr)
        throw runtime_error("cannot load model from " + model_path);

    for(mf_int i = 0; i < prob.nnz; i++)
    {
        mf_float r = mf_predict(model, prob.R[i].u, prob.R[i].v);
        f_out << r << endl;
    }

    switch(eval)
    {
        case RMSE:
        {
            auto rmse = calc_rmse(&prob, model);
            cout << fixed << setprecision(4) << "RMSE = " << rmse << endl;
            break;
        }
        case MAE:
        {
            auto mae = calc_mae(&prob, model);
            cout << fixed << setprecision(4) << "MAE = " << mae << endl;
            break;
        }
        case GKL:
        {
            auto gkl = calc_gkl(&prob, model);
            cout << fixed << setprecision(4) << "GKL = " << gkl << endl;
            break;
        }
        case LOGLOSS:
        {
            auto logloss = calc_logloss(&prob, model);
            cout << fixed << setprecision(4) << "LOGLOSS = " << logloss << endl;
            break;
        }
        case ACC:
        {
            auto acc = calc_accuracy(&prob, model);
            cout << fixed << setprecision(4) << "ACCURACY = " << acc << endl;
            break;
        }
        case ROW_AUC:
        {
            auto row_wise_auc = calc_auc(&prob, model, false);
            cout << fixed << setprecision(4) <<  "Row-wise AUC = " << row_wise_auc << endl;
            break;
        }
        case COL_AUC:
        {
            auto col_wise_auc = calc_auc(&prob, model, true);
            cout << fixed << setprecision(4) <<  "Colmn-wise AUC = " << col_wise_auc << endl;
            break;
        }
        case ROW_MPR:
        {
            auto row_wise_mpr = calc_mpr(&prob, model, false);
            cout << fixed << setprecision(4) <<  "Row-wise MPR = " << row_wise_mpr << endl;
            break;
        }
        case COL_MPR:
        {
            auto col_wise_mpr = calc_mpr(&prob, model, true);
            cout << fixed << setprecision(4) <<  "Column-wise MPR = " << col_wise_mpr << endl;
            break;
        }
        default:
        {
            throw invalid_argument("unknown evaluation criterion");
            break;
        }
    }
    mf_destroy_model(&model);
}

int main(int argc, char **argv)
{
    Option option;
    try
    {
        option = parse_option(argc, argv);
    }
    catch(invalid_argument &e)
    {
        cout << e.what() << endl;
        return 1;
    }

    try
    {
        predict(option.test_path, option.model_path, option.output_path, option.eval);
    }
    catch(runtime_error &e)
    {
        cout << e.what() << endl;
        return 1;
    }

    return 0;
}