main.py

from fastapi import FastAPI, Request, Body, Depends, Response
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
import numpy as np
from kmeans import KMeansVisualizer
from session import get_session, get_session_id
import io
import base64
import matplotlib.pyplot as plt

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

def get_plot_data(kmeans: KMeansVisualizer):
    fig, ax = plt.subplots(figsize=(10, 6))
    kmeans.plot(fig, ax)
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight', dpi=100)
    plt.close(fig)
    buf.seek(0)
    return {
        "plot": base64.b64encode(buf.getvalue()).decode(),
        "xlim": kmeans.xlim,
        "ylim": kmeans.ylim
    }

@app.get("/")
async def home(request: Request, session_id: str = Depends(get_session_id)):
    response = templates.TemplateResponse("index.html", {"request": request})
    response.set_cookie(key="session_id", value=session_id)
    return response

@app.post("/step")
async def step(kmeans: KMeansVisualizer = Depends(get_session)):
    try:
        if not kmeans.converged:
            kmeans.step()
        
        return {
            **get_plot_data(kmeans),
            "iteration": kmeans.iteration,
            "converged": kmeans.converged,
            "error": None
        }
    except Exception as e:
        return {
            "error": str(e),
            "iteration": kmeans.iteration,
            "converged": False
        }

@app.post("/clear")
async def clear(kmeans: KMeansVisualizer = Depends(get_session)):
    kmeans.clear_data()
    return {"success": True}

@app.post("/init")
async def initialize(kmeans: KMeansVisualizer = Depends(get_session)):
    kmeans.initialize_clients()
    return get_plot_data(kmeans)

@app.post("/select_start")
async def select_start(x: float, y: float, kmeans: KMeansVisualizer = Depends(get_session)):
    complete = kmeans.add_start_point(x, y)
    message = "Select next starting point" if not complete else "Starting points selected"
    return {
        **get_plot_data(kmeans),
        "complete": complete,
        "message": message
    }

@app.post("/init_clients")
async def init_clients(data: dict = Body(...), kmeans: KMeansVisualizer = Depends(get_session)):
    clients = data.get('clients', [])
    kmeans.initialize_custom_clients(
        np.array([[c['x'], c['y']] for c in clients]),
        [c['name'] for c in clients]
    )
    
    return get_plot_data(kmeans)

@app.post("/update_k")
async def update_k(data: dict = Body(...), kmeans: KMeansVisualizer = Depends(get_session)):
    kmeans.k = data.get('k', 2)
    kmeans.clear_data()
    
    return {
        "success": True,
        **get_plot_data(kmeans)
    }

@app.post("/step/{direction}")
async def step_direction(direction: str, kmeans: KMeansVisualizer = Depends(get_session)):
    try:
        if direction == "forward":
            if kmeans.current_step < len(kmeans.step_history) - 1:
                kmeans.load_state(kmeans.current_step + 1)
            else:
                kmeans.step()
        elif direction == "back":
            if kmeans.current_step > 0:
                kmeans.load_state(kmeans.current_step - 1)
        
        return {
            **get_plot_data(kmeans),
            "iteration": kmeans.iteration,
            "converged": kmeans.converged,
            "can_go_back": kmeans.current_step > 0,
            "can_go_forward": not kmeans.converged or kmeans.current_step < len(kmeans.step_history) - 1
        }
    except Exception as e:
        return {"error": str(e)}

@app.post("/elbow_method")
async def elbow_method(data: dict = Body(...)):
    k_min = data.get('k_min', 2)
    k_max = data.get('k_max', 5)
    num_runs = data.get('num_runs', 5) # Number of runs for each k to average out randomness

    k_values = range(k_min, k_max + 1)
    ssds = []

    for k in k_values:
        avg_ssd = 0
        for _ in range(num_runs):
            temp_kmeans = KMeansVisualizer() # Create a new instance for each run
            temp_kmeans.k = k
            temp_kmeans.client_data = temp_kmeans.initial_data.copy() # Use initial data for elbow
            temp_kmeans.client_names = temp_kmeans.initial_names.copy()
            
            # Run K-means until convergence
            temp_kmeans.init_centroids()
            while not temp_kmeans.converged and temp_kmeans.iteration < 100: # Max 100 iterations
                temp_kmeans.compute_distances()
                temp_kmeans.converged = temp_kmeans.update_centroids()
                temp_kmeans.iteration += 1
            avg_ssd += temp_kmeans.calculate_ssd()
        ssds.append(avg_ssd / num_runs)

    # Plot the elbow curve
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(list(k_values), ssds, marker='o')
    ax.set_xlabel('Number of Clusters (k)')
    ax.set_ylabel('Sum of Squared Distances (SSD)')
    ax.set_title('Elbow Method for Optimal k')
    ax.grid(True)

    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight', dpi=100)
    plt.close(fig)
    buf.seek(0)

    return {
        "plot": base64.b64encode(buf.getvalue()).decode(),
        "k_values": list(k_values),
        "ssds": ssds
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, port=8000)