# targeting.py
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score, average_precision_score

def build_dataset(users: pd.DataFrame, adopters: set[str]) -> tuple[pd.DataFrame, pd.Series]:
    """users has columns: user_id, age, logins_30d, features_used, tenure_days, team_size, industry.
       adopters is a set of user_ids that already converted."""
    df = users.copy()
    df["label"] = df["user_id"].isin(adopters).astype(int)
    y = df.pop("label")
    X = pd.get_dummies(df.drop(columns=["user_id"]), columns=["industry"], drop_first=True)
    return X, y

def train(X: pd.DataFrame, y: pd.Series):
    X_tr, X_val, y_tr, y_val = train_test_split(
        X, y, test_size=0.25, stratify=y, random_state=42
    )
    scaler = StandardScaler().fit(X_tr)
    X_tr_s = scaler.transform(X_tr); X_val_s = scaler.transform(X_val)

    # Baseline — logistic regression with class_weight for imbalance
    lr = LogisticRegression(class_weight="balanced", max_iter=500).fit(X_tr_s, y_tr)
    # Stronger — gradient boosting handles non-linear interactions
    gb = GradientBoostingClassifier(n_estimators=200, max_depth=3).fit(X_tr, y_tr)

    for name, model, X_eval in [("logreg", lr, X_val_s), ("gbm", gb, X_val)]:
        p = model.predict_proba(X_eval)[:, 1]
        print(f"{name}: AUC={roc_auc_score(y_val, p):.3f}  "
              f"PR-AUC={average_precision_score(y_val, p):.3f}")
    return gb, scaler

def score_and_rank(model, users: pd.DataFrame, adopters: set[str], top_k: int = 1000):
    """Score all non-adopters and return the top-K targets with lift."""
    non_adopters = users[~users["user_id"].isin(adopters)].copy()
    X = pd.get_dummies(non_adopters.drop(columns=["user_id"]),
                       columns=["industry"], drop_first=True)
    non_adopters["score"] = model.predict_proba(X)[:, 1]
    ranked = non_adopters.sort_values("score", ascending=False)

    base_rate = len(adopters) / len(users)
    top = ranked.head(top_k)
    # Lift = model's positive rate in top-K / random base rate
    # (true labels not known for non-adopters — use held-out to measure lift in practice)
    print(f"Base adoption rate: {base_rate:.3%} | Targeting top {top_k} users")
    return ranked[["user_id", "score"]]

# Split by signup date, not randomly
cutoff = "2025-06-01"
train_users = users[users["signup_date"] < cutoff]
test_users  = users[users["signup_date"] >= cutoff]

# Adopters in each cohort
train_adopters = adopter_events.query("event_date < @cutoff")["user_id"].unique()
test_adopters  = adopter_events.query("event_date >= @cutoff")["user_id"].unique()

def lift_chart(y_true, y_score, deciles=10):
    df = pd.DataFrame({"y": y_true, "p": y_score}).sort_values("p", ascending=False)
    df["decile"] = pd.qcut(df["p"].rank(method="first"), deciles, labels=False)
    base = df["y"].mean()
    table = df.groupby("decile")["y"].mean().rename("rate").to_frame()
    table["lift"] = table["rate"] / base
    return table.sort_index(ascending=False)