#include <nanobind/nanobind.h>
#include <nanobind/stl/set.h> 
#include <nanobind/stl/vector.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/tuple.h>
#include <nanobind/stl/pair.h>
#include <nanobind/ndarray.h>

#include <vector>
#include <string>
#include <tuple>
#include <cmath>
#include <limits>
#include <functional>
#include <unordered_map>
#include <random>
#include <algorithm>
#include <fstream> 
#include <iostream>
#include <sstream>
#include <iomanip>
#include <string>
#include <deque>
#include <queue>
#include <unordered_set>
#include <set>

#ifdef _OPENMP
#include <omp.h>
#endif

namespace nb = nanobind;
using namespace nb::literals;
using namespace std;

struct MinifyResult {
    bool                is_empty_language = false;
    int                 n_classes = 0;
    int                 new_initial = -1;
    vector<int>         eq_class;       // mapping: old_state_id -> new_class_id
    vector<vector<std::pair<int,int>>> class_trans; // mapping: new_class_id -> list of {sym, dest_class}
    vector<bool>        class_is_final; // mapping: new_class_id -> bool
};

class FastPartition {
public:
    vector<int> elements, location, block_of, block_start, block_size, marked_count;

    explicit FastPartition(int n) 
        : elements(n), location(n), block_of(n, 0), block_start(1, 0), block_size(1, n), marked_count(n, 0) {
        for (int i = 0; i < n; ++i) { elements[i] = i; location[i] = i; }
    }

    vector<pair<int, int>> refine(const vector<int>& movers) {
        vector<int> active_blocks;
        for (int s : movers) {
            int b = block_of[s];
            if (marked_count[b]++ == 0) active_blocks.push_back(b);
            int cur_idx = location[s], tgt_idx = block_start[b] + marked_count[b] - 1;
            swap(elements[cur_idx], elements[tgt_idx]);
            location[elements[cur_idx]] = cur_idx;
            location[elements[tgt_idx]] = tgt_idx;
        }

        vector<pair<int, int>> splits;
        for (int b : active_blocks) {
            int count = marked_count[b]; marked_count[b] = 0;
            if (count == block_size[b]) continue;
            int new_b = block_start.size();
            block_start.push_back(block_start[b]);
            block_size.push_back(count);
            block_start[b] += count;
            block_size[b] -= count;
            for (int i = 0; i < count; ++i) block_of[elements[block_start[new_b] + i]] = new_b;
            splits.emplace_back(new_b, b);
        }
        return splits;
    }
};

