import numpy as np
import matplotlib.pyplot as plt

A = np.array([[0.5, 0.5], [-0.5, 0.5]])
B = np.array([[0.5, -0.5], [0.5, 0.5]])

def blueFun(x):
    y = A @ x  # Matrix multiplication
    return y

def redFun(x):
    y = B @ x - np.array([[0.5], [0.5]])  # Matrix multiplication
    return y

# Initialize
x = np.array([[0.0], [0.0]])

yVecR = []  # Red points
yVecB = []  # Blue points

# Iteration with random selection
for _ in range(25000):
    if np.random.rand() < 0.5:  # Randomly choose blue or red function
        x = blueFun(x)
        yVecB.append(x)  # Append to blue points
    else:
        x = redFun(x)
        yVecR.append(x)  # Append to red points

# Convert lists to arrays for plotting
yVecR = np.hstack(yVecR).T if yVecR else np.empty((0, 2))
yVecB = np.hstack(yVecB).T if yVecB else np.empty((0, 2))

# Plotting the results
plt.scatter(yVecR[:, 0], yVecR[:, 1], s=0.1, color='red', label='Red Function')
plt.scatter(yVecB[:, 0], yVecB[:, 1], s=0.1, color='red', label='Blue Function')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.title("Levy C Curve with Red and Blue Points")
plt.show()