import numpy as np
import matplotlib.pyplot as plt

# Tell Matplotlib to use the system's LaTeX engine
plt.rcParams['text.usetex'] = True

# Add amsmath to the LaTeX preamble so pmatrix works
plt.rcParams['text.latex.preamble'] = r'\usepackage{amsmath}'

def plot_vector_field(A, x_range=(-5, 5), y_range=(-5, 5), density=1.5):
    """
    Takes a 2x2 numpy array A and plots its corresponding linear vector field Ax,
    including flow lines (streamlines) and eigenlines (if eigenvalues are real).
    """
    # 1. Create a dense grid for the streamplot (flow lines)
    x = np.linspace(x_range[0], x_range[1], 100)
    y = np.linspace(y_range[0], y_range[1], 100)
    X, Y = np.meshgrid(x, y)

    # 2. Calculate the vector field (U, V) = A * (X, Y)
    # U = a*X + b*Y
    # V = c*X + d*Y
    U = A[0, 0] * X + A[0, 1] * Y
    V = A[1, 0] * X + A[1, 1] * Y

    fig, ax = plt.subplots(figsize=(8, 8))

    # 3. Plot the flow lines
    # streamplot maps the trajectories of the vector field
    ax.streamplot(X, Y, U, V, color='cornflowerblue', density=density,
                  linewidth=1, arrowsize=1.5)

    # 4. Calculate eigenvalues and eigenvectors
    eigenvalues, eigenvectors = np.linalg.eig(A)

    # 5. Plot eigenlines (only if eigenvalues are real)
    colors = ['crimson', 'darkorange']
    has_real_eigenvalues = False

    for i, eigval in enumerate(eigenvalues):
        # Check if the imaginary part is effectively zero
        if np.isclose(eigval.imag, 0):
            has_real_eigenvalues = True
            real_eigval = eigval.real

            # The corresponding eigenvector is the i-th column
            v = eigenvectors[:, i].real

            # Create a line spanning across the plot area based on the eigenvector
            t = np.array([min(x_range[0], y_range[0]) * 3, max(x_range[1], y_range[1]) * 3])
            line_x = t * v[0]
            line_y = t * v[1]

# String version
#            ax.plot(line_x, line_y, color=colors[i % len(colors)], linestyle='--',
#                    linewidth=2.5, label=f'Eigenline (λ = {real_eigval:.2f})')

# LaTeX version
            ax.plot(line_x, line_y, color=colors[i % len(colors)], linestyle='--',
                    linewidth=2.5, label=f"Autoespacio ($\\lambda = {real_eigval:.2f}$)")

    # 6. Formatting the plot
    ax.set_xlim(x_range)
    ax.set_ylim(y_range)

    # Draw x and y axes passing through the origin
    ax.axhline(0, color='black', linewidth=1.2, zorder=3)
    ax.axvline(0, color='black', linewidth=1.2, zorder=3)

    ax.set_aspect('equal')
    ax.set_xlabel('x', fontsize=12)
    ax.set_ylabel('y', fontsize=12)

    # Format the matrix using LaTeX syntax
    # We use {{ and }} to escape the curly braces in the f-string
    # The :g formats the floats cleanly (e.g., 1 instead of 1.0)
    matrix_latex = f"\\begin{{pmatrix}} {A[0,0]:g} & {A[0,1]:g} \\\\ {A[1,0]:g} & {A[1,1]:g} \\end{{pmatrix}}"

    # Wrap the LaTeX string in $...$ for Matplotlib's math parser
    ax.set_title(f'$A = {matrix_latex}$. Autovalores: ${eigenvalues[0]:.2f}, {eigenvalues[1]:.2f}$', fontsize=16, pad=15)

#    # Format the matrix for the title (no LaTeX, just a string)
#    matrix_str = f"[[{A[0,0]}, {A[0,1]}],\n [{A[1,0]}, {A[1,1]}]]"
#    ax.set_title(f'Vector Field and Flow Lines for A =\n{matrix_str}', fontsize=14, pad=15)

    if has_real_eigenvalues:
        ax.legend(loc='upper left', framealpha=0.9)

    plt.grid(True, linestyle=':', alpha=0.6)
    plt.show()

if __name__ == "__main__":
    # --- Test Cases ---

    print("Close the current plot window to see the next example.")

    # Example 1: Source (Real, distinct, positive eigenvalues)
    A_source = np.array([[ 7.0,  -2.0],
                         [ 2.0, 2.0]])
    plot_vector_field(A_source)

    # Example 2: Sink (Real, distinct, negative eigenvalues)
    A_sink = np.array([[ -7.0,  2.0],
                         [ -2.0, -2.0]])
    plot_vector_field(A_sink)

    # Example 3: Saddle (Real eigenvalues with opposite signs)
    A_saddle = np.array([[ 1.0,  2.0],
                         [ 2.0, 1.0]])
    plot_vector_field(A_saddle)

    # Example 4: Stable Line (One eigenvalue zero, one negative)
    A_stline = np.array([[-1.0,  1.0],
                       [ 1.0, -1.0]])
    plot_vector_field(A_stline)

    # Example 5: Unstable Line (One eigenvalue zero, one positive)
    A_ustline = np.array([[1.0,  1.0],
                       [ 1.0, 1.0]])
    plot_vector_field(A_ustline)

    # Example 6: Radial source (Equal positive eigenvalues, diagonal)
    A_rsource = np.array([[1.0,  0.0],
                       [ 0.0, 1.0]])
    plot_vector_field(A_rsource)

    # Example 7: Radial sink (Equal negative eigenvalues, diagonal)
    A_rsink = np.array([[-1.0,  0.0],
                       [ 0.0, -1.0]])
    plot_vector_field(A_rsink)

    # Example 8: (Equal positive eigenvalues, non-diagonal)
    A_nondiagp = np.array([[2.0,  -1.0],
                       [ 1.0, 0.0]])
    plot_vector_field(A_nondiagp)

    # Example 9: (Equal negative eigenvalues, non-diagonal)
    A_nondiagn = np.array([[-2.0,  1.0],
                       [ -1.0, 0.0]])
    plot_vector_field(A_nondiagn)

    # Example 10: (Equal zero eigenvalues, non-diagonal)
    A_nondiag0 = np.array([[1.0,  -1.0],
                       [ 1.0, -1.0]])
    plot_vector_field(A_nondiag0)

    # Example 11: Spiral source (Complex eigenvalues, positive real part)
    A_spsource = np.array([[2.0,  4.0],
                       [ -2.0, 4.0]])
    plot_vector_field(A_spsource)

    # Example 12: Spiral sink (Complex eigenvalues, negative real part)
    A_spsink = np.array([[-2.0,  -4.0],
                       [ 2.0, -4.0]])
    plot_vector_field(A_spsink)

    # Example 13: Center (Complex eigenvalues, zero real part)
    A_center = np.array([[0.0,  1.0],
                       [ -1.0, 0.0]])
    plot_vector_field(A_center)