MinifyResult minify_dfa_cpp(
    int n_states,
    const vector<vector<std::pair<int,int>>>& adj,
    int initial,
    const vector<bool>& is_final_input
) {
    // Pruning (Forward and Backward Reachability)
    vector<bool> reachable(n_states, false), leads_to_final(n_states, false);
    queue<int> q; q.push(initial); reachable[initial] = true;
    while(!q.empty()){
        int u = q.front(); q.pop();
        for(auto& edge : adj[u]) if(!reachable[edge.second]) { reachable[edge.second] = true; q.push(edge.second); }
    }

    vector<vector<int>> rev_adj(n_states);
    for(int i=0; i<n_states; ++i) for(auto& edge : adj[i]) rev_adj[edge.second].push_back(i);
    for(int i=0; i<n_states; ++i) if(is_final_input[i] && reachable[i]) { leads_to_final[i] = true; q.push(i); }
    while(!q.empty()){
        int u = q.front(); q.pop();
        for(int v : rev_adj[u]) if(reachable[v] && !leads_to_final[v]) { leads_to_final[v] = true; q.push(v); }
    }

    vector<int> live_states;
    vector<int> global_to_live(n_states, -1);
    for(int i=0; i<n_states; ++i) if(reachable[i] && leads_to_final[i]) {
        global_to_live[i] = live_states.size();
        live_states.push_back(i);
    }
    
    if (global_to_live[initial] == -1) { MinifyResult r; r.is_empty_language = true; return r; }

    int N_live = live_states.size();
    vector<vector<pair<int, int>>> back(N_live);
    for(int u_live=0; u_live < N_live; ++u_live) {
        for(auto& edge : adj[live_states[u_live]]) {
            int v_live = global_to_live[edge.second];
            if(v_live != -1) back[v_live].push_back({u_live, edge.first});
        }
    }

    // Hopcrofts Algorithm
    FastPartition eq(N_live);
    vector<int> live_finals;
    for(int i=0; i<N_live; ++i) if(is_final_input[live_states[i]]) live_finals.push_back(i);
    auto first_split = eq.refine(live_finals);

    vector<int> worklist;
    vector<bool> in_worklist(N_live * 2, false); 
    auto add_to_wl = [&](int b) { if(!in_worklist[b]) { worklist.push_back(b); in_worklist[b] = true; }};

    if(first_split.empty()) add_to_wl(0);
    else {
        int b1 = first_split[0].first, b2 = first_split[0].second;
        add_to_wl(eq.block_size[b1] <= eq.block_size[b2] ? b1 : b2);
    }

    struct Move { int src; int sym; };
    vector<int> movers;
    vector<Move> all_moves;
    vector<int> splitter_elements;
    int head = 0;

    while(head < worklist.size()){
        int act = worklist[head++];
        in_worklist[act] = false;

        splitter_elements.clear();
        for(int i=0; i<eq.block_size[act]; ++i) 
            splitter_elements.push_back(eq.elements[eq.block_start[act] + i]);

        all_moves.clear();
        for(int v_live : splitter_elements) 
            for(auto& edge : back[v_live]) all_moves.push_back({edge.first, edge.second});

        sort(all_moves.begin(), all_moves.end(), [](const Move& a, const Move& b){ return a.sym < b.sym; });

        auto it = all_moves.begin();
        while(it != all_moves.end()){
            int cur_sym = it->sym;
            movers.clear();
            while(it != all_moves.end() && it->sym == cur_sym) { movers.push_back(it->src); ++it; }
            
            for(auto [YintX, YdiffX] : eq.refine(movers)) {
                if(in_worklist[YdiffX]) add_to_wl(YintX);
                else add_to_wl(eq.block_size[YintX] <= eq.block_size[YdiffX] ? YintX : YdiffX);
            }
        }
    }

    // Build Synchronized Result
    MinifyResult res;
    int num_blocks = eq.block_start.size();
    vector<int> block_to_res(num_blocks, -1);
    for(int i=0; i<num_blocks; ++i) if(eq.block_size[i] > 0) block_to_res[i] = res.n_classes++;

    res.eq_class.assign(n_states, -1);
    res.class_is_final.assign(res.n_classes, false);
    res.class_trans.assign(res.n_classes, {});
    res.new_initial = block_to_res[eq.block_of[global_to_live[initial]]];

    for(int i=0; i<N_live; ++i) {
        int u_global = live_states[i];
        int u_class = block_to_res[eq.block_of[i]];
        res.eq_class[u_global] = u_class;
        if(is_final_input[u_global]) res.class_is_final[u_class] = true;
    }

    for(int c=0; c < res.n_classes; ++c) {
        int rep_live = -1;
        for(int i=0; i<N_live; ++i) if(block_to_res[eq.block_of[i]] == c) { rep_live = i; break; }
        for(auto& edge : adj[live_states[rep_live]]) {
            int v_live = global_to_live[edge.second];
            if(v_live != -1) res.class_trans[c].push_back({edge.first, block_to_res[eq.block_of[v_live]]});
        }
    }
    return res;
}

// ---------------------------------------------------------------------------
//   branches : list[tuple[float, list[float]]]
//   trans    : list[list[tuple[int, float, int]]]
// ---------------------------------------------------------------------------

using PyBranch   = std::pair<double, std::vector<double>>;
using PyTrans    = std::tuple<int, double, int>;
using PyBranches = std::vector<PyBranch>;
using PyTransTable = std::vector<std::vector<PyTrans>>;
// ---------------------------------------------------------------------------

