SupplyAgent/supply_infra/scheduler/jobs/demand_pool/tree_weight.py

"""
类目树节点热度计算。

六维（ext_pop / plat_sust_pop / plat_ly_pop / recent_pop / real_rov_7d / real_vov_7d）各自独立；
后两维存 rov_diff / vov_diff（相对全局基线）：
1. 挂载点：将需求词级 avg/count 聚合为挂载点 avg/count
2. 树上任意节点：取其子树内全部「有数据」挂载点，
   score = sum(avg * count) / sum(count)，并保留 avg 与 count 供上层聚合
"""
from __future__ import annotations

import logging
from collections import defaultdict
from dataclasses import dataclass, field
from decimal import Decimal
from types import SimpleNamespace
from typing import Any

from supply_infra.db.repositories.category_tree_weight_repo import (
    CategoryTreeWeightRepository,
)
from supply_infra.db.repositories.demand_belong_category_repo import (
    DemandBelongCategoryRepository,
)
from supply_infra.db.repositories.demand_popularity_stats_repo import (
    DemandPopularityStatsRepository,
)
from supply_infra.db.repositories.global_tree_category_repo import (
    GlobalTreeCategoryRepository,
)
from supply_infra.scoring.ranking import rank_to_scores
from supply_infra.db.session import get_session

logger = logging.getLogger(__name__)

METRIC_KEYS: tuple[str, ...] = (
    "ext_pop",
    "plat_sust_pop",
    "plat_ly_pop",
    "recent_pop",
    "real_rov_7d",
    "real_vov_7d",
)

POP_DIM_KEYS: tuple[str, ...] = (
    "ext_pop",
    "plat_sust_pop",
    "plat_ly_pop",
    "recent_pop",
)

_UPSERT_BATCH = 500
_RANK_UPDATE_BATCH = 500


@dataclass
class DimStats:
    avg: float = 0.0
    count: int = 0


@dataclass
class HangFeatures:
    """挂载点（有需求词挂在该树上的节点）聚合特征。"""

    category_id: int
    word_count: int = 0
    dims: dict[str, DimStats] = field(default_factory=dict)


def _dec(value: float, places: int = 8) -> Decimal:
    return Decimal(str(round(float(value), places)))


def _optional_dec(value: float | None, places: int = 8) -> Decimal | None:
    if value is None:
        return None
    return _dec(value, places)


def _normalize_parent_id(parent_id: int | None) -> int | None:
    if parent_id is None or parent_id == 0:
        return None
    return parent_id


def _build_tree_index(
    categories: list[tuple[int, int | None]],
) -> tuple[set[int], dict[int | None, list[int]], dict[int, int | None]]:
    """categories: (category_id, parent_id) 列表。"""
    by_id: set[int] = set()
    children: dict[int | None, list[int]] = defaultdict(list)
    parent_of: dict[int, int | None] = {}
    for cid, parent_id in categories:
        by_id.add(cid)
        pid = _normalize_parent_id(parent_id)
        parent_of[cid] = pid
        children[pid].append(cid)
    for kids in children.values():
        kids.sort()
    return by_id, children, parent_of


def _aggregate_hang_features(
    belong_rows: list[tuple[int, int]],
    stats_by_belong_id: dict[int, Any],
) -> dict[int, HangFeatures]:
    """将词级 avg/count 按挂载类目聚合为挂载点 avg/count。"""
    acc: dict[int, dict[str, dict[str, float]]] = defaultdict(
        lambda: {dim: {"wsum": 0.0, "nsum": 0.0} for dim in METRIC_KEYS}
    )
    word_counts: dict[int, int] = defaultdict(int)

    for belong_id, category_id in belong_rows:
        word_counts[category_id] += 1
        stats = stats_by_belong_id.get(belong_id)
        if stats is None:
            continue
        for dim in METRIC_KEYS:
            n = int(getattr(stats, f"{dim}_count") or 0)
            if n <= 0:
                continue
            avg_raw = getattr(stats, f"{dim}_avg", None)
            if avg_raw is None:
                continue
            avg = float(avg_raw)
            cell = acc[category_id][dim]
            cell["wsum"] += avg * n
            cell["nsum"] += n

    features: dict[int, HangFeatures] = {}
    for category_id, wc in word_counts.items():
        dims: dict[str, DimStats] = {}
        for dim in METRIC_KEYS:
            cell = acc[category_id][dim]
            n = int(cell["nsum"])
            if n > 0:
                dims[dim] = DimStats(avg=cell["wsum"] / cell["nsum"], count=n)
            else:
                dims[dim] = DimStats()
        features[category_id] = HangFeatures(
            category_id=category_id,
            word_count=wc,
            dims=dims,
        )
    return features


