from collections import defaultdict
import copy
import random

import numpy as np

from .am_hmm          import HMM
from .am_causal_state import CausalState
from .am_transition   import Transition

from .am_random import uniform_dist, exp_uniform_blend, resolve_rng

def star_join(
	exit_symbol : str, 
	enter_symbols : list[str],
	machines : list[HMM],
	mode_residency_factor : float ) -> HMM :

	machine = HMM()

	isomorphic_groups = defaultdict(list)
	for i, m in enumerate( machines ) :
		if m.isoclass is not None :
			isomorphic_groups[ m.isoclass ].append( i )

	# since we are merging multuple machines which might have name collision
	# we need to rename the states to ensure uniqueness
	def rename_state( base_name : str, g : int ) :
		return f"{g}/{base_name}"

	def get_gid( idx : int, isoclass : int | None ) :
		return idx if isoclass is None else isoclass

	machine.set_alphabet( [ exit_symbol ] )

	for m in machines :
		machine.extend_alphabet( alphabet=m.alphabet )

	# create a connector state and connector state class

	connector_state = CausalState(
		name=f"/c", 
		classes=set({"connector"}) 
	)
	
	# initial states before adding each machines states
	machine.set_states( [ connector_state ] )
	machine.start_state = 0

	# number of machines (groups of states)
	n_groups = len( machines )

	# make sure we have enough symbols (otherwise connector can't be unifilar)
	if n_groups > len(enter_symbols) :
		raise Exception(
			f"Too few enter symbols given number of machines"
		)

	# for each given machine
	for m_idx, m in enumerate( machines ) : 

		# default to the index of the machine in the list
		m_gid = get_gid( m_idx, m.isoclass )

		# give the states from this machine a class name
		m_classes = { 
			f"m_{m_idx}", 
			f"isoclass_{m.isoclass}" 
		}

		added_states = []
		
		# create a state and extend our existing machine to include it
		for s_idx, state in enumerate( m.states ) :

			isomorphs=set()
			if m.isoclass is not None and m.isoclass in isomorphic_groups :

				for other_idx in isomorphic_groups[ m.isoclass ] :
					
					if other_idx == m_idx : 
						continue
					
					other_m = machines[ other_idx ]
					isomorphs.add(  
						rename_state( 
							other_m.states[ s_idx ].name, 
							get_gid( other_idx, other_m.isoclass ) )
					)

			added_states.append( 
				CausalState( 
					name=rename_state(state.name, m_gid),
					classes=( m_classes | state.classes ),
					isomorphs=isomorphs
				) 
			)

		machine.extend_states( added_states )

		added_transitions = []

		# add all of the transitions from the machine
		for tr in m.transitions :

			# get the names of the states for the transition
			origin_state_name = rename_state( m.states[ tr.origin_state_idx ].name, m_gid )
			target_state_name = rename_state( m.states[ tr.target_state_idx ].name, m_gid )

			# idx of the symbol remaped to this machines alphabet list
			new_symbol_idx = machine.symbol_idx_map[ m.alphabet[ tr.symbol_idx ] ]

			# create and add the new transition
			added_transitions.append( Transition(
				origin_state_idx=machine.state_idx_map[ origin_state_name ],
				target_state_idx=machine.state_idx_map[ target_state_name ],
				prob=tr.prob,
				symbol_idx=new_symbol_idx
			) )

		machine.extend_transitions( added_transitions )

		# Add connector transitions, and adjust transition probabilities to sum to 1

		# the name of the state that is the entry point to this group from the connector
		m_entry_state_name = rename_state( m.states[ m.start_state ].name, m_gid )

		# get the index of the entry state for this machine
		m_entry_state_idx = machine.state_idx_map[ m_entry_state_name ]


		# Get the within group transitions from m's entry state
		# ( the probabilities will need to be adjusted )
		transition_ids_from_m_entry = set()
		for i, tr in enumerate( machine.transitions ) : 
			if tr.origin_state_idx == m_entry_state_idx :
				transition_ids_from_m_entry.add( i )
		
		n_from_entry = len( transition_ids_from_m_entry )

		# Pr of staying in this group is distributed over the within group outgoing edges from the entry state 
		for i in transition_ids_from_m_entry :

			machine.transitions[ i ] = Transition(
				origin_state_idx=machine.transitions[ i ].origin_state_idx,
				target_state_idx=machine.transitions[ i ].target_state_idx,
				prob=mode_residency_factor / n_from_entry,
				symbol_idx=machine.transitions[ i ].symbol_idx
			)

		# from m's entry state back to connector
		escape_pr = 1.0 - mode_residency_factor

		machine.extend_transitions( transitions=[
			Transition(
				origin_state_idx=m_entry_state_idx,
				target_state_idx=machine.start_state,
				prob=escape_pr,
				symbol_idx=machine.symbol_idx_map[ exit_symbol ]
			)
		] )

		# from the connector to m's entry state
		machine.extend_alphabet( alphabet=[ enter_symbols[ m_idx ] ] )
		
		machine.extend_transitions( transitions=[
			Transition(
				origin_state_idx=machine.start_state,
				target_state_idx=m_entry_state_idx,
				prob=( 1.0 / n_groups ),
				symbol_idx=machine.symbol_idx_map[ enter_symbols[ m_idx ] ]
			)
		] )

	return machine

