import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from mpl_toolkits.mplot3d import Axes3D

plt.style.use('seaborn-v0_8')
np.random.seed(42)

# Creating matrices in NumPy

# Method 1: From a list of lists
A = np.array([[1, 2, 3],
              [4, 5, 6]])
print(f"Matrix A:")
print(A)
print(f"Shape: {A.shape} (2 rows, 3 columns)")
print(f"Size: {A.size} elements")

# Method 2: Using NumPy functions
B = np.zeros((3, 3))  # 3x3 matrix of zeros
C = np.ones((2, 4))   # 2x4 matrix of ones
D = np.eye(3)         # 3x3 identity matrix
E = np.random.rand(2, 3)  # 2x3 random matrix

print(f"\nZeros matrix B:")
print(B)
print(f"\nOnes matrix C:")
print(C)
print(f"\nIdentity matrix D:")
print(D)
print(f"\nRandom matrix E:")
print(E)

# Matrix indexing and slicing
M = np.array([[1, 2, 3, 4],
              [5, 6, 7, 8],
              [9, 10, 11, 12]])

print(f"Matrix M:")
print(M)
print()

# Access single element
print(f"Element at position (0,2): {M[0, 2]}")
print(f"Element at position (2,1): {M[2, 1]}")

# Access entire rows
print(f"\nFirst row: {M[0, :]}")
print(f"Second row: {M[1, :]}")

# Access entire columns
print(f"\nFirst column: {M[:, 0]}")
print(f"Third column: {M[:, 2]}")

# Submatrices
print(f"\nTop-left 2x2 submatrix:")
print(M[0:2, 0:2])

print(f"\nLast two rows, last two columns:")
print(M[1:, 2:])

# Matrix addition and scalar multiplication
A = np.array([[1, 2],
              [3, 4]])
B = np.array([[5, 6],
              [7, 8]])

print(f"Matrix A:")
print(A)
print(f"\nMatrix B:")
print(B)

# Matrix addition
C = A + B
print(f"\nA + B:")
print(C)

# Matrix subtraction
D = A - B
print(f"\nA - B:")
print(D)

# Scalar multiplication
E = 3 * A
print(f"\n3 * A:")
print(E)

# Element-wise operations
F = A * B  # Element-wise multiplication (Hadamard product)
print(f"\nA * B (element-wise multiplication):")
print(F)

# Visualize matrix operations
fig, axes = plt.subplots(2, 3, figsize=(12, 8))
matrices = [A, B, C, D, E, F]
titles = ['A', 'B', 'A + B', 'A - B', '3 x A', 'A * B']

for i, (mat, title) in enumerate(zip(matrices, titles)):
    row, col = i // 3, i % 3
    im = axes[row, col].imshow(mat, cmap='viridis', aspect='equal')
    axes[row, col].set_title(title)
    
    # Add text annotations
    for r in range(mat.shape[0]):
        for c in range(mat.shape[1]):
            axes[row, col].text(c, r, f'{mat[r,c]:.0f}', 
                               ha='center', va='center', 
                               color='white' if mat[r,c] < mat.max()/2 else 'black',
                               fontsize=14, fontweight='bold')
    
    # Add grid lines with black borders
    axes[row, col].set_xticks(np.arange(-0.5, mat.shape[1], 1), minor=True)
    axes[row, col].set_yticks(np.arange(-0.5, mat.shape[0], 1), minor=True)
    axes[row, col].grid(which='minor', color='black', linestyle='-', linewidth=2)
    axes[row, col].set_xticks([])
    axes[row, col].set_yticks([])

plt.tight_layout()
plt.show()

# Matrix multiplication
A = np.array([[1, 2, 3],
              [4, 5, 6]])
B = np.array([[7, 8],
              [9, 10],
              [11, 12]])

print(f"Matrix A (2x3):")
print(A)
print(f"\nMatrix B (3x2):")
print(B)

# Matrix multiplication using @
C = A @ B
print(f"\nA @ B (2x2):")
print(C)

# Alternative: using np.dot
C_alt = np.dot(A, B)
print(f"\nUsing np.dot (same result):")
print(C_alt)