def _subtree_hanging_ids(
    node_id: int,
    children: dict[int | None, list[int]],
    hang_ids: set[int],
    cache: dict[int, list[int]],
) -> list[int]:
    if node_id in cache:
        return cache[node_id]
    ids: list[int] = []
    if node_id in hang_ids:
        ids.append(node_id)
    for child in children.get(node_id, []):
        ids.extend(_subtree_hanging_ids(child, children, hang_ids, cache))
    cache[node_id] = ids
    return ids


def _rollup_avg_count(
    by_id: set[int],
    children: dict[int | None, list[int]],
    hangs: dict[int, HangFeatures],
) -> dict[int, dict[str, DimStats]]:
    """
    每个节点取其子树内全部有数据挂载点：
    avg = sum(avg * count) / sum(count)，count = sum(count)。
    """
    hang_ids = set(hangs)
    subtree_cache: dict[int, list[int]] = {}
    result: dict[int, dict[str, DimStats]] = {}

    for nid in by_id:
        hanging = _subtree_hanging_ids(nid, children, hang_ids, subtree_cache)
        dim_stats: dict[str, DimStats] = {}
        for dim in METRIC_KEYS:
            wsum = 0.0
            nsum = 0
            for hid in hanging:
                ds = hangs[hid].dims[dim]
                if ds.count <= 0:
                    continue
                wsum += ds.avg * ds.count
                nsum += ds.count
            if nsum > 0:
                dim_stats[dim] = DimStats(avg=wsum / nsum, count=nsum)
            else:
                dim_stats[dim] = DimStats()
        result[nid] = dim_stats
    return result


def _row_from_state(
    nid: int,
    biz_dt: str,
    hung_word_count: int,
    dim_stats: dict[str, DimStats],
) -> dict[str, Any]:
    row: dict[str, Any] = {
        "category_id": nid,
        "biz_dt": biz_dt,
        "hung_word_count": hung_word_count,
    }
    for dim in METRIC_KEYS:
        ds = dim_stats[dim]
        row[f"{dim}_count"] = int(ds.count)
        row[f"{dim}_avg"] = _optional_dec(ds.avg) if ds.count > 0 else None
    return row


def _materialize_stats_row(stats: Any) -> SimpleNamespace:
    """在 session 内抽出标量，避免 DetachedInstanceError。"""
    payload: dict[str, Any] = {"demand_category_id": int(stats.demand_category_id)}
    for dim in METRIC_KEYS:
        count = int(getattr(stats, f"{dim}_count") or 0)
        payload[f"{dim}_count"] = count
        if count > 0:
            avg_raw = getattr(stats, f"{dim}_avg", None)
            payload[f"{dim}_avg"] = float(avg_raw) if avg_raw is not None else None
        else:
            payload[f"{dim}_avg"] = None
    return SimpleNamespace(**payload)


def _materialize_weight_row(row: Any) -> dict[str, Any]:
    """在 session 内抽出标量，避免 DetachedInstanceError。"""
    payload: dict[str, Any] = {
        "category_id": int(row.category_id),
        "biz_dt": str(row.biz_dt),
    }
    for dim in POP_DIM_KEYS:
        count = int(getattr(row, f"{dim}_count", 0) or 0)
        payload[f"{dim}_count"] = count
        if count > 0:
            avg_raw = getattr(row, f"{dim}_avg", None)
            payload[f"{dim}_avg"] = float(avg_raw) if avg_raw is not None else None
        else:
            payload[f"{dim}_avg"] = None
    return payload


