#include <cstdlib>

#include <fstream>

#include <iostream>

#include <sstream>

#include <stdexcept>

#include <string>

#include <unordered_map>

#include <utility>

#include <disco.hpp>

#include <pgvector/pqxx.hpp>

#include <pqxx/pqxx>

std::string convert_to_utf8(const std::string& str) {

std::stringstream buf;

for (const unsigned char v : str) {

if (v < 128) {

buf << v;

} else {

buf << static_cast<unsigned char>(195) << static_cast<unsigned char>(v - 64);

}

}

return buf.str();

}

disco::Dataset<int, std::string> load_movielens(const std::string& path) {

std::string line;

// read movies

std::unordered_map<std::string, std::string> movies;

std::ifstream movies_file(path + "/u.item");

if (!movies_file.is_open()) {

throw std::runtime_error{"Could not open file"};

}

while (std::getline(movies_file, line)) {

std::string::size_type n = line.find('|');

std::string::size_type n2 = line.find('|', n + 1);

movies.emplace(line.substr(0, n), convert_to_utf8(line.substr(n + 1, n2 - n - 1)));

}

// read ratings and create dataset

disco::Dataset<int, std::string> data;

std::ifstream ratings_file(path + "/u.data");

if (!ratings_file.is_open()) {

throw std::runtime_error{"Could not open file"};

}

while (std::getline(ratings_file, line)) {

std::string::size_type n = line.find('\t');

std::string::size_type n2 = line.find('\t', n + 1);

std::string::size_type n3 = line.find('\t', n2 + 1);

data.push(

std::stoi(line.substr(0, n)),

movies.at(line.substr(n + 1, n2 - n - 1)),

std::stof(line.substr(n2 + 1, n3 - n2 - 1))

);

}

return data;

}

int main() {

// https://grouplens.org/datasets/movielens/100k/

char* movielens_path = std::getenv("MOVIELENS_100K_PATH");

if (!movielens_path) {

std::cout << "Set MOVIELENS_100K_PATH" << std::endl;

return 1;

}

pqxx::connection conn{"dbname=pgvector_example"};

pqxx::nontransaction tx{conn};

tx.exec("CREATE EXTENSION IF NOT EXISTS vector");

tx.exec("DROP TABLE IF EXISTS users");

tx.exec("DROP TABLE IF EXISTS movies");

tx.exec("CREATE TABLE users (id integer PRIMARY KEY, factors vector(20))");

tx.exec("CREATE TABLE movies (name text PRIMARY KEY, factors vector(20))");

disco::Dataset<int, std::string> data = load_movielens(movielens_path);

auto recommender = disco::Recommender<int, std::string>::fit_explicit(data, {.factors = 20});

for (const auto& user_id : recommender.user_ids()) {

pgvector::Vector factors{*recommender.user_factors(user_id)};

tx.exec("INSERT INTO users (id, factors) VALUES ($1, $2)", pqxx::params{user_id, factors});

}

for (const auto& item_id : recommender.item_ids()) {

pgvector::Vector factors{*recommender.item_factors(item_id)};

tx.exec(

"INSERT INTO movies (name, factors) VALUES ($1, $2)", pqxx::params{item_id, factors}

);

}

std::string movie{"Star Wars (1977)"};

std::cout << "Item-based recommendations for " << movie << std::endl;

pqxx::result result = tx.exec(

"SELECT name FROM movies WHERE name != $1 ORDER BY factors <=> (SELECT factors FROM movies WHERE name = $1) LIMIT 5",

pqxx::params{movie}

);

for (const auto& row : result) {

std::cout << "- " << row[0].as<std::string>() << std::endl;

}

int user_id = 123;

std::cout << std::endl << "User-based recommendations for " << user_id << std::endl;

result = tx.exec(

"SELECT name FROM movies ORDER BY factors <#> (SELECT factors FROM users WHERE id = $1) LIMIT 5",

pqxx::params{user_id}

);

for (const auto& row : result) {

std::cout << "- " << row[0].as<std::string>() << std::endl;

}

return 0;

}