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

logger = CodedFlowsLogger(name="Clu. HDBSCAN", level=logging.INFO)


def _validate_numerical_data(data):
    """
    Validates if the input data contains only numerical (integer, float) or boolean values.
    """
    if sparse.issparse(data):
        if not (
            np.issubdtype(data.dtype, np.number) or np.issubdtype(data.dtype, np.bool_)
        ):
            raise TypeError(
                f"Sparse matrix contains unsupported data type: {data.dtype}."
            )
        return
    if isinstance(data, np.ndarray):
        if not (
            np.issubdtype(data.dtype, np.number) or np.issubdtype(data.dtype, np.bool_)
        ):
            raise TypeError(
                f"NumPy array contains unsupported data type: {data.dtype}."
            )
        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.")
        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.")
        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: {type(data)}")


def clu_hdbscan(
    data: DataFrame, options=None
) -> Tuple[DataFrame, Union[DataFrame, Dict]]:
    options = options or {}
    verbose = options.get("verbose", True)
    min_cluster_size = options.get("min_cluster_size", 5)
    min_samples = options.get("min_samples", 5)
    cluster_selection_epsilon = options.get("cluster_selection_epsilon", 0.0)
    metric = options.get("metric", "euclidean").lower()
    cluster_selection_method = options.get("cluster_selection_method", "eom").lower()
    algorithm = options.get("algorithm", "auto").lower().replace("-", "_")
    allow_single_cluster = options.get("allow_single_cluster", False)
    cluster_col_name = options.get("cluster_column", "cluster_id")
    metrics_as = options.get("metrics_as", "Dataframe")
    AUTOMATED_THRESHOLD = 10000
    Clustered_data = None
    Metrics = None
    try:
        verbose and logger.info(
            f"Initializing Clu. HDBSCAN (min_cluster={min_cluster_size}, min_samples={min_samples})."
        )
        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 numerical requirements...")
        _validate_numerical_data(data)
        n_samples = data.shape[0]
        verbose and logger.info(f"Processing {n_samples} samples.")
        Model = HDBSCAN(
            min_cluster_size=min_cluster_size,
            min_samples=min_samples,
            cluster_selection_epsilon=cluster_selection_epsilon,
            metric=metric,
            cluster_selection_method=cluster_selection_method,
            algorithm=algorithm,
            allow_single_cluster=allow_single_cluster,
            copy=True,
            n_jobs=-1,
        )
        verbose and logger.info("Fitting model and predicting labels...")
        labels = Model.fit_predict(data)
        n_clusters_found = len(set(labels)) - (1 if -1 in labels else 0)
        n_noise = list(labels).count(-1)
        verbose and logger.info(
            f"Found {n_clusters_found} clusters. Noise points: {n_noise}."
        )
        unique_labels = set(labels)
        if len(unique_labels) < 2:
            verbose and logger.warning(
                "Not enough clusters (or only noise) found to calculate metrics. Returning NaNs."
            )
            (s_score, ch_score, db_score) = (np.nan, np.nan, np.nan)
        else:
            verbose and logger.info("Computing Intrinsic metrics.")
            if n_samples > AUTOMATED_THRESHOLD:
                s_score = silhouette_score(
                    data, labels, sample_size=min(n_samples, 20000), random_state=42
                )
                verbose and logger.info(f"Silhouette Score  : {s_score:.4f} (Sampled)")
            else:
                s_score = silhouette_score(data, labels)
                verbose and logger.info(f"Silhouette Score  : {s_score:.4f}")
            ch_score = calinski_harabasz_score(data, labels)
            verbose and logger.info(f"Calinski-Harabasz : {ch_score:.4f}")
            db_score = davies_bouldin_score(data, 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] = labels
            verbose and logger.info("Results assigned in-place to Pandas DataFrame.")
        elif is_polars:
            Clustered_data = data.with_columns(
                pl.Series(name=cluster_col_name, values=labels)
            )
            verbose and logger.info("Results attached to Polars DataFrame.")
    except Exception as e:
        verbose and logger.error(f"Clu. HDBSCAN operation failed: {e}")
        raise
    return (Clustered_data, Metrics)