// kmeans.ts - K-Means clustering algorithm

export interface Point {
    x: number;
    y: number;
}

export interface Cluster {
    centroid: Point;
    points: Point[];
}

export interface KMeansResult {
    clusters: Cluster[];
    iterations: number;
    converged: boolean;
}

export class KMeans {
    private readonly k: number;
    private readonly maxIterations: number;
    private readonly data: Point[];
    private clusters: Cluster[] = [];

    constructor(data: Point[], k: number, maxIterations: number = 50) {
        this.k = k;
        this.maxIterations = maxIterations;
        this.data = data;
    }

    private static euclideanDistance(p1: Point, p2: Point): number {
        const dx = p2.x - p1.x;
        const dy = p2.y - p1.y;
        return Math.sqrt(dx * dx + dy * dy);
    }

    private initializeCentroids(): void {
        const centroids: Point[] = [];
        const dataCopy = [...this.data];
        for (let i = 0; i < this.k && dataCopy.length > 0; i++) {
            const randomIndex = Math.floor(Math.random() * dataCopy.length);
            const centroid = { ...dataCopy[randomIndex] };
            centroids.push(centroid);
            dataCopy.splice(randomIndex, 1);
        }
        this.clusters = centroids.map(c => ({ centroid: c, points: [] }));
    }

    private assignClusters(pointAssignments: number[]): boolean {
        let hasChanged = false;

        for (const cluster of this.clusters) {
            cluster.points = [];
        }

        this.data.forEach((point, pointIndex) => {
            let minDistance = Infinity;
            let closestClusterIndex = -1;

            this.clusters.forEach((cluster, clusterIndex) => {
                const distance = KMeans.euclideanDistance(point, cluster.centroid);
                if (distance < minDistance) {
                    minDistance = distance;
                    closestClusterIndex = clusterIndex;
                }
            });

            if (pointAssignments[pointIndex] !== closestClusterIndex) {
                hasChanged = true;
                pointAssignments[pointIndex] = closestClusterIndex;
            }
            
            if (closestClusterIndex !== -1) {
               this.clusters[closestClusterIndex].points.push(point);
            }
        });

        return hasChanged;
    }

    private updateCentroids(): void {
        for (const cluster of this.clusters) {
            if (cluster.points.length === 0) continue;

            const sumX = cluster.points.reduce((sum, p) => sum + p.x, 0);
            const sumY = cluster.points.reduce((sum, p) => sum + p.y, 0);

            cluster.centroid.x = sumX / cluster.points.length;
            cluster.centroid.y = sumY / cluster.points.length;
        }
    }

    public run(): KMeansResult {
        this.initializeCentroids();
        
        const pointAssignments = new Array(this.data.length).fill(-1);

        let iterations = 0;
        let converged = false;
        
        for (let i = 0; i < this.maxIterations; i++) {
            iterations = i + 1;
            const hasChanged = this.assignClusters(pointAssignments);
            this.updateCentroids();

            if (!hasChanged) {
                converged = true;
                break;
            }
        }
        
        return {
            clusters: this.clusters,
            iterations,
            converged
        };
    }
}