from __future__ import annotations

import copy
from fractions import Fraction
import gc
import warnings
from collections import deque

import numpy as np
import scipy.sparse as sp_sparse
import scipy.sparse.linalg as sp_linalg

from .am_fast.distance import jensenshannondivergence_cpu as af_jensenshannondivergence_cpu

# --------- Conditionally use cupy & cupyx if available else numpy and scipy

try : 
	import cupy as cp
	CUPY_AVAILABLE = True
except ImportError :
	CUPY_AVAILABLE = False

if CUPY_AVAILABLE :
	import cupyx.scipy.sparse as cpx_sparse
	import cupyx.scipy.sparse.linalg as cpx_linalg
	from .am_fast.distance import jensenshannondivergence_gpu as af_jensenshannondivergence_device
else :
	cp = None
	cpx_sparse = None
	cpx_linalg = None

def to_cpu( array ):
    return array.get() if CUPY_AVAILABLE and hasattr( array, 'get' ) else array

# --------- Conditionally use cagra if available else hnswlib

if CUPY_AVAILABLE :
	try :
		from cuvs.neighbors import cagra
		CAGRA_AVAILABLE = True
	except ImportError :
		CAGRA_AVAILABLE = False

if not CAGRA_AVAILABLE :
	import hnswlib
else :
	hnswlib = None

# -------------------------------------------------------------

from .am_causal_state import CausalState
from .am_transition import Transition

