"""Customer and sales analytics models.

Provides standard e-commerce analytics wrapped as functions:
    - RFM segmentation (Recency, Frequency, Monetary)
    - Customer LTV estimation
    - Cohort retention analysis
    - Revenue decomposition

Uses scikit-learn for clustering and polars for data processing.

Usage via provider:
    from platoon.learning.providers import get_provider
    analytics = get_provider("analytics")
    segments = analytics.rfm_segment(orders)

Or direct:
    from platoon.learning.analytics import rfm_segment, customer_ltv
"""

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple

from platoon.learning.providers import ModelProvider, register_provider


# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------


@dataclass
class RFMResult:
    """Output of RFM segmentation."""

    customers: List[Dict[str, Any]]  # Each has: customer_id, recency, frequency, monetary, segment, rfm_score
    segments: Dict[str, int]  # segment_name -> count
    n_clusters: int
    method: str


@dataclass
class LTVResult:
    """Output of customer lifetime value estimation."""

    customers: List[Dict[str, Any]]  # Each has: customer_id, historical_ltv, predicted_ltv, tenure_days
    avg_ltv: float
    median_ltv: float
    total_ltv: float


@dataclass
class CohortResult:
    """Output of cohort retention analysis."""

    retention_matrix: List[List[float]]  # [cohort_index][period] = retention %
    cohort_labels: List[str]  # e.g. ["2024-01", "2024-02", ...]
    period_labels: List[str]  # e.g. ["M+0", "M+1", ...]
    cohort_sizes: List[int]


# ---------------------------------------------------------------------------
# Implementations
# ---------------------------------------------------------------------------


def rfm_segment(
    orders: List[Dict],
    customer_id_key: str = "customer_id",
    date_key: str = "order_date",
    amount_key: str = "line_total",
    reference_date: Optional[str] = None,
    n_clusters: int = 4,
) -> RFMResult:
    """Segment customers by Recency, Frequency, Monetary value.

    Computes RFM scores per customer, then clusters with KMeans.

    Args:
        orders: List of order dicts (flat — one row per line item)
        customer_id_key: Key for customer identifier
        date_key: Key for order date (ISO string)
        amount_key: Key for monetary value
        reference_date: "now" for RFM recency calc (default: max date in data)
        n_clusters: Number of KMeans clusters (default 4)

    Returns:
        RFMResult with per-customer RFM scores and segment labels
    """
    import polars as pl
    from sklearn.cluster import KMeans
    from sklearn.preprocessing import StandardScaler

    df = pl.DataFrame(orders)
    df = df.with_columns(pl.col(date_key).str.to_date().alias("date"))
    df = df.with_columns(pl.col(amount_key).cast(pl.Float64).alias("amount"))

    if reference_date:
        ref = datetime.fromisoformat(reference_date)
    else:
        ref = df["date"].max()

    # Compute RFM per customer
    rfm = (
        df.group_by(customer_id_key)
        .agg(
            recency=((pl.lit(ref) - pl.col("date").max()).dt.total_days()),
            frequency=pl.col("date").n_unique(),
            monetary=pl.col("amount").sum(),
        )
    )

    rfm_data = rfm.to_dicts()
    if not rfm_data:
        return RFMResult(customers=[], segments={}, n_clusters=0, method="kmeans")

    # Cluster with KMeans
    features = [[r["recency"], r["frequency"], r["monetary"]] for r in rfm_data]
    scaler = StandardScaler()
    scaled = scaler.fit_transform(features)

    n_clusters = min(n_clusters, len(rfm_data))
    km = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
    labels = km.fit_predict(scaled)

    # Name segments by cluster characteristics
    segment_names = _name_rfm_segments(rfm_data, labels, n_clusters)

    customers = []
    for i, row in enumerate(rfm_data):
        customers.append({
            "customer_id": row[customer_id_key],
            "recency": row["recency"],
            "frequency": row["frequency"],
            "monetary": round(row["monetary"], 2),
            "segment": segment_names[labels[i]],
            "cluster": int(labels[i]),
        })

    segments = {}
    for c in customers:
        seg = c["segment"]
        segments[seg] = segments.get(seg, 0) + 1

    return RFMResult(
        customers=customers,
        segments=segments,
        n_clusters=n_clusters,
        method="kmeans",
    )


