diff --git a/src/api/PartitionManager.cc b/src/api/PartitionManager.cc
index 2d7e43e70..6d1958a69 100644
--- a/src/api/PartitionManager.cc
+++ b/src/api/PartitionManager.cc
@@ -90,6 +90,28 @@ vec<PTRef> PartitionManager::getPartitions(ipartitions_t const & mask) const {
     return res;
 }
 
+PTRef PartitionManager::getPartition(int partitionIndex) const {
+    if (partitionIndex < 0) { return PTRef_Undef; }
+    return partitionInfo.getTopLevelFormula(static_cast<unsigned>(partitionIndex));
+}
+
+PTRef PartitionManager::getPartitionForInternal(PTRef term) const {
+    // we ignore if the term is actually top-level ..
+    int const partitionIndex = getPartitionIndex(term);
+    return getPartition(partitionIndex);
+}
+
+PTRef PartitionManager::getInternalPartition(int partitionIndex) const {
+    if (partitionIndex < 0) { return PTRef_Undef; }
+    return partitionInfo.getInternalFormula(static_cast<unsigned>(partitionIndex));
+}
+
+PTRef PartitionManager::getInternalPartitionFor(PTRef term) const {
+    // we ignore if the term is actually internal ..
+    int const partitionIndex = getPartitionIndex(term);
+    return getInternalPartition(partitionIndex);
+}
+
 void PartitionManager::addClauseClassMask(CRef c, ipartitions_t const & toadd) {
     partitionInfo.addClausePartition(c, toadd);
 }