class MSP:
	"""
	Mixed-State Presentation (MSP) for computing 'exact' intrinsic complexity.
 
	Computes E (excess entropy),  S (synchronization information), and T (transient information) 
	using spectral decomposition of the fundamental matrix.
 
	See
		Exact Complexity: The Spectral Decomposition of Intrinsic Computation
		https://arxiv.org/abs/1309.3792
	"""

	def __init__(
		self,
		states,
		mixed_state_vectors,
		transitions,
		alphabet,
		start_state_idx: int   = 0,
		gmres_rtol:      float = 1e-10,
		gmres_maxiter:   int   = 1000,
		eps:             float = 1e-12,
	):
		self.states              = states
		self.mixed_state_vectors = np.array(mixed_state_vectors)
		self.transitions         = transitions
		self.alphabet            = alphabet
		self.start_state_idx     = start_state_idx
		self.gmres_rtol          = gmres_rtol
		self.gmres_maxiter       = gmres_maxiter
		self.EPS                 = eps
		self._cache              = {}
 
		M = len(states)
		if self.mixed_state_vectors.ndim != 2 or self.mixed_state_vectors.shape[0] != M:
			raise ValueError(
				f"mixed_state_vectors must have shape (M={M}, N_causal), "
				f"got {self.mixed_state_vectors.shape}."
			)
		if not (0 <= start_state_idx < M):
			raise ValueError(
				f"start_state_idx={start_state_idx} is out of range for M={M} states."
			)
		if gmres_rtol <= 0:
			raise ValueError(f"gmres_rtol must be positive, got {gmres_rtol}.")
		if eps <= 0:
			raise ValueError(f"eps must be positive, got {eps}.")
 
	def _build_W(self) -> sp_sparse.csr_matrix:
		"""
		Build and return the sparse row-stochastic transition matrix W (M × M).
		Rows are re-normalised to absorb floating-point drift.
		"""
		if "W" in self._cache:
			return self._cache["W"]
 
		M    = len(self.states)
		
		rows = [tr.origin_state_idx for tr in self.transitions]
		cols = [tr.target_state_idx for tr in self.transitions]
		data = [tr.prob              for tr in self.transitions]
		
		W    = sp_sparse.csr_matrix((data, (rows, cols)), shape=(M, M))
 
		row_sums = np.array(W.sum(axis=1)).ravel()
		
		if np.any(row_sums < self.EPS):
			bad = np.where(row_sums < self.EPS)[0].tolist()
			raise ValueError(
				f"States at indices {bad} have zero total outgoing probability. "
				"Check that every MSP state has at least one outgoing transition."
			)
		
		W = (sp_sparse.diags(1.0 / row_sums) @ W).tocsr()
 
		self._cache["W"] = W
		return W
 
	def _compute_pi_W(self, W) -> np.ndarray:
		"""
		Solve for the stationary distribution of W via GPU GMRES, then free
		all intermediate GPU arrays before returning the CPU result.
		"""
		if "pi_W" in self._cache:
			return self._cache["pi_W"]
 
		xp = cp if CUPY_AVAILABLE else np
		
		xp_sparse = cpx_sparse if CUPY_AVAILABLE else sp_sparse
		xp_linalg = cpx_linalg if CUPY_AVAILABLE else sp_linalg

		W_device = cpx_sparse.csr_matrix(W)
		M     = W_device.shape[0]
 
		I_device  = xp_sparse.eye(M, format="csr", dtype=W_device.dtype)
		A_device  = (W_device.T - I_device).tocsr()
 
		ones_row = xp_sparse.csr_matrix(xp.ones((1, M), dtype=W_device.dtype))
		A_aug    = xp_sparse.vstack([A_device[:-1, :], ones_row]).tocsr()
 
		b_device = xp.zeros(M, dtype=W_device.dtype)
		b_device[-1] = 1.0
 
		diag_vals = A_aug.diagonal()
		diag_vals[diag_vals == 0] = 1.0
		M_inv = xp_sparse.diags(1.0 / diag_vals)
 
		pi_W_device, info = xp_linalg.gmres(
			A_aug, b_device,
			M=M_inv,
			restart=100,
			maxiter=1000,
		)
 
		if info > 0:
			print(f"Warning: GMRES did not converge after {info} iterations")
		elif info < 0:
			raise ValueError("GMRES failed due to illegal input or breakdown.")
 
		pi_W_device = xp.clip(pi_W_device, 0.0, None)
		s = pi_W_device.sum()
		
		if s < self.EPS:
			raise ValueError("Stationary distribution collapsed.")
		
		pi_W_device /= s
 
		# Download before freeing
		result = to_cpu( pi_W_device )
 
		# -- Free every GPU object created in this method ----------------------
		del W_device, I_device, A_device, ones_row, A_aug, b_device, M_inv, diag_vals, pi_W_device

		if CUPY_AVAILABLE :
			cp.get_default_memory_pool().free_all_blocks()
			cp.get_default_pinned_memory_pool().free_all_blocks()
		
		gc.collect()
 
		self._cache["pi_W"] = result

		return result
 
	def _fundamental(
		self,
		W:     sp_sparse.csr_matrix,
		pi_W:  np.ndarray,
	) -> tuple[sp_linalg.LinearOperator, np.ndarray, np.ndarray, LinearOperator]:
		"""
		Compute the fundamental matrix
		"""
		if "fundamental" in self._cache:
			return self._cache["fundamental"]
 
		M  = W.shape[0]
		WT = W.T.tocsc()
 
		e0          = np.zeros(M)
		e0[self.start_state_idx] = 1.0
 
		def Q_T_matvec(v: np.ndarray) -> np.ndarray:
			return v - WT @ v + pi_W * v.sum()
 
		Q_LO = sp_linalg.LinearOperator((M, M), matvec=Q_T_matvec, dtype=float)
 
		W_diag  = np.array(W.diagonal())
		M_diag  = 1.0 - W_diag + pi_W
		M_diag  = np.where(np.abs(M_diag) > self.EPS, M_diag, 1.0)
		precond = sp_linalg.LinearOperator((M, M), matvec=lambda v: v / M_diag, dtype=float)
 
		Z_row, exit_code = sp_linalg.lgmres(
			Q_LO,
			e0,
			M=precond,
			rtol=self.gmres_rtol,
			atol=self.gmres_rtol,
			maxiter=self.gmres_maxiter,
		)
		if exit_code != 0:
			raise RuntimeError(
				f"GMRES failed on the fundamental matrix solve (exit code {exit_code}). "
				"Try increasing gmres_maxiter, loosening gmres_rtol, or verify that "
				"the transition matrix is irreducible."
			)
 
		fund_row = Z_row - pi_W
 
		self._cache["fundamental"] = (Q_LO, Z_row, fund_row, precond)
		return Q_LO, Z_row, fund_row, precond
 
	def _get_H_WA(self) -> np.ndarray:
		"""
		Per-state Shannon entropy of the output symbol distribution,
		marginalised over target states. Rows are normalised before use.
		"""
		if "H_WA" in self._cache:
			return self._cache["H_WA"]
 
		M, A       = len(self.states), len(self.alphabet)
		state_sym  = np.zeros((M, A))
		for tr in self.transitions:
			state_sym[tr.origin_state_idx, tr.symbol_idx] += tr.prob
 
		sym_row_sums = state_sym.sum(axis=1, keepdims=True)
		sym_row_sums = np.where(sym_row_sums > self.EPS, sym_row_sums, 1.0)
		state_sym   /= sym_row_sums
 
		safe     = state_sym > self.EPS
		log_vals = np.where(safe, np.log2(np.where(safe, state_sym, 1.0)), 0.0)
		H_WA     = -np.sum(state_sym * log_vals, axis=1)
 
		self._cache["H_WA"] = H_WA
		return H_WA
 
	def _get_H_eta(self) -> np.ndarray:
		"""
		Per-state Shannon entropy of the mixed-state (belief) vector.
		Rows are normalised before entropy computation.
		"""
		if "H_eta" in self._cache:
			return self._cache["H_eta"]
 
		msv      = self.mixed_state_vectors.copy().astype(float)
		row_sums = msv.sum(axis=1, keepdims=True)
		row_sums = np.where(row_sums > self.EPS, row_sums, 1.0)
		msv     /= row_sums
 
		safe     = msv > self.EPS
		log_vals = np.where(safe, np.log2(np.where(safe, msv, 1.0)), 0.0)
		H_eta    = -np.sum(msv * log_vals, axis=1)
 
		self._cache["H_eta"] = H_eta
		return H_eta
 
	def get_E_S_T(self) -> tuple[float, float, float]:

		try:
			W    = self._build_W()
			pi_W = self._compute_pi_W(W)   # GPU memory freed inside _compute_pi_W
			H_WA  = self._get_H_WA()
			H_eta = self._get_H_eta()
 
			Q_LO, Z_row, fund_row, precond = self._fundamental(W, pi_W)
 
			E = float(fund_row @ H_WA)
			S = float(fund_row @ H_eta)
 
			ZZ_row, exit_code = sp_linalg.lgmres(
				Q_LO,
				Z_row,
				x0=Z_row,
				M=precond,
				rtol=self.gmres_rtol,
				atol=self.gmres_rtol,
				maxiter=self.gmres_maxiter,
			)
			if exit_code != 0:
				raise RuntimeError(
					f"GMRES failed on the Z² solve for T (exit code {exit_code}). "
					"Try increasing gmres_maxiter or loosening gmres_rtol."
				)
 
			t_row = ZZ_row - pi_W
			T     = float(t_row @ H_WA)

			return E, S, T
 
		finally:

			if CUPY_AVAILABLE :
				cp.get_default_memory_pool().free_all_blocks()
				cp.get_default_pinned_memory_pool().free_all_blocks()

			gc.collect()
 
	def clear_cache(self) -> None:
		self._cache.clear()


