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)
    """

    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):
        # Convert NumPy array if needed
        if isinstance(matrix_in, np.ndarray):
            matrix_in = sp.Matrix(matrix_in)

        # Compute raw eigen data: list of (eigenval, multiplicity, [eigenvecs])
        raw_data = sorted(matrix_in.eigenvects(), key=lambda x: x[0], reverse=reverse)
        
        vals_clean = []
        vecs_clean = []

        for eigval, mult, vects in raw_data:
            for i in range(mult):
                v_raw = vects[i]
                
                # Normalize vector by its largest element for cleaner display
                max_val = max(v_raw, key=lambda x: abs(x.evalf()))
                v_norm = v_raw / max_val

                # Apply numerical formatting if requested
                if evalf:
                    val_processed = _safe_eval(eigval)
                    vec_processed = v_norm.applyfunc(_safe_eval)
                else:
                    val_processed = eigval
                    vec_processed = v_norm
                
                vals_clean.append(val_processed)
                vecs_clean.append(vec_processed)
                
        return vals_clean, vecs_clean

    # Render Output
    def _render_display(vals, vecs):
        # Header
        if show == "both":
            ipy.display(ipy.Math(f"\\text{{Eigenvalues and Eigenvectors of }} {name}:"))
            
            if output == "default":
                for i, (val, vec) in enumerate(zip(vals, vecs)):
                    latex_str = f"\\lambda_{{{i+1}}} = {latex(val)}, v_{{{i+1}}} = {latex(vec)} \\qquad"
                    ipy.display(ipy.Math(latex_str))
            
            elif output == "compact":
                for i, (val, vec) in enumerate(zip(vals, vecs)):
                    # Transpose vector for compact row view
                    latex_str = f"\\lambda_{{{i+1}}} = {latex(val)}, \\quad v_{{{i+1}}} = {latex(vec.T)}"
                    ipy.display(ipy.Math(latex_str))
            
            elif output == "maple":
                # Vector of eigenvalues
                ipy.display(ipy.Math(f"\\lambda = {latex(sp.Matrix(vals))}"))
                # Matrix of eigenvectors
                V_matrix = sp.Matrix.hstack(*vecs)
                ipy.display(ipy.Math(f"v = {latex(V_matrix)}"))

        elif show == "eigvals":
            ipy.display(ipy.Math(f"\\text{{Eigenvalues of }} {name}: \\quad \\lambda = {latex(sp.Matrix(vals))}"))

        elif show == "eigvecs":
            V_matrix = sp.Matrix.hstack(*vecs)
            ipy.display(ipy.Math(f"\\text{{Eigenvectors of }} {name}: \\quad v = {latex(V_matrix)}"))

    # Main Execution Flow
    
    # Validation
    show_opt = show.lower()
    out_opt = output.lower()
    if show_opt not in ["both", "eigvals", "eigvecs", "none"]:
        raise ValueError(f"Invalid 'show' option: {show}")
    if out_opt not in ["default", "compact", "maple"]:
        raise ValueError(f"Invalid 'output' option: {output}")

    # Computation
    eig_vals, eig_vecs = _compute_eigen_data(A)

    # Display
    if show_opt != "none":
        _render_display(eig_vals, eig_vecs)

    # Return Data
    if return_data:
        return sp.Matrix(eig_vals), sp.Matrix.hstack(*eig_vecs)
    
    return None