import numpy as np
import matplotlib.pyplot as plt

from sklearn.cross_decomposition import PLSRegression

# Generate random data for input and output spaces
input_space = (np.random.randn(100, 2) @ np.array([[20, 0], [0, 3.0]]))
output_space = np.random.randn(100, 2)
output_space[:,0]=input_space[:,0] 
output_space=output_space @ np.array([[3, 3], [8, 20.0]])

# Perform PLS with 2 components
pls = PLSRegression(n_components=2)
pls.fit(input_space, output_space)

# Calculate the PLS direction vectors
input_pls = pls.x_weights_
output_pls = pls.y_weights_

# Create the scatterplots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))

# Plot input space
ax1.scatter(input_space[:, 0], input_space[:, 1], c='blue', alpha=0.5)
v1=ax1.quiver(0, 0, *input_pls[:, 0], scale=3, color='red', width=0.01, headwidth=3, headlength=4 )

ax1.set_title('Input Space')
ax1.set_xlabel('$x_1$')
ax1.set_ylabel('$x_2$')
ax1.axis('equal')
ax1.set_xlim(ax1.get_xlim()[0] - 1, ax1.get_xlim()[1] + 1)
ax1.set_ylim(ax1.get_ylim()[0] - 1, ax1.get_ylim()[1] + 1)

# Plot output space
ax2.scatter(output_space[:, 0], output_space[:, 1], c='green', alpha=0.5)
ax2.quiver(0, 0, *output_pls[:, 0], scale=3, color='red', width=0.01, headwidth=5, headlength=4)
ax2.set_title('Output Space')
ax2.set_xlabel('$y_1$')
ax2.set_ylabel('$y_2$')
ax2.axis('equal')
ax2.set_xlim(ax2.get_xlim()[0] - 1, ax2.get_xlim()[1] + 1)
ax2.set_ylim(ax2.get_ylim()[0] - 1, ax2.get_ylim()[1] + 1)

plt.show()