from pathlib import Path
from typing import Dict, Any, Optional
import pandas as pd
import numpy as np
from flaml.automl.automl import AutoML
from sklearn.metrics import accuracy_score, r2_score

from mlpipeline.logging.logger import get_logger
from mlpipeline.utils.common import save_pickle, load_pickle

logger = get_logger(__name__)


class FLAMLTrainer:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.automl: Optional[AutoML] = None
        self.task = config.get('task', 'classification')
    
    def train(self, train_data: pd.DataFrame, target_column: str, model_path: Path) -> Dict[str, float]:
        logger.info("Starting FLAML training")
        
        X_train = train_data.drop(columns=[target_column])
        y_train = train_data[target_column]
        
        self.automl = AutoML()
        
        settings = {
            'time_budget': self.config.get('time_budget', 600),
            'metric': self.config.get('metric', 'auto'),
            'task': self.task,
            'estimator_list': self.config.get('estimator_list', ['lgbm', 'xgboost', 'rf']),
            'n_jobs': self.config.get('n_jobs', -1),
            'verbose': self.config.get('verbose', 1),
            'early_stop': self.config.get('early_stop', True),
        }
        
        self.automl.fit(X_train=X_train, y_train=y_train, **settings)
        
        y_pred = self.automl.predict(X_train)
        
        if self.task == 'classification':
            score = accuracy_score(y_train, y_pred)
            metric_name = 'accuracy'
        else:
            score = r2_score(y_train, y_pred)
            metric_name = 'r2_score'
        
        save_pickle(model_path / 'model.pkl', self.automl)
        
        metrics = {
            metric_name: float(score),
            'best_loss': float(self.automl.best_loss),
        }
        
        logger.info(f"FLAML training completed. Best {metric_name}: {score}")
        return metrics
    
    def predict(self, data: pd.DataFrame) -> np.ndarray:
        if self.automl is None:
            raise ValueError("Model not trained. Call train() first.")
        return self.automl.predict(data)
    
    def load(self, model_path: Path):
        logger.info(f"Loading FLAML model from {model_path}")
        self.automl = load_pickle(model_path / 'model.pkl')
        return self