fashion_ai/service.py

from __future__ import annotations

import os
from typing import Any

import numpy as np

from .classifier import FashionClassifier
from .encoder import FashionItemEncoder
from .ranker import NeuralOutfitScorer
from .retriever import OutfitCandidateRetriever
from .schemas import (
    EncodedWardrobeItem,
    OutfitCandidate,
    RecommendationContext,
    WeatherContext,
)


DEFAULT_TOP_K = int(os.getenv("FASHION_RECOMMEND_TOP_K", "5"))
DEFAULT_CANDIDATE_POOL = int(os.getenv("FASHION_CANDIDATE_POOL", "24"))
DEFAULT_MAX_BEAM = int(os.getenv("FASHION_MAX_BEAM", "64"))
DEFAULT_DIVERSITY_LAMBDA = float(os.getenv("FASHION_DIVERSITY_LAMBDA", "0.28"))

_SERVICE_SINGLETON: MultimodalOutfitRecommendationService | None = None


class MultimodalOutfitRecommendationService:
    """Multimodal retrieval + ranking service for outfit recommendations."""

    def __init__(
        self,
        encoder: FashionItemEncoder | None = None,
        retriever: OutfitCandidateRetriever | None = None,
        scorer: NeuralOutfitScorer | None = None,
        classifier: FashionClassifier | None = None,
        top_k: int = DEFAULT_TOP_K,
        candidate_pool: int = DEFAULT_CANDIDATE_POOL,
        max_beam: int = DEFAULT_MAX_BEAM,
        diversity_lambda: float = DEFAULT_DIVERSITY_LAMBDA,
    ) -> None:
        self.encoder = encoder or FashionItemEncoder()
        self.retriever = retriever or OutfitCandidateRetriever(
            self.encoder,
            slot_pool_size=candidate_pool,
        )
        self.scorer = scorer or NeuralOutfitScorer(d_model=self.encoder.embedding_dim)
        self.classifier = classifier or FashionClassifier()
        self.top_k = top_k
        self.candidate_pool = candidate_pool
        self.max_beam = max_beam
        self.diversity_lambda = diversity_lambda

    def recommend(
        self,
        wardrobe_items: list[dict[str, Any]],
        occasion: str = "casual",
        top_selected: dict[str, Any] | None = None,
        bottom_selected: dict[str, Any] | None = None,
        other_selected: dict[str, Any] | None = None,
        weather: dict[str, Any] | None = None,
        user_profile: dict[str, Any] | None = None,
        region: str = "global",
        top_k: int | None = None,
        candidate_pool: int | None = None,
        diversity_lambda: float | None = None,
    ) -> dict[str, Any]:
        context = self._build_context(occasion, weather, user_profile, region)
        top_k = top_k or self.top_k
        candidate_pool = candidate_pool or self.candidate_pool
        diversity_lambda = self.diversity_lambda if diversity_lambda is None else diversity_lambda

        encoded_items = self.retriever.encode_wardrobe(wardrobe_items)
        slot_buckets = self.retriever.split_by_slot(encoded_items)
        if not encoded_items:
            return self._empty_payload(occasion, "A", "Your wardrobe is empty. Add garments to get outfit recommendations.")
        if not slot_buckets["top"] or not slot_buckets["bottom"]:
            return self._empty_payload(
                occasion,
                "A",
                "You need at least one topwear and one bottomwear item to generate outfits.",
            )

        case_name = self._resolve_case_name(top_selected, bottom_selected)
        retrieved = self.retriever.retrieve(
            encoded_items=encoded_items,
            context=context,
            locked_top=top_selected,
            locked_bottom=bottom_selected,
            locked_other=other_selected,
            candidate_pool=candidate_pool,
        )
        context_vector = self.encoder.encode_context(context)
        user_vector = self._build_user_vector(context, encoded_items)

        candidates = self._assemble_candidates(retrieved)
        scored = self.scorer.score_candidates(candidates, context_vector, user_vector, context)
        diversified = self._diversify(scored, top_k=top_k, diversity_lambda=diversity_lambda)
        payloads = [
            self._candidate_to_payload(candidate, rank=index + 1)
            for index, candidate in enumerate(diversified)
        ]

        selected_outfit_score = None
        improved_recommendations: list[dict[str, Any]] = []
        recommendations = payloads

        if case_name == "D" and top_selected and bottom_selected:
            selected_candidates = self._assemble_candidates(
                {
                    "top": [self.encoder.encode_item(top_selected)],
                    "bottom": [self.encoder.encode_item(bottom_selected)],
                    "shoes": [],
                    "accessory": [],
                    "unknown": [],
                }
            )
            selected_scored = self.scorer.score_candidates(
                selected_candidates,
                context_vector,
                user_vector,
                context,
            )
            best_selected = self._diversify(selected_scored, top_k=1, diversity_lambda=0.0)
            selected_outfit_score = self._candidate_to_payload(best_selected[0], rank=1) if best_selected else None
            improved_recommendations = payloads[:top_k]
            recommendations = []

        return {
            "occasion": occasion,
            "case": case_name,
            "selected_outfit_score": selected_outfit_score,
            "recommendations": recommendations,
            "improved_recommendations": improved_recommendations,
            "total_combinations_checked": len(candidates),
            "notice": (
                "Selected outfit may not be the best fit for this occasion."
                if selected_outfit_score and selected_outfit_score.get("score", 0) < 45
                else None
            ),
            "engine_version": (
                f"fashion-mm-v1::{self.encoder.backend_name}"
                f"::{ 'trained-transformer' if self.scorer.is_trained else 'zero-shot-ranker' }"
            ),
        }

    def _assemble_candidates(
        self,
        slot_candidates: dict[str, list[EncodedWardrobeItem]],
    ) -> list[OutfitCandidate]:
        tops = slot_candidates.get("top") or []
        bottoms = slot_candidates.get("bottom") or []

        candidates: list[OutfitCandidate] = []
        for top in tops[: self.candidate_pool]:
            for bottom in bottoms[: self.candidate_pool]:
                if top.item.get("id") == bottom.item.get("id"):
                    continue

                candidates.append(
                    OutfitCandidate(
                        top=top,
                        bottom=bottom,
                        shoes=None,
                        accessory=None,
                    )
                )
                if len(candidates) >= self.max_beam:
                    return candidates
        return candidates

    def _diversify(
        self,
        scored: list[OutfitCandidate],
        top_k: int,
        diversity_lambda: float,
    ) -> list[OutfitCandidate]:
        selected: list[OutfitCandidate] = []
        remaining = list(scored)
        while remaining and len(selected) < top_k:
            best_index = 0
            best_score = -1e9
            for index, candidate in enumerate(remaining):
                relevance = candidate.score / 100.0
                redundancy = 0.0
                if selected:
                    redundancy = max(self._candidate_similarity(candidate, prev) for prev in selected)
                mmr_score = (1.0 - diversity_lambda) * relevance - diversity_lambda * redundancy
                if mmr_score > best_score:
                    best_index = index
                    best_score = mmr_score
            selected.append(remaining.pop(best_index))
        return selected

    def _build_user_vector(
        self,
        context: RecommendationContext,
        encoded_items: list[EncodedWardrobeItem],
    ) -> np.ndarray:
        liked_ids = set(str(item_id) for item_id in context.user_profile.get("liked_item_ids", []) if item_id)
        disliked_ids = set(str(item_id) for item_id in context.user_profile.get("disliked_item_ids", []) if item_id)

        vectors = []
        for encoded_item in encoded_items:
            item_id = str(encoded_item.item.get("id") or "")
            if item_id in liked_ids:
                vectors.append(encoded_item.vector)
            if item_id in disliked_ids:
                vectors.append(-encoded_item.vector)

        if vectors:
            merged = np.mean(np.stack(vectors, axis=0), axis=0)
            norm = float(np.linalg.norm(merged))
            if norm > 1e-8:
                return merged / norm
        return self.encoder.encode_text("User prefers versatile, well coordinated, context-appropriate outfits")

    @staticmethod
    def _candidate_to_payload(candidate: OutfitCandidate, rank: int) -> dict[str, Any]:
        payload = {
            "rank": rank,
            "score": candidate.score,
            "breakdown": candidate.breakdown,
            "reason": candidate.reason,
            "tip": candidate.tip,
            "top": MultimodalOutfitRecommendationService._item_payload(candidate.top),
            "bottom": MultimodalOutfitRecommendationService._item_payload(candidate.bottom),
        }
        return payload

    @staticmethod
    def _item_payload(encoded_item: EncodedWardrobeItem) -> dict[str, Any]:
        return {
            "id": encoded_item.item.get("id"),
            "category": encoded_item.item.get("category"),
            "color": encoded_item.item.get("color"),
            "image_url": encoded_item.item.get("image_url", ""),
        }

    @staticmethod
    def _build_context(
        occasion: str,
        weather: dict[str, Any] | None,
        user_profile: dict[str, Any] | None,
        region: str,
    ) -> RecommendationContext:
        weather = weather or {}
        return RecommendationContext(
            occasion=occasion or "casual",
            weather=WeatherContext(
                season=str(weather.get("season") or "all-season"),
                temperature_c=weather.get("temperature_c"),
                is_rainy=weather.get("is_rainy"),
            ),
            region=region or "global",
            user_profile=user_profile or {},
        )

    @staticmethod
    def _resolve_case_name(
        top_selected: dict[str, Any] | None,
        bottom_selected: dict[str, Any] | None,
    ) -> str:
        if top_selected and bottom_selected:
            return "D"
        if top_selected:
            return "B"
        if bottom_selected:
            return "C"
        return "A"

    @staticmethod
    def _empty_payload(occasion: str, case_name: str, notice: str) -> dict[str, Any]:
        return {
            "occasion": occasion,
            "case": case_name,
            "selected_outfit_score": None,
            "recommendations": [],
            "improved_recommendations": [],
            "total_combinations_checked": 0,
            "notice": notice,
            "engine_version": "fashion-mm-v1::empty",
        }

    def score_outfit(
        self,
        top: dict[str, Any],
        bottom: dict[str, Any],
        other: dict[str, Any] | None = None,
        occasion: str = "casual",
        weather: dict[str, Any] | None = None,
        user_profile: dict[str, Any] | None = None,
        region: str = "global",
    ) -> dict[str, Any]:
        context = self._build_context(occasion, weather, user_profile, region)
        other_encoded = self.encoder.encode_item(other) if other else None
        candidate = OutfitCandidate(
            top=self.encoder.encode_item(top),
            bottom=self.encoder.encode_item(bottom),
            shoes=other_encoded if other_encoded and other_encoded.slot == "shoes" else None,
            accessory=other_encoded if other_encoded and other_encoded.slot != "shoes" else None,
        )
        scored = self.scorer.score_candidates(
            candidates=[candidate],
            context_vector=self.encoder.encode_context(context),
            user_vector=self._build_user_vector(context, candidate.slot_items()),
            context=context,
        )
        best = scored[0]
        return {
            "score": best.score,
            "breakdown": best.breakdown,
            "reason": best.reason,
            "tip": best.tip,
            "engine_version": (
                f"fashion-mm-v1::{self.encoder.backend_name}"
                f"::{ 'trained-transformer' if self.scorer.is_trained else 'zero-shot-ranker' }"
            ),
        }

    @staticmethod
    def _candidate_similarity(left: OutfitCandidate, right: OutfitCandidate) -> float:
        left_vec = np.mean(np.stack([item.vector for item in left.slot_items()], axis=0), axis=0)
        right_vec = np.mean(np.stack([item.vector for item in right.slot_items()], axis=0), axis=0)
        left_norm = float(np.linalg.norm(left_vec))
        right_norm = float(np.linalg.norm(right_vec))
        if left_norm < 1e-8 or right_norm < 1e-8:
            return 0.0
        return float(np.dot(left_vec / left_norm, right_vec / right_norm))

    def classify_item(self, item: dict[str, Any]) -> dict[str, Any]:
        """
        Classify a fashion item using the integrated classifier.
        
        Uses NVIDIA model as primary, HuggingFace as fallback.
        
        Args:
            item: Wardrobe item dict with metadata and/or image_url
            
        Returns:
            Classification result with category, confidence, attributes
        """
        return self.classifier.classify_item(item)

    def match_items(
        self,
        item1: dict[str, Any],
        item2: dict[str, Any],
        match_threshold: float = 0.5,
    ) -> dict[str, Any]:
        """
        Determine if two fashion items match well together.
        
        Uses NVIDIA model as primary, HuggingFace as fallback.
        
        Args:
            item1: First wardrobe item
            item2: Second wardrobe item
            match_threshold: Confidence threshold for match (0-1)
            
        Returns:
            Dict with match result, score, reason, and compatibility breakdown
        """
        return self.classifier.match_items(item1, item2, match_threshold)


def get_recommendation_service() -> MultimodalOutfitRecommendationService:
    global _SERVICE_SINGLETON
    if _SERVICE_SINGLETON is None:
        _SERVICE_SINGLETON = MultimodalOutfitRecommendationService()
    return _SERVICE_SINGLETON