import logging
import pandas as pd
import polars as pl
import numpy as np
from scipy import sparse
from sklearn.cluster import KMeans, MiniBatchKMeans
from sklearn.metrics import (
    silhouette_score,
    calinski_harabasz_score,
    davies_bouldin_score,
)
from coded_flows.types import DataFrame, Tuple, Any, Union, Dict
from coded_flows.utils import CodedFlowsLogger

logger = CodedFlowsLogger(name="Clu. K-Means", level=logging.INFO)


def _validate_numerical_data(data):
    """
    Validates if the input data contains only numerical or boolean values.
    """
    if sparse.issparse(data):
        if not (
            np.issubdtype(data.dtype, np.number) or np.issubdtype(data.dtype, np.bool_)
        ):
            raise TypeError("Sparse matrix contains unsupported data type.")
        return
    if isinstance(data, np.ndarray):
        if not (
            np.issubdtype(data.dtype, np.number) or np.issubdtype(data.dtype, np.bool_)
        ):
            raise TypeError("NumPy array contains unsupported data type.")
        return
    if isinstance(data, (pd.DataFrame, pd.Series)):
        if isinstance(data, pd.Series):
            if not (
                pd.api.types.is_numeric_dtype(data) or pd.api.types.is_bool_dtype(data)
            ):
                raise TypeError(
                    f"Pandas Series '{data.name}' is not numerical or boolean."
                )
        else:
            numeric_cols = data.select_dtypes(include=["number", "bool"])
            if numeric_cols.shape[1] != data.shape[1]:
                invalid_cols = list(set(data.columns) - set(numeric_cols.columns))
                raise TypeError(
                    f"Pandas DataFrame contains non-numerical columns: {invalid_cols}."
                )
        return
    if isinstance(data, (pl.DataFrame, pl.Series)):
        if isinstance(data, pl.Series):
            if not (data.dtype.is_numeric() or data.dtype == pl.Boolean):
                raise TypeError(
                    f"Polars Series '{data.name}' is not numerical or boolean."
                )
        else:
            for col_name, dtype in zip(data.columns, data.dtypes):
                if not (dtype.is_numeric() or dtype == pl.Boolean):
                    raise TypeError(
                        f"Polars DataFrame contains non-numerical column: '{col_name}'."
                    )
        return
    raise ValueError(f"Unsupported data type for validation: {type(data)}")


def clu_kmeans(
    data: DataFrame, options=None
) -> Tuple[Any, DataFrame, Union[DataFrame, Dict]]:
    options = options or {}
    verbose = options.get("verbose", True)
    n_clusters = options.get("n_clusters", 2)
    cluster_col_name = options.get("cluster_column", "cluster_id")
    metrics_as = options.get("metrics_as", "Dataframe")
    random_state = options.get("random_state", 42)
    AUTOMATED_THRESHOLD = 10000
    Clustered_data = None
    Model = None
    Metrics = None
    try:
        verbose and logger.info(f"Starting Smart K-Means with K={n_clusters}.")
        is_pandas = isinstance(data, pd.DataFrame)
        is_polars = isinstance(data, pl.DataFrame)
        if not (is_pandas or is_polars):
            raise ValueError("Input data must be a pandas or polars DataFrame.")
        verbose and logger.info("Validating input data types...")
        _validate_numerical_data(data)
        n_samples = data.shape[0]
        verbose and logger.info(f"Dataset contains {n_samples} samples.")
        if n_samples >= AUTOMATED_THRESHOLD:
            verbose and logger.info(
                f"Sample size >= {AUTOMATED_THRESHOLD}. Using MiniBatchKMeans."
            )
            Model = MiniBatchKMeans(
                n_clusters=n_clusters,
                random_state=random_state,
                batch_size=256,
                n_init="auto",
            )
        else:
            verbose and logger.info(
                f"Sample size < {AUTOMATED_THRESHOLD}. Using standard KMeans."
            )
            Model = KMeans(
                n_clusters=n_clusters, random_state=random_state, n_init="auto"
            )
        verbose and logger.info("Fitting model...")
        cluster_labels = Model.fit_predict(data)
        metric_names = []
        metric_values = []
        verbose and logger.info("Computing Intrinsic metrics.")
        if n_samples > AUTOMATED_THRESHOLD:
            s_score = silhouette_score(
                data,
                cluster_labels,
                sample_size=min(n_samples, 20000),
                random_state=random_state,
            )
            verbose and logger.info(f"Silhouette Score  : {s_score:.4f} (Sampled)")
        else:
            s_score = silhouette_score(data, cluster_labels)
            verbose and logger.info(f"Silhouette Score  : {s_score:.4f}")
        ch_score = calinski_harabasz_score(data, cluster_labels)
        verbose and logger.info(f"Calinski-Harabasz : {ch_score:.4f}")
        db_score = davies_bouldin_score(data, cluster_labels)
        verbose and logger.info(f"Davies-Bouldin    : {db_score:.4f}")
        metric_names = ["Silhouette Score", "Calinski-Harabasz", "Davies-Bouldin"]
        metric_values = [s_score, ch_score, db_score]
        if metrics_as == "Dataframe":
            Metrics = pd.DataFrame({"Metric": metric_names, "Value": metric_values})
        else:
            Metrics = dict(zip(metric_names, metric_values))
        if is_pandas:
            Clustered_data = data.copy()
            Clustered_data[cluster_col_name] = cluster_labels
            verbose and logger.info("Appended clusters to Pandas DataFrame.")
        elif is_polars:
            Clustered_data = data.with_columns(
                pl.Series(name=cluster_col_name, values=cluster_labels)
            )
            verbose and logger.info("Appended clusters to Polars DataFrame.")
    except Exception as e:
        verbose and logger.error(f"Error during K-Means process: {e}")
        raise
    return (Model, Clustered_data, Metrics)