bimap.hpp

#include <set>
#include <vector>
#include <utility>
#include <iterator>
#include <ranges>
#include <unordered_set>
#include <algorithm>

namespace ojl {
  template <typename T>
  struct bimap_left {
      T val;
  };

  template <typename T>
  struct bimap_right {
      T val;
  };


  template <typename L, typename R>
  struct CompareLeft {
      bool
      operator()(const std::pair<L, R> &lhs, const std::pair<L, R> &rhs) const
      {
          return lhs.first < rhs.first || (lhs.first == rhs.first && lhs.second < rhs.second);
      }

      bool
      operator()(const L &lhs, const std::pair<L, R> &rhs) const
      {
          return lhs < rhs.first;
      }

      bool
      operator()(const std::pair<L, R> &lhs, const L &rhs) const 
      {
          return lhs.first < rhs;
      }

      using is_transparent = L;
  };


  template <typename L, typename R>
  struct CompareRight {
      bool
      operator()(const std::pair<L, R> &lhs, const std::pair<L, R> &rhs) const
      {
          return lhs.second < rhs.second || (lhs.second == rhs.second && lhs.first < rhs.first);
      }

      bool
      operator()(const R &lhs, const std::pair<L, R> &rhs) const 
      {
          return lhs < rhs.second;
      }

      bool
      operator()(const std::pair<L, R> &lhs, const R &rhs) const 
      {
          return lhs.second < rhs;
      }

      using is_transparent = L;
  };


  //TODO: maybe add unkeyed versions of member functions for when L != R
  // (replacing left with L, right with R)

  template <typename L, typename R>
  struct bimap {
      using left = bimap_left<L>;
      using right = bimap_right<R>;

      using relation = std::pair<L,R>;

      using l_relation_set  = std::set<relation, CompareLeft<L,R>>;
      using r_relation_set  = std::set<relation, CompareRight<L,R>>;

      using relation_iterator = typename l_relation_set::iterator;

      //TODO: see if the iterator types are the same

      std::set<L>  left_elements;
      std::set<R>  right_elements;
      l_relation_set l_relations;
      r_relation_set r_relations;

      //TODO: badly named? Maybe should make the sets private if I have accessors...?
      const l_relation_set&
      relations_ref()
      {
          return l_relations;
      }

      std::vector<relation>
      all_relations_vec()
      {
          std::vector<relation> rel;
          for (auto it = l_relations.begin(); it != l_relations.end(); it++){
              rel.push_back(*it);
          }
          return rel;
      }


      // Insert elements
      void
      insert(left l_key)
      {
          left_elements.insert(l_key.val);
      }

      void
      insert(right r_key)
      {
          right_elements.insert(r_key.val);
      }

      void
      l_insert(L l)
      {
          left_elements.insert(l);
      }

      void
      r_insert(R r)
      {
          right_elements.insert(r);
      }

      // Insert relations
      void
      insert(relation rel)
      {
          left_elements.insert(rel.first);
          right_elements.insert(rel.second);
          l_relations.insert(rel);
          r_relations.insert(rel);
      }

      void
      insert(const L &l, const R &r)
      {
          left_elements.insert(l);
          right_elements.insert(r);
          l_relations.insert({l,r});
          r_relations.insert({l,r});
      }

      // Erase relations
      void
      erase(const relation &rel)
      {
          r_relations.erase(rel);
          l_relations.erase(rel);
      }

      void
      erase(const L &l, const R &r)
      {
          relation rel{l,r};
          r_relations.erase(rel);
          l_relations.erase(rel);
      }


      // Erase elements
      void
      erase(left l_key)
      {
          left_elements.erase(l_key.val);
          //Find all relations in l_relations and remove them from r_relations
          auto &[lb, ub] = l_relations.equal_range(l_key.val);
          for (auto it = lb; it < ub; it++){
              r_relations.erase(*it);
          }
          l_relations.erase(l_key.val);
      }

      void
      erase(right r_key)
      {
          right_elements.erase(r_key.val);
          auto &[lb, ub] = r_relations.equal_range(r_key.val);
          for (auto it = lb; it < ub; it++){
              l_relations.erase(*it);
          }
          r_relations.erase(r_key.val);
      }

      void
      l_erase(L l)
      {
          left_elements.erase(l);
          auto &[lb, ub] = l_relations.equal_range(l);
          for (auto it = lb; it < ub; it++){
              r_relations.erase(*it);
          }
          l_relations.erase(l);
      }