def compute_msp_exact(
	T_x : list[ list[ list[ Fraction ] ] ],
	pi  : tuple[Fraction],
	n_states : int,
	alphabet : list[str],
	exact_state_cap: int = 1000,
	verbose: bool = True,
):

	n_symbols  = len(alphabet)

	def _apply(mu: tuple[Fraction, ...], x: int) :

		Tx = T_x[x]
		mu_Tx = [Fraction(0)] * n_states
		for i, w in enumerate(mu):
			if w == Fraction(0):
				continue
			row = Tx[i]
			for j in range(n_states):
				if row[j]:
					mu_Tx[j] += w * row[j]

		prob = sum(mu_Tx)
		if prob == Fraction(0):
			return None, Fraction(0)

		mu_next = tuple(v / prob for v in mu_Tx)
		return mu_next, prob

	belief_vectors: list[tuple[Fraction, ...]] = []
	seen_states:    dict[tuple, int]           = {}
	msp_transitions: list                      = []

	def _register(mu: tuple[Fraction, ...]) -> int:
		idx = len(belief_vectors)
		belief_vectors.append(mu)
		seen_states[mu] = idx
		return idx

	start_mu = tuple(pi)
	_register(start_mu)
	frontier = deque([0])

	while frontier and len(belief_vectors) < exact_state_cap:
		
		idx = frontier.popleft()
		mu  = belief_vectors[idx]

		for x in range(n_symbols):

			mu_next, prob = _apply(mu, x)

			if mu_next is None:
				continue

			if mu_next not in seen_states:
				next_idx = _register(mu_next)
				frontier.append(next_idx)
			else:
				next_idx = seen_states[mu_next]

			msp_transitions.append(
				Transition(
					origin_state_idx = idx,
					target_state_idx = next_idx,
					prob             = float(prob),
					symbol_idx       = x,
					pq               = prob
				)
			)

	if frontier:
		raise RuntimeError( f"Exact state cap {exact_state_cap} exceded." )

	if verbose:
		print(f"Exact MSP found: {len(belief_vectors)} states.")

	msp = MSP(
		states              = [CausalState(name=f"MS_{i}") for i in range(n_belief_vectors)],
		mixed_state_vectors = mixed_state_vectors,
		transitions         = msp_transitions,
		alphabet            = alphabet.copy(),
	)

	return msp