// Helper to pass std::vector memory to Python without copying
template <typename T>
nb::ndarray<nb::numpy, nb::c_contig, T, nb::shape<1>> as_ndarray_ref(std::vector<T>&& vec) {

    if (vec.empty()) {
        return nb::ndarray<nb::numpy, nb::c_contig, T, nb::shape<1>>(nullptr, {0});
    }
    
    auto* ptr = new std::vector<T>(std::move(vec));
    
    nb::capsule owner(ptr, [](void* p) noexcept {
        delete static_cast<std::vector<T>*>(p);
    });
    
    return nb::ndarray<nb::numpy, nb::c_contig, T, nb::shape<1>>(
        ptr->data(),
        { ptr->size() },
        owner
    );
}

template <typename T>
nb::ndarray<nb::numpy, nb::c_contig, T, nb::shape<1>> as_ndarray_cpy(const std::vector<T>& vec) {
    if (vec.empty()) {
        return nb::ndarray<nb::numpy, nb::c_contig, T, nb::shape<1>>(nullptr, {0});
    }
    auto* ptr = new std::vector<T>(vec);
    
    nb::capsule owner(ptr, [](void* p) noexcept {
        delete static_cast<std::vector<T>*>(p);
    });
    
    return nb::ndarray<nb::numpy, nb::c_contig, T, nb::shape<1>>(
        ptr->data(),
        { ptr->size() },
        owner
    );
}

// Return type for generate: (symbol_indices, state_indices).
// state_indices is empty when include_states=false.
using GenerateResult = std::pair<std::vector<std::uint16_t>, std::vector<std::uint16_t>>;

struct Transition {
    int    sym_idx;
    double prob;
    int    target;
};

struct Branch {
    double              weight;
    std::vector<double> state_vec;
    Branch() {}
    Branch(double weight, std::vector<double> state_vec)
        : weight(weight), state_vec(std::move(state_vec)) {}
};

struct SparseEntry {
    int index;
    double val;
};

struct SparseBranch {
    double weight;
    std::vector<SparseEntry> state_vec;
    size_t cached_hash;
};

// ---------------------------------------------------------------------------
//  Hashing helpers
// ---------------------------------------------------------------------------
 
inline void hash_mix(size_t& seed, int idx, int64_t q) {
    size_t h = std::hash<int>{}(idx) ^ (std::hash<int64_t>{}(q) + 0x9e3779b9ULL);
    seed ^= h + 0x9e3779b9ULL + (seed << 6) + (seed >> 2);
}
 
size_t compute_sparse_hash(const std::vector<SparseEntry>& v, double scale) {
    size_t seed = 0;
    for (const auto& e : v) {
        int64_t q = static_cast<int64_t>(std::round(e.val * scale));
        if (q == 0) continue;
        hash_mix(seed, e.index, q);
    }
    return seed;
}

struct MemoizedHasher {
    size_t operator()(const SparseBranch* b) const {
        return b->cached_hash;
    }
};

struct ExactEquality {
    double scale;
    explicit ExactEquality(double s) : scale(s) {}
    bool operator()(const SparseBranch* a, const SparseBranch* b) const {
        if (a->state_vec.size() != b->state_vec.size()) return false;
        for (size_t i = 0; i < a->state_vec.size(); ++i) {
            if (a->state_vec[i].index != b->state_vec[i].index) return false;
            auto qa = static_cast<int64_t>(std::round(a->state_vec[i].val * scale));
            auto qb = static_cast<int64_t>(std::round(b->state_vec[i].val * scale));
            if (qa != qb) return false;
        }
        return true;
    }
};

std::vector<SparseBranch> merge_branches_exact(
    std::vector<SparseBranch>& branches, 
    double scale,
    double eps_weight)
{
    if (branches.empty()) return {};
 
    ExactEquality eq(scale);

    std::unordered_map<SparseBranch*, SparseBranch*, MemoizedHasher, ExactEquality>
        merged(branches.size(), MemoizedHasher{}, eq);
    
    merged.reserve(branches.size());
 
    for (auto& branch : branches) {
        if (branch.weight < eps_weight) continue;
        auto it = merged.find(&branch);
        if (it != merged.end())
            it->second->weight += branch.weight;
        else
            merged[&branch] = &branch;
    }
 
    std::vector<SparseBranch> out;
    out.reserve(merged.size());
    for (auto& [key, val] : merged)
        out.push_back(std::move(*val));
    return out;
}

