from __future__ import annotations

from abc import abstractmethod
from typing import override
import gc
import copy
from collections import deque
from collections import defaultdict
import json
from fractions import Fraction
from pathlib import Path
import warnings
import time
from dataclasses import dataclass, asdict, field

import networkx as nx

from automata.fa.dfa import DFA

import sympy
import numpy as np

from .am_machine import Machine
from .am_symbol import Symbol

from   .am_msp import MSP, compute_msp, compute_msp_exact
from   .am_solve import solve_for_pi, solve_for_pi_fractional
from . import am_vis
from . import am_fast

from .am_causal_state import CausalState
from .am_transition   import Transition

from .am_fast.distance import jensenshannondivergence_gpu as af_jensenshannondivergence
from .am_fast.distance import jensenshannondivergence_cpu as af_jensenshannondivergence_cpu

class HMM(Machine) :

	@staticmethod
	def isomorphic_shift(
		input_symbol_indices: np.ndarray,
		input_state_indices:  np.ndarray,
		m : HMM,
		shift : int = 1
	) -> dict[str, np.ndarray]:

		if not any( state.isomorphs for state in m.states ):
			raise ValueError("HMM has no states with isomorphs")

		inputs = np.asarray(input_symbol_indices)
		states = np.asarray(input_state_indices)

		n_states = len(m.states)

		tr_sym_table = np.full((n_states, n_states), -1, dtype=np.int32)

		for tr in m.transitions:
			tr_sym_table[ tr.origin_state_idx, tr.target_state_idx ] = tr.symbol_idx

		# Build isomorph remapping: identity by default, overridden where isomorphs exist
		iso_table = np.arange(n_states, dtype=np.int32)

		for i, state in enumerate(m.states):
			if state.isomorphs:
				isomorph       = sorted(state.isomorphs)[0]
				iso_table[ i ] = m.state_idx_map[isomorph]

		for i, state in enumerate(m.states):
			if state.isomorphs:

				# extend the isomorph list to include the identity
				isormorphs_with_identity = sorted( [ i ] + [ m.state_idx_map[iso] for iso in state.isomorphs ] )

				# find the identity index
				pos = isormorphs_with_identity.index( i )

				# cyclical shift 
				iso_table[ i ] = isormorphs_with_identity[ ( pos + shift ) % len( isormorphs_with_identity ) ]

		origins = states[:-1]
		targets = states[1:]

		out_origins = iso_table[origins]
		out_targets = iso_table[targets]

		inv_sym = tr_sym_table[ out_origins, out_targets ]
		inv_sts = np.empty(states.size, dtype=states.dtype)

		inv_sts[:-1] = out_origins
		inv_sts[-1]  = out_targets[-1]

		return {
			"symbol_index": inv_sym.astype(inputs.dtype),
			"state_index":  inv_sts,
		}

	def __init__( 
		self,
		states : list[CausalState] | None = None,
		transitions : list[Transition] | None = None,
		start_state : int = 0,
		alphabet : list[str] | None = None,
		isoclasses : dict[str,set[str]] | None = None,
		name : str = "",
		description : str = "" ) : 

		self.alphabet    = alphabet or []
		self.states      = states or []
		self.transitions = transitions or []

		self.set_alphabet( self.alphabet )
		self.set_states( self.states )
		self.set_transitions( self.transitions )

		self.start_state = start_state
		self.isoclasses  = isoclasses or {}

		self.name = name
		self.description = description

		# to be depreciated
		self.isoclass = None

		# --- derived --------

		self.complexity : dict[str, any] = {}

		self.symbol_idx_map : dict[str,int] = {}
		self.state_idx_map  : dict[str,int] = {}

		self.pi_fractional = None
		self.pi = None
		self.T  = None
		self.msp : MSP = None
		self.reverse_am : HMM = None
		self.is_q_weighted = False
		self.is_minimal = False

		# --- const ----------

		self.EPS = 1e-12

	#-------------------------------------------------------#
	#                Overrides / Getters                    #
	#-------------------------------------------------------#

	@override
	def get_states(self) -> list[CausalStates] : 
		return self.states

	@override
	def get_transitions(self) -> list[Transition] : 
		return self.transitions

	@override
	def get_alphabet(self) -> list[Symbol]:
		return [ Symbol(a) for a in self.alphabet ]

	#-------------------------------------------------------#
	#                     Serialization                     #
	#-------------------------------------------------------#


	def to_dict(self) :
		return {
			"name"            : self.name,
			"description"     : self.description,
			"states"          : [ asdict(state)      for state      in self.states      ],
			"transitions"     : [ asdict(transition) for transition in self.transitions ],
			"alphabet"        : self.alphabet
		}

	def from_dict( self, config ) :

		self.name            = config.name 
		self.description     = config.description 
		
		self.set_states( states=[ 
			CausalState( 
				name=state[ "name" ],
				classes=set( state[ "classes" ] )
			)
			for state in config.states
		] )

		self.set_transitions( transitions=[ 
			Transition( 
				origin_state_idx=tr[ "origin_state_idx" ],
				target_state_idx=tr[ "target_state_idx" ],
				prob=tr[ "prob" ],
				symbol_idx=tr[ "symbol_idx" ]
			)
			for tr in config.transitions
		] )

		self.set_alphabet( alphabet=config.alphabet ) 

	def save_config(
		self, 
		path : Path, 
		with_complexity : bool = False, 
		with_block_convergence :  bool = False,
		with_structural_properties : bool = False,
		with_causal_properties : bool = False ) :

		config = self.to_dict()

		if with_complexity :
			
			complexity = self.get_complexities( 
				with_block_convergence=with_block_convergence,
				with_reversal=False # not implemented
			)

			config[ "complexity" ] = complexity

		config[ "structural_properties" ] = {
			"unifilar"              : self.is_unifilar(),
			"row_stochastic"        : self.is_row_stochastic(),
			"strongly_connected"    : self.is_strongly_connected(),
			"aperiodic"             : self.is_aperiodic(),
			"minimal"               : self.is_minimal_as_dfa( topological_only=False ),
			"topologically_minimal" : self.is_minimal_as_dfa( topological_only=True ),
			"is_epsilon_machine"    : self.is_epsilon_machine()
		}

		# Not implemented
		# config[ "causal_properties" ] = {
		# 	"k_X"                        : self.k_X(),
		# 	"R"                          : self.R(),
		# 	"causally_reversible"        : self.causally_reversible(),
		# 	"microscopically_reversible" : self.microscopically_reversible(),
		# }

		with open( path / "am_config.json", "w", encoding="utf-8" ) as f :
			json.dump( config, f, ensure_ascii=False, indent=2, default=list )

	def from_file( self, path : Path ) :
		with open( Path / "am_config.json", "r" ) as f:
			config = json.load(f)
		self.from_dict()

	#-------------------------------------------------------#
	#             Setters and State Management              #
	#-------------------------------------------------------#
	
	def _invalidate(self) :

		"""
		Invalidate

		Reset all derived properties so they will be recomputed when requested later
		"""

		self.complexity = {}
		self.T = None
		self.T_x = None
		self.pi = None
		self.pi_fractional = None
		self.msp = None 
		self.reverse_am = None
		self.is_q_weighted = False
		self.is_minimal = False

	def set_states( self, states : list[CausalState] ) :
		self._invalidate()
		self.states = states.copy()
		self.state_idx_map = {}
		for idx, state in enumerate( self.states ) :
			self.state_idx_map[ state.name ] = idx

	def set_alphabet( self, alphabet : list[str] ) :

		self._invalidate()

		old_alphabet = self.alphabet.copy()

		aSet = set()
		aSet.update( alphabet )

		self.alphabet = sorted(list(aSet))

		self.symbol_idx_map = {}
		for idx, symbol in enumerate( self.alphabet ) :
			self.symbol_idx_map[ symbol ] = idx

		for i, tr in enumerate( self.transitions ) :
			symbol = old_alphabet[ tr.symbol_idx ]
			self.transitions[ i ] = Transition(
				origin_state_idx=tr.origin_state_idx,
				target_state_idx=tr.target_state_idx,
				prob=tr.prob,
				symbol_idx=self.symbol_idx_map[ symbol ]
			)

	def set_transitions( self, transitions : list[Transition] ) :
		self._invalidate()
		self.transitions = transitions.copy()

	def extend_states( self, states : list[CausalState] ) :
		self.set_states( self.states + states )

	def extend_alphabet( self, alphabet : list[str] ) :
		self.set_alphabet( self.alphabet + alphabet  )

	def extend_transitions( self, transitions : list[Transition] ) :
		self.set_transitions( self.transitions + transitions  )

	def get_complexity_measure_if_exists(self, measure ) :
		m = self.complexity.get( measure, None )
		return m

	def set_complexity_measure(self, measure, value ) :
		self.complexity[ measure ] = value

	#-------------------------------------------------------#
	#                   Get and Compute                     #
	#-------------------------------------------------------#

	def get_complexities( 
		self, 
		with_block_convergence=False, 
		with_reversal=False ) :

		directly_calculable = [
			self.C_mu,
			self.h_mu,
			self.H_1,
			self.rho_mu
		]

		requires_block_convergence = [
			self.E, 
			self.T_inf,
			self.S,
			self.Chi
		]

		requires_reversal = [
			self.q_mu,
			self.r_mu,
			self.b_mu
		]
			
		complexities = { m.__name__ : m() for m in directly_calculable }

		if with_block_convergence :

			complexities |= { m.__name__ : m() for m in requires_block_convergence }

			for key in [ 'H_L', 'T_L', 'h_mu_L', 'H_sync' ] :
				if key in self.complexity :
					complexities[ key ] = self.complexity[ key ]

		if with_reversal :
			complexities |= { m.__name__ : m() for m in requires_reversal }

		return complexities

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

	def get_metadata(self) :
		return {
			"name" : self.name,
			'complexity' : self.complexity,
			"description" : self.description
		}

	def get_transition_matrix(self) :

		if self.T  is not None :
			return self.T

		n_states = len( self.states )
		T = np.zeros((n_states, n_states))

		for tr in self.transitions :    
			T[ tr.origin_state_idx, tr.target_state_idx  ] = tr.prob

		self.T = T

		return self.T

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

	def get_T_X(self) :

		if self.T_x  is not None :
			return self.T_x

		n_states  = len( self.states )
		n_symbols = len( self.alphabet )

		T_x = [ np.zeros((n_states, n_states)) for _ in range( n_symbols ) ]

		for tr in self.transitions :
			T_x[ tr.symbol_idx ][tr.origin_state_idx, tr.target_state_idx] = tr.prob

		self.T_x = T_x
		return self.T_x

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

	def get_msp_qw(
		self,
		exact_state_cap: int = 1000,
		verbose: bool = True,
	):
		if self.msp is not None:
			return self.msp

		try : 

			print( "\nTrying to Compute Mixed State Presentation using Exact Fractions\n" )

			self.msp = compute_msp_exact(
				T_x=self.get_Tx_fractional(),
				pi=self.get_fractional_stationary_distribution(),
				n_states=len(self.states),
				alphabet=self.alphabet,
				exact_state_cap=1000,
				bool = True
			)

			return self.msp 

		except RuntimeError as e :
			warnings.warn( f"Exact msp failed: {e} Falling back to msp approximation." )

		return self.get_msp()

	def get_msp(
		self,
		exact_state_cap: int = 175_000,
		jsd_eps:         float = 1e-7,
		k_ann:           int   = 50,
		verbose                = True,
	) :

		if self.msp is not None:
			return self.msp
	 
		T_x = self.get_T_X()
		pi  = self.get_stationary_distribution()

		T_stacked      = np.stack(T_x)
		n_symbols      = len(self.alphabet)
		n_input_states = T_stacked.shape[1]

		print( "\nComputing Mixed State Presentation\n" )

		self.msp = compute_msp( 
			T_x=T_x,
			pi=pi,
			n_states=len(self.states),
			alphabet=self.alphabet,
			exact_state_cap=exact_state_cap,
			verbose=verbose
		)

		return self.msp

	def get_reverse_am(self) :

		if self.reverse_am is not None:
			return self.reverse_am

		pi = self.get_stationary_distribution()
		self.reverse_am = copy.deepcopy(self)
		
		new_transitions = []
		for tr in self.transitions:
			i = tr.target_state_idx
			j = tr.origin_state_idx
			
			p_reversed = (pi[j] * tr.prob) / pi[i]
			
			new_transitions.append(
				Transition(
					origin_state_idx=i,
					target_state_idx=j,
					prob=p_reversed,
					symbol_idx=tr.symbol_idx
				)
			)

		self.reverse_am.set_transitions(new_transitions)

		if self.reverse_am.is_epsilon_machine():
			return self.reverse_am

		rmsp = self.reverse_am.get_msp_qw( exact_state_cap=len(self.states)*4 )

		self.reverse_am.set_states( rmsp.states )
		self.reverse_am.set_transitions( rmsp.transitions )
		self.reverse_am.msp = rmsp
		self.reverse_am.start_state = 0

		self.reverse_am.collapse_to_largest_strongly_connected_subgraph()
		self.reverse_am.minimize()

		return self.reverse_am

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

	def get_Tx_fractional(self) -> list[ list[ list[ Fraction ] ] ] :

		self.to_q_weighted()

		n_states  = len( self.states )
		n_symbols = len( self.alphabet )

		T_x = []

		for x in range( n_symbols ) :
			T_x.append( [] )
			for i in range( n_states ) :
				T_x[ x ].append( [ 0 for _ in range( n_states ) ] )

		for tr in self.transitions :
			T_x[ tr.symbol_idx ][ tr.origin_state_idx ][ tr.target_state_idx ] = tr.pq

		return T_x

	def get_T_sympy( self ) :

		self.to_q_weighted()

		n = len( self.states )
		T = sympy.zeros( n, n )

		for tr in self.transitions :
			T[ tr.origin_state_idx, tr.target_state_idx ] = tr.pq

		return T

	def get_fractional_stationary_distribution(self) :

		T = self.get_T_sympy()

		if self.pi_fractional is not None :
			return self.pi_fractional

		G = self.as_digraph()

		if not nx.is_strongly_connected(G):
			raise ValueError( "Single stationary distribution requires strongly connected HMM." )

		self.pi_fractional = solve_for_pi_fractional( T )

		return self.pi_fractional

	def get_stationary_distribution(self):

		if self.pi is not None :
			return self.pi

		G = self.as_digraph()
		
		if not nx.is_strongly_connected(G):
			raise ValueError( "Single stationary distribution requires strongly connected HMM." )

		T = self.get_transition_matrix()
		return solve_for_pi( T )		

	#-------------------------------------------------------#
	#                 Complexity Measures                   #
	#-------------------------------------------------------#
	
	def C_mu( self ) :

		"""
		c_mu (statistical complexity, or forecasting complexity)

		Interpretations
			- The amount of historical information a process stores.
			- The amount of structure in a process.
		"""

		m = self.get_complexity_measure_if_exists( "C_mu" )

		if m is not None :
			return m

		pi = self.get_stationary_distribution()

		h = 0
		for i, pr in enumerate( pi ) :
			
			if pr < self.EPS :
				continue

			h += -pr * np.log2( pr )

		self.set_complexity_measure( "C_mu", h )

		return h

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

	def h_mu( self ) :

		"""
		h_mu (entropy rate)

		Interpretation
			- Intrinsic Randomness in the process.
		"""

		m = self.get_complexity_measure_if_exists( "h_mu" )

		if m is not None :
			return m

		T  = self.get_transition_matrix()
		pi = self.get_stationary_distribution()

		n_states = pi.size

		h = 0
		for i, pr in enumerate( pi ) :

			if pr < self.EPS :
				continue

			row_entropy = 0
			for j in range( len( pi ) ) :

				if T[ i, j ]  < self.EPS :
					continue

				row_entropy -= T[ i, j ] * np.log2( T[ i, j ] )

			h += pr * row_entropy

		self.set_complexity_measure( "h_mu", h )

		return h

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

	def H_1(self):

		"""
		H_1 (single-symbol entropy, or marginal entropy)
		Interpretations:
			- How uncertain you are about a single measurement with no context.
			- The total information content per symbol before seeing any neighbors.
			- The maximum possible entropy rate, achieved by a memoryless process.
		"""

		m = self.get_complexity_measure_if_exists("H_1")
		if m is not None:
			return m

		pi  = self.get_stationary_distribution()
		T_X = self.get_T_X()  # dict: symbol -> matrix

		h = 0.0
		for T_x in T_X:
			# Pr(x) = sum_i pi[i] * sum_j T^(x)[i,j]
			p_sym = 0.0
			for i, pr in enumerate(pi):
				if pr < self.EPS:
					continue
				p_sym += pr * T_x[i, :].sum()

			if p_sym < self.EPS:
				continue
			h -= p_sym * np.log2(p_sym)

		self.set_complexity_measure("H_1", h)
		return h

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

	def rho_mu(self) :
		
		"""
		rho_mu (anticipated information, or single-symbol redundancy)
		Interpretations:
			- How much the past reduces your uncertainty about the next symbol.
			- The portion of each symbol's information that was already predictable.
			- The difference between naive per-symbol uncertainty and true randomness.
			- Equal to H_1 - h_mu: the gap closed by knowing context.
		"""

		m = self.get_complexity_measure_if_exists("rho_mu")
		
		if m is not None:
			return m

		rho = self.H_1() - self.h_mu()
		
		self.set_complexity_measure("rho_mu", rho)
		
		return rho

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

	def block_convergence( self ) :

		"""
		block_convergence
			- Estimates a range of complexity measures using block entropy convergence.
			- E, S, T, T_L, H_L, h_mu_L, H_sync
		"""

		trs = [ [] for _ in range( len( self.states ) ) ]
		for tr in self.transitions :
			trs[ tr.origin_state_idx ].append( ( 
				tr.symbol_idx, 
				float( tr.prob ),
				tr.target_state_idx ) )

		pi = self.get_stationary_distribution()

		state_dist = [ float( pi[ i ] ) for i in range( len( self.states ) ) ]
		branches = [(1.0, list(state_dist))]

		print( "\nComputing Block Entropy\n" )

		C = am_fast.block_entropy_convergence(
			h_mu            = self.h_mu(),
			n_states        = len( self.states ),
			n_symbols       = len( self.alphabet ),
			convergence_tol = 1e-6,
			precision       = 10,
			eps             = 1e-25,
			branches        = branches,
			trans           = trs,
			max_branches    = 30_000_000
		)

		print( "Done\n" )

		self.set_complexity_measure( f"E",          C.E )
		self.set_complexity_measure( f"S",          C.S )
		self.set_complexity_measure( f"T_inf",      C.T )
		self.set_complexity_measure( f"T_L",        C.T_L.tolist() )
		self.set_complexity_measure( f"H_L",        C.H_L.tolist() )
		self.set_complexity_measure( f"h_mu_L",  C.h_mu_L.tolist() )
		self.set_complexity_measure( f"H_sync",  C.H_sync.tolist() )

		return C

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

	def E( self ) :

		"""
		E (excess entropy)

		Interpretations:
			- Mutual information from past to future I(Past;Future).
			- How much total information from the past reduces uncertainty in the future
			- Amount of (apparent) stored information

		Exact Complexity: The Spectral Decomposition of Intrinsic Computation
		https://arxiv.org/abs/1309.3792
		"""

		m = self.get_complexity_measure_if_exists( "E" )

		if m is not None :
			return m

		try : 
			msp = self.get_msp()
			E, S, T = msp.get_E_S_T()
			self.set_complexity_measure( "E", E )
			self.set_complexity_measure( "S", S )
			self.set_complexity_measure( "T_inf", T )
			
		except Exception as e :

			print( f"MSP failed {e}" )
			exit()

			C = self.block_convergence()	
			E = C.E
			self.set_complexity_measure( "E", E )

		return E

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

	def S( self ) :

		"""
		S (synchronization information)

		Interpretation
			-synchronization information

		Exact Complexity: The Spectral Decomposition of Intrinsic Computation
		https://arxiv.org/abs/1309.3792
		"""

		m = self.get_complexity_measure_if_exists( "S" )

		if m is not None :
			return m

		try : 
			msp = self.get_msp()
			E, S, T = msp.get_E_S_T()
			self.set_complexity_measure( "E", E )
			self.set_complexity_measure( "S", S )
			self.set_complexity_measure( "T_inf", T )

		except Exception as e :
			print( f"{e} \nFalling back to iterative estimation.")
			exit()

			C = self.block_convergence()	
			S = C.S
			self.set_complexity_measure( "S", S )

		return S

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

	def T_inf( self ) :

		"""
		T_inf (transient information)

		Interpretation
			-Transient information

		Exact Complexity: The Spectral Decomposition of Intrinsic Computation
		https://arxiv.org/abs/1309.3792
		"""

		m = self.get_complexity_measure_if_exists( "T_inf" )

		if m is not None :
			return m

		try : 
			msp = self.get_msp()
			E, S, T = msp.get_E_S_T()
			self.set_complexity_measure( "E", E )
			self.set_complexity_measure( "S", S )
			self.set_complexity_measure( "T_inf", T )

		except Exception as e :
			print( f"{e} \nFalling back to iterative estimation.")
			exit()
			C = self.block_convergence()	
			T_inf = C.T
			self.set_complexity_measure( "T_inf", T_inf )

		return T_inf

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

	def Chi( self ) :

		"""
		Chi (crypticity)

		Interpretation
			- Difference between internal stored information and apparent information to an observer
			- How muching information is hiding in the system
		"""

		m = self.get_complexity_measure_if_exists( "Chi" )

		if m is not None :
			return m

		Chi = self.C_mu() - self.E()

		self.set_complexity_measure( "Chi", Chi )

		return Chi

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

	def k_X( self ) :

		"""
		k_x (cryptic order)

		Interpretation
			
		"""

		m = self.get_complexity_measure_if_exists( "k_X" )

		if m is not None :
			return m

		raise NotImplementedError("Not implemented")

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

	def R( self ) :

		"""
		R (Markov order)
		Interpretations:
			- The smallest L such that h_mu(L) = h_mu.
			- The context length beyond which additional history gives no predictive benefit.
			- The memory length of the process — how far back correlations extend.
			- The ideal context window length for a predictor.
			- R = 0 means the process is IID (memoryless).
			- R = infinity means the process has unbounded memory (non-finitary).
			— only report if convergence was clean, otherwise None
		"""

		m = self.get_complexity_measure_if_exists( "R" )

		if m is not None :
			return m

		raise NotImplementedError("Not implemented")

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

	def b_mu( self ) :

		"""
		b_mu - requires reverse HMM


		Interpretation
			
		"""

		m = self.get_complexity_measure_if_exists( "b_mu" )

		if m is not None :
			return m

		raise NotImplementedError("Not implemented")

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

	def r_mu( self ) :

		"""
		r_mu - requires reverse HMM


		Interpretation
			
		"""

		m = self.get_complexity_measure_if_exists( "r_mu" )

		if m is not None :
			return m

		raise NotImplementedError("Not implemented")

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

	def q_mu( self ) :

		"""
		q_mu 

		Interpretation
		"""

		m = self.get_complexity_measure_if_exists( "q_mu" )

		if m is not None :
			return m

		raise NotImplementedError("Not implemented")

	#-------------------------------------------------------#
	#                      Properties                       #
	#-------------------------------------------------------#

	def causally_reversible( self ) :

		"""
		causally_reversible
		"""

		m = self.get_complexity_measure_if_exists( "causally_reversible" )

		if m is not None :
			return m

		raise NotImplementedError("Not implemented")

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

	def microscopically_reversible( self ) :

		"""
		microscopically_reversible
		"""

		m = self.get_complexity_measure_if_exists( "microscopically_reversible" )

		if m is not None :
			return m

		raise NotImplementedError("Not implemented")

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

	def is_row_stochastic(self) :

		"""
		is_row_stochastic
			-All states have outgoing transition probabilities which sum 1
		"""

		sums = np.zeros( len( self.states ) )
		for tr in self.transitions :
			sums[ tr.origin_state_idx ] += tr.prob
		return np.allclose( sums, 1.0 )

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

	def is_unifilar(self) :

		"""
		is_unifilar
			-All states have no more than one outgoing transition per-symbol
		"""

		symbol_trs = np.full( ( len( self.states ), len( self.alphabet) ), -1 )

		for tr in self.transitions : 
		
			if symbol_trs[ tr.origin_state_idx, tr.symbol_idx ] == -1 :
				symbol_trs[ tr.origin_state_idx, tr.symbol_idx ] = tr.target_state_idx
		
			elif symbol_trs[ tr.origin_state_idx, tr.symbol_idx ] != tr.target_state_idx :
				return False
		
		return True

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

	def is_strongly_connected(self) :

		"""
		is_strongly_connected
			-Every state is reachable from every other state
		"""

		return nx.is_strongly_connected( self.as_digraph() )

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

	def is_aperiodic(self) :

		"""
		is_aperiodic
			-"A strongly connected directed graph is aperiodic 
				if there is no integer k > 1 that divides the length of every cycle in the graph."
		"""

		return nx.is_aperiodic( self.as_digraph() )

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

	def is_minimal_as_dfa( self, topological_only : bool, verbose=True ) :

		with_probs = not topological_only

		# Construct the DFA
		dfa = self.as_dfa( with_probs=with_probs )

		# Minimize the DFA
		#dfa = dfa.minify(retain_names=True)
		dfa = am_fast.minify_cpp( dfa, retain_names=True )

		# check we have minimal number of states
		if len( dfa.states ) != len( self.states ) :
			if verbose : 
				print( f"Not minimal reduces from {len( self.states )} to {len( dfa.states )} states" )
			return False

		return True

	def is_topological_epsilon_machine( self, verbose=True ) :

		"""
		is_topological_epsilon_machine
			-see algorithm 2 in Enumerating Finitary Processes, Johnson et. al, 2024
		"""

		if not ( self.is_unifilar() and self.is_strongly_connected() ) :
			if verbose : 
				print( f"Either non unifilar or not strongly connected" )
			return False
		else :
			return self.is_minimal_as_dfa( topological_only=True, verbose=verbose )

	def is_epsilon_machine( self, verbose=True ) :

		if not ( self.is_unifilar() and self.is_strongly_connected() ) :
			if verbose : 
				print( f"Either non unifilar or not strongly connected" )
			return False
		else :
			return self.is_minimal_as_dfa( topological_only=False, verbose=verbose )

	#-------------------------------------------------------#
	#                      Properties                       #
	#-------------------------------------------------------#

	def minimize(self, retain_names: bool = True):

		"""
		minimize
			-Simplify to minimal form, using Myhill Nerode
		"""

		if self.is_minimal :
			return

		start = time.perf_counter()

		if not self.is_unifilar():
			raise ValueError(
				"DFA minimization is not valid for non-unifilar HMMs"
			)

		was_strongly_connected = self.is_strongly_connected()

		was_row_stochastic = self.is_row_stochastic()
		n_states_before = len(self.states)

		dfa = self.as_dfa(with_probs=True)

		#min_dfa = self.as_dfa(with_probs=True).minify(retain_names=True)
		min_dfa = am_fast.minify_cpp( dfa, retain_names=True )

		# Build lookup from original state index -> CausalState object
		orig_state   = {i: s for i, s in enumerate(self.states)}
		eq_list      = list(min_dfa.states)

		start_eq = min_dfa.initial_state
		
		# Separate the start state, then sort the rest by the 
		# smallest original state index inside each equivalence class.
		other_eqs = [eq for eq in eq_list if eq != start_eq]
		other_eqs.sort(key=lambda eq: min(eq))

		# Recombine so start eq comes first, followed by the sorted remaining classes
		eq_list = [start_eq] + other_eqs
		# ----------------------------------------------------------

		# Recompute eq_to_idx with the new ordering
		eq_to_idx = {eq: i for i, eq in enumerate(eq_list)}

		# new_start is now guaranteed to be 0
		new_start = 0

		# Map each original state index -> its equivalence class
		# Guard: minify() silently drops unreachable states
		orig_to_eq = {s: eq for eq in min_dfa.states for s in eq}

		# Build lookup from original state index -> its transitions
		orig_trs = defaultdict(list)
		for t in self.transitions:
			orig_trs[t.origin_state_idx].append(t)

		new_trs = []
		for eq in min_dfa.states:
			rep        = next(iter(eq))
			origin_idx = eq_to_idx[eq]
			for t in orig_trs[rep]:
				target_eq  = orig_to_eq[t.target_state_idx]
				target_idx = eq_to_idx[target_eq]
				new_trs.append(Transition(
					origin_state_idx = origin_idx,
					target_state_idx = target_idx,
					prob             = t.prob,
					symbol_idx       = t.symbol_idx,
				))

		members_list = [[orig_state[i] for i in sorted(eq)] for eq in eq_list]  # sorted for determinism

		# Compute new names
		if retain_names:
			new_names = [
				"{" + ",".join(str(m.name) for m in members) + "}" if len(members) > 1
				else members[0].name
				for members in members_list
			]
		else:
			new_names = [str(j) for j in range(len(eq_list))]

		old_name_to_new_name = {
			m.name: new_names[j]
			for j, members in enumerate(members_list)
			for m in members
		}

		# Build the new states, preserving classes and isomorphs regardless of naming
		new_states = []
		for j, (eq, members, name) in enumerate(zip(eq_list, members_list, new_names)):
			classes   = set().union(*(m.classes for m in members))
			isomorphs = {
				old_name_to_new_name.get(iso, iso)
				for m in members
				for iso in m.isomorphs
				if old_name_to_new_name.get(iso, iso) != name
			}
			new_states.append(CausalState(
				name      = name,
				classes   = classes,
				isomorphs = isomorphs,
			))

		self.set_states(new_states)
		self.set_transitions(new_trs)
		self.start_state = new_start

		if n_states_before == len(new_states) :
			print( f"{n_states_before} state HMM was already minimal." )
		else :
			print( f"Minimized from {n_states_before} to {len(new_states)}" )

		if not ( was_strongly_connected ==  self.is_strongly_connected() ) :
			raise RuntimeError(
				f"Minimization broke strongly connected"
			)

		if not ( was_row_stochastic ==  self.is_row_stochastic() ) :
			raise RuntimeError(
				f"Minimization broke row stochasticity"
			)

		self.is_minimal = True

	#-------------------------------------------------------#
	#                      Modifiers                        #
	#-------------------------------------------------------#


	def collapse_to_largest_strongly_connected_subgraph( self, rename_states=True ) :

		# get equivalent networkx graph
		G = self.as_digraph()

		# if already strongly connected, nothing to do
		if not nx.is_strongly_connected( G ) :

			start = time.perf_counter()
			subgraph_nodes = list( nx.strongly_connected_components( G ) )

			# decompose into strongly connected components and sort by length
			# subgraph_nodes = list(nx.strongly_connected_components( G ))
			subgraph_nodes.sort(key=len)
			component_state_set = subgraph_nodes[-1]

			# Take the largest strongly connected component (as list of state names)
			component_states = sorted( list( component_state_set ) )

			# make temporary copies of the old transitions and states
			old_transitions = [
				Transition(
					origin_state_idx=tr.origin_state_idx,
					target_state_idx=tr.target_state_idx,
					prob=tr.prob,
					symbol_idx=tr.symbol_idx
				)

				for tr in self.transitions
			]

			old_states = [
				CausalState(
					name=s.name,
					classes=s.classes,
					isomorphs=s.isomorphs
				)
				for s in self.states
			]

			self.set_states(
				states=[ 
					state
					for i, state in enumerate( old_states ) if i in component_state_set
				]
			)

			# we will build new transition list based on those belonging to the component
			self.set_transitions( transitions= [] )

			# for tracking which new transitions leave each state
			transitions_from_state = { state : set() for state in component_states }
			new_transitions = []

			for tr in old_transitions :

				origin_state_name = old_states[ tr.origin_state_idx ].name
				target_state_name = old_states[ tr.target_state_idx ].name

				# skip transitions that connect separate strongly connected components
				if not ( tr.origin_state_idx in component_state_set and tr.target_state_idx in component_state_set ) :
					continue

				# track transitions (by index in new transitions list) that leave this state
				transitions_from_state[ tr.origin_state_idx ].add( len( new_transitions ) )

				my_origin_state_idx = self.state_idx_map[ origin_state_name ]
				my_target_state_idx = self.state_idx_map[ target_state_name ]

				new_transitions.append( 
					Transition(
						origin_state_idx=my_origin_state_idx,
						target_state_idx=my_target_state_idx,
						prob=tr.prob,
						symbol_idx=tr.symbol_idx
					)  )

			self.set_transitions( transitions=new_transitions )

			# if we removed an outgoing transition from a state, we need to distribute its probability 
			# among the remaining outgoing transitions from the state
			for state in component_states :
				
				# get the set of transitions leaving this state
				state_trs = transitions_from_state[ state ]

				# sum the probabilities of the outgoing transitions from the state
				p_sum = np.sum( [ self.transitions[ i ].prob for i in state_trs ] )

				# how much probability is missing
				diff = 1.0 - p_sum

				# if significant difference
				if abs( diff ) > self.EPS : 

					# calculate how much of the difference each transition gets
					adjustment = diff / len( state_trs )
					
					# update the transitions
					for i in state_trs :

						# transition with origional probability
						tr = self.transitions[ i ]

						# adjusted probability
						self.transitions[ i ] = Transition(
							origin_state_idx=tr.origin_state_idx, 
							target_state_idx=tr.target_state_idx, 
							prob=tr.prob + adjustment, 
							symbol_idx=tr.symbol_idx )

			if rename_states :
				self.set_states( [
					CausalState( name=f"{i}" )
					for i, s in enumerate( self.states )
				] )


	def to_q_weighted( self, denominator_limit=1000 ) :

		"""
		to_q_weighted
			-Snaps edge probabilities to exact Fractions, and assigns Fraction values to Transition.pq
		"""

		if self.is_q_weighted :
			return

		if not self.is_row_stochastic() :
			raise ValueError( "Cannot convert to q-weighted because not row stochastic" )

		t_from = [[] for _ in range(len(self.states))]

		for i, tr in enumerate( self.transitions ) :
			t_from[ tr.origin_state_idx ].append( i )

		new_transitions = []
		for t_list in t_from :
			
			if not t_list : 
				continue

			p_q_sum = Fraction(0,1)
			p_qs = []

			for t_idx in t_list :
			
				p_q = Fraction( self.transitions[ t_idx ].prob ).limit_denominator( denominator_limit )
				p_q_sum += p_q
				p_qs.append( p_q )

			if p_q_sum != Fraction(1,1) :
				
				max_pq_i = np.argmax( p_qs )
				max_oq = p_qs[ max_pq_i ]

				diff = p_q_sum - Fraction(1,1)

				# recurse with higher resolution
				if diff > max_oq :
					return self.to_q_weighted( denominator_limit*10 )
				else :
					p_qs[ max_pq_i ] -= diff

			for i, t_idx in enumerate( t_list ) :	
				new_transitions.append( 
					Transition( 
						origin_state_idx=self.transitions[  t_idx ].origin_state_idx,
						target_state_idx=self.transitions[  t_idx ].target_state_idx,
						prob=float(p_qs[ i ]),
						symbol_idx=self.transitions[  t_idx ].symbol_idx,
						pq=p_qs[ i ]
					)
				)

		self.set_transitions( new_transitions )
		self.is_q_weighted = True

	#-------------------------------------------------------#
	#                    Data Generation                    #
	#-------------------------------------------------------#

	def generate_data(
		self,
		file_prefix: str,
		n_gen: int,
		include_states: bool,
		isomorphic_shifts : set[int]=None,
		random_seed : int=42 ) -> dict[any] : 

		trs = [ [] for _ in range( len( self.states ) ) ]
		for tr in self.transitions :
			trs[ tr.origin_state_idx ].append( ( 
				tr.symbol_idx, 
				float( tr.prob ),
				tr.target_state_idx ) )

		data = am_fast.generate_data(
			n_gen=n_gen,
			start_state=self.start_state,
			transitions=trs,
			alphabet=sorted(list(self.alphabet)),
			include_states=include_states,
			random_seed=random_seed
		)

		if isomorphic_shifts is not None :

			if not include_states :
				raise ValueError( "Isomorphic inversion requires include_states=True" )

			data[ "isomorphic_shifts" ] = {}

			for shift in isomorphic_shifts :

				shifted = HMM.isomorphic_shift(
					symbol_indices=data[ "symbol_index" ], 
					state_indices=data[ "state_index" ], 
					m=self,
					shift=shift
				)

				data[ "isomorphic_shifts" ][ shift ] = {
					"symbol_index" : shifted[ "symbol_index" ],
					"state_index"  : shifted[ "state_index" ]
				}

		am_fast.save_data(
			data=data,
			file_prefix=file_prefix,
			alphabet=sorted(list(self.alphabet)),
			n_states=len( self.states ),
			start_state=self.start_state,
			random_seed=random_seed,
			metadata=self.get_metadata() )

		return data

	#-------------------------------------------------------#
	#                 Basic Visualization                   #
	#-------------------------------------------------------#

	def draw_graph(
		self,
		engine     : str  = 'dot',
		output_dir : Path = Path('.'),
		show       : bool = True
	) -> None :

		G = self.as_digraph()

		subgraphs = None if nx.is_strongly_connected( G ) else list( nx.strongly_connected_components( G ) )
		
		am_vis.draw_graph( 
			self, 
			output_dir=output_dir, 
			title="am_graph", 
			view=show, 
			subgraphs=subgraphs, 
			engine=engine )

	#-------------------------------------------------------#
	#                   Alternate Forms                     #
	#-------------------------------------------------------#


	def as_digraph( self ) :

		"""
		Build nx DiGraph based on the HMM, without edge probabilities or symbols
		"""

		G = nx.DiGraph()
		G.add_nodes_from( [ i for i, s in enumerate( self.states ) ] )

		for tr in self.transitions :
			G.add_edge( tr.origin_state_idx, tr.target_state_idx )

		return G

	def as_dfa( self, with_probs : bool ) :

		precision=8

		def edge_label( symb, prob ) :
			return f"({symb},{round(prob, precision)})"

		# Build states, symbols, and transitions 
		dfa_states  = { i for i, _ in enumerate( self.states ) }
			
		if not with_probs :
			dfa_symbols = set( { t.symbol_idx for t in self.transitions } )
		else : 
			dfa_symbols = set( { edge_label( t.symbol_idx, t.prob ) for t in self.transitions } )

		dfa_transitions = defaultdict(dict)

		if not with_probs :
			for t in self.transitions :
				dfa_transitions[ t.origin_state_idx ][ t.symbol_idx ] = t.target_state_idx
		else :
			for t in self.transitions :
				dfa_transitions[ t.origin_state_idx ][ edge_label( t.symbol_idx, t.prob ) ] = t.target_state_idx

		# Construct the DFA
		return DFA(
			states=dfa_states,
			input_symbols=dfa_symbols,
			transitions=dfa_transitions,
			initial_state=self.start_state,
			allow_partial=True,
			final_states={ 
				s for s in dfa_states
			}
		)
