#pragma once
#include <vector>
#include <map>
#include <math.h>
#include <string.h>
#include <queue>
#include <set>
#include <algorithm>

#include "utils.h"


class KnowledgeGraph
{
public:
    KnowledgeGraph();
    ~KnowledgeGraph();

    int n_entity, n_relation;
    std::vector<Triplet> triplets;
    std::vector<Triplet> valid_triplets;
    std::vector<Triplet> test_triplets;
    std::set<Triplet> all_triplets;
    std::map<std::string, int> e2id, r2id;
    std::map<int, std::string> id2e, id2r;
    inline int n_relation_without_reverse() {
        return this->n_relation;
    }
    inline int n_relation_with_reverse() {
        return this->n_relation * 2;
    }

    void clear_triplets();
    void load_data(std::string path);
    void load_valid_data(std::string path);
    void load_test_data(std::string path);
};