import numpy as np
from ran_array_complete import RANArray as RAN, real_ran, const_ran


# NOTE: For iterative algorithms, keep b = 1 everywhere.
# Use real_ran() or const_ran() to create matrices and vectors.
# Use A @ B, not np.dot(A, B), for RAN matrix multiplication.

np.random.seed(42)


# ============================================================
# 1. Basic matrix multiplication
# ============================================================
A = real_ran(np.random.rand(40, 30))
B = real_ran(np.random.rand(30, 20))

C = A @ B
print("A @ B shape:", C.shape)
print("Collapsed shape:", C.collapse().shape)

# Verify it matches ordinary numpy
expected = A.collapse() @ B.collapse()
print("Matches numpy matmul:", np.allclose(C.collapse(), expected))


# ============================================================
# 2. Iteration: power iteration (dominant eigenvector)
# ============================================================
print("\n--- Power iteration ---")

# Make a symmetric positive definite RAN matrix
M = real_ran(np.random.rand(3, 3))
M = M @ M.T

x = real_ran(np.random.rand(3, 1))

for step in range(30):
    x = M @ x

    # Normalize
    norm = np.sqrt(np.sum(x.collapse() ** 2))
    x = x / norm

    # Rayleigh quotient x^T M x / x^T x
    rayleigh = (x.T @ (M @ x)).collapse().item()

    if step % 5 == 0:
        print(f"step {step}: Rayleigh quotient = {rayleigh:.6f}")


# ============================================================
# 3. Iteration: linear regression via gradient descent
# ============================================================
print("\n--- Linear regression gradient descent ---")

N, D = 100, 5
X = real_ran(np.random.randn(N, D))
w_true = real_ran(np.random.randn(D))
y = X @ w_true + 0.1 * real_ran(np.random.randn(N))

w = real_ran(np.zeros(D))
lr = 0.01

for step in range(1000):
    pred = X @ w
    err = pred - y
    grad = X.T @ err
    w = w - lr * grad

    if step % 100 == 0:
        loss = np.mean(err.collapse() ** 2)
        print(f"step {step}: loss = {loss:.6f}")

print("\nFinal weights vs true weights:")
print("  estimated:", w.collapse())
print("  true:     ", w_true.collapse())