import numpy as np
import sklearn
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler


customer_names = np.array(
    [
        "trial-01",
        "trial-02",
        "trial-03",
        "steady-01",
        "steady-02",
        "steady-03",
        "premium-01",
        "premium-02",
        "premium-03",
    ]
)

features = np.array(
    [
        [1.2, 18.0],
        [1.8, 22.0],
        [2.0, 24.0],
        [5.8, 47.0],
        [6.4, 52.0],
        [6.9, 49.0],
        [10.2, 85.0],
        [10.8, 91.0],
        [11.4, 88.0],
    ],
    dtype=float,
)

scaler = StandardScaler()
features_scaled = scaler.fit_transform(features)

kmeans = KMeans(n_clusters=3, n_init=10, random_state=42)
labels = kmeans.fit_predict(features_scaled)
centers = scaler.inverse_transform(kmeans.cluster_centers_)

print(f"scikit-learn {sklearn.__version__}")
print("Cluster assignments:")
for customer, label in zip(customer_names, labels):
    print(f"- {customer}: cluster {label}")

print("\nCluster centers in original units:")
for cluster_id, center in enumerate(centers):
    visits, spend = center
    print(f"- cluster {cluster_id}: visits={visits:.1f}, avg_order=${spend:.2f}")

print(f"\nInertia: {kmeans.inertia_:.3f}")
print(f"Iterations: {kmeans.n_iter_}")

new_rows = np.array(
    [
        [2.3, 23.0],
        [10.5, 89.0],
    ],
    dtype=float,
)
new_labels = kmeans.predict(scaler.transform(new_rows))

print("\nNew row clusters:")
for row, label in zip(new_rows, new_labels):
    visits, spend = row
    print(f"- visits={visits:.1f}, avg_order=${spend:.2f}: cluster {label}")