import mglearn
import matplotlib.pyplot as plt
import numpy as np


# generate dataset
# Datos
longitud_cm = [
    5.2, 5.8, 6.5, 7.1, 7.9, 10, 9.3, 10.5, 11.2, 12.4,
    10.7, 12.0, 13.6, 14.3, 14.9, 15.5, 15.2, 13.8, 6.7, 8.4,
    3.1, 3.4, 3.6, 4.2, 4.5, 5.0, 5.3, 5.8, 6.1, 7.0,
    6.5, 7.4, 8.0, 9.2, 9.5, 9.8, 10.0, 8.6, 7.8, 6.4
]
ancho_cm = [
    2.0, 2.5, 3.0, 3.2, 3.5, 4.0, 3.8, 4.2, 4.5, 4.0,
    3.9, 4.1, 1.8, 5.0, 5.2, 5.5, 5.3, 4.7, 2.3, 3.1,
    1.5, 1.7, 1.8, 2.0, 2.2, 2.5, 2.6, 2.8, 3.0, 3.2,
    2.9, 3.1, 3.4, 3.6, 3.8, 4.0, 4.1, 3.3, 3.0, 3.7
]

# Crear X como pares de (longitud_cm, ancho_cm)
X = np.array(list(zip(longitud_cm, ancho_cm)))

# Crear etiquetas: 0 para Estafiate y 1 para Albahaca
y = np.array([0] * 20 + [1] * 20)

# plot dataset
mglearn.discrete_scatter(X[:, 0], X[:, 1], y)
plt.legend(["Estafiate", "Albaca"], loc=4)
plt.xlabel("Longitud (cm)")
plt.ylabel("Ancho (cm)")
plt.title("Clasificación de Estafiate y Albahaca")
print("X.shape: {}".format(X.shape))

plt.show()