import numpy as np
import scipy as sp
import sympy
from fractions import Fraction

def solve_for_pi( T : np.ndarray, EPS : float = 1e-12 ) -> np.ndarray :

	n = T.shape[0]

	# helper ---------------------------------------------

	def normalize_pi( pi ) :
		if np.any( pi < -1e-8 ):
			raise ValueError("Significant negative values in solve result")

		pi = np.clip( pi, a_min=0.0, a_max=None )
		s = pi.sum()

		if abs(s) < EPS :
			raise Exception("Degenerate stationary distribution (sum ≈ 0)")

		return pi / s

	# least squares method --------------------------------

	try:
		A = T.T - np.eye(n)
		A = np.vstack([np.eye(n) - T.T, np.ones((1, n))])
		b = np.zeros(n + 1)
		b[-1] = 1.0
		pi, *_ = np.linalg.lstsq(A, b, rcond=None)
		return normalize_pi(pi)

	except (np.linalg.LinAlgError, ValueError):
		warnings.warn( "Falling back to linalg.solve" )

	# Solve method ----------------------------------------

	try:
		A = T.T - np.eye(n)
		A[-1, :] = 1.0
		b = np.zeros(n)
		b[-1] = 1.0
		pi = np.linalg.solve(A, b)
		return normalize_pi( pi )

	except (np.linalg.LinAlgError, ValueError):
		warnings.warn( "Falling back to nullspace method" )

	# Nullspace fallback -------------------------------

	try:
		A = T.T - np.eye(n)
		ns = sp.linalg.null_space( A )

		if ns.shape[1] != 1:
			raise ValueError("Unexpected nullspace dimension")

		pi = ns[:, 0]
		pi *= np.sign(pi[np.argmax(np.abs(pi))])
		return normalize_pi( pi )

	except (np.linalg.LinAlgError, ValueError):
		warnings.warn( "Falling back to eigen method" )

	# Eigen fallback ----------------------------------

	eigvals, eigvecs = linalg.eig( T.T )
	idx = np.argmin(np.abs(eigvals - 1))
	pi = np.real_if_close(eigvecs[:, idx])
	pi = np.real(pi)
	pi *= np.sign(pi[np.argmax(np.abs(pi))])

	return normalize_pi( pi )

def solve_for_pi_fractional( T : sympy.Matrix ) -> tuple[Fraction] :

	n = P.shape[0]

	# Setup the system (T.T - I)
	A = T.transpose() - sympy.Matrix.eye(n)

	# Replace the last row with 1s to enforce sum(pi) = 1
	for j in range(n):
		A[n-1, j] = 1

	# Setup the target vector [0, 0, ..., 1]
	b = sympy.Matrix([0] * (n - 1) + [1])

	# 4. Solve exactly
	pi_exact = A.LUsolve(b)

	pi_exact = ( Fraction( int(x.p), int(x.q) ) for x in pi_exact )

	return tuple(pi_exact)