"""
This module provides the Clustering operator that operate on embeddings.

Operators:
- `Clustering`: Performs the clustering on examples in embedding space.

NOTE: The above operators take embeddings as input. Refer uptrain.operators.language.embedding and uptrain.operators.embs to learn more.

"""

from __future__ import annotations
import typing as t

from loguru import logger
import numpy as np
import polars as pl

if t.TYPE_CHECKING:
    from uptrain.framework import Settings
from uptrain.operators.base import (
    ColumnOp,
    register_op,
    TYPE_TABLE_OUTPUT,
)
from uptrain.utilities import lazy_load_dep

nltk = lazy_load_dep("nltk", "nltk")


@register_op
class Clustering(ColumnOp):
    """
    Operator for performing Kmeans clustering.

    Attributes:
        algorithm (str): algorithm to perform clustering.
        n_clusters (int): number of clusters.
        col_in (str):  The name of the column in the DataFrame where clustering needs to be performed.
        col_out (str):  The name of the column in the DataFrame to output the assigned cluster index.
        col_out_dist (str): The name of the column in the DataFrame to output the euclidean distance from its cluster centroid.
        col_aggs (list[str]): Optional, can be used to specify name of columns to aggregate by and run individual clustering, ex: if you want separate clustering for seperate organisations
    Example:
        ```
        import polars as pl
        from uptrain.operators import Clustering

        # Create an instance of the UMAP operator
        op = Clustering(
                algorithm="kmeans",
                n_clusters=40,
                col_in = 'umap_embedding',
                col_out = 'cluster_index'
            )

        # Set up the operator
        op.setup()

        # Run the operator on the input data
        input_data = pl.DataFrame(...)
        output = op.run(input_data)

        # Get the output DataFrame
        assigned_clusters = output["output"]
        ```
    """

    algorithm: str = "kmeans"
    n_clusters: int = 40
    col_in: str = "umap_embedding"
    col_out: str = "cluster_index"
    col_out_dist: str = "cluster_index_distance"
    col_aggs: list[str] = []
    min_samples_each_cluster: int = 50

    def setup(self, settings: Settings):
        return self

    def run(self, data: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
        self.cluster_centroids = {}
        unique_agg_keys = ["default"]

        if len(self.col_aggs):
            # First aggregate by col_aggs
            agg_data = data.groupby(self.col_aggs).agg(
                [pl.col(self.col_in).count().alias("num_rows_" + self.col_in)]
            )
            agg_data = agg_data.drop("num_rows_" + self.col_in)
            unique_agg_keys = agg_data.to_dicts()

        res_data_arr = []
        for unique_agg_key in unique_agg_keys:
            cond = True
            if isinstance(unique_agg_key, dict):
                unique_agg_key = dict(
                    [(key, unique_agg_key[key]) for key in sorted(unique_agg_key)]
                )
                for key, val in unique_agg_key.items():
                    cond = cond & (data[key] == val)
            data_subset = data.filter(cond)

            n_clusters = max(
                1,
                min(
                    self.n_clusters,
                    int(len(data_subset) / self.min_samples_each_cluster),
                ),
            )

            if self.algorithm == "kmeans":
                algorithm_obj = nltk.cluster.KMeansClusterer(
                    n_clusters,
                    distance=nltk.cluster.util.cosine_distance,
                    avoid_empty_clusters=True,
                )
            else:
                raise Exception(f"{self.algorithm} is not supported yet.")

            embeddings = np.asarray(data_subset[self.col_in])
            assigned_clusters = algorithm_obj.cluster(embeddings, assign_clusters=True)
            scores = []
            for index in range(len(embeddings)):
                scores.append(
                    np.linalg.norm(
                        embeddings[index]
                        - algorithm_obj.means()[assigned_clusters[index]]
                    )
                )

            data_subset = data_subset.with_columns(
                [
                    pl.Series(assigned_clusters).alias(self.col_out),
                    pl.Series(scores).alias(self.col_out_dist),
                    pl.Series([str(unique_agg_key)] * len(data_subset)).alias(
                        "_unique_agg_key_for_clustering"
                    ),
                ]
            )
            res_data_arr.append(data_subset)
            unique_assigned_clusters = np.unique(np.array(assigned_clusters))
            all_means = [
                list(x.astype(np.float64)) for x in list(algorithm_obj.means())
            ]
            assigned_means = []
            for clus_idx in unique_assigned_clusters:
                assigned_means.append([round(y, 8) for y in all_means[clus_idx]])
            self.cluster_centroids[str(unique_agg_key)] = assigned_means

        return {"output": pl.concat(res_data_arr)}
