import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans

def generate_data(n_samples=300):
    data, _ = make_blobs(n_samples=n_samples, centers=4, cluster_std=1.0, random_state=42)
    return data

def plot_clusters(data, k):
    kmeans = KMeans(n_clusters=k)
    y_kmeans = kmeans.fit_predict(data)

    plt.scatter(data[:, 0], data[:, 1], c=y_kmeans, s=50, cmap='viridis')
    centers = kmeans.cluster_centers_
    plt.scatter(centers[:, 0], centers[:, 1], c='red', s=200, alpha=0.75, marker='X')
    plt.title(f'K-means Clustering with k={k}')
    plt.xlabel('Feature 1')
    plt.ylabel('Feature 2')
    plt_fig = plt
    return plt_fig

def main():
    st.title("K-means Clustering Simulator")
    st.write("This is a simple simulator to visualize how k-means clustering works.")

    # 데이터 생성 및 추가
    if 'data' not in st.session_state:
        st.session_state.data = np.array([]).reshape(0, 2)

    add_point = st.button("Add Random Data Point")
    if add_point:
        new_point = np.array([[np.random.uniform(0, 100), np.random.uniform(0, 100)]])
        st.session_state.data = np.vstack([st.session_state.data, new_point])

    x_coord = st.number_input("X Coordinate", min_value=0.0, max_value=100.0, value=50.0)
    y_coord = st.number_input("Y Coordinate", min_value=0.0, max_value=100.0, value=50.0)
    add_custom_point = st.button("Add Custom Data Point")
    if add_custom_point:
        st.session_state.data = np.vstack([st.session_state.data, [x_coord, y_coord]])

    st.write("Here is the data:")
    plt.scatter(st.session_state.data[:, 0], st.session_state.data[:, 1], s=50, cmap='viridis')
    st.pyplot()

    k = st.slider("Select the number of clusters (k)", 1, 10, 4)
    st.write(f"You selected k={k}")

    plt_fig = plot_clusters(st.session_state.data, k)
    st.pyplot(plt_fig)

if __name__ == '__main__':
    main()