def _name_rfm_segments(
    rfm_data: List[Dict], labels, n_clusters: int
) -> Dict[int, str]:
    """Assign descriptive names to clusters based on RFM means."""
    cluster_stats = {}
    for i in range(n_clusters):
        members = [rfm_data[j] for j in range(len(rfm_data)) if labels[j] == i]
        if not members:
            cluster_stats[i] = {"recency": 999, "frequency": 0, "monetary": 0}
            continue
        cluster_stats[i] = {
            "recency": sum(m["recency"] for m in members) / len(members),
            "frequency": sum(m["frequency"] for m in members) / len(members),
            "monetary": sum(m["monetary"] for m in members) / len(members),
        }

    # Rank clusters
    ranked = sorted(cluster_stats.items(),
                    key=lambda x: (-x[1]["monetary"], -x[1]["frequency"], x[1]["recency"]))

    name_pool = ["Champions", "Loyal", "Potential", "At Risk", "Hibernating", "Lost"]
    names = {}
    for rank, (cluster_id, _) in enumerate(ranked):
        names[cluster_id] = name_pool[min(rank, len(name_pool) - 1)]
    return names


def customer_ltv(
    orders: List[Dict],
    customer_id_key: str = "customer_id",
    date_key: str = "order_date",
    amount_key: str = "line_total",
    reference_date: Optional[str] = None,
    projection_days: int = 365,
) -> LTVResult:
    """Estimate customer lifetime value.

    Simple method: LTV = (avg_order_value * purchase_frequency) * projected_lifespan.
    Uses historical data to compute per-customer metrics, then projects forward.

    Args:
        orders: Flat order records
        customer_id_key: Customer ID field
        date_key: Order date field
        amount_key: Order amount field
        reference_date: "now" for calculations
        projection_days: Days to project LTV forward

    Returns:
        LTVResult with per-customer LTV estimates
    """
    import polars as pl

    df = pl.DataFrame(orders)
    df = df.with_columns(
        pl.col(date_key).str.to_date().alias("date"),
        pl.col(amount_key).cast(pl.Float64).alias("amount"),
    )

    if reference_date:
        ref = datetime.fromisoformat(reference_date)
    else:
        ref = df["date"].max()

    # Per-customer stats
    stats = (
        df.group_by(customer_id_key)
        .agg(
            total_spent=pl.col("amount").sum(),
            order_count=pl.col("date").n_unique(),
            first_order=pl.col("date").min(),
            last_order=pl.col("date").max(),
        )
    ).to_dicts()

    customers = []
    ltvs = []

    for row in stats:
        tenure = (ref - row["first_order"]).days
        tenure = max(tenure, 1)
        order_count = row["order_count"]
        total = row["total_spent"]

        avg_order_value = total / max(order_count, 1)
        purchase_rate = order_count / (tenure / 365)  # orders per year
        historical_ltv = total
        predicted_ltv = avg_order_value * purchase_rate * (projection_days / 365)

        customers.append({
            "customer_id": row[customer_id_key],
            "historical_ltv": round(historical_ltv, 2),
            "predicted_ltv": round(predicted_ltv, 2),
            "avg_order_value": round(avg_order_value, 2),
            "purchase_rate_annual": round(purchase_rate, 2),
            "tenure_days": tenure,
            "order_count": order_count,
        })
        ltvs.append(predicted_ltv)

    ltvs.sort()
    avg = sum(ltvs) / max(len(ltvs), 1)
    median = ltvs[len(ltvs) // 2] if ltvs else 0

    return LTVResult(
        customers=customers,
        avg_ltv=round(avg, 2),
        median_ltv=round(median, 2),
        total_ltv=round(sum(ltvs), 2),
    )


def cohort_retention(
    orders: List[Dict],
    customer_id_key: str = "customer_id",
    date_key: str = "order_date",
    period: str = "month",
) -> CohortResult:
    """Compute cohort retention matrix.

    Groups customers by their first purchase month (cohort), then tracks
    what percentage return in each subsequent month.

    Args:
        orders: Flat order records
        customer_id_key: Customer ID field
        date_key: Order date field (ISO string)
        period: Cohort period — "month" or "week"

    Returns:
        CohortResult with retention matrix and labels
    """
    import polars as pl

    df = pl.DataFrame(orders)
    df = df.with_columns(pl.col(date_key).str.to_date().alias("date"))

    if period == "month":
        df = df.with_columns(
            pl.col("date").dt.truncate("1mo").alias("period")
        )
    else:
        df = df.with_columns(
            pl.col("date").dt.truncate("1w").alias("period")
        )

    # First purchase period per customer
    first_purchase = (
        df.group_by(customer_id_key)
        .agg(cohort=pl.col("period").min())
    )

    df = df.join(first_purchase, on=customer_id_key)

    # Unique customer-period pairs
    activity = df.select(customer_id_key, "cohort", "period").unique()

    # Compute period offset
    if period == "month":
        activity = activity.with_columns(
            ((pl.col("period") - pl.col("cohort")).dt.total_days() / 30).cast(pl.Int32).alias("period_offset")
        )
    else:
        activity = activity.with_columns(
            ((pl.col("period") - pl.col("cohort")).dt.total_days() / 7).cast(pl.Int32).alias("period_offset")
        )

    # Build retention counts
    cohort_groups = activity.group_by("cohort", "period_offset").agg(
        customers=pl.col(customer_id_key).n_unique()
    ).sort("cohort", "period_offset")

    # Cohort sizes (period_offset == 0)
    cohort_sizes_df = cohort_groups.filter(pl.col("period_offset") == 0)

    cohorts = sorted(cohort_sizes_df["cohort"].to_list())
    max_offset = activity["period_offset"].max() if len(activity) > 0 else 0

    # Build matrix
    cohort_size_map = {}
    for row in cohort_sizes_df.to_dicts():
        cohort_size_map[row["cohort"]] = row["customers"]

    counts_map = {}
    for row in cohort_groups.to_dicts():
        counts_map[(row["cohort"], row["period_offset"])] = row["customers"]

    retention_matrix = []
    cohort_labels = []
    cohort_sizes = []

    for cohort in cohorts:
        size = cohort_size_map.get(cohort, 0)
        cohort_sizes.append(size)
        cohort_labels.append(str(cohort))
        row = []
        for offset in range(max_offset + 1):
            count = counts_map.get((cohort, offset), 0)
            pct = round(count / size * 100, 1) if size > 0 else 0.0
            row.append(pct)
        retention_matrix.append(row)

    prefix = "M" if period == "month" else "W"
    period_labels = [f"{prefix}+{i}" for i in range(max_offset + 1)]

    return CohortResult(
        retention_matrix=retention_matrix,
        cohort_labels=cohort_labels,
        period_labels=period_labels,
        cohort_sizes=cohort_sizes,
    )


# ---------------------------------------------------------------------------
# Analytics provider
# ---------------------------------------------------------------------------


@register_provider
class AnalyticsProvider(ModelProvider):
    """Customer and sales analytics — sklearn + polars backed."""

    name = "analytics_provider"
    domain = "analytics"
    backend = "sklearn"

    @classmethod
    def is_available(cls) -> bool:
        try:
            import sklearn
            import polars
            return True
        except ImportError:
            return False

    def rfm_segment(self, orders: List[Dict], **kwargs) -> RFMResult:
        return rfm_segment(orders, **kwargs)

    def customer_ltv(self, orders: List[Dict], **kwargs) -> LTVResult:
        return customer_ltv(orders, **kwargs)

    def cohort_retention(self, orders: List[Dict], **kwargs) -> CohortResult:
        return cohort_retention(orders, **kwargs)
