from pathlib import Path
from typing import Dict, Any, Optional, Tuple
import pandas as pd
from autogluon.tabular import TabularPredictor

from mlpipeline.logging.logger import get_logger

logger = get_logger(__name__)


class AutoGluonTrainer:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.predictor: Optional[TabularPredictor] = None
    
    def train(self, train_data: pd.DataFrame, target_column: str, model_path: Path) -> Dict[str, float]:
        logger.info("Starting AutoGluon training")
        
        self.predictor = TabularPredictor(
            label=target_column,
            path=str(model_path),
            eval_metric=self.config.get('eval_metric'),
            verbosity=self.config.get('verbosity', 2),
        )
        
        self.predictor.fit(
            train_data=train_data,
            time_limit=self.config.get('time_limit', 600),
            presets=self.config.get('presets', 'medium_quality'),
            num_bag_folds=self.config.get('num_bag_folds', 5),
            num_stack_levels=self.config.get('num_stack_levels', 1),
        )
        
        leaderboard = self.predictor.leaderboard(silent=True)
        best_model = leaderboard.iloc[0]
        
        # Get feature importance if available
        try:
            feature_importance = self.predictor.feature_importance(data=train_data)
        except:
            feature_importance = None
        
        metrics = {
            'validation_accuracy': float(best_model['score_val']),
            'score': float(best_model['score_val']),  # Keep for backward compatibility
            'score_test': float(best_model.get('score_test', 0.0)),
            'fit_time': float(best_model.get('fit_time', 0.0)),
            'pred_time_val': float(best_model.get('pred_time_val', 0.0)),
            'num_models_trained': len(leaderboard),
            'best_model_name': str(best_model['model']),
        }
        
        logger.info(f"AutoGluon training completed. Best score: {metrics['score']}")
        return metrics, feature_importance
    
    def predict(self, data: pd.DataFrame) -> pd.Series:
        if self.predictor is None:
            raise ValueError("Model not trained. Call train() first.")
        return self.predictor.predict(data)
    
    def load_model(self, model_path: Path):
        logger.info(f"Loading AutoGluon model from {model_path}")
        self.predictor = TabularPredictor.load(str(model_path))
        return self