#include "directed_graph.h"
#include <unordered_set>
#include <unordered_map>
#include <string>
#include <fstream>

DirectedGraph::DirectedGraph(){
    // nothing to do
}

DirectedGraph::~DirectedGraph(){
    // nothing to do
}

DirectedGraph::DirectedGraph(const DirectedGraph& g){
    this->_vertices = g._vertices;
    this->_inbounds.clear();
    this->_outbounds.clear();
    this->_weights.clear();
    for(std::pair<const int, std::unordered_set<int>> p : g._inbounds){
        this->_inbounds[p.first] = p.second;
    }
    for(std::pair<const int, std::unordered_set<int>> p : g._outbounds){
        this->_outbounds[p.first] = p.second;
    }
    for(std::pair<const std::pair<int, int>, int> p : g._weights){
        this->_weights[p.first] = p.second;
    }
}

DirectedGraph DirectedGraph::operator=(const DirectedGraph& g){
    this->_vertices = g._vertices;
    this->_inbounds.clear();
    this->_outbounds.clear();
    this->_weights.clear();
    for(std::pair<const int, std::unordered_set<int>> p : g._inbounds){
        this->_inbounds[p.first] = p.second;
    }
    for(std::pair<const int, std::unordered_set<int>> p : g._outbounds){
        this->_outbounds[p.first] = p.second;
    }
    for(std::pair<const std::pair<int, int>, int> p : g._weights){
        this->_weights[p.first] = p.second;
    }
    return *this;
}

DirectedGraph DirectedGraph::reindex() const{
    std::unordered_map<int, int> old_to_new;
    int i = this->numVertices()-1;
    for(int v : this->_vertices){
        old_to_new[v] = i;
        i--;
    }
    DirectedGraph new_g;
    for(int v : this->_vertices){
        new_g.addVertex(old_to_new[v]);
    }
    for(std::pair<const std::pair<int, int>, int> p : this->_weights){
        new_g.addEdge(old_to_new[p.first.first], old_to_new[p.first.second], p.second);
    }
    return new_g;
}

int DirectedGraph::numVertices() const{
    return this->_vertices.size();
}

int DirectedGraph::numEdges() const{
    return this->_weights.size();
}

std::unordered_set<int> DirectedGraph::vertices() const{
    return this->_vertices;
}

std::unordered_map<std::pair<int, int>, int> DirectedGraph::edges() const{
    std::unordered_map<std::pair<int, int>, int> weights;
    for(std::pair<const std::pair<int, int>, int> p : this->_weights){
        weights[p.first] = p.second;
    }
    return weights;
}

bool DirectedGraph::isEdge(int v1, int v2) const{
    if(this->_vertices.find(v1) == this->_vertices.end() || this->_vertices.find(v2) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    return this->_weights.find(std::make_pair(v1, v2)) != this->_weights.end();
}

int DirectedGraph::inDegree(int v) const{
    if(this->_vertices.find(v) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    return this->_inbounds.at(v).size();
}

int DirectedGraph::outDegree(int v) const{
    if(this->_vertices.find(v) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    return this->_outbounds.at(v).size();
}

std::unordered_set<int> DirectedGraph::inbounds(int v) const{
    if(this->_vertices.find(v) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    return this->_inbounds.at(v);
}

std::unordered_set<int> DirectedGraph::outbounds(int v) const{
    if(this->_vertices.find(v) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    return this->_outbounds.at(v);
}

int DirectedGraph::weight(int v1, int v2) const{
    if(this->_vertices.find(v1) == this->_vertices.end() || this->_vertices.find(v2) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    if(this->_weights.find(std::make_pair(v1, v2)) == this->_weights.end())
        throw std::invalid_argument("Edge not found");
    return this->_weights.at(std::make_pair(v1, v2));
}

void DirectedGraph::weight(int v1, int v2, int w){
    if(this->_vertices.find(v1) == this->_vertices.end() || this->_vertices.find(v2) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    if(this->_weights.find(std::make_pair(v1, v2)) == this->_weights.end())
        throw std::invalid_argument("Edge not found");
    this->_weights[std::make_pair(v1, v2)] = w;
}

void DirectedGraph::addVertex(int v){
    if(this->_vertices.find(v) != this->_vertices.end())
        throw std::invalid_argument("Vertex already exists");
    this->_vertices.insert(v);
    this->_inbounds[v] = std::unordered_set<int>();
    this->_outbounds[v] = std::unordered_set<int>();
}

void DirectedGraph::removeVertex(int v){
    if(this->_vertices.find(v) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    this->_vertices.erase(v);
    for(int i : this->_inbounds[v]){
        this->_outbounds[i].erase(v);
        this->_weights.erase(std::make_pair(i, v));
    }
    for(int i : this->_outbounds[v]){
        this->_inbounds[i].erase(v);
        this->_weights.erase(std::make_pair(v, i));
    }
    this->_inbounds.erase(v);
    this->_outbounds.erase(v);
}

void DirectedGraph::addEdge(int v1, int v2, int w){
    if(this->_vertices.find(v1) == this->_vertices.end() || this->_vertices.find(v2) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    if(this->_weights.find(std::make_pair(v1, v2)) != this->_weights.end())
        throw std::invalid_argument("Edge already exists");
    this->_weights[std::make_pair(v1, v2)] = w;
    this->_inbounds[v2].insert(v1);
    this->_outbounds[v1].insert(v2);
}

void DirectedGraph::removeEdge(int v1, int v2){
    if(this->_vertices.find(v1) == this->_vertices.end() || this->_vertices.find(v2) == this->_vertices.end())
        throw std::invalid_argument("Vertex not found");
    if(this->_weights.find(std::make_pair(v1, v2)) == this->_weights.end())
        throw std::invalid_argument("Edge not found");
    this->_weights.erase(std::make_pair(v1, v2));
    this->_inbounds[v2].erase(v1);
    this->_outbounds[v1].erase(v2);
}

DirectedGraph read_from_file(std::string filename){
    std::ifstream file(filename);
    if(!file.is_open())
        throw std::invalid_argument("File not found");
    DirectedGraph g;
    int num_vertices=0;
    file >> num_vertices;
    for(int i = 0; i < num_vertices; i++){
        g.addVertex(i);
    }
    int num_edges=0;
    file >> num_edges;
    int v1, v2, w;
    for(int i = 0; i < num_edges; i++){
        file >> v1 >> v2 >> w;
        g.addEdge(v1, v2, w);
    }
    file.close();
    return g;
}

void write_to_file(const DirectedGraph& g, std::string filename){
    DirectedGraph new_g = g.reindex();
    std::ofstream file(filename);
    if(!file.is_open())
        throw std::invalid_argument("File could not be opened");
    file << new_g.numVertices() << " " << new_g.numEdges() << std::endl;
    for(std::pair<const std::pair<int, int>, int> p : new_g.edges()){
        file << p.first.first << " " << p.first.second << " " << p.second << std::endl;
    }
    file.close();
}

DirectedGraph random_graph(int num_vertices, int num_edges){
    DirectedGraph g;
    for(int i = 0; i < num_vertices; i++){
        g.addVertex(i);
    }
    for(int i = 0; i < num_edges; i++){
        int v1 = rand() % num_vertices;
        int v2 = rand() % num_vertices;
        int w = rand() % 100;
        if (g.isEdge(v1, v2)){
            i--;
            continue;
        }
        g.addEdge(v1, v2, w);
    }
    return g;
}