# Manual calculation for verification
print(f"\nManual calculation:")
print(f"C[0,0] = A[0,:] · B[:,0] = {A[0,:]} · {B[:,0]} = {np.dot(A[0,:], B[:,0])}")
print(f"C[0,1] = A[0,:] · B[:,1] = {A[0,:]} · {B[:,1]} = {np.dot(A[0,:], B[:,1])}")
print(f"C[1,0] = A[1,:] · B[:,0] = {A[1,:]} · {B[:,0]} = {np.dot(A[1,:], B[:,0])}")
print(f"C[1,1] = A[1,:] · B[:,1] = {A[1,:]} · {B[:,1]} = {np.dot(A[1,:], B[:,1])}")

# Dimension compatibility check
print(f"\nDimension compatibility:")
print(f"A.shape = {A.shape}, B.shape = {B.shape}")
print(f"A has {A.shape[1]} columns, B has {B.shape[0]} rows")
print(f"Since {A.shape[1]} == {B.shape[0]}, multiplication is valid")
print(f"Result shape: ({A.shape[0]}, {B.shape[1]}) = {C.shape}")

# Matrix-vector multiplication
A = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])
x = np.array([1, 2, 1])

print(f"Matrix A:")
print(A)
print(f"\nVector x: {x}")

# Matrix-vector multiplication
y = A @ x
print(f"\nA @ x = {y}")

# Show row-wise computation
print(f"\nRow-wise computation:")
for i in range(A.shape[0]):
    row_product = np.dot(A[i, :], x)
    print(f"Row {i}: {A[i, :]} · {x} = {row_product}")

# Visualize as linear transformation
# Let's use a 2D example for better visualization
A_2d = np.array([[2, 1],
                 [1, 2]])
x_2d = np.array([1, 1])
y_2d = A_2d @ x_2d

plt.figure(figsize=(10, 5))

# Plot original and transformed vectors
plt.subplot(1, 2, 1)
plt.arrow(0, 0, x_2d[0], x_2d[1], head_width=0.1, head_length=0.1, 
          fc='blue', ec='blue', label='Original x')
plt.arrow(0, 0, y_2d[0], y_2d[1], head_width=0.1, head_length=0.1, 
          fc='red', ec='red', label='Transformed Ax')
plt.grid(True, alpha=0.3)
plt.axis('equal')
plt.legend()
plt.title('Linear Transformation')
plt.xlabel('x')
plt.ylabel('y')

# Show the computation
plt.subplot(1, 2, 2)
computation_text = f"""Matrix A:
{A_2d}

Vector x: {x_2d}

A @ x = {y_2d}

Computation:
y[0] = {A_2d[0,0]}x{x_2d[0]} + {A_2d[0,1]}x{x_2d[1]} = {y_2d[0]}
y[1] = {A_2d[1,0]}x{x_2d[0]} + {A_2d[1,1]}x{x_2d[1]} = {y_2d[1]}"""
plt.text(0.1, 0.5, computation_text, transform=plt.gca().transAxes, 
         fontsize=10, verticalalignment='center', fontfamily='monospace')
plt.axis('off')
plt.title('Computation Details')

plt.tight_layout()
plt.show()

# Special matrices
print("Special Matrices:")
print("=" * 30)

# Identity matrix
I = np.eye(3)
print(f"3x3 Identity matrix:")
print(I)

# Test identity property
A = np.random.randint(1, 10, (3, 3))
AI = A @ I
IA = I @ A
print(f"\nOriginal matrix A:")
print(A)
print(f"\nA @ I:")
print(AI)
print(f"\nI @ A:")
print(IA)
print(f"\nAre they equal? {np.allclose(A, AI) and np.allclose(A, IA)}")

# Diagonal matrix
diagonal_values = [2, 3, 5]
D = np.diag(diagonal_values)
print(f"\nDiagonal matrix with values {diagonal_values}:")
print(D)

# Symmetric matrix
# Create a symmetric matrix
B = np.random.rand(3, 3)
S = B + B.T  # Adding a matrix to its transpose makes it symmetric
print(f"\nSymmetric matrix S:")
print(S)
print(f"\nS transpose:")
print(S.T)
print(f"\nIs S symmetric? {np.allclose(S, S.T)}")