struct ComplexityMeasures {
    double              E;
    double              T;
    double              S;
    std::vector<double> H_L;
    std::vector<double> T_L;
    std::vector<double> h_mu_L;
    std::vector<double> H_sync;
    bool                converged;
};

ComplexityMeasures block_entropy_convergence_cpp(
    const double h_mu,
    const int    n_states,
    const int    n_symbols,
    const double convergence_tol,
    const double precision,
    const double eps,
    const PyBranches&   py_branches,
    const PyTransTable& py_trans,
    const size_t max_branches
) {
    const double prune_tol    = 1e-6;
    const double eps_weight   = std::numeric_limits<double>::epsilon();
    const double fine_scale   = std::pow(10.0, precision);
    const int    MAX_L        = 1000; 

    // Transition Table
    std::vector<int> trans_offsets(n_states + 1, 0);
    std::vector<Transition> flat_trans;
    flat_trans.reserve(n_states * n_symbols);
    for (int i = 0; i < n_states; ++i) {
        trans_offsets[i] = static_cast<int>(flat_trans.size());
        for (const auto& [sym, prob, target] : py_trans[i])
            flat_trans.push_back({sym, prob, target});
    }
    trans_offsets[n_states] = static_cast<int>(flat_trans.size());

    // Initial State Setup
    std::vector<SparseBranch> branches;
    auto state_entropy = [](const std::vector<SparseEntry>& sv) {
        double h = 0.0;
        for (const auto& e : sv) if (e.val > 0.0) h -= e.val * std::log2(e.val);
        return h;
    };

    for (const auto& [w, sv] : py_branches) {
        std::vector<SparseEntry> sparse_v;
        for (int i = 0; i < n_states; ++i) if (sv[i] >= eps) sparse_v.push_back({i, sv[i]});
        branches.push_back({w, sparse_v, compute_sparse_hash(sparse_v, fine_scale)});
    }

    ComplexityMeasures result;
    result.converged = false;
    double block_entropy = 0.0, kahan_c = 0.0, prev_excess = std::numeric_limits<double>::infinity();
    double T = 0.0, T_kahan_c = 0.0, H_1 = 0.0;
    double S_accum = 0.0, S_kahan_c = 0.0;

    // L=0 synchronization term
    double H_sync_0 = 0.0;
    for (const auto& b : branches) H_sync_0 += b.weight * state_entropy(b.state_vec);
    S_accum = H_sync_0; 
    result.H_sync.push_back(H_sync_0);

    // Working buffers
    std::vector<double> sym_probs(n_symbols, 0.0), next_dense(n_symbols * n_states, 0.0);
    std::vector<SparseEntry> scratch; scratch.reserve(n_states);

    // Main Convergence Loop
    for (int step = 0; step < MAX_L; ++step) {
        int L = step + 1;
        std::deque<SparseBranch> merged_deque;
        ExactEquality eq(fine_scale);
        std::unordered_map<SparseBranch*, SparseBranch*, MemoizedHasher, ExactEquality> 
            merge_map(0, MemoizedHasher{}, eq);

        double conditional_entropy = 0.0; 

        // Generate next mixed states
        for (const auto& branch : branches) {
            std::fill(sym_probs.begin(), sym_probs.end(), 0.0);
            std::fill(next_dense.begin(), next_dense.end(), 0.0);

            for (const auto& entry : branch.state_vec) {
                for (int t = trans_offsets[entry.index]; t < trans_offsets[entry.index + 1]; ++t) {
                    const auto& tr = flat_trans[t];
                    double move_prob = entry.val * tr.prob;
                    sym_probs[tr.sym_idx] += move_prob;
                    next_dense[tr.sym_idx * n_states + tr.target] += move_prob;
                }
            }

            for (int s = 0; s < n_symbols; ++s) {
                double p_s = sym_probs[s];
                if (p_s < eps) continue;

                conditional_entropy -= branch.weight * p_s * std::log2(p_s);
                scratch.clear();
                for (int i = 0; i < n_states; ++i) {
                    double val = next_dense[s * n_states + i] / p_s;
                    if (val >= eps) scratch.push_back({i, val});
                }

                double new_weight = branch.weight * p_s;
                if (new_weight < eps_weight) continue;

                SparseBranch probe{new_weight, scratch, compute_sparse_hash(scratch, fine_scale)};
                auto it = merge_map.find(&probe);
                if (it != merge_map.end()) it->second->weight += new_weight;
                else {
                    merged_deque.push_back(std::move(probe));
                    merge_map[&merged_deque.back()] = &merged_deque.back();
                }
            }
        }

        // Accumulate Metrics for Step L
        {
            double y = conditional_entropy - kahan_c;
            double t = block_entropy + y;
            kahan_c = (t - block_entropy) - y;
            block_entropy = t;
        }

        if (step == 0) H_1 = block_entropy;

        double H_sync_L = 0.0;
        for (const auto& b : merged_deque) H_sync_L += b.weight * state_entropy(b.state_vec);
        
        { // Kahan accumulation for S
            double y = H_sync_L - S_kahan_c;
            double t = S_accum + y;
            S_kahan_c = (t - S_accum) - y;
            S_accum = t;
        }

        double term = static_cast<double>(L) * (conditional_entropy - h_mu);
        { // Kahan accumulation for T
            double y = term - T_kahan_c;
            double t = T + y;
            T_kahan_c = (t - T) - y;
            T = t;
        }

        result.H_L.push_back(block_entropy);
        result.T_L.push_back(L * H_1 - block_entropy);
        result.h_mu_L.push_back(conditional_entropy);
        result.H_sync.push_back(H_sync_L);

        // Convergence Check
        double excess = block_entropy - L * h_mu;

        if (H_sync_L < 1e-5 && std::abs(excess - prev_excess) < convergence_tol) {
            result.E = std::max(0.0, excess);
            result.T = std::max(0.0, T);
            result.S = std::max(0.0, S_accum);
            result.converged = true;
            return result;
        }

        prev_excess = excess;

        branches.assign(std::make_move_iterator(merged_deque.begin()), 
                        std::make_move_iterator(merged_deque.end()));

        std::cout <<  "Branches: " << branches.size() << ",  E=" << excess << ", H_s=" << H_sync_L << std::endl;

        if(branches.size() >= max_branches) {
            std::cout << "Warning: max branches reached before convergence." << std::endl;
            break;
        }
    }

    result.E = std::max(0.0, prev_excess);
    result.T = std::max(0.0, T);
    result.S = std::max(0.0, S_accum);
    return result;
}

