#pragma once
#include <algorithm>
#include <cassert>
#include <numeric>
#include <vector>

class Edge {
  public:
    unsigned int u, v;

    bool operator==(const Edge &e) const { return u == e.u && v == e.v; }
};

class DiDegree {
  public:
    unsigned int in;
    unsigned int out;
};

typedef std::vector<unsigned int> DegreeSequence;
typedef std::vector<DiDegree> DiDegreeSequence;

class Graph {
  public:
    Graph() {}

    Graph(unsigned int n) { adj.resize(n); }

    ~Graph() {}

    void resize(unsigned int n) {
        if (n < adj.size()) {
            edges.clear();
        }
        adj.resize(n);
    }

    unsigned int edgeCount() const { return edges.size(); }

    const Edge &getEdge(unsigned int i) const { return edges[i]; }

    bool createFromDegreeSequence(const DegreeSequence &d) {
        // Havel-Hakimi algorithm

        unsigned int n = d.size();

        // degree, vertex index
        std::vector<std::pair<unsigned int, unsigned int>> degrees(n);
        for (unsigned int i = 0; i < n; ++i) {
            degrees[i].first = d[i];
            degrees[i].second = i;
        }

        edges.clear();
        adj.resize(n);
        while (!degrees.empty()) {
            std::sort(degrees.begin(), degrees.end());
            // Highest degree is at back of the vector
            // Take it out
            unsigned int degree = degrees.back().first;
            unsigned int u = degrees.back().second;
            degrees.pop_back();
            if (degree > degrees.size()) {
                edges.clear();
                adj.clear();
                return false;
            }
            // Now loop over the last 'degree' entries of degrees
            auto rit = degrees.rbegin();
            for (unsigned int i = 0; i < degree; ++i) {
                if (rit->first == 0 || !addEdge({u, rit->second})) {
                    edges.clear();
                    adj.clear();
                    return false;
                }
                rit->first--;
                ++rit;
            }
        }
        return true;
    }

    DegreeSequence getDegreeSequence() const {
        DegreeSequence d(adj.size());
        std::transform(adj.begin(), adj.end(), d.begin(),
                       [](const auto &u) { return u.size(); });
        return d;
    }

    // Assumes valid vertex indices
    bool hasEdge(const Edge &e) const {
        for (unsigned int v : adj[e.u]) {
            if (v == e.v)
                return true;
        }
        return false;
    }

    bool addEdge(const Edge &e) {
        if (e.u >= adj.size() || e.v >= adj.size())
            return false;
        if (hasEdge(e))
            return false;
        edges.push_back(e);
        adj[e.u].push_back(e.v);
        adj[e.v].push_back(e.u);
        return true;
    }

    // There are two possible edge exchanges
    // switchType indicates which one is desired
    // Returns false if the switch is not possible
    bool exchangeEdges(const Edge &e1, const Edge &e2, bool switchType) {
        // The new edges configuration is one of these two
        // A) e1.u - e2.u and e1.v - e2.v
        // B) e1.u - e2.v and e1.v - e2.u
        // First check if the move is possible
        if (switchType) {
            if (hasEdge({e1.u, e2.u}) || hasEdge({e1.v, e2.v}))
                return false; // conflicting edges
        } else {
            if (hasEdge({e1.u, e2.v}) || hasEdge({e1.v, e2.u}))
                return false; // conflicting edges
        }

        // Find the edges in the adjacency lists
        unsigned int i1, j1, i2, j2;
        for (i1 = 0; i1 < adj[e1.u].size(); ++i1) {
            if (adj[e1.u][i1] == e1.v)
                break;
        }
        for (j1 = 0; j1 < adj[e1.v].size(); ++j1) {
            if (adj[e1.v][j1] == e1.u)
                break;
        }
        for (i2 = 0; i2 < adj[e2.u].size(); ++i2) {
            if (adj[e2.u][i2] == e2.v)
                break;
        }
        for (j2 = 0; j2 < adj[e2.v].size(); ++j2) {
            if (adj[e2.v][j2] == e2.u)
                break;
        }

        // Remove the old edges
        bool removedOne = false;
        for (auto iter = edges.begin(); iter != edges.end();) {
            if (*iter == e1) {
                iter = edges.erase(iter);
                if (removedOne)
                    break;
                removedOne = true;
            } else if (*iter == e2) {
                iter = edges.erase(iter);
                if (removedOne)
                    break;
                removedOne = true;
            } else {
                ++iter;
            }
        }

        // Add the new edges
        if (switchType) {
            adj[e1.u][i1] = e2.u;
            adj[e1.v][j1] = e2.v;
            adj[e2.u][i2] = e1.u;
            adj[e2.v][j2] = e1.v;
            edges.push_back({e1.u, e2.u});
            edges.push_back({e1.v, e2.v});
        } else {
            adj[e1.u][i1] = e2.v;
            adj[e1.v][j1] = e2.u;
            adj[e2.u][i2] = e1.v;
            adj[e2.v][j2] = e1.u;
            edges.push_back({e1.u, e2.v});
            edges.push_back({e1.v, e2.u});
        }
        return true;
    }

    int countTriangles() const {
        int triangles = 0;
        for (auto& v : adj) {
            for (unsigned int i = 0; i < v.size(); ++i) {
                for (unsigned int j = i + 1; j < v.size(); ++j) {
                    if (hasEdge({v[i], v[j]})) {
                        ++triangles;
                    }
                }
            }
        }
        assert(triangles % 3 == 0);
        return triangles / 3;
    }

  private:
    // Graph is saved in two formats for speed
    // The two should be kept consistent at all times
    std::vector<std::vector<unsigned int>> adj;
    std::vector<Edge> edges;
};