def _build_rank_score_rows(
    rows: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """根据已写入的权重行计算各节点排名分。"""
    dim_scores: dict[str, dict[int, float]] = {}
    for dim in POP_DIM_KEYS:
        candidates: list[tuple[int, float]] = []
        for row in rows:
            count = int(row[f"{dim}_count"])
            if count <= 0:
                continue
            avg_raw = row[f"{dim}_avg"]
            if avg_raw is None:
                continue
            avg = float(avg_raw)
            candidates.append((int(row["category_id"]), avg))
        dim_scores[dim] = rank_to_scores(candidates)

    updates: list[dict[str, Any]] = []
    for row in rows:
        category_id = int(row["category_id"])
        biz_dt = str(row["biz_dt"])
        score_values = {
            f"{dim}_score": dim_scores[dim].get(category_id)
            for dim in POP_DIM_KEYS
        }
        non_null_scores = [v for v in score_values.values() if v is not None]
        total = sum(non_null_scores) if non_null_scores else None
        updates.append(
            {
                "category_id": category_id,
                "biz_dt": biz_dt,
                **{
                    col: _optional_dec(score_values[col])
                    for col in (
                        "ext_pop_score",
                        "plat_sust_pop_score",
                        "plat_ly_pop_score",
                        "recent_pop_score",
                    )
                },
                "total_score": _optional_dec(total),
            }
        )
    return updates


def _update_category_tree_rank_scores(biz_dt: str) -> dict[str, Any]:
    """在 category_tree_weight 全部写入后，按 biz_dt 更新四维排名分与 total_score。"""
    with get_session() as session:
        repo = CategoryTreeWeightRepository(session)
        rows = [
            _materialize_weight_row(row)
            for row in repo.list_by_biz_dt(biz_dt)
        ]

    if not rows:
        logger.info("Category tree rank scores: no rows, skip biz_dt=%s", biz_dt)
        return {"biz_dt": biz_dt, "nodes": 0, "updated": 0}

    updates = _build_rank_score_rows(rows)
    updated = 0
    with get_session() as session:
        repo = CategoryTreeWeightRepository(session)
        for i in range(0, len(updates), _RANK_UPDATE_BATCH):
            updated += repo.update_rank_scores(updates[i : i + _RANK_UPDATE_BATCH])

    result = {
        "biz_dt": biz_dt,
        "nodes": len(updates),
        "updated": updated,
    }
    logger.info("Category tree rank scores completed: %s", result)
    return result


def compute_category_tree_weight(biz_dt: str) -> dict[str, Any]:
    """按 biz_dt 计算整棵树各维度加权平均分并写入 category_tree_weight。"""
    with get_session() as session:
        categories = [
            (int(c.id), c.parent_id)
            for c in GlobalTreeCategoryRepository(session).list_active_categories()
        ]
        belong_rows_raw = [
            (int(b.id), int(b.category_id))
            for b in DemandBelongCategoryRepository(session).list_active()
            if b.category_id is not None
        ]
        stats_by_belong = {
            int(s.demand_category_id): _materialize_stats_row(s)
            for s in DemandPopularityStatsRepository(session).list_by_biz_dt(biz_dt)
        }

    if not categories:
        logger.info("Category tree weight: no categories, skip biz_dt=%s", biz_dt)
        return {"biz_dt": biz_dt, "nodes": 0, "hanging": 0, "upserted": 0}

    by_id, children, _parent_of = _build_tree_index(categories)
    belong_rows = [
        (belong_id, category_id)
        for belong_id, category_id in belong_rows_raw
        if category_id in by_id
    ]

    hangs = _aggregate_hang_features(belong_rows, stats_by_belong)
    rolled = _rollup_avg_count(by_id, children, hangs)

    rows: list[dict[str, Any]] = []
    for nid in by_id:
        feat = hangs.get(nid)
        rows.append(
            _row_from_state(
                nid,
                biz_dt,
                feat.word_count if feat else 0,
                rolled[nid],
            )
        )

    upserted = 0
    with get_session() as session:
        repo = CategoryTreeWeightRepository(session)
        for i in range(0, len(rows), _UPSERT_BATCH):
            upserted += repo.upsert_rows(rows[i : i + _UPSERT_BATCH])

    result = {
        "biz_dt": biz_dt,
        "nodes": len(rows),
        "hanging": len(hangs),
        "upserted": upserted,
    }
    logger.info("Category tree weight completed: %s", result)

    rank_stats = _update_category_tree_rank_scores(biz_dt)
    result["rank_scores"] = rank_stats
    return result