GenerateResult generate(
    const uint64_t      n_gen,
    const int           start_state_index,
    const PyTransTable& transitions,    // [state][tr] = (sym_idx, prob, target)
    const bool          include_states,  // whether to populate the state vector
    const uint64_t      random_seed
)
{
    const size_t n_states = transitions.size();
 
    // Build one discrete_distribution<int> per state.
    // The sampled integer is the intra-state transition index.
    std::vector<std::discrete_distribution<int>> dists;
    dists.reserve(n_states);

    for (const auto& state_trans : transitions) {
        std::vector<double> weights;
        weights.reserve(state_trans.size());
        for (const auto& [sym_idx, prob, target] : state_trans) {
            weights.push_back(prob);
        }
        dists.emplace_back(weights.begin(), weights.end());
    }
 
    // std::mt19937_64 rng{ std::random_device{}() }; // for random seed
    std::mt19937_64 rng{ random_seed };

    const auto sz = static_cast<size_t>(n_gen);
    
    int cur = start_state_index;
    
    std::vector<std::uint16_t> states;

    std::vector<std::uint16_t> symbols( sz );
    symbols.shrink_to_fit();

    if( include_states ) {
        
        states.resize( sz + 1 );
        states.shrink_to_fit();

        for (uint64_t i = 0; i < n_gen; ++i) {
            states[ i ] = static_cast<std::uint16_t>( cur );
            const int chosen                    = dists[cur](rng);
            const auto& [sym_idx, prob, target] = transitions[cur][chosen];
            symbols[ i ] = static_cast<std::uint16_t>( sym_idx );
            cur = target;
        }

        states[n_gen] = static_cast<std::uint16_t>( cur );  // final target state
    }
    else 
    {
        for (uint64_t i = 0; i < n_gen; ++i) {
            const int chosen                    = dists[cur](rng);
            const auto& [sym_idx, prob, target] = transitions[cur][chosen];
            symbols[ i ] = static_cast<std::uint16_t>( sym_idx );
            cur = target;
        }         
    }
 
    return { std::move(symbols), std::move(states) };
}

