from pathlib import Path
from typing import Dict, Any, Optional
import pandas as pd
from pycaret.classification import setup as classification_setup, compare_models as classification_compare, finalize_model as classification_finalize, save_model as classification_save, load_model as classification_load
from pycaret.regression import setup as regression_setup, compare_models as regression_compare, finalize_model as regression_finalize, save_model as regression_save, load_model as regression_load

from mlpipeline.logging.logger import get_logger

logger = get_logger(__name__)


class PyCaretTrainer:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.model: Optional[Any] = None
        self.is_classification = None
    
    def train(self, train_data: pd.DataFrame, target_column: str, model_path: Path) -> Dict[str, float]:
        logger.info("Starting PyCaret training")
        
        if train_data[target_column].dtype == 'object' or train_data[target_column].nunique() < 20:
            self.is_classification = True
            setup_fn = classification_setup
            compare_fn = classification_compare
            finalize_fn = classification_finalize
            save_fn = classification_save
        else:
            self.is_classification = False
            setup_fn = regression_setup
            compare_fn = regression_compare
            finalize_fn = regression_finalize
            save_fn = regression_save
        
        exp = setup_fn(
            data=train_data,
            target=target_column,
            session_id=self.config.get('session_id', 42),
            fold=self.config.get('fold', 5),
            verbose=self.config.get('verbose', False),
            use_gpu=self.config.get('use_gpu', False),
        )
        
        best_model = compare_fn(
            n_select=self.config.get('n_select', 5),
            verbose=self.config.get('verbose', False),
        )
        
        if self.config.get('tuning', {}).get('enabled', True):
            from pycaret.classification import tune_model as classification_tune
            from pycaret.regression import tune_model as regression_tune
            tune_fn = classification_tune if self.is_classification else regression_tune
            
            best_model = tune_fn(
                best_model,
                n_iter=self.config.get('tuning', {}).get('n_iter', 10),
                optimize=self.config.get('tuning', {}).get('optimize', 'Accuracy'),
            )
        
        self.model = finalize_fn(best_model)
        
        save_fn(self.model, str(model_path / 'model'))
        
        from pycaret.classification import pull as classification_pull
        from pycaret.regression import pull as regression_pull
        pull_fn = classification_pull if self.is_classification else regression_pull
        
        results = pull_fn()
        
        metrics = {
            'score': float(results.iloc[0]['Mean']) if not results.empty else 0.0,
        }
        
        logger.info(f"PyCaret training completed. Score: {metrics['score']}")
        return metrics
    
    def predict(self, data: pd.DataFrame) -> pd.Series:
        if self.model is None:
            raise ValueError("Model not trained. Call train() first.")
        
        from pycaret.classification import predict_model as classification_predict
        from pycaret.regression import predict_model as regression_predict
        predict_fn = classification_predict if self.is_classification else regression_predict
        
        predictions = predict_fn(self.model, data=data)
        return predictions.iloc[:, -1]
    
    def load(self, model_path: Path):
        logger.info(f"Loading PyCaret model from {model_path}")
        load_fn = classification_load if self.is_classification else regression_load
        self.model = load_fn(str(model_path / 'model'))
        return self