"""Shared split + metric definitions. Every model is scored through here."""
from __future__ import annotations

import numpy as np
import pandas as pd
from sklearn.metrics import average_precision_score, brier_score_loss, roc_auc_score

# Host-disjoint: hosts in buckets >= TEST_HOST_BUCKET never appear in training.
TEST_HOST_BUCKET = 80

# Temporal: a training row's label window must END before the first test anchor,
# so no training label observes the test period.
SPLITS = {
    7: dict(train_anchors=range(0, 10), val_anchors=range(10, 11), test_anchors=range(17, 22)),
    14: dict(train_anchors=range(0, 1), val_anchors=range(1, 2), test_anchors=range(15, 16)),
}

NUMERIC = [
    "price_usd", "calls_30d", "payers_30d", "calls_per_payer", "days_since_last_call",
    "days_since_last_update", "days_listed_observed", "first_seen_censored", "is_new_listing",
    "present_7d_ago", "payers_delta_7d", "calls_delta_7d", "path_depth", "has_path_param",
    "desc_len", "host_listings", "host_active_rate_7d", "host_one_call_per_payer_share",
    "host_median_payers", "host_networks",
]


def split(df: pd.DataFrame, h: int) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    lab = f"paid_{h}d"
    df = df[df[lab] >= 0]
    s = SPLITS[h]
    train_host = df.host_bucket < TEST_HOST_BUCKET
    tr = df[df.anchor_idx.isin(s["train_anchors"]) & train_host]
    va = df[df.anchor_idx.isin(s["val_anchors"]) & train_host]
    te = df[df.anchor_idx.isin(s["test_anchors"]) & ~train_host]
    return tr, va, te


def ece(p: np.ndarray, y: np.ndarray, bins: int = 15) -> float:
    edges = np.linspace(0, 1, bins + 1)
    idx = np.clip(np.digitize(p, edges) - 1, 0, bins - 1)
    return float(sum(abs(p[idx == b].mean() - y[idx == b].mean()) * (idx == b).mean()
                     for b in range(bins) if (idx == b).any()))


def metrics(y: np.ndarray, p: np.ndarray) -> dict:
    y = np.asarray(y).astype(int)
    p = np.clip(np.asarray(p, float), 1e-6, 1 - 1e-6)
    return dict(n=int(len(y)), base_rate=round(float(y.mean()), 4),
                auc=round(float(roc_auc_score(y, p)), 4) if 0 < y.mean() < 1 else None,
                ap=round(float(average_precision_score(y, p)), 4),
                brier=round(float(brier_score_loss(y, p)), 4), ece=round(ece(p, y), 4))


def segments(te: pd.DataFrame) -> dict[str, pd.Series]:
    """Named boolean masks over the test set."""
    return {
        "all": pd.Series(True, index=te.index),
        "new_listings": te.is_new_listing == 1,
        "small_hosts(<=20)": te.host_listings <= 20,
        "stale(>7d since call)": te.days_since_last_call > 7,
    }


def report(te: pd.DataFrame, p: np.ndarray, h: int) -> dict:
    y = te[f"paid_{h}d"].to_numpy()
    p = np.asarray(p)
    out = {}
    for name, m in segments(te).items():
        m = m.to_numpy()
        if m.sum() >= 50 and 0 < y[m].mean() < 1:
            out[name] = metrics(y[m], p[m])
    return out