      void
      r_erase(R r)
      {
          right_elements.erase(r);
          auto &[lb, ub] = l_relations.equal_range(r);
          for (auto it = lb; it < ub; it++){
              r_relations.erase(*it);
          }
          r_relations.erase(r);
      }

      // Contains
      bool contains(relation rel)
      {
          return l_relations.contains(rel);
      }

      bool contains(L l, R r)
      {
          return l_relations.contains({l,r});
      }


      bool contains(left l_key)
      {
          return l_relations.contains(l_key.val);
      }

      bool contains(right r_key)
      {
          return r_relations.contains(r_key.val);
      }

      bool l_contains(L l)
      {
          return l_relations.contains(l);
      }

      bool r_contains(R r)
      {
          return r_relations.contains(r);
      }

      
      // Count
      size_t
      count(left l_key)
      {
          return l_relations.count(l_key.val);
      }

      size_t
      count(right r_key)
      {
          return r_relations.count(r_key.val);
      }

      size_t
      l_count(L l)
      {
          return l_relations.count(l);
      }

      size_t
      r_count(R r)
      {
          return r_relations.count(r);
      }

      // Min elements
      std::optional<L>
      l_min_element()
      {
          auto it = std::min_element(left_elements.begin(), left_elements.end());
          if (it != left_elements.end()){
              return *it;
          }
          return std::nullopt;
      }

      std::optional<R>
      r_min_element()
      {
          auto it = std::min_element(right_elements.begin(), right_elements.end());
          if (it != right_elements.end()){
              return *it;
          }
          return std::nullopt;
      }

      // Max elements
      std::optional<L>
      l_max_element()
      {
          auto it = std::max_element(left_elements.begin(), left_elements.end());
          if (it != left_elements.end()){
              return *it;
          }
          return std::nullopt;
      }

      std::optional<R>
      r_max_element()
      {
          auto it = std::max_element(right_elements.begin(), right_elements.end());
          if (it != right_elements.end()){
              return *it;
          }
          return std::nullopt;
      }

      // Min mapped elements
      std::optional<R>
      min_mapped_element(left l_key)
      {
          auto it = std::ranges::min_element(l_relations | std::views::filter([&l_key](relation rel){return rel.first == l_key.val;}));
          if (it != l_relations.end()){
              return it->second;
          }
          return std::nullopt;

      }

      std::optional<R>
      min_mapped_element(right r_key)
      {
          auto it = std::ranges::min_element(l_relations | std::views::filter([&r_key](relation rel){return rel.second == r_key.val;}));
          if (it != l_relations.end()){
              return it->first;
          }
          return std::nullopt;

      }

      std::optional<R>
      l_min_mapped_element(L l)
      {
          auto it = std::ranges::min_element(l_relations | std::views::filter([&l](relation rel){return rel.first == l;}));
          if (it != l_relations.end()){
              return it->second;
          }
          return std::nullopt;

      }

      std::optional<R>
      r_min_mapped_element(R r)
      {
          auto it = std::ranges::min_element(l_relations | std::views::filter([&r](relation rel){return rel.second == r;}));
          if (it != l_relations.end()){
              return it->first;
          }
          return std::nullopt;

      }

      // Max mapped elements
      std::optional<R>
      max_mapped_element(left l_key)
      {
          auto it = std::ranges::max_element(l_relations | std::views::filter([&l_key](relation rel){return rel.first == l_key.val;}));
          if (it != l_relations.end()){
              return it->second;
          }
          return std::nullopt;

      }

      std::optional<R>
      max_mapped_element(right r_key)
      {
          auto it = std::ranges::max_element(l_relations | std::views::filter([&r_key](relation rel){return rel.second == r_key.val;}));
          if (it != l_relations.end()){
              return it->first;
          }
          return std::nullopt;

      }

      std::optional<R>
      l_max_mapped_element(L l)
      {
          auto it = std::ranges::max_element(l_relations | std::views::filter([&l](relation rel){return rel.first == l;}));
          if (it != l_relations.end()){
              return it->second;
          }
          return std::nullopt;

      }

      std::optional<R>
      r_max_mapped_element(R r)
      {
          auto it = std::ranges::max_element(l_relations | std::views::filter([&r](relation rel){return rel.second == r;}));
          if (it != l_relations.end()){
              return it->first;
          }
          return std::nullopt;

      }

