"""
ML — Customer segmentation and clustering algorithms.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np
import pandas as pd
from sklearn.cluster import DBSCAN, KMeans
from sklearn.metrics import davies_bouldin_score, silhouette_score
from sklearn.preprocessing import StandardScaler

from bizanalyticsai.core.dataframe import BizFrame


@dataclass
class ClusterResult:
    """Results from a clustering operation."""

    labels:         list[int]
    n_clusters:     int
    silhouette:     float
    davies_bouldin: float
    cluster_sizes:  dict[int, int]
    cluster_stats:  BizFrame
    algorithm:      str

    def summary(self) -> str:
        lines = [
            f"Algorithm: {self.algorithm} | Clusters: {self.n_clusters}",
            f"Silhouette Score: {self.silhouette:.4f} (higher is better, max 1.0)",
            f"Davies-Bouldin: {self.davies_bouldin:.4f} (lower is better)",
            f"Cluster sizes: {self.cluster_sizes}",
        ]
        return "\n".join(lines)


def cluster(
    df: pd.DataFrame,
    cols: list[str] | None = None,
    n_clusters: int = 4,
    algorithm: Literal["kmeans", "dbscan"] = "kmeans",
    scale: bool = True,
    random_state: int = 42,
    **kwargs,
) -> BizFrame:
    """
    Cluster rows into segments.

    Args:
        df:          Input DataFrame.
        cols:        Numeric columns to use. Defaults to all numeric.
        n_clusters:  Number of clusters (KMeans only).
        algorithm:   'kmeans' or 'dbscan'.
        scale:       Standardize features before clustering.
        random_state: Random seed for reproducibility.
        **kwargs:    Extra args passed to the sklearn algorithm.

    Returns:
        BizFrame with a 'cluster' column added.
    """
    num_cols = cols or list(df.select_dtypes(include="number").columns)
    X = df[num_cols].fillna(df[num_cols].median()).values

    if scale:
        X = StandardScaler().fit_transform(X)

    if algorithm == "kmeans":
        model = KMeans(n_clusters=n_clusters, random_state=random_state, n_init="auto", **kwargs)
    elif algorithm == "dbscan":
        eps = kwargs.pop("eps", 0.5)
        min_samples = kwargs.pop("min_samples", 5)
        model = DBSCAN(eps=eps, min_samples=min_samples, **kwargs)
    else:
        raise ValueError(f"Unknown algorithm: {algorithm}")

    labels = model.fit_predict(X)
    out = df.copy()
    out["cluster"] = labels
    return BizFrame(out)


def cluster_full(
    df: pd.DataFrame,
    cols: list[str] | None = None,
    n_clusters: int = 4,
    algorithm: Literal["kmeans", "dbscan"] = "kmeans",
    scale: bool = True,
    random_state: int = 42,
    **kwargs,
) -> ClusterResult:
    """
    Like cluster(), but returns a full ClusterResult with metrics and stats.
    """
    num_cols = cols or list(df.select_dtypes(include="number").columns)
    X_raw = df[num_cols].fillna(df[num_cols].median()).values
    X = StandardScaler().fit_transform(X_raw) if scale else X_raw

    if algorithm == "kmeans":
        model = KMeans(n_clusters=n_clusters, random_state=random_state, n_init="auto", **kwargs)
    else:
        eps = kwargs.pop("eps", 0.5)
        min_samples = kwargs.pop("min_samples", 5)
        model = DBSCAN(eps=eps, min_samples=min_samples, **kwargs)

    labels = model.fit_predict(X)

    valid_mask = labels != -1
    sil   = silhouette_score(X[valid_mask], labels[valid_mask]) if valid_mask.sum() > 1 else 0.0
    db    = davies_bouldin_score(X[valid_mask], labels[valid_mask]) if valid_mask.sum() > 1 else float("inf")

    out = df.copy()
    out["cluster"] = labels

    # Per-cluster statistics
    stats = out.groupby("cluster")[num_cols].agg(["mean", "median", "std", "count"])
    stats.columns = ["_".join(c) for c in stats.columns]

    cluster_sizes = dict(pd.Series(labels).value_counts().sort_index())

    return ClusterResult(
        labels=list(labels),
        n_clusters=len(set(l for l in labels if l != -1)),
        silhouette=round(float(sil), 4),
        davies_bouldin=round(float(db), 4),
        cluster_sizes={int(k): int(v) for k, v in cluster_sizes.items()},
        cluster_stats=BizFrame(stats),
        algorithm=algorithm,
    )


def find_optimal_k(
    df: pd.DataFrame,
    cols: list[str] | None = None,
    k_range: tuple[int, int] = (2, 10),
    scale: bool = True,
) -> BizFrame:
    """
    Find optimal number of clusters using Elbow + Silhouette methods.

    Args:
        df:      Input DataFrame.
        cols:    Numeric columns to use.
        k_range: (min_k, max_k) to evaluate.
        scale:   Standardize features.

    Returns:
        BizFrame with columns: k, inertia, silhouette, davies_bouldin.
    """
    num_cols = cols or list(df.select_dtypes(include="number").columns)
    X = df[num_cols].fillna(df[num_cols].median()).values
    if scale:
        X = StandardScaler().fit_transform(X)

    rows = []
    for k in range(k_range[0], k_range[1] + 1):
        model = KMeans(n_clusters=k, random_state=42, n_init="auto")
        labels = model.fit_predict(X)
        rows.append({
            "k":             k,
            "inertia":       round(float(model.inertia_), 2),
            "silhouette":    round(float(silhouette_score(X, labels)), 4),
            "davies_bouldin": round(float(davies_bouldin_score(X, labels)), 4),
        })

    result = BizFrame(rows)
    best_k = int(result.loc[result["silhouette"].idxmax(), "k"])
    result["recommended"] = result["k"] == best_k
    return result


def rfm_segments(
    df: pd.DataFrame,
    customer_col: str,
    date_col: str,
    revenue_col: str,
    reference_date: str | None = None,
) -> BizFrame:
    """
    RFM (Recency, Frequency, Monetary) segmentation.

    Scores each customer 1–5 on each dimension and assigns a segment label.

    Args:
        df:             Transaction DataFrame.
        customer_col:   Customer ID column.
        date_col:       Transaction date column.
        revenue_col:    Revenue/amount column.
        reference_date: Reference date for recency (default: max date in df).

    Returns:
        BizFrame with recency_days, frequency, monetary, rfm_score, segment.
    """
    out = df.copy()
    out[date_col] = pd.to_datetime(out[date_col])
    ref = pd.Timestamp(reference_date) if reference_date else out[date_col].max()

    rfm = out.groupby(customer_col).agg(
        recency  =(date_col,    lambda x: (ref - x.max()).days),
        frequency=(customer_col,"count"),
        monetary =(revenue_col, "sum"),
    ).reset_index()

    # Score 1–5 (5 = best)
    rfm["r_score"] = pd.qcut(rfm["recency"],   5, labels=[5, 4, 3, 2, 1], duplicates="drop").astype(int)
    rfm["f_score"] = pd.qcut(rfm["frequency"], 5, labels=[1, 2, 3, 4, 5], duplicates="drop").astype(int)
    rfm["m_score"] = pd.qcut(rfm["monetary"],  5, labels=[1, 2, 3, 4, 5], duplicates="drop").astype(int)
    rfm["rfm_score"] = rfm["r_score"] + rfm["f_score"] + rfm["m_score"]

    def _label(row) -> str:
        r, f, m = row["r_score"], row["f_score"], row["m_score"]
        if r >= 4 and f >= 4 and m >= 4: return "Champions"
        if r >= 4 and f >= 3:            return "Loyal Customers"
        if r >= 4 and f <= 2:            return "Recent Customers"
        if r <= 2 and f >= 4:            return "At Risk"
        if r <= 2 and f <= 2:            return "Lost"
        if m >= 4:                       return "Big Spenders"
        return "Potential Loyalists"

    rfm["segment"] = rfm.apply(_label, axis=1)
    return BizFrame(rfm)