# Visualize these matrices
fig, axes = plt.subplots(1, 3, figsize=(12, 4))
matrices = [I, D, S]
titles = ['Identity Matrix', 'Diagonal Matrix', 'Symmetric Matrix']

for i, (mat, title) in enumerate(zip(matrices, titles)):
    im = axes[i].imshow(mat, cmap='Greys', aspect='equal')
    axes[i].set_title(title, fontsize=16)
    
    # Add text annotations
    for r in range(mat.shape[0]):
        for c in range(mat.shape[1]):
            # Use white text on dark background, black text on light background
            text_color = 'white' if mat[r,c] > (mat.min() + mat.max()) / 2 else 'black'
            axes[i].text(c, r, f'{mat[r,c]:.1f}', 
                        ha='center', va='center', fontsize=14,
                        color=text_color)
    
    # Add border and set white background
    axes[i].set_xticks([])
    axes[i].set_yticks([])
    axes[i].set_facecolor('white')
    for spine in axes[i].spines.values():
        spine.set_edgecolor('black')
        spine.set_linewidth(2)

plt.tight_layout()
plt.show()

# Matrix transpose
A = np.array([[1, 2, 3],
              [4, 5, 6]])

print(f"Original matrix A (2x3):")
print(A)

# Transpose
A_T = A.T
print(f"\nTranspose A^T (3x2):")
print(A_T)

# Alternative method
A_T_alt = np.transpose(A)
print(f"\nUsing np.transpose (same result):")
print(A_T_alt)

# Verify properties
print(f"\nVerifying properties:")
print(f"(A^T)^T == A? {np.array_equal(A_T.T, A)}")

# For matrix multiplication property
B = np.array([[1, 2],
              [3, 4],
              [5, 6]])
AB = A @ B
AB_T = AB.T
BT_AT = B.T @ A.T

print(f"\nMatrix B (3x2):")
print(B)
print(f"\nAB (2x2):")
print(AB)
print(f"\n(AB)^T:")
print(AB_T)
print(f"\nB^T @ A^T:")
print(BT_AT)
print(f"\n(AB)^T == B^T A^T? {np.allclose(AB_T, BT_AT)}")

# Practice Exercise: Matrix Operations
print("Practice Exercises:")
print("=" * 30)

# Create random matrices for practice
np.random.seed(123)  # For reproducible results
A = np.random.randint(1, 6, (3, 3))
B = np.random.randint(1, 6, (3, 3))
x = np.random.randint(1, 6, 3)

print(f"Matrix A:")
print(A)
print(f"\nMatrix B:")
print(B)
print(f"\nVector x: {x}")

# Exercise 1: Basic operations
print(f"\nExercise 1: Basic Operations")
print(f"A + B =")
print(A + B)

print(f"\nA - B =")
print(A - B)

print(f"\n2 * A =")
print(2 * A)

# Exercise 2: Matrix multiplication
print(f"\nExercise 2: Matrix Multiplication")
print(f"A @ B =")
print(A @ B)

print(f"\nB @ A =")
print(B @ A)
print(f"\nNote: A @ B ≠ B @ A (matrix multiplication is not commutative!)")

# Exercise 3: Matrix-vector multiplication
print(f"\nExercise 3: Matrix-Vector Multiplication")
y = A @ x
print(f"A @ x = {y}")

# Exercise 4: Transpose properties
print(f"\nExercise 4: Transpose Properties")
AT = A.T
BT = B.T
AB = A @ B
AB_T = AB.T
BT_AT = BT @ AT

print(f"(A @ B)^T =")
print(AB_T)
print(f"\nB^T @ A^T =")
print(BT_AT)
print(f"\nAre they equal? {np.allclose(AB_T, BT_AT)}")

# Exercise 5: Check if matrix is symmetric
print(f"\nExercise 5: Symmetry Check")
S = A @ A.T  # This creates a symmetric matrix
print(f"S = A @ A^T =")
print(S)
print(f"\nIs S symmetric? {np.allclose(S, S.T)}")
print(f"\nS^T =")
print(S.T)