app/services/recommender.py

import numpy as np
import torch
import contextlib
from app.core.ml_manager import ml_manager

class RecommenderService:
    
    @staticmethod
    def topk_indices(scores: np.ndarray, k: int) -> np.ndarray:
        k = min(k, len(scores))
        if k <= 0: return np.empty((0,), dtype=np.int64)
        idx = np.argpartition(scores, -k)[-k:]
        return idx[np.argsort(scores[idx])[::-1]]

    @staticmethod
    def get_recommendations(interactions: list, top_k: int = 20):
        mgr = ml_manager
        
        history = np.array([x["itemIndex"] for x in interactions], dtype=np.int64)
        ratings = np.array([x["rating"] for x in interactions], dtype=np.float32)
        
        # 0. Cold Start Fallback
        if len(history) == 0:
            top_pop = RecommenderService.topk_indices(mgr.scalars[:, 7], top_k).tolist()
            return [{"item_index": idx, "badges": ["Trending"], "score": 1.0} for idx in top_pop]
            
        pos_idx = history[ratings >= 3.0]
        neg_idx = history[ratings <= 2.5]
        
        # === STAGE 1: CANDIDATE GENERATION ===
        
        ease_scores = mgr.B_ease[pos_idx].sum(axis=0) if len(pos_idx) > 0 else np.zeros(mgr.num_items, dtype=np.float32)
        ease_scores[history] = -np.inf
        top_ease = RecommenderService.topk_indices(ease_scores, mgr.TOP_EASE_K)
        
        if len(pos_idx) > 0:
            user_text = mgr.text_embs_norm[pos_idx].mean(axis=0)
            text_scores = mgr.text_embs_norm @ (user_text / (np.linalg.norm(user_text) + 1e-12))
        else: text_scores = np.zeros(mgr.num_items, dtype=np.float32)
        text_scores[history] = -np.inf
        top_text = RecommenderService.topk_indices(text_scores, mgr.TOP_TEXT_K)
        
        recent_pos = pos_idx[-5:] if len(pos_idx) > 0 else pos_idx
        if len(recent_pos) > 0:
            rec_text = mgr.text_embs_norm[recent_pos].mean(axis=0)
            recent_scores = mgr.text_embs_norm @ (rec_text / (np.linalg.norm(rec_text) + 1e-12))
        else: recent_scores = np.zeros(mgr.num_items, dtype=np.float32)
        recent_scores[history] = -np.inf
        top_recent = RecommenderService.topk_indices(recent_scores, mgr.TOP_RECENT_K)
        
        if len(pos_idx) >= mgr.MIN_HISTORY_FOR_TWO_TOWER:
            with torch.no_grad():
                l_items = torch.tensor([x + 1 for x in pos_idx[-100:]], dtype=torch.long, device=mgr.device).unsqueeze(0)
                l_rats = torch.tensor([int(r) for r in ratings[ratings >= 3.0][-100:]], dtype=torch.long, device=mgr.device).unsqueeze(0)
                s_items = torch.tensor([x + 1 for x in pos_idx[-10:]], dtype=torch.long, device=mgr.device).unsqueeze(0)
                s_rats = torch.tensor([int(r) for r in ratings[ratings >= 3.0][-10:]], dtype=torch.long, device=mgr.device).unsqueeze(0)
                
                user_vec = mgr.two_tower.forward_user(l_items, l_rats, s_items, s_rats).cpu().numpy()[0]
                tt_scores = mgr.tt_item_corpus @ user_vec
                
            # Safely pad history check against tt_scores bounds
            tt_scores[history[history < mgr.num_items]] = -np.inf
            top_tt = RecommenderService.topk_indices(tt_scores, mgr.TOP_TT_K)
        else: 
            tt_scores = np.zeros(mgr.num_items, dtype=np.float32)
            # FIX: Return an empty array so no badges are falsely awarded!
            top_tt = np.array([], dtype=np.int64)
        
        # === STAGE 1.5: SAFE CANDIDATE POOLING ===
        # Ensure we don't accidentally pass out-of-bounds dynamic IDs to the ML models
        top_ease = top_ease[top_ease < mgr.num_items]
        top_text = top_text[top_text < mgr.num_items]
        top_recent = top_recent[top_recent < mgr.num_items]
        top_tt = top_tt[top_tt < mgr.num_items]

        # === UNION & FEATURE ENGINEERING ===
        # To prevent starving the Two-Tower model, we ensure proportional representation
        # if the total unique count exceeds SLATE_SIZE.
        candidates = np.unique(np.concatenate([top_ease, top_text, top_recent, top_tt])).astype(np.int64)
        
        if len(candidates) > mgr.SLATE_SIZE:
            # Deterministic truncation: take an equal chunk from the TOP of each retriever's list
            # We divide SLATE_SIZE (e.g., 160) by our 4 sources = 40 from each
            chunk = mgr.SLATE_SIZE // 4
            balanced_pool = np.concatenate([
                top_ease[:chunk], 
                top_text[:chunk], 
                top_recent[:chunk], 
                top_tt[:chunk]
            ])
            # Fill any remaining slots up to SLATE_SIZE with the next best overall candidates
            remaining_slots = mgr.SLATE_SIZE - len(np.unique(balanced_pool))
            fillers = np.setdiff1d(candidates, balanced_pool)[:remaining_slots]
            
            candidates = np.unique(np.concatenate([balanced_pool, fillers])).astype(np.int64)

        K = len(candidates)
        
        # 🛠️ THE FIX: Safe v_diff generation avoiding NoneType crashes
        v_diff_len = mgr.tag_scores.shape[1] if mgr.tag_scores is not None else 783
        v_diff = np.zeros(v_diff_len, dtype=np.float32)
        
        if len(pos_idx) > 0 and len(neg_idx) > 0 and mgr.tag_scores is not None:
            v_diff = np.maximum(0, mgr.tag_scores[neg_idx].mean(axis=0) - mgr.tag_scores[pos_idx].mean(axis=0))

        in_ease = np.isin(candidates, top_ease).astype(np.float32)
        in_text = np.isin(candidates, top_text).astype(np.float32)
        in_tt = np.isin(candidates, top_tt).astype(np.float32)
        
        tab_feats = np.column_stack([
            ease_scores[candidates], 
            text_scores[candidates], 
            tt_scores[candidates[:len(candidates)]], # Safe slice
            recent_scores[candidates],
            (mgr.tag_scores[candidates] @ v_diff) if mgr.tag_scores is not None else np.zeros(K),
            mgr.scalars[candidates, 7],
            np.abs(mgr.scalars[history, 2].mean() - mgr.scalars[candidates, 2]) if len(history) else np.zeros(K),
            (in_ease + in_text + in_tt),
            in_ease, in_text, in_tt
        ]).astype(np.float32)

        median = np.median(tab_feats, axis=0)
        iqr = np.clip(np.percentile(tab_feats, 75, axis=0) - np.percentile(tab_feats, 25, axis=0), 1e-6, None)
        tab_feats = np.clip((tab_feats - median) / iqr, -10.0, 10.0)

        franchise_ids = mgr.scalars[candidates, 4].astype(np.int64)
        
        # === STAGE 2: SLATE RANKER ===
        pad_len = mgr.SLATE_SIZE - K
        if pad_len > 0:
            tab_feats = np.vstack([tab_feats, np.zeros((pad_len, 11), dtype=np.float32)])
            franchise_ids = np.concatenate([franchise_ids, np.zeros(pad_len, dtype=np.int64)])
            padded_cands = np.concatenate([candidates + 1, np.zeros(pad_len, dtype=np.int64)])
        else:
            padded_cands = candidates + 1

        valid_mask = np.zeros(mgr.SLATE_SIZE, dtype=bool)
        valid_mask[:K] = True

        autocast_ctx = torch.autocast(device_type=mgr.device.type) if mgr.device.type == 'cuda' else contextlib.nullcontext()
        
        with torch.no_grad(), autocast_ctx:
            t_tab = torch.tensor(
                tab_feats,
                dtype=torch.float32,
                device=mgr.device
            ).unsqueeze(0)
            
            t_f = torch.tensor(
                franchise_ids,
                dtype=torch.long,
                device=mgr.device
            ).unsqueeze(0)

            t_cands = torch.tensor(
                padded_cands,
                dtype=torch.long,
                device=mgr.device
            ).unsqueeze(0)

            t_mask = torch.tensor(
                valid_mask,
                dtype=torch.bool,
                device=mgr.device
            ).unsqueeze(0)

            logits = mgr.ranker(t_tab, t_f, t_cands, t_mask).squeeze(0)
            logits = logits.masked_fill(~torch.tensor(valid_mask, device=mgr.device), -float('inf'))
            sorted_idx = torch.argsort(logits, descending=True).cpu().numpy()

        results = []
        for loc_idx in sorted_idx[:top_k]:
            item_idx = padded_cands[loc_idx] - 1
            if item_idx == -1: continue
            
            badges = []
            if in_ease[loc_idx]: badges.append("Collaborative")
            if in_text[loc_idx]: badges.append("Semantic")
            if in_tt[loc_idx]: badges.append("Deep Two-Tower")
                
            results.append({
                "item_index": int(item_idx),
                "badges": badges,
                "score": float(logits[loc_idx].item())
            })

        return results