def star(
	exit_symbol          : str,
	enter_symbols        : list[str],
	normal_symbols       : list[str],
	n_modes              : int = 7,
	n_isomorphic         : int = 2,
	randomness           : float = 0.3,
	connectedness        : float = 0.5,
	residency_factor     : float = 0.5,
	n_normal_symbols     : int = 4,
	t_states_per_machine : int = 17 )  -> HMM :

	if len( normal_symbols ) < n_normal_symbols*n_isomorphic :
		raise ValueError( "Must have at least n_normal_symbols*n_isomorphic normal symbols" )

	if len( enter_symbols ) < n_modes*n_isomorphic :
		raise ValueError( "Must have at least n_modes*n_isomorphic enter symbols" )

	alphabet     = [ f"{normal_symbols[i]}" for i in range( 0, n_normal_symbols            ) ]
	iso_alphabet = [ f"{normal_symbols[i]}" for i in range( n_normal_symbols, n_normal_symbols*2  ) ]

	random_machines = []

	for i in range( n_modes ) :

		m = random_machine( 
			n_states=t_states_per_machine, 
			symbols=alphabet, 
			randomness=randomness,
			connectedness=connectedness ) 
		
		m.collapse_to_largest_strongly_connected_subgraph()
		m_iso = isomorphic_to( m, alphabet=iso_alphabet )

		m.isoclass     = f"{i}"
		m_iso.isoclass = f"{i}"

		for j, state in enumerate( m.states ) : 
			m.states[ j ].add_isomorph( m_iso.states[ j ].name )
			m_iso.states[ j ].add_isomorph( m.states[ j ].name )

		random_machines.append( m )
		random_machines.append( m_iso )

	mode_machine = star_join(  
		exit_symbol=exit_symbol, 
		enter_symbols=enter_symbols,
		machines=random_machines,
		mode_residency_factor=residency_factor
	)

	return mode_machine


def isomorphic_to( 
	m : HMM, 
	alphabet : list[str],
	decorator : str = '@' ) -> HMM :

	# make sure there are enough symbols
	if len( alphabet ) < len( m.alphabet ) :
		raise ValueError( "Not enough symbols in the alphabet" )

	# take the as much of them as needed
	alphabet_used = alphabet[ 0 : len( m.alphabet ) ]

	states = [ 
		CausalState( 
			name=f"{s.name}{decorator}",
			classes=copy.deepcopy( s.classes )
		)
		for s in m.states
	]

	return HMM(
		states=states,
		transitions=copy.deepcopy( m.transitions ),
		start_state=0,
		alphabet=alphabet_used
	)

def random_machine( 
	n_states : int, 
	symbols  : list[str], 
	connectedness : float,
	randomness : float,
	random_seed : int | None = None )  -> HMM  :

	if random_seed is not None:
		py_rng = random.Random(random_seed)
		np_rng = np.random.default_rng(random_seed)
	else:
		py_rng = random
		np_rng = resolve_rng(None)

	states=[
		CausalState( name=f"{i}" )
		for i in range( n_states  )
	] 

	n_symbols = len( symbols )
	transitions = []

	for state_idx, state in enumerate( states ) : 

		n_transitions = sum( py_rng.random() < connectedness for _ in range( n_symbols - 1 ) ) + 1
		transition_to = py_rng.sample( range( n_states ), n_transitions )

		transition_probabilities = exp_uniform_blend( 
			n=n_transitions, 
			alpha=randomness,
			np_rng=np_rng )

		transition_symbols_indices = py_rng.sample( range( n_symbols ), n_transitions )

		for i, p in enumerate( transition_probabilities ) :
			transitions.append(
				Transition(
					origin_state_idx=state_idx,
					target_state_idx=transition_to[ i ],
					prob=p,
					symbol_idx=transition_symbols_indices[ i ]
				)
			) 

	return HMM( 
		states=states,
		transitions=transitions,
		start_state=0,
		alphabet=symbols.copy()
	)