diff --git a/src/api/PartitionManager.h b/src/api/PartitionManager.h
index bd62e7a67..9c2bb345b 100644
--- a/src/api/PartitionManager.h
+++ b/src/api/PartitionManager.h
@@ -49,11 +49,17 @@ class PartitionManager {
     void invalidatePartitions(ipartitions_t const & toinvalidate);
 
     inline std::vector<PTRef> getPartitions() const { return partitionInfo.getTopLevelFormulas(); }
+    inline std::vector<PTRef> getInternalPartitions() const { return partitionInfo.getInternalFormulas(); }
 
     vec<PTRef> getPartitions(ipartitions_t const &) const;
 
     unsigned getNofPartitions() const { return partitionInfo.getNoOfPartitions(); }
 
+    PTRef getPartition(int partitionIndex) const;
+    PTRef getPartitionForInternal(PTRef) const;
+    PTRef getInternalPartition(int partitionIndex) const;
+    PTRef getInternalPartitionFor(PTRef) const;
+
     void transferPartitionMembership(PTRef old, PTRef new_ptref) {
         if (new_ptref == old) { return; }
         this->addIPartitions(new_ptref, getIPartitions(old));
diff --git a/src/cnfizers/TermMapper.h b/src/cnfizers/TermMapper.h
index 6d21d7ef5..e5d69a219 100644
--- a/src/cnfizers/TermMapper.h
+++ b/src/cnfizers/TermMapper.h
@@ -51,6 +51,7 @@ class TermMapper {
     // Returns the literal corresponding to the term. The connection must already exist.
     Lit getLit(PTRef) const;
     // Test if the given term already has an assigned SAT variable
+    // Only *internal* terms that were given to the SMT solver are tracked, not necessarily top-level formulas!
     bool hasLit(PTRef tr) const {
         Var v = var_Undef;
         peekVar(toPositive(tr), v);
diff --git a/src/common/FlaPartitionMap.cc b/src/common/FlaPartitionMap.cc
index 8d24f8645..614781db2 100644
--- a/src/common/FlaPartitionMap.cc
+++ b/src/common/FlaPartitionMap.cc
@@ -1,34 +1,53 @@
-//
-// Created by Martin Blicha on 06.10.18.
-//
+/*
+ *  Copyright (c) 2018, Martin Blicha <martin.blicha@gmail.com>
+ *                2025, Tomas Kolarik <tomaqa@gmail.com>
+ *
+ *  SPDX-License-Identifier: MIT
+ *
+ */
 
 #include "FlaPartitionMap.h"
 
 namespace opensmt {
-std::vector<PTRef> FlaPartitionMap::get_top_level_flas() const {
+std::vector<PTRef> FlaPartitionMap::getTopLevelFlas() const {
     std::vector<PTRef> res;
-    res.reserve(top_level_flas.size());
-    for (auto entry : top_level_flas) {
+    res.reserve(topLevelFlaToIdxMap.size());
+    for (auto entry : topLevelFlaToIdxMap) {
+        res.push_back(entry.first);
+    }
+    return res;
+}
+
+std::vector<PTRef> FlaPartitionMap::getInternalFlas() const {
+    std::vector<PTRef> res;
+    res.reserve(topLevelFlaToIdxMap.size());
+    for (auto entry : topLevelFlaToIdxMap) {
         res.push_back(entry.first);
     }
     return res;
 }
 
 void FlaPartitionMap::transferPartitionMembership(PTRef old, PTRef new_ptref) {
-    auto it = top_level_flas.find(old);
-    if (it != top_level_flas.end()) {
-        store_other_fla_index(new_ptref, it->second);
+    if (auto it = internalFlaToIdxMap.find(old); it != internalFlaToIdxMap.end()) {
+        unsigned int const idx = it->second;
+        storeInternalFlaIndex(new_ptref, idx);
+        internalFlaToIdxMap.erase(it);
+        assert(internalIdxToFlaMap.contains(idx) and internalIdxToFlaMap.at(idx) == new_ptref);
         return;
     }
-    auto other_it = other_flas.find(old);
-    if (other_it != other_flas.end()) { store_other_fla_index(new_ptref, other_it->second); }
+
+    if (auto it = topLevelFlaToIdxMap.find(old); it != topLevelFlaToIdxMap.end()) {
+        storeInternalFlaIndex(new_ptref, it->second);
+    }
 }
 
 int FlaPartitionMap::getPartitionIndex(PTRef ref) const {
-    auto it = top_level_flas.find(ref);
-    if (it != top_level_flas.end()) { return static_cast<int>(it->second); }
-    auto other_it = other_flas.find(ref);
-    if (other_it != other_flas.end()) { return static_cast<int>(other_it->second); }
+    if (auto it = topLevelFlaToIdxMap.find(ref); it != topLevelFlaToIdxMap.end()) {
+        return static_cast<int>(it->second);
+    }
+    if (auto it = internalFlaToIdxMap.find(ref); it != internalFlaToIdxMap.end()) {
+        return static_cast<int>(it->second);
+    }
     return -1;
 }
 } // namespace opensmt
diff --git a/src/common/FlaPartitionMap.h b/src/common/FlaPartitionMap.h
index 9c9fa61f5..f33ba0876 100644
--- a/src/common/FlaPartitionMap.h
+++ b/src/common/FlaPartitionMap.h
@@ -1,29 +1,67 @@
-//
-// Created by Martin Blicha on 06.10.18.
-//
+/*
+ *  Copyright (c) 2018, Martin Blicha <martin.blicha@gmail.com>
+ *                2025, Tomas Kolarik <tomaqa@gmail.com>
+ *
+ *  SPDX-License-Identifier: MIT
+ *
+ */
 
 #ifndef OPENSMT_FLAPARTITIONMAP_H
 #define OPENSMT_FLAPARTITIONMAP_H
 
 #include <pterms/PTRef.h>
 
+#include <cassert>
 #include <map>
+#include <unordered_map>
 #include <vector>
 
 namespace opensmt {
 class FlaPartitionMap {
 public:
-    void store_top_level_fla_index(PTRef fla, unsigned int idx) { top_level_flas[fla] = idx; }
-    void store_other_fla_index(PTRef fla, unsigned int idx) { other_flas[fla] = idx; }
-    std::vector<PTRef> get_top_level_flas() const;
-    unsigned getNoOfPartitions() const { return get_top_level_flas().size(); }
+    inline void storeTopLevelFlaIndex(PTRef fla, unsigned int idx);
+    inline void storeInternalFlaIndex(PTRef fla, unsigned int idx);
+    std::vector<PTRef> getTopLevelFlas() const;
+    std::vector<PTRef> getInternalFlas() const;
+    inline PTRef getTopLevelFla(unsigned int idx) const noexcept;
+    inline PTRef getInternalFla(unsigned int idx) const noexcept;
+    inline unsigned getNoOfPartitions() const;
     void transferPartitionMembership(PTRef old, PTRef new_ptref);
     int getPartitionIndex(PTRef ref) const;
 
 private:
-    std::map<PTRef, unsigned int> top_level_flas;
-    std::map<PTRef, unsigned int> other_flas;
+    std::map<PTRef, unsigned int> topLevelFlaToIdxMap;
+    std::map<PTRef, unsigned int> internalFlaToIdxMap;
+    std::unordered_map<unsigned int, PTRef> topLevelIdxToFlaMap;
+    std::unordered_map<unsigned int, PTRef> internalIdxToFlaMap;
 };
+
+void FlaPartitionMap::storeTopLevelFlaIndex(PTRef fla, unsigned int idx) {
+    topLevelFlaToIdxMap[fla] = idx;
+    topLevelIdxToFlaMap[idx] = fla;
+}
+
+void FlaPartitionMap::storeInternalFlaIndex(PTRef fla, unsigned int idx) {
+    internalFlaToIdxMap[fla] = idx;
+    internalIdxToFlaMap[idx] = fla;
+}
+
+PTRef FlaPartitionMap::getTopLevelFla(unsigned int idx) const noexcept {
+    if (auto it = topLevelIdxToFlaMap.find(idx); it != topLevelIdxToFlaMap.end()) { return it->second; }
+    return PTRef_Undef;
+}
+
+PTRef FlaPartitionMap::getInternalFla(unsigned int idx) const noexcept {
+    if (auto it = internalIdxToFlaMap.find(idx); it != internalIdxToFlaMap.end()) { return it->second; }
+    return PTRef_Undef;
+}
+
+unsigned FlaPartitionMap::getNoOfPartitions() const {
+    assert(getTopLevelFlas().size() == topLevelFlaToIdxMap.size());
+    assert(topLevelFlaToIdxMap.size() <= topLevelIdxToFlaMap.size());
+    return topLevelFlaToIdxMap.size();
+}
+
 } // namespace opensmt
 
 #endif // OPENSMT_FLAPARTITIONMAP_H
diff --git a/src/common/PartitionInfo.cc b/src/common/PartitionInfo.cc
index 2cd6bccb7..c7b236b9a 100644
--- a/src/common/PartitionInfo.cc
+++ b/src/common/PartitionInfo.cc
@@ -6,7 +6,7 @@
 
 namespace opensmt {
 void PartitionInfo::assignTopLevelPartitionIndex(unsigned int n, PTRef tr) {
-    flaPartitionMap.store_top_level_fla_index(tr, n);
+    flaPartitionMap.storeTopLevelFlaIndex(tr, n);
     ipartitions_t p = 0;
     setbit(p, n);
     addIPartitions(tr, p);
diff --git a/src/common/PartitionInfo.h b/src/common/PartitionInfo.h
index 18906d722..3d2bc395a 100644
--- a/src/common/PartitionInfo.h
+++ b/src/common/PartitionInfo.h
@@ -28,7 +28,10 @@ class PartitionInfo {
     ipartitions_t const & getClausePartitions(CRef) const;
     void addClausePartition(CRef c, ipartitions_t const & p);
 
-    inline std::vector<PTRef> getTopLevelFormulas() const { return flaPartitionMap.get_top_level_flas(); }
+    inline std::vector<PTRef> getTopLevelFormulas() const { return flaPartitionMap.getTopLevelFlas(); }
+    inline std::vector<PTRef> getInternalFormulas() const { return flaPartitionMap.getInternalFlas(); }
+    inline PTRef getTopLevelFormula(unsigned int idx) const { return flaPartitionMap.getTopLevelFla(idx); }
+    inline PTRef getInternalFormula(unsigned int idx) const { return flaPartitionMap.getInternalFla(idx); }
     inline unsigned int getNoOfPartitions() const { return flaPartitionMap.getNoOfPartitions(); }
     inline void transferPartitionMembership(PTRef o, PTRef n) {
         return flaPartitionMap.transferPartitionMembership(o, n);