/** @file
 * @brief SimplifiedDBN class - the Simplified DBN click model.
 */
/* Copyright (C) 2017 Vivek Pal
 *
 * 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 2 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
 * <https://www.gnu.org/licenses/>.
 */

#include <config.h>

#include "simplifieddbn.h"

#include <algorithm>
#include <iomanip>
#include <fstream>
#include <sstream>
#include <string>
#include <utility>
#include <vector>

using namespace std;

#define QID 0
#define DOCIDS 1
#define CLICKS 2

string
SimplifiedDBN::name()
{
    return "SimplifiedDBN";
}

static vector<string>
get_docid_list(const string &str_docids)
{
    vector<string> docids;
    string docid;
    for (size_t j = 0; j <= str_docids.length(); ++j) {
        char ch = str_docids[j];
        if (ch != ',' && ch != '\0') {
            docid += ch;
        } else {
            docids.push_back(docid);
            docid.clear();
        }
    }
    return docids;
}

static vector<int>
get_click_list(const string &str_clicks)
{
    vector<int> clicks;
    string clickstring, clickcount;
    int click = 0;

    for (size_t j = 0; j <= str_clicks.length(); ++j) {
        char ch = str_clicks[j];
        if (ch != ',' && ch != '\0') {
            // Get clickstring of the form "docid:click_count".
            clickstring += str_clicks[j];

            // Get click count string.
            size_t delimiter_pos = clickstring.find(':');
            UNSIGNED_OVERFLOW_OK(++delimiter_pos);
            clickcount = clickstring.substr(delimiter_pos);

            // Convert click count string to an integer.
            stringstream ss(clickcount);
            ss >> click;
        } else {
            clicks.push_back(click);
            clickstring.clear();
        }
    }
    return clicks;
}

vector<Session>
SimplifiedDBN::build_sessions(const string &logfile)
{
    ifstream file;
    file.open(logfile, ios::in);

    if (!file) {
        throw runtime_error("Couldn't open file: " + logfile);
    }

    string line;

    // Skip the first line.
    getline(file, line);

    vector<Session> sessions;

    // Start reading file from the second line.
    while (getline(file, line)) {
        istringstream ss(line);

        vector<string> row_data;

        while (ss >> std::ws) {
            string column_element;
            if (ss.peek() == '"') {
                int pos = ss.tellg();
                ss.seekg(pos + 1);
                char ch;
                while (ss.get(ch)) {
                    if (ch == '"')
                        break;
                    column_element += ch;
                }
            } else {
                if (ss.peek() == ',') {
                    int pos = ss.tellg();
                    ss.seekg(pos + 1);
                }
                getline(ss, column_element, ',');
            }
            row_data.push_back(column_element);
        }

        string qid = row_data[0];
        string query = row_data[1];
        string docids = row_data[2];
        string clicks = row_data[4];

        Session s;
        s.create_session(qid, docids, clicks);
        sessions.push_back(s);
    }

    file.close();
    return sessions;
}

void
SimplifiedDBN::train(const vector<Session> &sessions)
{
    map<string, map<string, map<int, map<int, double>>>> doc_rel_fractions;

    for (auto&& session : sessions) {
        string qid = session.get_qid();

        vector<string> docids = get_docid_list(session.get_docids());

        vector<int> clicks = get_click_list(session.get_clicks());

        int last_clicked_pos = clicks.size() - 1;

        for (size_t j = 0; j < clicks.size(); ++j)
            if (clicks[j] != 0)
                last_clicked_pos = j;

        // Initialise some values.
        for (int k = 0; k <= last_clicked_pos; ++k) {
            doc_rel_fractions[qid][docids[k]][PARAM_ATTR_PROB][0] = 1.0;
            doc_rel_fractions[qid][docids[k]][PARAM_ATTR_PROB][1] = 1.0;
            doc_rel_fractions[qid][docids[k]][PARAM_SAT_PROB][0] = 1.0;
            doc_rel_fractions[qid][docids[k]][PARAM_SAT_PROB][1] = 1.0;
        }

        for (int k = 0; k <= last_clicked_pos; ++k) {
            if (clicks[k] != 0) {
                doc_rel_fractions[qid][docids[k]][PARAM_ATTR_PROB][1] += 1;
                if (int(k) == last_clicked_pos)
                    doc_rel_fractions[qid][docids[k]][PARAM_SAT_PROB][1] += 1;
                else
                    doc_rel_fractions[qid][docids[k]][PARAM_SAT_PROB][0] += 1;
            } else {
                doc_rel_fractions[qid][docids[k]][PARAM_ATTR_PROB][0] += 1;
            }
        }
    }

    for (auto i = doc_rel_fractions.begin(); i != doc_rel_fractions.end();
        ++i) {
        string qid = i->first;
        for (auto&& j : i->second) {
            doc_relevances[qid][j.first][PARAM_ATTR_PROB] =
                j.second[PARAM_ATTR_PROB][1] /
                (j.second[PARAM_ATTR_PROB][1] +
                j.second[PARAM_ATTR_PROB][0]);
            doc_relevances[qid][j.first][PARAM_SAT_PROB] =
                j.second[PARAM_SAT_PROB][1] /
                (j.second[PARAM_SAT_PROB][1] +
                j.second[PARAM_SAT_PROB][0]);
        }
    }
}

vector<pair<string, double>>
SimplifiedDBN::get_predicted_relevances(const Session &session)
{
    vector<pair<string, double>> docid_relevances;

    vector<string> docids = get_docid_list(session.get_docids());

    for (size_t i = 0; i < docids.size(); ++i) {
        double attr_prob =
                doc_relevances[session.get_qid()][docids[i]][PARAM_ATTR_PROB];
        double sat_prob =
                doc_relevances[session.get_qid()][docids[i]][PARAM_SAT_PROB];
        docid_relevances.push_back(make_pair(docids[i], attr_prob * sat_prob));
    }
    return docid_relevances;
}
