import random
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from copy import deepcopy

from .am_create import star, star_join, isomorphic_to
import random

class Generator :

	def __init__(self, random_seed=None ) -> None :
		
		if random_seed :
		
			np.random.seed( random_seed ) 
			random.seed( random_seed )

		self.vocabulary = { 
			
			# 1. Printable ASCII: space (32) through tilde (126) → 95 chars
			"ASCII" : [chr(i) for i in range(32, 127)],

			# 2. Greek capitals Α–Τ (19), skipping U+03A2 (undefined slot)
			"greek_upper" : [chr(i) for i in range(0x0391, 0x03A5) if i != 0x03A2][:19],

			# 3. Greek smalls α–τ (19), skipping U+03C2 (ς, final-sigma variant)
			"greek_lower" : [chr(i) for i in range(0x03B1, 0x03C5) if i != 0x03C2][:19],

			# 4. Geometric Shapes block: U+25A0–U+25FF → 96 chars
			"geometric" : [chr(i) for i in range(0x25A0, 0x2600)]
		}

		self.exit_symbol = self.vocabulary[ "geometric" ][ 0 ]

	def generate( 
		self, n_states : int, 
		n_symbols, 
		n_modes : int,
		mode_alphas : list[float],
		residency_factor : float,
		levels : int,
		output_dir : Path ) :
		
		# 0-9 ASCII + lowercase letters + uppercase letters
		normal_symbols  = [ chr(i) for i in range( 48,  58 ) ]
		normal_symbols += [ chr(i) for i in range( 97, 123 ) ]
		normal_symbols += [ chr(i) for i in range( 65,  91 ) ]
		enter_symbols = self.vocabulary[ "greek_lower" ] + self.vocabulary[ "greek_upper" ]

		n_gen = n_modes*2

		# n_gen^levels leaves — correct for a full n-ary tree
		HMMs = [
		    star(
		        exit_symbol=self.exit_symbol,
		        enter_symbols=enter_symbols[ 0:n_modes*2 ],
		        normal_symbols=normal_symbols,
		        n_modes=n_modes,
		        n_isomorphic=2,
		        randomness=0.3,
		        connectedness=0.5,
		        residency_factor=residency_factor,
		        n_normal_symbols=n_symbols,
		        t_states_per_HMM=n_states
		    )
		    for i in range( n_gen**levels )
		]

		es_per_lev = n_gen
		es_offset = n_modes*2

		# Reduce bottom-up, levels-1 times:
		#   n_gen^levels → n_gen^(levels-1) → ... → n_gen
		for i in range( levels - 1 ) :
		    next_level = []
		    x = 0

		    es = enter_symbols[ es_offset + es_per_lev*i : es_offset+es_per_lev*(i+1) ]

		    for k in range( len(HMMs) // n_gen ) :  # shrinks each pass
		        m = star_join(
		            exit_symbol=self.exit_symbol,
		            enter_symbols=es,
		            HMMs=HMMs[ x:x+n_gen ],
		            mode_residency_factor=0.5
		        )
		        next_level.append( m )
		        x += n_gen
		    HMMs = deepcopy( next_level )

		# HMMs now has exactly n_gen elements — join into root
		return star_join(
		    exit_symbol=self.exit_symbol,
		    enter_symbols=normal_symbols,
		    HMMs=HMMs,
		    mode_residency_factor=0.5
		)

if __name__ == "__main__":

	n_modes = 1
	alpha = 0.3
	residency_factor = 0.5
	n_symbols = 4
	n_states = 27
	random_seed=42

	generator = Generator( random_seed=random_seed )

	path = Path("../data/aM_7-mode-iso" )
	path.mkdir(parents=True, exist_ok=True)

	aM = generator.generate( 
		n_states=n_states, 
		n_symbols=n_symbols, 
		n_modes=n_modes, 
		mode_alphas=[ alpha ]*n_modes,
		residency_factor=residency_factor,
		levels=3,
		output_dir=path )

	aM.minimize(retain_names=True)
	print( aM.is_epsilon_HMM() )

	aM.draw_graph(output_dir="./")

	print( f"h_mu : {aM.h_mu()}" )
	print( f"C_mu : {aM.C_mu()}" )
	# print( f"E    : {aM.E()}" )

	# aM.to_q_weighted()
	# print( aM.is_epsilon_HMM( verbose=True) )
	# aM.draw_graph(output_dir="./")

	# p = aM.get_fractional_stationary_distribution()

	# rm = aM.get_reverse()

	# mspq = aM.get_msp_qw( exact_state_cap=5000 )
	# print( "Done" )

	# print( aM.is_topological_epsilon_HMM() )
	# print( aM.is_epsilon_HMM() )
	# aM.draw_graph(output_dir="./")

	# aM.minimize(retain_names=True)
	# print( aM.is_topological_epsilon_HMM() )
	# print( aM.is_epsilon_HMM() )
	# aM.draw_graph(output_dir="./")

	# aM.save_config( path, with_complexity=True, with_block_convergence=True )

	# data = aM.generate_data( 
	# 	file_prefix=path / "train", 
	# 	n_gen=1_000_000_000, 
	# 	include_states=True,
	# 	isomorphic_shifts=set({1}),
	# 	random_seed=random_seed )

	# aM.generate_data( 
	# 	file_prefix=path / "test", 
	# 	n_gen=250_000_000, 
	# 	include_states=True,
	# 	isomorphic_shifts=set({1}),
	# 	random_seed=random_seed )


