
import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
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 Logistic Regression with Randomized Search.
        """
        if param_distributions is None:
            param_distributions = {
                'C': np.logspace(-4, 4, 20),
                'penalty': ['l2'], # 'l1', 'elasticnet' require solver tuning, keeping simple
                'solver': ['lbfgs', 'liblinear']
            }
            
        lr = LogisticRegression(random_state=self.random_state, max_iter=1000)
        
        search = RandomizedSearchCV(
            estimator=lr,
            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, y)
        
        self.model = search.best_estimator_
        self.best_params = search.best_params_
        self.cv_results = pd.DataFrame(search.cv_results_)
        
        return self.model
        
    def set_coefficients(self, coefs, intercept):
        """
        Manually set coefficients for the model.
        coefs: array-like of shape (n_features,)
        intercept: float
        """
        if self.model is None:
            # Initialize a dummy model if not trained
            self.model = LogisticRegression(random_state=self.random_state)
            # Need to initialize classes_ etc. for predict_proba to work
            # This is tricky without fitting.
            # Best practice: Fit on dummy data or ensure train() is called first.
            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]
