"""
Cohort — Retention cohort analysis and revenue cohort heatmaps.
"""

from __future__ import annotations

from typing import Literal

import numpy as np
import pandas as pd

from bizanalyticsai.core.dataframe import BizFrame


def cohort_retention(
    df: pd.DataFrame,
    customer_col: str,
    date_col: str,
    period: str = "M",
) -> BizFrame:
    """
    Build a cohort retention table.

    Each row = a cohort (first seen period).
    Each column = period offset (0, 1, 2, …).
    Values = % of original cohort still active.

    Args:
        df:           DataFrame with one row per customer-period activity.
        customer_col: Customer ID column.
        date_col:     Activity date column.
        period:       Grouping period ('M' = month, 'Q' = quarter, 'W' = week).

    Returns:
        BizFrame where index = cohort period, columns = period offsets (0, 1, 2, …).
    """
    out = df[[customer_col, date_col]].copy()
    out[date_col] = pd.to_datetime(out[date_col])
    out["_period"] = out[date_col].dt.to_period(period)

    # First period each customer appeared
    first_period = out.groupby(customer_col)["_period"].min().rename("_cohort")
    out = out.join(first_period, on=customer_col)

    # Period offset
    out["_offset"] = (out["_period"] - out["_cohort"]).apply(lambda x: x.n if hasattr(x, "n") else int(x))

    # Count unique customers per cohort-offset
    cohort_size = out.groupby("_cohort")[customer_col].nunique()
    cohort_data = out.groupby(["_cohort", "_offset"])[customer_col].nunique().unstack(fill_value=0)

    # Retention as % of cohort size
    retention = cohort_data.divide(cohort_size, axis=0) * 100
    retention.columns = [f"Period {c}" if c > 0 else "Period 0 (base)" for c in retention.columns]
    retention.index   = [str(i) for i in retention.index]

    return BizFrame(retention.round(1))


def cohort_revenue(
    df: pd.DataFrame,
    customer_col: str,
    date_col: str,
    revenue_col: str,
    period: str = "M",
    metric: Literal["total", "avg", "cumulative"] = "total",
) -> BizFrame:
    """
    Build a revenue cohort table.

    Args:
        df:           DataFrame with customer_col, date_col, revenue_col.
        customer_col: Customer ID column.
        date_col:     Date column.
        revenue_col:  Revenue column.
        period:       Grouping period.
        metric:       'total', 'avg', or 'cumulative' revenue per cohort-offset.

    Returns:
        BizFrame cohort revenue matrix.
    """
    out = df[[customer_col, date_col, revenue_col]].copy()
    out[date_col] = pd.to_datetime(out[date_col])
    out["_period"] = out[date_col].dt.to_period(period)

    first_period = out.groupby(customer_col)["_period"].min().rename("_cohort")
    out = out.join(first_period, on=customer_col)
    out["_offset"] = (out["_period"] - out["_cohort"]).apply(lambda x: x.n if hasattr(x, "n") else int(x))

    if metric == "total":
        pivot = out.groupby(["_cohort", "_offset"])[revenue_col].sum().unstack(fill_value=0)
    elif metric == "avg":
        pivot = out.groupby(["_cohort", "_offset"])[revenue_col].mean().unstack(fill_value=0)
    elif metric == "cumulative":
        pivot = out.groupby(["_cohort", "_offset"])[revenue_col].sum().unstack(fill_value=0)
        pivot = pivot.cumsum(axis=1)

    pivot.columns = [f"Period {c}" if c > 0 else "Period 0 (base)" for c in pivot.columns]
    pivot.index   = [str(i) for i in pivot.index]

    return BizFrame(pivot.round(2))


def ltv_by_cohort(
    df: pd.DataFrame,
    customer_col: str,
    date_col: str,
    revenue_col: str,
    period: str = "M",
    max_periods: int = 24,
) -> BizFrame:
    """
    Compute average cumulative LTV per cohort across periods.

    Returns BizFrame with cohort | period_N_ltv columns representing
    the average cumulative revenue per customer by month N.
    """
    rev = cohort_revenue(df, customer_col, date_col, revenue_col, period, metric="cumulative")
    cohort_sizes = (
        df.copy()
        .assign(_period=pd.to_datetime(df[date_col]).dt.to_period(period))
        .pipe(lambda d: d.join(d.groupby(customer_col)["_period"].min().rename("_cohort"), on=customer_col))
        .groupby("_cohort")[customer_col].nunique()
    )
    cohort_sizes.index = [str(i) for i in cohort_sizes.index]

    ltv_table = rev.divide(cohort_sizes, axis=0)
    ltv_table = ltv_table.iloc[:, :max_periods]
    return BizFrame(ltv_table.round(2))


def cohort_summary(
    df: pd.DataFrame,
    customer_col: str,
    date_col: str,
    revenue_col: str | None = None,
    period: str = "M",
) -> dict:
    """
    High-level cohort summary statistics.

    Returns:
        Dict with avg_retention_1m, avg_retention_3m, avg_retention_12m,
        best_cohort, worst_cohort, avg_ltv_12m.
    """
    retention = cohort_retention(df, customer_col, date_col, period)
    cols = list(retention.columns)

    def _avg_at(offset: int) -> float:
        col = f"Period {offset}"
        if col not in cols:
            return np.nan
        return float(retention[col].mean())

    avg_1m  = _avg_at(1)
    avg_3m  = _avg_at(3)
    avg_12m = _avg_at(12)

    col_1 = "Period 1"
    if col_1 in cols:
        best_cohort  = str(retention[col_1].idxmax())
        worst_cohort = str(retention[col_1].idxmin())
    else:
        best_cohort = worst_cohort = "N/A"

    result: dict = {
        "avg_retention_1m":  round(avg_1m, 1) if not np.isnan(avg_1m) else None,
        "avg_retention_3m":  round(avg_3m, 1) if not np.isnan(avg_3m) else None,
        "avg_retention_12m": round(avg_12m, 1) if not np.isnan(avg_12m) else None,
        "best_cohort":       best_cohort,
        "worst_cohort":      worst_cohort,
        "n_cohorts":         len(retention),
    }

    if revenue_col:
        ltv = ltv_by_cohort(df, customer_col, date_col, revenue_col, period)
        if "Period 12" in ltv.columns:
            result["avg_ltv_12m"] = round(float(ltv["Period 12"].mean()), 2)

    return result
