app/core/ml_manager.py

import numpy as np
import torch
import faiss
import pandas as pd
from pathlib import Path
import logging
import json

logger = logging.getLogger(__name__)

class MLManager:
    _instance = None
    
    MIN_HISTORY_FOR_TWO_TOWER = 5
    TOP_EASE_K = 100
    TOP_TEXT_K = 40
    TOP_RECENT_K = 20
    TOP_TT_K = 50
    SLATE_SIZE = 160

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(MLManager, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def initialize(self, data_dir: Path):
        if self._initialized: return

        logger.info("Initializing ML Manager...")
        
        def get_path(filename):
            paths = [
                data_dir / filename,
                data_dir / "raw" / filename,
                data_dir / "processed" / filename,
                data_dir / "ranking" / filename,
                data_dir / "kg" / "node_features" / filename,
                data_dir / "exports" / filename,
            ]
            for p in paths:
                if p.exists(): return p
            return data_dir / filename

        # 1. Load EASE Matrix
        self.ease_matrix = np.load(get_path("ease_item_item_matrix.npy"), mmap_mode='r')
        self.B_ease = self.ease_matrix 

        # 2. Load movies.npy
        movies_array = np.load(get_path("movies.npy")).astype(np.float32)
        self.num_items = movies_array.shape[0]
        
        self.plot_embeddings = movies_array[:, :1024]
        self.scalars = movies_array[:, 1024:1034]

        self.d_semantic = 1024
        plot_embeddings_norm = self.plot_embeddings.copy()
        faiss.normalize_L2(plot_embeddings_norm)
        
        self.faiss_index = faiss.IndexFlatIP(self.d_semantic)
        self.faiss_index.add(plot_embeddings_norm)
        self.text_embs_norm = plot_embeddings_norm 

        # 3. Load Static Metadata directly from item_features.parquet
        logger.info("Loading item_features.parquet...")
        
        # 🛠️ THE FIX: Added columns required for the Detail Modal
        columns_to_load = [
            "item_idx", "tmdb_id", "title", "poster_path", 
            "overview", "release_date", "runtime", "genres", "backdrop_path"
        ]
        
        # Fallback in case your parquet is missing some of these columns
        try:
            self.movie_meta = pd.read_parquet(get_path("item_features.parquet"), columns=columns_to_load)
        except ValueError:
            logger.warning("Some rich metadata columns missing. Loading basic columns.")
            self.movie_meta = pd.read_parquet(get_path("item_features.parquet"))

        # Drop rows where item_idx is missing, then set it as index
        self.movie_meta = self.movie_meta.dropna(subset=["item_idx"])
        self.movie_meta["item_idx"] = self.movie_meta["item_idx"].astype(int)
        self.movie_meta = self.movie_meta.set_index("item_idx")

        # 4. Load PyTorch Models
        from app.core.models import TriModalTwoTower, SemanticSlateRanker
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        gnn_embs = np.load(get_path("film_embeddings.npy"))
        scalars_norm = (self.scalars - np.mean(self.scalars, axis=0)) / (np.std(self.scalars, axis=0) + 1e-6)

        self.two_tower = TriModalTwoTower(self.num_items, self.plot_embeddings, gnn_embs, scalars_norm).to(self.device)
        self.two_tower.load_state_dict(
            torch.load(get_path("two_tower_best.pth"), map_location=self.device, weights_only=True), 
            strict=False # Ignores missing registered buffers
        )
        self.two_tower.eval()

        # Initialize Ranker (Must be exactly 11 features!)
        padded_movie_vectors = np.vstack([np.zeros((1, 1024), dtype=np.float32), plot_embeddings_norm])
        num_franchises = int(np.max(self.scalars[:, 4]) + 2) 
        continuous_dim = 11 # EXACT FIX: 8 Numeric + 3 Binary

        self.ranker = SemanticSlateRanker(
            continuous_dim, num_items=padded_movie_vectors.shape[0],
            num_franchises=num_franchises, pretrained_text_embs=padded_movie_vectors
        ).to(self.device)
        self.ranker.load_state_dict(
            torch.load(get_path("best_semantic_slate_ranker.pt"), map_location=self.device, weights_only=True),
            strict=False
        )
        self.ranker.eval()

        logger.info("Precomputing Two-Tower Item Vectors...")
        with torch.no_grad():
            all_item_ids = torch.arange(1, self.num_items + 1).to(self.device)
            self.tt_item_corpus = self.two_tower.forward_item(all_item_ids).cpu().numpy()

        # 5. Load Cluster Ontology & Tag Genome
        with open(get_path("cluster_ontology.json"), "r", encoding="utf-8") as f:
            self.ontology = json.load(f)

        self.fc_keys = list(self.ontology.keys())
        self.fc_names = {k: self.ontology[k]["label"] for k in self.fc_keys}
        self.fc_matrix = np.array([self.ontology[k]["prototype_embedding"] for k in self.fc_keys], dtype=np.float32)
        faiss.normalize_L2(self.fc_matrix)

        axis_mapping = {
            "Adrenaline & Spectacle":["FC02", "FC03", "FC10", "FC25", "FC26", "FC29"],
            "Dark & Gritty":["FC04", "FC08", "FC11", "FC18"],
            "Heartwarming & Joy":["FC06", "FC14", "FC16"],
            "Cerebral & Surreal":["FC07", "FC12", "FC15"],
            "Grounded & Historical":["FC01", "FC05", "FC09", "FC17", "FC21"],
            "Intimate & Human":["FC19", "FC20", "FC22", "FC23", "FC24", "FC28", "FC30"],
        }
        self.macro_axes_names = list(axis_mapping.keys())
        macro_vectors =[np.mean([self.ontology[fc]["prototype_embedding"] for fc in fcs if fc in self.ontology], axis=0).astype(np.float32) for fcs in axis_mapping.values()]
        self.macro_matrix = np.array(macro_vectors, dtype=np.float32)
        faiss.normalize_L2(self.macro_matrix)

        # Try loading the tag genome files
        try:
            self.tag_scores = np.load(get_path("film_tag_scores.npy"), mmap_mode="r")
            with open(get_path("kg_tag_index.json"), "r") as f:
                # Force keys to int and values to string during inversion
                raw_index = json.load(f)
                self.tag_col_to_id = {int(v): str(k) for k, v in raw_index.items()}
                
            with open(get_path("tags.json"), "r", encoding="utf-8") as f:
                self.tag_id_to_name = {
                    str(t["id"]): t["tag"] 
                    for t in [json.loads(line) for line in f if line.strip()]
                }
        except Exception as e:
            logger.error(f"⚠️ TAG GENOME FAILED TO LOAD: {e}")
            self.tag_scores = None

        # 6. Load Browse Categories
        try:
            with open(get_path("browse_index.json"), "r", encoding="utf-8") as f:
                self.browse_data = json.load(f)
            logger.info("Browse Categories loaded successfully.")
        except Exception as e:
            logger.error(f"⚠️ BROWSE DATA FAILED TO LOAD: {e}")
            self.browse_data = {}
            
        self._initialized = True
        logger.info("ML Manager fully initialized.")

ml_manager = MLManager()