"""
Model Trainer
=============

This module provides functionality for training machine learning models.
"""

import logging
import time
from pathlib import Path
from typing import Dict, Any, Optional
import joblib
import json
import pandas as pd
import numpy as np
from sklearn.pipeline import Pipeline

from .model_builder import ModelBuilder

logger = logging.getLogger(__name__)


class ModelTrainer:
    """
    Trains and saves machine learning models for sentiment analysis.
    """
    
    LABEL_MAP = {0: "negative", 2: "neutral", 4: "positive"}
    
    def __init__(self, model_params: Optional[Dict[str, Any]] = None):
        """
        Initialize the model trainer.
        
        Args:
            model_params: Optional dictionary of model hyperparameters
        """
        self.model_params = model_params or {}
        self.builder = ModelBuilder(**self.model_params)
        self.model: Optional[Pipeline] = None
        self.training_info: Dict[str, Any] = {}
    
    def train(
        self,
        X_train: list,
        y_train: np.ndarray,
        verbose: bool = True
    ) -> Pipeline:
        """
        Train the model on the provided data.
        
        Args:
            X_train: List of training texts
            y_train: Training labels
            verbose: Whether to print training information
            
        Returns:
            Trained model pipeline
        """
        logger.info(f"Training model on {len(X_train)} samples...")
        
        # Build model pipeline
        self.model = self.builder.build_tfidf_svm_pipeline()
        
        # Train model
        start_time = time.time()
        self.model.fit(X_train, y_train)
        training_time = time.time() - start_time
        
        # Store training information
        self.training_info = {
            "training_samples": len(X_train),
            "training_time_seconds": training_time,
            "model_params": self.model_params,
            "label_distribution": pd.Series(y_train).value_counts().to_dict()
        }
        
        if verbose:
            logger.info(f"Training completed in {training_time:.2f} seconds")
            logger.info(f"Label distribution: {self.training_info['label_distribution']}")
        
        return self.model
    
    def save_model(
        self,
        output_dir: Path,
        model_name: str = "tfidf_svm_sentiment140"
    ):
        """
        Save the trained model and metadata.
        
        Args:
            output_dir: Directory to save the model
            model_name: Base name for the model files
        """
        if self.model is None:
            raise ValueError("No trained model to save. Train a model first.")
        
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # Save model
        model_path = output_dir / f"{model_name}.joblib"
        joblib.dump(self.model, model_path)
        logger.info(f"Model saved to {model_path}")
        
        # Save label map
        labelmap_path = output_dir / "label_map.json"
        with open(labelmap_path, "w", encoding="utf-8") as f:
            json.dump(self.LABEL_MAP, f, ensure_ascii=False, indent=2)
        logger.info(f"Label map saved to {labelmap_path}")
        
        # Save training info
        info_path = output_dir / "training_info.json"
        with open(info_path, "w", encoding="utf-8") as f:
            json.dump(self.training_info, f, indent=2, ensure_ascii=False)
        logger.info(f"Training info saved to {info_path}")
    
    def load_model(self, model_path: Path) -> Pipeline:
        """
        Load a trained model from disk.
        
        Args:
            model_path: Path to the model file
            
        Returns:
            Loaded model pipeline
        """
        logger.info(f"Loading model from {model_path}")
        self.model = joblib.load(model_path)
        logger.info("Model loaded successfully")
        return self.model
    
    def predict(self, texts: list) -> np.ndarray:
        """
        Make predictions on new texts.
        
        Args:
            texts: List of texts to predict
            
        Returns:
            Array of predicted labels
        """
        if self.model is None:
            raise ValueError("No model loaded. Train or load a model first.")
        
        return self.model.predict(texts)
    
    def predict_proba(self, texts: list) -> np.ndarray:
        """
        Get prediction probabilities (if supported by the model).
        
        Args:
            texts: List of texts to predict
            
        Returns:
            Array of prediction probabilities
        """
        if self.model is None:
            raise ValueError("No model loaded. Train or load a model first.")
        
        if hasattr(self.model, "predict_proba"):
            return self.model.predict_proba(texts)
        elif hasattr(self.model, "decision_function"):
            # For LinearSVC, use decision function
            return self.model.decision_function(texts)
        else:
            raise NotImplementedError("Model does not support probability predictions")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Example usage
    from pathlib import Path
    import pandas as pd
    
    # Load sample data
    train_path = Path("data/sentiment140/train_split.csv")
    if train_path.exists():
        train_df = pd.read_csv(train_path, nrows=1000)  # Use subset for testing
        
        # Initialize trainer
        trainer = ModelTrainer(model_params={
            "c_value": 1.0,
            "min_df": 5,
            "max_df": 0.9,
            "ngram_max": 2
        })
        
        # Train model
        trainer.train(
            X_train=train_df["text"].tolist(),
            y_train=train_df["target"].to_numpy()
        )
        
        # Save model
        output_dir = Path("artifacts_sentiment140")
        trainer.save_model(output_dir)
        
        print("\nTraining completed successfully!")
    else:
        print(f"Train data not found at {train_path}")