// Strongly conn. components port of networkx (Tarjan's algorithm)
std::vector<std::set<int>> strongly_connected_components_cpp(
    const std::vector<std::vector<int>>& transitions) 
{
    const size_t N = transitions.size();
    
    std::vector<int> preorder(N, -1);
    std::vector<int> lowlink(N, -1);
    std::vector<bool> scc_found(N, false);
    
    std::vector<size_t> next_neighbor_idx(N, 0);
    
    std::vector<int> scc_stack;
    std::vector<std::set<int>> res;
    int counter = 0;

    for (size_t start_node = 0; start_node < N; ++start_node) {
        if (scc_found[start_node] || preorder[start_node] != -1) continue;

        std::vector<int> dfs_stack = {(int)start_node};

        while (!dfs_stack.empty()) {
            int v = dfs_stack.back();

            if (preorder[v] == -1) {
                preorder[v] = ++counter;
            }

            bool pushed = false;
            for (; next_neighbor_idx[v] < transitions[v].size(); ++next_neighbor_idx[v]) {
                int w = transitions[v][next_neighbor_idx[v]];
                if (preorder[w] == -1) {
                    dfs_stack.push_back(w);
                    pushed = true;
                    next_neighbor_idx[v]++; 
                    break;
                }
            }

            if (pushed) continue;

            lowlink[v] = preorder[v];
            for (int w : transitions[v]) {
                if (!scc_found[w]) {
                    if (preorder[w] > preorder[v]) {
                        lowlink[v] = std::min(lowlink[v], lowlink[w]);
                    } else {
                        lowlink[v] = std::min(lowlink[v], preorder[w]);
                    }
                }
            }

            dfs_stack.pop_back();

            if (lowlink[v] == preorder[v]) {
                std::set<int> scc = {v};
                while (!scc_stack.empty() && preorder[scc_stack.back()] > preorder[v]) {
                    int k = scc_stack.back();
                    scc_stack.pop_back();
                    scc.insert(k);
                }
                for (int node : scc) scc_found[node] = true;
                res.push_back(std::move(scc));
            } else {
                scc_stack.push_back(v);
            }
        }
    }
    return res;
}

// ---------------------------------------------------------------------------
//                                Bindings
// ---------------------------------------------------------------------------

