import numpy                as    np
import sympy                as    sp
from   sympy                import latex
import IPython.display      as    ipy

def display_eigen(A, 
                # -- Labeling & Formatting --
                name="Matrix", evalf=False, prec=5, 
                # --- Output Control ---
                show="both", output="default", return_data=False, reverse=False):

    """
    Computes and displays the eigenvalues and eigenvectors of a matrix A using internal sub-functions.

    Parameters:
        A           : Square matrix (SymPy Matrix or NumPy ndarray)
        name        : Label used for the matrix (default="Matrix")
        evalf       : If True, show decimal values (default=False)
        prec        : Accuracy of decimal evaluation (default: 5)
        show        : What to display ("both", "eigvals", "eigvecs", or "none") (default="both")
        output      : How to display ("default", "compact", "Maple") (default="default")
        return_data : If True, returns eigenvalues and eigenvectors (default=False)
        reverse     : If True, sorts eigenvalues in descending order (default=False)

    Returns (if return_data=True):
        - sp.Matrix: Column vector of eigenvalues
        - sp.Matrix: Matrix of eigenvectors (stacked columns)
    """

    # Normalize inputs
    show = show.lower()
    output = output.lower()
    
    threshold = 10**(-10)

    # Safe Evaluation
    def _safe_eval(x):
        try:
            val = x.evalf(prec)
            if abs(val) < threshold:
                return sp.Integer(0)
            return round(float(val), prec)
        except (TypeError, ValueError):
            return x

    # Compute & Process Data
    def _compute_eigen_data(matrix_in):
        if isinstance(matrix_in, np.ndarray):
            matrix_in = sp.Matrix(matrix_in)

        try:
            # Sort by magnitude if evalf, else by real part
            key_func = (lambda x: abs(x[0].evalf())) if evalf else (lambda x: sp.re(x[0]))
            raw_data = sorted(matrix_in.eigenvects(), key=key_func, reverse=reverse)
        except:
            raw_data = matrix_in.eigenvects()
        
        vals_clean = []
        vecs_clean = []

        for eigval, mult, vects in raw_data:
            for i in range(mult):
                if i >= len(vects): break
                v_raw = vects[i]
                
                # Normalize vector
                try:
                    max_val = max(v_raw, key=lambda x: abs(x.evalf()))
                    v_norm = v_raw / max_val if max_val != 0 else v_raw
                except:
                    v_norm = v_raw

                if evalf:
                    vals_clean.append(_safe_eval(eigval))
                    vecs_clean.append(v_norm.applyfunc(_safe_eval))
                else:
                    vals_clean.append(eigval)
                    vecs_clean.append(v_norm)
                
        return vals_clean, vecs_clean

    # Render Output
    def _render_display(vals, vecs):
        # Formatting helper
        def _fmt(x):
            return latex(x).replace(r'\\', r'\\[5pt]')

        if show in ["both", "eigvals", "eigvecs"]:
            title_str = f"\\text{{Eigenvalues and Eigenvectors of }} {name}:"
            ipy.display(ipy.Math(rf"\large {title_str}"))

        if show == "both":
            if output == "default":
                for i, (val, vec) in enumerate(zip(vals, vecs)):
                    latex_str = rf"\large \lambda_{{{i+1}}} = {_fmt(val)}, \quad v_{{{i+1}}} = {_fmt(vec)}"
                    ipy.display(ipy.Math(latex_str))
            
            elif output == "compact":
                for i, (val, vec) in enumerate(zip(vals, vecs)):
                    latex_str = rf"\large \lambda_{{{i+1}}} = {_fmt(val)}, \quad v_{{{i+1}}} = {_fmt(vec.T)}"
                    ipy.display(ipy.Math(latex_str))
            
            elif output == "maple":
                L_matrix = sp.Matrix(vals)
                V_matrix = sp.Matrix.hstack(*vecs) if vecs else sp.Matrix([])
                ipy.display(ipy.Math(rf"\large \lambda = {_fmt(L_matrix)}"))
                ipy.display(ipy.Math(rf"\large V = {_fmt(V_matrix)}"))

        elif show == "eigvals":
            L_matrix = sp.Matrix(vals)
            ipy.display(ipy.Math(rf"\large \lambda = {_fmt(L_matrix)}"))

        elif show == "eigvecs":
            V_matrix = sp.Matrix.hstack(*vecs) if vecs else sp.Matrix([])
            ipy.display(ipy.Math(rf"\large V = {_fmt(V_matrix)}"))

    # Validation
    if show not in ["both", "eigvals", "eigvecs", "none"]:
        raise ValueError(f"Invalid 'show' option: {show}")
    if output not in ["default", "compact", "maple"]:
        raise ValueError(f"Invalid 'output' option: {output}")

    # Execution
    eig_vals, eig_vecs = _compute_eigen_data(A)

    if show != "none":
        _render_display(eig_vals, eig_vecs)

    if return_data:
        return sp.Matrix(eig_vals), (sp.Matrix.hstack(*eig_vecs) if eig_vecs else sp.Matrix([]))
    
    return None