ml_models.py

"""
ML Model Loader and Utilities
Handles loading and using the conflict prediction model and package embeddings.
"""

import json
import pickle
from pathlib import Path
from typing import Dict, List, Tuple, Optional
import numpy as np
from packaging.requirements import Requirement


class ConflictPredictor:
    """Load and use the conflict prediction model."""
    
    def __init__(self, model_path: Optional[Path] = None):
        """Initialize the conflict predictor."""
        if model_path is None:
            model_path = Path(__file__).parent / "models" / "conflict_predictor.pkl"
        
        self.model = None
        self.model_path = model_path
        
        if model_path.exists():
            try:
                with open(model_path, 'rb') as f:
                    self.model = pickle.load(f)
                print(f"✅ Loaded conflict prediction model from {model_path}")
            except Exception as e:
                print(f"⚠️ Could not load conflict prediction model: {e}")
        else:
            print(f"⚠️ Conflict prediction model not found at {model_path}")
    
    def extract_features(self, requirements_text: str) -> np.ndarray:
        """Extract features from requirements text (same as training)."""
        features = []
        
        packages = {}
        lines = requirements_text.strip().split('\n')
        num_packages = 0
        has_pins = 0
        version_specificity = []
        
        for line in lines:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            
            try:
                req = Requirement(line)
                pkg_name = req.name.lower()
                specifier = str(req.specifier) if req.specifier else ''
                
                if pkg_name in packages:
                    features.append(1)  # has_duplicate flag
                else:
                    packages[pkg_name] = specifier
                    num_packages += 1
                    
                    if specifier:
                        has_pins += 1
                        if '==' in specifier:
                            version_specificity.append(3)
                        elif '>=' in specifier or '<=' in specifier:
                            version_specificity.append(2)
                        else:
                            version_specificity.append(1)
                    else:
                        version_specificity.append(0)
            except:
                pass
        
        feature_vec = []
        feature_vec.append(min(num_packages / 20.0, 1.0))
        feature_vec.append(has_pins / max(num_packages, 1))
        feature_vec.append(np.mean(version_specificity) / 3.0 if version_specificity else 0)
        feature_vec.append(1 if len(packages) < num_packages else 0)
        
        common_packages = [
            'torch', 'pytorch-lightning', 'tensorflow', 'keras', 'fastapi', 'pydantic',
            'numpy', 'pandas', 'scipy', 'scikit-learn', 'matplotlib', 'seaborn',
            'requests', 'httpx', 'sqlalchemy', 'alembic', 'uvicorn', 'starlette',
            'langchain', 'openai', 'chromadb', 'redis', 'celery', 'gunicorn',
            'pillow', 'opencv-python', 'beautifulsoup4', 'scrapy', 'plotly', 'jax'
        ]
        
        for pkg in common_packages:
            feature_vec.append(1 if pkg in packages else 0)
        
        has_torch = 'torch' in packages
        has_pl = 'pytorch-lightning' in packages
        has_tf = 'tensorflow' in packages
        has_keras = 'keras' in packages
        has_fastapi = 'fastapi' in packages
        has_pydantic = 'pydantic' in packages
        
        feature_vec.append(1 if (has_torch and has_pl) else 0)
        feature_vec.append(1 if (has_tf and has_keras) else 0)
        feature_vec.append(1 if (has_fastapi and has_pydantic) else 0)
        
        return np.array(feature_vec)
    
    def predict(self, requirements_text: str) -> Tuple[bool, float]:
        """
        Predict if requirements have conflicts.
        
        Returns:
            (has_conflict, confidence_score)
        """
        if self.model is None:
            return False, 0.0
        
        try:
            features = self.extract_features(requirements_text)
            features = features.reshape(1, -1)
            
            prediction = self.model.predict(features)[0]
            probability = self.model.predict_proba(features)[0]
            
            has_conflict = bool(prediction)
            confidence = float(probability[1] if has_conflict else probability[0])
            
            return has_conflict, confidence
        except Exception as e:
            print(f"Error in conflict prediction: {e}")
            return False, 0.0


class PackageEmbeddings:
    """Load and use package embeddings for similarity matching."""
    
    def __init__(self, embeddings_path: Optional[Path] = None):
        """Initialize package embeddings."""
        if embeddings_path is None:
            embeddings_path = Path(__file__).parent / "models" / "package_embeddings.json"
        
        self.embeddings = {}
        self.embeddings_path = embeddings_path
        self.model = None
        
        if embeddings_path.exists():
            try:
                with open(embeddings_path, 'r') as f:
                    self.embeddings = json.load(f)
                print(f"✅ Loaded {len(self.embeddings)} package embeddings from {embeddings_path}")
            except Exception as e:
                print(f"⚠️ Could not load embeddings: {e}")
        else:
            print(f"⚠️ Embeddings not found at {embeddings_path}")
    
    def _load_model(self):
        """Lazy load the sentence transformer model."""
        if self.model is None:
            try:
                from sentence_transformers import SentenceTransformer
                self.model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
            except ImportError:
                print("⚠️ sentence-transformers not available, embedding similarity disabled")
                return None
        return self.model
    
    def get_embedding(self, package_name: str) -> Optional[np.ndarray]:
        """Get embedding for a package (from cache or compute on-the-fly)."""
        package_lower = package_name.lower()
        
        # Check cache first
        if package_lower in self.embeddings:
            return np.array(self.embeddings[package_lower])
        
        # Compute on-the-fly if model available
        model = self._load_model()
        if model is not None:
            embedding = model.encode([package_name])[0]
            # Cache it
            self.embeddings[package_lower] = embedding.tolist()
            return embedding
        
        return None
    
    def find_similar(self, package_name: str, top_k: int = 5, threshold: float = 0.6) -> List[Tuple[str, float]]:
        """
        Find similar packages using cosine similarity.
        
        Returns:
            List of (package_name, similarity_score) tuples
        """
        query_emb = self.get_embedding(package_name)
        if query_emb is None:
            return []
        
        similarities = []
        
        for pkg, emb in self.embeddings.items():
            if pkg == package_name.lower():
                continue
            
            emb_array = np.array(emb)
            # Cosine similarity
            similarity = np.dot(query_emb, emb_array) / (
                np.linalg.norm(query_emb) * np.linalg.norm(emb_array)
            )
            
            if similarity >= threshold:
                similarities.append((pkg, float(similarity)))
        
        # Sort by similarity and return top_k
        similarities.sort(key=lambda x: x[1], reverse=True)
        return similarities[:top_k]
    
    def get_best_match(self, package_name: str, threshold: float = 0.7) -> Optional[str]:
        """Get the best matching package name."""
        similar = self.find_similar(package_name, top_k=1, threshold=threshold)
        if similar:
            return similar[0][0]
        return None