// from https://github.com/wjakob/dset, changed to add connected component
// computation
//
// Copyright (c) 2015 Wenzel Jakob <wenzel@inf.ethz.ch>
//
// This software is provided 'as-is', without any express or implied
// warranty.  In no event will the authors be held liable for any damages
// arising from the use of this software.
//
// Permission is granted to anyone to use this software for any purpose,
// including commercial applications, and to alter it and redistribute it
// freely, subject to the following restrictions:
//
// 1. The origin of this software must not be misrepresented; you must not
// claim that you wrote the original software. If you use this software
// in a product, an acknowledgment in the product documentation would be
// appreciated but is not required.
// 2. Altered source versions must be plainly marked as such, and must not be
// misrepresented as being the original software.
// 3. This notice may not be removed or altered from any source distribution.
//
#pragma once
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <unordered_map>
#include <vector>

class DisjointSets {
 public:
  DisjointSets(uint32_t size) : mData(size) {
    for (uint32_t i = 0; i < size; ++i) mData[i] = (uint32_t)i;
  }

  uint32_t find(uint32_t id) const {
    while (id != parent(id)) {
      uint64_t value = mData[id];
      uint32_t new_parent = parent((uint32_t)value);
      uint64_t new_value = (value & 0xFFFFFFFF00000000ULL) | new_parent;
      /* Try to update parent (may fail, that's ok) */
      if (value != new_value) mData[id].compare_exchange_weak(value, new_value);
      id = new_parent;
    }
    return id;
  }

  bool same(uint32_t id1, uint32_t id2) const {
    for (;;) {
      id1 = find(id1);
      id2 = find(id2);
      if (id1 == id2) return true;
      if (parent(id1) == id1) return false;
    }
  }

  uint32_t unite(uint32_t id1, uint32_t id2) {
    for (;;) {
      id1 = find(id1);
      id2 = find(id2);

      if (id1 == id2) return id1;

      uint32_t r1 = rank(id1), r2 = rank(id2);

      if (r1 > r2 || (r1 == r2 && id1 < id2)) {
        std::swap(r1, r2);
        std::swap(id1, id2);
      }

      uint64_t oldEntry = ((uint64_t)r1 << 32) | id1;
      uint64_t newEntry = ((uint64_t)r1 << 32) | id2;

      if (!mData[id1].compare_exchange_strong(oldEntry, newEntry)) continue;

      if (r1 == r2) {
        oldEntry = ((uint64_t)r2 << 32) | id2;
        newEntry = ((uint64_t)(r2 + 1) << 32) | id2;
        /* Try to update the rank (may fail, retry if rank = 0) */
        if (!mData[id2].compare_exchange_strong(oldEntry, newEntry) && r2 == 0)
          continue;
      }

      break;
    }
    return id2;
  }

  uint32_t size() const { return (uint32_t)mData.size(); }

  uint32_t rank(uint32_t id) const {
    return ((uint32_t)(mData[id] >> 32)) & 0x7FFFFFFFu;
  }

  uint32_t parent(uint32_t id) const { return (uint32_t)mData[id]; }

  int connectedComponents(std::vector<int>& components) {
    components.resize(mData.size());
    int lonelyNodes = 0;
    std::unordered_map<uint32_t, int> toLabel;
    for (size_t i = 0; i < mData.size(); ++i) {
      // we optimize for connected component of size 1
      // no need to put them into the hashmap
      auto iParent = find(i);
      if (rank(iParent) == 0) {
        components[i] = static_cast<int>(toLabel.size()) + lonelyNodes++;
        continue;
      }
      auto iter = toLabel.find(iParent);
      if (iter == toLabel.end()) {
        auto s = static_cast<uint32_t>(toLabel.size()) + lonelyNodes;
        toLabel.insert(std::make_pair(iParent, s));
        components[i] = s;
      } else {
        components[i] = iter->second;
      }
    }
    return toLabel.size() + lonelyNodes;
  }

  mutable std::vector<std::atomic<uint64_t>> mData;
};