import pandas as pd
import numpy as np
from sklearn.linear_model import SGDClassifier
from sklearn.model_selection import RandomizedSearchCV
from .config import Config

class ModelTrainer:
    def __init__(self, random_state=None):
        self.random_state = random_state if random_state else Config.RANDOM_STATE
        self.model = None
        self.best_params = None
        self.cv_results = None
        
    def train(self, X, y, param_distributions=None, n_iter=10, cv=5, scoring='roc_auc'):
        """
        Train SGD Classifier with Randomized Search.
        SGDClassifier supports incremental learning via partial_fit().
        """
        # Handle NaN values - SGDClassifier doesn't support NaN
        X_clean = X.fillna(0)
        
        if param_distributions is None:
            param_distributions = {
                'alpha': np.logspace(-6, 0, 20),
                'learning_rate': ['constant', 'adaptive', 'invscaling'],
                'eta0': [0.01, 0.1, 1.0],
                'penalty': ['l2'],
                'loss': [Config.SGD_LOSS],
                'max_iter': [1000, 2000],
                'tol': [1e-3, 1e-4]
            }
            
        sgd = SGDClassifier(
            loss=Config.SGD_LOSS,
            penalty='l2',
            random_state=self.random_state,
            warm_start=False
        )
        
        search = RandomizedSearchCV(
            estimator=sgd,
            param_distributions=param_distributions,
            n_iter=n_iter,
            cv=cv,
            scoring=scoring,
            random_state=self.random_state,
            n_jobs=1,  # Avoid joblib parallel issues on Windows or small data
            verbose=0
        )
        
        search.fit(X_clean, y)
        
        # Validate that we have samples to train
        if X_clean.shape[0] == 0:
            raise ValueError("No samples available for training after cleaning")
        
        self.model = search.best_estimator_
        self.best_params = search.best_params_
        self.cv_results = pd.DataFrame(search.cv_results_)
        
        return self.model
    
    def partial_fit(self, X, y, classes=None, sample_weight=None):
        """
        Incremental learning: Update model with new data.
        This is useful for online learning scenarios where data arrives in batches.
        
        Parameters:
        -----------
        X : array-like of shape (n_samples, n_features)
            Training data
        y : array-like of shape (n_samples,)
            Target values
        classes : array-like of shape (n_classes,), default=None
            All classes for the model. Required on first call to partial_fit.
        sample_weight : array-like of shape (n_samples,), default=None
            Sample weights
            
        Returns:
        --------
        self : object
            Returns the instance itself.
        """
        if self.model is None:
            self.model = SGDClassifier(
                loss=Config.SGD_LOSS,
                penalty='l2',
                random_state=self.random_state,
                warm_start=False
            )
        
        # Handle NaN values
        X_clean = X.fillna(0) if hasattr(X, 'fillna') else X
        
        # If classes not provided, try to infer from y if this is first call (which shouldn't happen usually in partial_fit if not init)
        # But if model is new, we need classes.
        if classes is None and not hasattr(self.model, 'classes_'):
            classes = np.unique(y)
            
        # Use partial_fit for incremental learning
        self.model.partial_fit(X_clean, y, classes=classes, sample_weight=sample_weight)
        
        return self
        
    def set_coefficients(self, coefs, intercept):
        """
        Manually set coefficients for model.
        coefs: array-like of shape (n_features,)
        intercept: float
        """
        if self.model is None:
            self.model = SGDClassifier(
                loss=Config.SGD_LOSS,
                penalty='l2',
                random_state=self.random_state
            )
            raise ValueError("Model must be trained/initialized first.")
            
        if hasattr(self.model, 'best_estimator_'):
             estimator = self.model.best_estimator_
        else:
             estimator = self.model
             
        estimator.coef_ = np.array([coefs])
        estimator.intercept_ = np.array([intercept])
        
    def predict_proba(self, X):
        if self.model is None:
            raise ValueError("Model not trained.")
        return self.model.predict_proba(X)[:, 1]