"""
Inference script for Linear Regression Text Classification Model
This script demonstrates how to load and use the trained model for predictions.
"""

import joblib
import numpy as np
from typing import Union, List


class TextClassifier:
    """Wrapper class for the Linear Regression text classification model."""
    
    def __init__(self, model_path: str):
        """
        Initialize the classifier by loading the model.
        
        Args:
            model_path: Path to the saved joblib model file
        """
        self.model = joblib.load(model_path)
        self.class_names = {
            0: "Strictly Necessary",
            1: "Functionality", 
            2: "Analytics",
            3: "Advertising/Tracking"
        }
    
    def predict(self, features: np.ndarray) -> np.ndarray:
        """
        Make predictions on input features.
        
        Args:
            features: Preprocessed TF-IDF features (numpy array)
            
        Returns:
            Array of predicted class labels
        """
        predictions = self.model.predict(features)
        return predictions.astype(int)
    
    def predict_single(self, features: np.ndarray) -> int:
        """
        Make a prediction for a single sample.
        
        Args:
            features: Preprocessed TF-IDF features for one sample
            
        Returns:
            Predicted class label
        """
        if len(features.shape) == 1:
            features = features.reshape(1, -1)
        prediction = self.model.predict(features)[0]
        return int(prediction)
    
    def get_class_name(self, class_id: int) -> str:
        """
        Get the name of a class given its ID.
        
        Args:
            class_id: The numeric class identifier
            
        Returns:
            The name/description of the class
        """
        return self.class_names.get(class_id, f"Unknown Class {class_id}")


def load_model(model_path: str = "Linear Regression/LR_TFIDF+NAME.joblib") -> TextClassifier:
    """
    Load the trained model from disk.
    
    Args:
        model_path: Path to the model file
        
    Returns:
        TextClassifier instance
    """
    return TextClassifier(model_path)


def main():
    """Example usage of the model."""
    # Load the model
    print("Loading model...")
    classifier = load_model()
    print("Model loaded successfully!")
    
    # Example: Create dummy features for demonstration
    # In practice, you would use your TF-IDF vectorizer to transform text
    print("\nExample prediction (using random features for demonstration):")
    dummy_features = np.random.randn(1, 100)  # Replace with actual TF-IDF features
    
    prediction = classifier.predict_single(dummy_features)
    class_name = classifier.get_class_name(prediction)
    
    print(f"Predicted class: {prediction}")
    print(f"Class name: {class_name}")
    
    # Batch prediction example
    print("\nBatch prediction example:")
    batch_features = np.random.randn(5, 100)  # 5 samples
    batch_predictions = classifier.predict(batch_features)
    
    print(f"Predictions for {len(batch_predictions)} samples:")
    for i, pred in enumerate(batch_predictions):
        print(f"  Sample {i+1}: Class {pred} ({classifier.get_class_name(pred)})")


if __name__ == "__main__":
    main()