      // Equal range
      std::pair<relation_iterator, relation_iterator>
      equal_range(left l_key)
      {
          return l_relations.equal_range(l_key.val);
      }

      std::pair<relation_iterator, relation_iterator>
      equal_range(right r_key)
      {
          return r_relations.equal_range(r_key.val);
      }

      std::pair<relation_iterator, relation_iterator>
      l_equal_range(L l)
      {
          return equal_range(left{l});
      }

      std::pair<relation_iterator, relation_iterator>
      r_equal_range(R r)
      {
          return equal_range(right{r});
      }

      // Convenience function for getting a vector out
      std::vector<R>
      mapped_vector(left l_key)
      {
          std::vector<R> v;
          auto [lb, ub] = equal_range(l_key);
          v.reserve(std::distance(lb,ub));
          for (auto it = lb; it != ub; it++){
              v.push_back(it->second);
          }
          return v;
      }

      std::vector<L>
      mapped_vector(right r_key)
      {
          std::vector<L> v;
          auto [lb, ub] = equal_range(r_key);
          v.reserve(std::distance(lb,ub));
          for (auto it = lb; it != ub; it++){
              v.push_back(it->first);
          }
          return v;
      }

      // Convenience function for getting a set
      std::vector<R>
      mapped_set(left l_key)
      {
          std::set<R> s;
          auto [lb, ub] = equal_range(l_key);
          for (auto it = lb; it != ub; it++){
              s.insert(it->second);
          }
          return s;
      }

      std::vector<L>
      mapped_set(right r_key)
      {
          std::set<L> s;
          auto [lb, ub] = equal_range(r_key);
          s.insert(ub - lb);
          for (auto it = lb; it != ub; it++){
              s.insert(it->first);
          }
          return s;
      }

      std::vector<R>
      l_mapped_vector(L l)
      {
          return mapped_vector(left{l});
      }

      std::vector<L>
      r_mapped_vector(R r)
      {
          return mapped_vector(right{r});
      }

      std::vector<R>
      l_mapped_set(L l)
      {
          return mapped_set(left{l});
      }

      std::vector<L>
      r_mapped_set(R r)
      {
          return mapped_set(right{r});
      }

      // A complement of this bimap (in the set of possible relations between left_elements and right_elements)
      bimap<L,R>
      complement()
      {
          bimap<L,R> ret;
          // Insert elements
          for (const auto &l : left_elements){
              ret.l_insert(l);
          }
          for (const auto &r : right_elements){
              ret.r_insert(r);
          }

          // Insert relations
          for (const auto &l : left_elements){
              for (const auto &r : right_elements){
                  relation rel{l,r};
                  if (!this->contains(rel)){
                      ret.insert(rel);
                  }
              }
          }
          return ret;
      }

      bimap<L,R>
      complement(std::set<L> left_elems, std::set<R> right_elems)
      {
          bimap<L,R> ret;
          // Insert elements
          for (const auto &l : left_elems){
              ret.l_insert(l);
          }
          for (const auto &r : right_elems){
              ret.r_insert(r);
          }

          // Insert relations
          for (const auto &l : left_elems){
              for (const auto &r : right_elems){
                  relation rel{l,r};
                  if (!this->contains(rel)){
                      ret.insert(rel);
                  }
              }
          }
          return ret;
      }



      bimap<L,R>
      complement(std::vector<L> left_elems, std::vector<R> right_elems)
      {
          bimap<L,R> ret;
          // Insert elements
          for (const auto &l : left_elems){
              ret.l_insert(l);
          }
          for (const auto &r : right_elems){
              ret.r_insert(r);
          }

          // Insert relations
          for (const auto &l : left_elems){
              for (const auto &r : right_elems){
                  relation rel{l,r};
                  if (!this->contains(rel)){
                      ret.insert(rel);
                  }
              }
          }
          return ret;
      }


  };

  // Equality comparison
  template<typename L, typename R>
  bool
  operator==(const bimap<L,R> &lhs, const bimap<L,R> &rhs)
  {
      return lhs.left_elements  == rhs.left_elements 
          && lhs.right_elements == rhs.right_elements
          && lhs.l_relations    == rhs.l_relations;
     
      if (lhs.left_elements != rhs.left_elements || lhs.right_elements != rhs.right_elements) {
          return false;
      }
      return lhs.l_relations == rhs.l_relations;
  }

}