fashion_ai/ranker.py

from __future__ import annotations

import math
import os

import numpy as np
import torch
import torch.nn as nn

from .schemas import OutfitCandidate, RecommendationContext


class OutfitCompatibilityRanker(nn.Module):
    """
    Transformer ranker.

    Input:
      outfit_tokens  [B, 6, D] = [CONTEXT, USER, TOP, BOTTOM, SHOES, ACCESSORY]
      attention_mask [B, 6]

    Output:
      logits [B, 1]
    """

    def __init__(
        self,
        d_model: int = 512,
        n_layers: int = 4,
        n_heads: int = 8,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.d_model = d_model
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
        self.slot_embedding = nn.Embedding(7, d_model)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.head = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model // 2, 1),
        )
        nn.init.trunc_normal_(self.cls_token, std=0.02)

    def forward(
        self,
        outfit_tokens: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        if outfit_tokens.ndim != 3:
            raise ValueError("outfit_tokens must have shape [B, S, D]")
        if attention_mask.ndim != 2:
            raise ValueError("attention_mask must have shape [B, S]")

        batch_size, seq_len, _ = outfit_tokens.shape
        cls = self.cls_token.expand(batch_size, 1, self.d_model)
        tokens = torch.cat([cls, outfit_tokens], dim=1)
        slot_ids = torch.arange(seq_len + 1, device=tokens.device).unsqueeze(0)
        tokens = tokens + self.slot_embedding(slot_ids)

        cls_mask = torch.ones((batch_size, 1), device=attention_mask.device, dtype=attention_mask.dtype)
        mask = torch.cat([cls_mask, attention_mask], dim=1)
        encoded = self.encoder(tokens, src_key_padding_mask=mask == 0)
        return self.head(encoded[:, 0, :])


class NeuralOutfitScorer:
    """
    Uses a trained transformer checkpoint when available.
    Otherwise falls back to zero-shot geometric scoring over multimodal
    embeddings so the endpoint stays usable before fine-tuning.
    """

    def __init__(
        self,
        d_model: int = 512,
        checkpoint_path: str | None = None,
        device: str | None = None,
    ) -> None:
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model = OutfitCompatibilityRanker(d_model=d_model).to(self.device).eval()
        self.checkpoint_path = checkpoint_path or os.getenv("FASHION_RANKER_CHECKPOINT")
        self.is_trained = False
        self._load_checkpoint_if_available()

    def score_candidates(
        self,
        candidates: list[OutfitCandidate],
        context_vector: np.ndarray,
        user_vector: np.ndarray,
        context: RecommendationContext,
    ) -> list[OutfitCandidate]:
        if not candidates:
            return []
        if self.is_trained:
            return self._score_with_transformer(candidates, context_vector, user_vector, context)
        return self._score_zero_shot(candidates, context_vector, user_vector, context)

    def _score_with_transformer(
        self,
        candidates: list[OutfitCandidate],
        context_vector: np.ndarray,
        user_vector: np.ndarray,
        context: RecommendationContext,
    ) -> list[OutfitCandidate]:
        token_batch = []
        mask_batch = []
        for candidate in candidates:
            vectors = [
                context_vector,
                user_vector,
                candidate.top.vector,
                candidate.bottom.vector,
                candidate.shoes.vector if candidate.shoes is not None else np.zeros_like(context_vector),
                candidate.accessory.vector if candidate.accessory is not None else np.zeros_like(context_vector),
            ]
            mask = [
                1,
                1,
                1,
                1,
                1 if candidate.shoes is not None else 0,
                1 if candidate.accessory is not None else 0,
            ]
            token_batch.append(np.stack(vectors, axis=0))
            mask_batch.append(mask)

        with torch.inference_mode():
            logits = self.model(
                torch.tensor(np.stack(token_batch), dtype=torch.float32, device=self.device),
                torch.tensor(np.asarray(mask_batch), dtype=torch.long, device=self.device),
            ).squeeze(-1)
            probs = torch.sigmoid(logits).detach().cpu().numpy()

        return self._finalize_candidates(candidates, probs, context_vector, user_vector, context)

    def _score_zero_shot(
        self,
        candidates: list[OutfitCandidate],
        context_vector: np.ndarray,
        user_vector: np.ndarray,
        context: RecommendationContext,
    ) -> list[OutfitCandidate]:
        scores = []
        for candidate in candidates:
            item_vectors = [slot_item.vector for slot_item in candidate.slot_items()]
            outfit_centroid = self._normalize(np.mean(np.stack(item_vectors, axis=0), axis=0))
            context_alignment = self._cosine(outfit_centroid, context_vector)
            user_alignment = self._cosine(outfit_centroid, user_vector)
            pairwise_cohesion = self._pairwise_cohesion(item_vectors)
            slot_coverage = math.log1p(len(item_vectors)) / math.log1p(4)

            score = np.mean(
                np.asarray(
                    [
                        self._to_unit_interval(context_alignment),
                        self._to_unit_interval(user_alignment),
                        self._to_unit_interval(pairwise_cohesion),
                        slot_coverage,
                    ],
                    dtype=np.float32,
                )
            )
            scores.append(float(np.clip(score, 0.0, 1.0)))

        return self._finalize_candidates(candidates, scores, context_vector, user_vector, context)

    def _finalize_candidates(
        self,
        candidates: list[OutfitCandidate],
        probs: list[float] | np.ndarray,
        context_vector: np.ndarray,
        user_vector: np.ndarray,
        context: RecommendationContext,
    ) -> list[OutfitCandidate]:
        scored = []
        for candidate, prob in zip(candidates, probs, strict=False):
            candidate.score = round(float(prob) * 100.0, 2)
            candidate.breakdown = self._build_breakdown(candidate, context_vector, user_vector, context)
            candidate.reason = self._build_reason(candidate, context)
            candidate.tip = self._build_tip(candidate)
            scored.append(candidate)
        scored.sort(key=lambda item: item.score, reverse=True)
        return scored

    def _load_checkpoint_if_available(self) -> None:
        if not self.checkpoint_path or not os.path.isfile(self.checkpoint_path):
            return
        try:
            payload = torch.load(self.checkpoint_path, map_location=self.device)
            state_dict = payload.get("model_state_dict", payload) if isinstance(payload, dict) else payload
            self.model.load_state_dict(state_dict, strict=False)
            self.model.eval()
            self.is_trained = True
        except Exception:
            self.is_trained = False

    def _build_breakdown(
        self,
        candidate: OutfitCandidate,
        context_vector: np.ndarray,
        user_vector: np.ndarray,
        context: RecommendationContext,
    ) -> dict[str, float]:
        item_vectors = [slot_item.vector for slot_item in candidate.slot_items()]
        outfit_centroid = self._normalize(np.mean(np.stack(item_vectors, axis=0), axis=0))
        context_alignment = self._to_score(self._cosine(outfit_centroid, context_vector))
        user_affinity = self._to_score(self._cosine(outfit_centroid, user_vector))
        visual_cohesion = self._to_score(self._pairwise_cohesion(item_vectors))
        top_bottom_compat = self._to_score(self._cosine(candidate.top.vector, candidate.bottom.vector))
        occasion_fit = self._occasion_slot_fit(candidate, context.occasion)

        return {
            "color": round((visual_cohesion + top_bottom_compat) / 2.0, 2),
            "style": round((context_alignment + visual_cohesion) / 2.0, 2),
            "occasion": round((occasion_fit + context_alignment) / 2.0, 2),
            "context_alignment": round(context_alignment, 2),
            "user_affinity": round(user_affinity, 2),
            "visual_cohesion": round(visual_cohesion, 2),
        }

    def _build_reason(self, candidate: OutfitCandidate, context: RecommendationContext) -> str:
        parts = [
            f"{candidate.top.item.get('color', 'Unknown')} {candidate.top.item.get('category', 'Topwear')}",
            f"{candidate.bottom.item.get('color', 'Unknown')} {candidate.bottom.item.get('category', 'Bottomwear')}",
        ]
        if candidate.shoes is not None:
            parts.append(
                f"{candidate.shoes.item.get('color', 'Unknown')} {candidate.shoes.item.get('category', 'Footwear')}"
            )
        if candidate.accessory is not None:
            parts.append(
                f"{candidate.accessory.item.get('color', 'Unknown')} {candidate.accessory.item.get('category', 'Accessory')}"
            )
        return (
            f"Learned multimodal embeddings rate {' + '.join(parts)} as a coherent "
            f"combination for {context.occasion or 'casual'} context."
        )

    @staticmethod
    def _build_tip(candidate: OutfitCandidate) -> str:
        if candidate.score >= 85:
            return "Strong outfit match. Keep accessories minimal so the silhouette stays clean."
        if candidate.score >= 70:
            return "Good base outfit. Add one tonal accessory to reinforce the palette."
        return "This outfit is acceptable, but one slot can be swapped for stronger style alignment."

    @staticmethod
    def _occasion_slot_fit(candidate: OutfitCandidate, occasion: str) -> float:
        occ = str(occasion or "casual").lower()
        texts = " ".join(slot.metadata_text.lower() for slot in candidate.slot_items())
        if occ in texts:
            return 95.0
        if occ in {"formal", "interview", "business", "office"} and any(
            token in texts for token in ["shirt", "blazer", "trouser", "loafer"]
        ):
            return 88.0
        if occ in {"party", "festive", "wedding"} and any(
            token in texts for token in ["silk", "embroidered", "dress", "kurta"]
        ):
            return 88.0
        if occ in {"sports", "gym", "active"} and any(
            token in texts for token in ["sneaker", "jogger", "tee", "hoodie"]
        ):
            return 85.0
        return 72.0

    @staticmethod
    def _pairwise_cohesion(vectors: list[np.ndarray]) -> float:
        if len(vectors) < 2:
            return 0.0
        scores = []
        for left_index in range(len(vectors)):
            for right_index in range(left_index + 1, len(vectors)):
                scores.append(NeuralOutfitScorer._cosine(vectors[left_index], vectors[right_index]))
        return float(np.mean(np.asarray(scores, dtype=np.float32)))

    @staticmethod
    def _cosine(left: np.ndarray, right: np.ndarray) -> float:
        left_vec = NeuralOutfitScorer._normalize(left)
        right_vec = NeuralOutfitScorer._normalize(right)
        return float(np.dot(left_vec, right_vec))

    @staticmethod
    def _to_score(value: float) -> float:
        return round(100.0 * NeuralOutfitScorer._to_unit_interval(value), 2)

    @staticmethod
    def _to_unit_interval(value: float) -> float:
        return float(np.clip((value + 1.0) / 2.0, 0.0, 1.0))

    @staticmethod
    def _normalize(vector: np.ndarray) -> np.ndarray:
        arr = np.asarray(vector, dtype=np.float32).reshape(-1)
        norm = float(np.linalg.norm(arr))
        if norm < 1e-8:
            return np.zeros_like(arr, dtype=np.float32)
        return arr / norm