def compute_msp(
	T_x : list[np.ndarray],
	pi : np.ndarray,
	n_states : int,
	alphabet : list[str],
	exact_state_cap: int = 175_000,
	jsd_eps: float = 1e-7,
	k_ann: int   = 50,
	EPS : float = 1e-12,
	verbose = True,
) :

	xp = cp if CUPY_AVAILABLE else np

	n_input_states = T_x[0].shape[0]
	n_symbols      = len(alphabet)
	
	T_flat_device = xp.asarray(
		np.stack(T_x).transpose(1, 0, 2).reshape(n_input_states, n_symbols * n_input_states)
	)
 
	# -- Sparse Belief Vectors -------------------------------------------

	sparse_belief_vector_indices: list[np.ndarray] = []
	sparse_belief_vector_values: list[np.ndarray] = []
 
	def _to_sparse(v: np.ndarray):

		non_zero = np.where( np.abs( v ) > 1e-20 )[ 0 ]
		return non_zero.astype(np.int32), v[ non_zero ]
 
	def _to_dense( sparse_indices: np.ndarray, sparse_values: np.ndarray ):
		dense = np.zeros( n_input_states )
		dense[ sparse_indices ] = sparse_values
		return dense

	def _batch_to_dense( indicies: list[int] ) -> np.ndarray:
		
		out = np.zeros((len(indicies), n_input_states), dtype=np.float64)
		for row, i in enumerate(indicies):
			out[row, sparse_belief_vector_indices[i]] = sparse_belief_vector_values[i]
		
		return out
 
	# -- Transitions ----------------------------------------------

	t_origin: list[int]   = []
	t_target: list[int]   = []
	t_prob:   list[float] = []
	t_symbol: list[int]   = []
 
	# -- Sparse dual bucket coarse hash ----------------------------------

	# Since we use dual bins 
	# Some x, y in p, q needs to differ by more than 1/(2*coarse_scale) to not share a bin
	# Thus minimal total variation between non-bin sharing distributions is 
	# To ensure jsd(p,q) < jsd_eps -> p,q share a bin, we use Pinsker's inequality
	# 1 /(2*coarse_scale) <= 2*sqrt(2*jsd_eps)
	# 1 /( 4*sqrt(2*jsd_eps) ) <= coarse_scale
	
	coarse_scale = 1 / ( 4*np.sqrt( 2 * jsd_eps ) )
 
	def _coarse_keys( sparse_indices: np.ndarray, sparse_values: np.ndarray) -> list[bytes]:
		scaled = _to_dense( sparse_indices, sparse_values*coarse_scale )
		return [ 
			np.round(scaled).astype(np.int32).tobytes(), 
			np.round(scaled + 0.5).astype(np.int32).tobytes() 
		]

	buckets: dict[bytes, list[int]] = {}
 
	# Jenson Shannon Divergence for sparse vectors
	def _jsd_sparse(
		ai: np.ndarray, av: np.ndarray,
		bi: np.ndarray, bv: np.ndarray,
	) -> float:

		# get sorted union of the [non-zero] indices over both vectors
		all_i = np.union1d( ai, bi )
	
		# form dense vectors
		p = np.zeros( len( all_i ), dtype=np.float64 )
		q = np.zeros( len( all_i ), dtype=np.float64 )

		# np.searchsorted( all_i, ai ) finds position in all_i where each ai is 
		p[ np.searchsorted( all_i, ai ) ] = av
		q[ np.searchsorted( all_i, bi ) ] = bv

		# Jenson Shannon divergence (JDS)
		return af_jensenshannondivergence_cpu( p, q )

	# Check if an existing belief vector is within jsd_eps 
	def _lookup( sparse_indices: np.ndarray, sparse_values: np.ndarray ) -> int:
		
		keys = _coarse_keys( sparse_indices, sparse_values )
		candidates = list( { idx for key in keys for idx in buckets.get( key, [] ) } )

		if not candidates :
			return -1

		dense = _to_dense( sparse_indices, sparse_values ).reshape( 1, n_input_states )
		candidate_vectors = _batch_to_dense( candidates )

		distances = af_jensenshannondivergence_cpu( candidate_vectors, dense, axis=1 )
		min_idx  = np.argmin( distances )
		min_dist = distances[ min_idx ] 

		if min_dist < jsd_eps :
			return candidates[ min_idx ]
		
		return -1

	# Map belief vector to course hash buckets for subsequent fast lookup
	def _register( sparse_indices: np.ndarray, sparse_values: np.ndarray, idx: int):
		for key in _coarse_keys( sparse_indices, sparse_values ):
			buckets.setdefault( key, [] ).append( idx ) 

	pi_indicies, pi_values = _to_sparse( pi )
	sparse_belief_vector_indices.append( pi_indicies )
	sparse_belief_vector_values.append( pi_values )
	n_belief_vectors = 1

	_register( pi_indicies, pi_values, 0 )

	frontier = deque([0])

	FRONTIER_CHUNK = 256
 
	try:
		while frontier and n_belief_vectors < exact_state_cap:
 
			batch_indices = []
			while frontier and len(batch_indices) < FRONTIER_CHUNK:
				batch_indices.append(frontier.popleft())
 
			B = len(batch_indices)
 
			mus_device       = xp.asarray(_batch_to_dense( batch_indices ) )
			flat_device      = mus_device @ T_flat_device
			all_mu_Tx_device = flat_device.reshape( B, n_symbols, n_input_states )
			probs_device     = all_mu_Tx_device.sum(axis=2)
 
			vb_device, vx_device = xp.where(probs_device > EPS)
 
			if vb_device.size == 0:
				del mus_device, flat_device, all_mu_Tx_device, probs_device, vb_device, vx_device
				continue
 
			vp_device      = probs_device[vb_device, vx_device]
			mu_next_device = all_mu_Tx_device[vb_device, vx_device] / vp_device[:, None]
 
			mu_next_cpu = to_cpu( mu_next_device )

			vp_cpu      = to_cpu( vp_device      )
			vb_cpu      = to_cpu( vb_device      )
			vx_cpu      = to_cpu( vx_device      )
 
			# Free GPU arrays for this batch immediately after download
			del mus_device, flat_device, all_mu_Tx_device, probs_device
			del vb_device, vx_device, vp_device, mu_next_device
 
			if not np.isfinite(mu_next_cpu).all():
				raise ValueError("Non-finite belief vectors in BFS batch")
 
			for q in range( len( vb_cpu ) ):
				
				indicies, values = _to_sparse(mu_next_cpu[q])
				orig = int( batch_indices[ int( vb_cpu[ q ] ) ] )
				idx = _lookup( indicies, values )
 
				if idx == -1:
					
					idx = n_belief_vectors
					sparse_belief_vector_indices.append( indicies )
					sparse_belief_vector_values.append( values )
					_register( indicies, values, idx )

					n_belief_vectors += 1
					frontier.append(idx)
 
				t_origin.append(orig)
				t_target.append(idx)
				t_prob.append(float(vp_cpu[q]))
				t_symbol.append(int(vx_cpu[q]))

		# -- Close the graph by mapping each dangling state to the nearest closed state

		if frontier:
 
			if verbose:
				print(f"Max exact states ({exact_state_cap}) exceded ({n_belief_vectors}). Closing graph...")
 
			dense_all = np.zeros((n_belief_vectors, n_input_states), dtype=np.float32)

			for i in range(n_belief_vectors):
				dense_all[i, sparse_belief_vector_indices[i]] = sparse_belief_vector_values[i].astype(np.float32)
	
			ixp = xp if CAGRA_AVAILABLE else np

			if CAGRA_AVAILABLE:

				bv_device_f32 = xp.asarray(dense_all)
				del dense_all

				if verbose : 
					print("Building CAGRA index...")

				index = cagra.build(
					cagra.IndexParams(graph_degree=32, metric='sqeuclidean'),
					bv_device_f32,
				)

			else :

				if verbose : 
					print("Building hnswlib index...")

				T_flat_device = to_cpu( T_flat_device )
				bv_device_f32 = dense_all
				index = hnswlib.Index(space='l2', dim=bv_device_f32[0].size)
				index.init_index(max_elements=len(bv_device_f32), ef_construction=500, M=50)
				index.set_ef(50)
				index.add_items(bv_device_f32)

			if verbose:
				print("Index built.")
 
			closure_buckets: dict[bytes, list[tuple]] = {}
			max_error = 0.0
			avg_error = 0.0
			n_mapped  = 0
 
			def _closure_lookup( indicies: np.ndarray, values: np.ndarray) -> int:
				seen = set()
				for key in _coarse_keys( indicies, values):
					for entry in closure_buckets.get(key, []):
						eid = id(entry)
						if eid not in seen:
							seen.add(eid)
							s_idx, s_val, target = entry
							if _jsd_sparse(indicies, values, s_idx, s_val) < jsd_eps:
								return target
				return -1
 
			def _closure_register(indicies: np.ndarray, values: np.ndarray, target: int):
				entry = (indicies, values, target)
				for key in _coarse_keys(indicies, values):
					closure_buckets.setdefault(key, []).append(entry)
 
			CHUNK = 10_000
			frontier_list = list(frontier)

			for start in range(0, len(frontier_list), CHUNK):

				chunk_idxs = frontier_list[start : start + CHUNK]
				C  = len(chunk_idxs)

				if CAGRA_AVAILABLE :
					chunk_device  = cp.asarray(_batch_to_dense(chunk_idxs))
				else :
					chunk_device  = _batch_to_dense(chunk_idxs)

				flat_device   = chunk_device @ T_flat_device

				mu_Tx_device  = flat_device.reshape(C, n_symbols, n_input_states)

				probs_device  = mu_Tx_device.sum(axis=2)
 
				# Free intermediates we no longer need
				del chunk_device, flat_device

				vf_device, vx_device = ixp.where(probs_device > EPS)
 
				if vf_device.size == 0:
					del mu_Tx_device, probs_device, vf_device, vx_device
					continue
 
				vp_device   = probs_device[vf_device, vx_device]
				mu_v_device = mu_Tx_device[vf_device, vx_device] / vp_device[:, None]
 
				# Free now-consumed GPU arrays
				del mu_Tx_device, probs_device

				mu_v_cpu = to_cpu( mu_v_device )
				vp_cpu   = to_cpu( vp_device )
				vf_cpu   = to_cpu( vf_device )
				vx_cpu   = to_cpu( vx_device )

				del vf_device, vx_device, vp_device
 
				# float32 copy for CAGRA; free float64 immediately after cast
				unknown_device = mu_v_device.astype(ixp.float32)
				del mu_v_device

				N            = len(vf_cpu)
				orig_arr     = np.array([chunk_idxs[f] for f in vf_cpu], dtype=np.int32)
				next_idx     = np.full(N, -1, dtype=np.int32)
				unknown_mask = np.ones(N, dtype=bool)

				sparse_cache = [_to_sparse(mu_v_cpu[q]) for q in range(N)]

				for q in range(N):
					indicies, values = sparse_cache[q]
					target = _closure_lookup(indicies, values)
					if target != -1:
						next_idx[q]     = target
						unknown_mask[q] = False
 
				if unknown_mask.any():

					unknown_rows = unknown_device[ixp.asarray(unknown_mask)]
 
					if CAGRA_AVAILABLE :

						if verbose:
							print(f"CARGA search for {unknown_mask.sum()} unknown states...")
						
						_, indices_device = cagra.search(
							cagra.SearchParams(itopk_size=128),
							index, unknown_rows, k_ann,
						)

					else :
						if verbose :
							print(f"hnswlib search for {unknown_mask.sum()} unknown states...")
						
						indices_device, _ = index.knn_query( to_cpu( unknown_rows ), k=k_ann )

					output_indices = ixp.asarray( indices_device )
 
					min_dist = ixp.full(unknown_rows.shape[0], ixp.inf, dtype=ixp.float32)
					min_ks   = ixp.zeros(unknown_rows.shape[0], dtype=ixp.int32)

					for idx_k in range(output_indices.shape[1]):
						sv   = bv_device_f32[output_indices[:, idx_k], :]

						if CAGRA_AVAILABLE :
							d = af_jensenshannondivergence_device(unknown_rows, sv, axis=1)
						else :
							d = af_jensenshannondivergence_cpu(unknown_rows, sv, axis=1)

						mask = d < min_dist
						min_dist = ixp.where(mask, d, min_dist)
						min_ks   = ixp.where(mask, idx_k, min_ks)

					max_dist  = to_cpu( ixp.max(min_dist) )
					mean_dist = to_cpu( ixp.mean(min_dist) )
 
					n_mapped  += N
					avg_error += mean_dist
					max_error  = max(max_dist, max_error)

					resolved   = to_cpu( output_indices[ixp.arange(unknown_rows.shape[0]), min_ks] )
					unknown_qs = np.where(unknown_mask)[0]
 
					for i, q in enumerate(unknown_qs):
						nz_idx, nz_val = sparse_cache[q]
						next_idx[q]    = int(resolved[i])
						_closure_register(nz_idx, nz_val, int(resolved[i]))
 
					# Free ANN-related GPU arrays
					del unknown_rows, indices_device, output_indices
					del min_dist, min_ks
 
				del unknown_device
 
				for q in range(N):
					t_origin.append(int(orig_arr[q]))
					t_target.append(int(next_idx[q]))
					t_prob.append(float(vp_cpu[q]))
					t_symbol.append(int(vx_cpu[q]))
 
			# -- Free CAGRA index and float32 belief matrix --------------------

			# CAGRA memory lives outside CuPy's pool and is not released by free_all_blocks()
			del index
			del bv_device_f32
 
			if CUPY_AVAILABLE :

				cp.get_default_memory_pool().free_all_blocks()
				cp.get_default_pinned_memory_pool().free_all_blocks()

			gc.collect()
 
			if verbose:
				print(f"Closed: {n_mapped} dangling states")
				print(f"Avg JSD state closure error: {(1.0 / np.log(2)) * avg_error / n_mapped} bits")
				print(f"Max JSD state closure error: {(1.0 / np.log(2)) * max_error} bits")
 
		# ------------------ Create the MSP --------------------#
		
		msp_transitions = [
			Transition(
				origin_state_idx = t_origin[i],
				target_state_idx = t_target[i],
				prob             = t_prob[i],
				symbol_idx       = t_symbol[i],
			)
			for i in range(len(t_origin))
		]
 
		mixed_state_vectors = list(_batch_to_dense(list(range(n_belief_vectors))))

 
		msp = MSP(
			states              = [CausalState(name=f"MS_{i}") for i in range(n_belief_vectors)],
			mixed_state_vectors = mixed_state_vectors,
			transitions         = msp_transitions,
			alphabet            = alphabet.copy(),
		)
 
		if verbose:
			print("Done.\n")
 
		return msp
 
	finally:

		# ------------------ Cleanup --------------------#

		if 'T_flat_device' in dir():
			del T_flat_device

		if 'index' in locals():
			del index
		
		if 'bv_device_f32' in locals():
			del bv_device_f32

		if CUPY_AVAILABLE :

			cp.get_default_memory_pool().free_all_blocks()
			cp.get_default_pinned_memory_pool().free_all_blocks()

		gc.collect()