NB_MODULE(_am_fast, m) {

    m.doc() = "HMM utilities (nanobind)";
 
    nb::class_<ComplexityMeasures>(m, "ComplexityMeasures")
        .def_ro("E",              &ComplexityMeasures::E)
        .def_ro("T",              &ComplexityMeasures::T)
        .def_ro("S",              &ComplexityMeasures::S)
        .def_ro("converged",      &ComplexityMeasures::converged)
        .def_prop_ro("H_L",    [](ComplexityMeasures& r) { return as_ndarray_cpy( r.H_L ); }, nb::rv_policy::move )
        .def_prop_ro("T_L",    [](ComplexityMeasures& r) { return as_ndarray_cpy( r.T_L ); }, nb::rv_policy::move )
        .def_prop_ro("h_mu_L", [](ComplexityMeasures& r) { return as_ndarray_cpy( r.h_mu_L ); }, nb::rv_policy::move )
        .def_prop_ro("H_sync", [](ComplexityMeasures& r) { return as_ndarray_cpy( r.H_sync ); }, nb::rv_policy::move )
        .def("__repr__", [](const ComplexityMeasures& c) {

                // Helper lambda to format vectors as [1.23, 4.56, ..., 8.90, 0.12]
                auto fmt_vec = [](const std::vector<double>& v) -> std::string {
                    if (v.empty()) return "[]";
                    std::ostringstream oss;
                    oss << std::fixed << std::setprecision(4) << "[";
                    if (v.size() <= 4) {
                        for (size_t i = 0; i < v.size(); ++i) {
                            oss << v[i] << (i == v.size() - 1 ? "" : ", ");
                        }
                    } else {
                        oss << v[0] << ", " << v[1] << ", ..., " 
                            << v[v.size() - 2] << ", " << v[v.size() - 1];
                    }
                    oss << "]";
                    return oss.str();
                };

                std::ostringstream out;
                out << "ComplexityMeasures(\n"
                    << "  E=" << std::fixed << std::setprecision(6) << c.E << ",\n"
                    << "  T=" << c.T << ",\n"
                    << "  S=" << c.S << ",\n"
                    << "  converged=" << (c.converged ? "True" : "False") << ",\n"
                    << "  H_L="     << fmt_vec(c.H_L) << ",\n"
                    << "  T_L="     << fmt_vec(c.T_L) << ",\n"
                    << "  h_mu_L="  << fmt_vec(c.h_mu_L) << ",\n"
                    << "  H_sync="  << fmt_vec(c.H_sync) << "\n"
                    << ")";
                return out.str();
            });

    m.def(
        "block_entropy_convergence_cpp",
        [](const double h_mu,
           const int n_states,
           const int n_symbols,
           const double convergence_tol,
           const double precision,
           const double eps,
           const PyBranches& py_branches,
           const PyTransTable& py_trans,
           const size_t max_branches) {

            return block_entropy_convergence_cpp(
                h_mu, 
                n_states, 
                n_symbols,
                convergence_tol, 
                precision, 
                eps,
                py_branches, 
                py_trans, 
                max_branches
            );
        },
        "h_mu"_a, "n_states"_a, "n_symbols"_a,
        "convergence_tol"_a, "precision"_a, "eps"_a,
        "branches"_a, "trans"_a, "max_branches"_a
    );

    m.def(
        "generate_cpp",
        [](const uint64_t n_gen, 
            const int start_state_index, 
            const PyTransTable& transitions, 
            const bool include_states, 
            const uint64_t random_seed ) 
        {
            auto [symbols, states] = generate(n_gen, start_state_index, transitions, include_states, random_seed );
            return std::make_pair(
                as_ndarray_ref(std::move(symbols)),
                as_ndarray_ref(std::move(states))
            );
        },
        "n_gen"_a, 
        "start_state_index"_a, 
        "transitions"_a,
        "include_states"_a,
        "random_seed"_a,
        "Walk the HMM for n_gen steps.\n"
        "Returns a zero-copy tuple of (symbol_indices, state_indices)."
    );

    m.def("strongly_connected_components_cpp",
        &strongly_connected_components_cpp,
        "Computes strongly connected components of DiGraph, ported from networkx"
    );

    nb::class_<MinifyResult>(m, "MinifyResult")
        .def_ro("is_empty_language", &MinifyResult::is_empty_language,
            "True if the minimised DFA accepts no strings.")
        .def_ro("n_classes",         &MinifyResult::n_classes,
            "Number of equivalence classes (= states in the minimised DFA).")
        .def_ro("new_initial",       &MinifyResult::new_initial,
            "Class id of the initial state.")
        .def_ro("eq_class",          &MinifyResult::eq_class,
            "eq_class[old_idx] = new class id, or -1 for trap-equivalent states.")
        .def_ro("class_trans",       &MinifyResult::class_trans,
            "class_trans[class] = list of DFTransition objects.")
        .def_ro("class_is_final",    &MinifyResult::class_is_final,
            "class_is_final[class] = True if the class is an accepting state.");

    m.def("minify_dfa_cpp", 
        &minify_dfa_cpp,
        "n_states"_a,
        "adj"_a,
        "initial"_a,
        "is_final_input"_a,
        "DFA Minimization using Sparse Hopcroft"
    );
}
