
import pandas as pd
import numpy as np
import joblib
from sklearn.model_selection import train_test_split

from .config import Config
from .eda import FeatureAnalyzer
from .binning import BinningProcess
from .selection import FeatureSelector
from .modeling import ModelTrainer
from .evaluation import ModelEvaluator
from .scorecard import ScorecardTransformer
from .reporting import ReportGenerator
from .plot_utils import plot_binning

class AutoScore:
    def __init__(self, random_state=42):
        self.random_state = random_state
        self.eda = FeatureAnalyzer()
        self.binning = None
        self.selector = FeatureSelector()
        self.trainer = ModelTrainer(random_state=random_state)
        self.evaluator = ModelEvaluator()
        self.sc_transformer = None
        self.reporter = ReportGenerator()
        
        # Intermediate results
        self.data_summary = None
        self.iv_df = None
        self.selected_features = []
        self.model_metrics = {}
        self.scorecard_df = None
        
    def fit(self, X, y, 
            test_size=0.3,
            date_col='create_time',
            oot_split_date=None,
            excluded_features=None,
            special_values=None,
            monotonicity_constraints=None,
            max_bins=10,
            min_bin_pct=0.05,
            iv_threshold=0.02,
            selection_method='stepwise',
            pdo=20, base_score=600, base_odds=50,
            custom_binning_rules=None,
            final_features=None):
        """
        Execute the full scorecard development pipeline.
        """
        # Preprocessing: Exclude features
        if excluded_features:
            X = X.drop(columns=[c for c in excluded_features if c in X.columns])
            
        # 1. Split Data (Train/Test/OOT)
        X_oot, y_oot = None, None
        
        if date_col in X.columns and oot_split_date:
            print(f"Splitting OOT data based on {date_col} > {oot_split_date}")
            mask_oot = X[date_col] > oot_split_date
            X_oot = X[mask_oot].copy()
            y_oot = y[mask_oot].copy()
            X = X[~mask_oot].copy()
            y = y[~mask_oot].copy()
            
            # Drop date column from modeling
            X_oot = X_oot.drop(columns=[date_col])
            X = X.drop(columns=[date_col])
        elif date_col in X.columns:
             X = X.drop(columns=[date_col])
        
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=self.random_state, stratify=y
        )
        
        # 2. EDA
        print("Running EDA...")
        self.data_summary = self.eda.analyze(X_train)
        
        # 3. Binning
        print("Running Binning...")
        self.binning = BinningProcess(max_bins=max_bins, min_bin_pct=min_bin_pct, special_values=special_values)
        
        if custom_binning_rules:
            print("Using custom binning rules...")
            self.binning.set_rules(custom_binning_rules)
            # Recalculate stats for the current training data using these rules
            self.binning.recalculate_stats(X_train, y_train)
        else:
            self.binning.fit(X_train, y_train, monotonicity_constraints=monotonicity_constraints)
        
        self.iv_df = self.binning.iv_df
        
        # Transform to WOE
        X_train_woe = self.binning.transform(X_train)
        X_test_woe = self.binning.transform(X_test)
        X_oot_woe = self.binning.transform(X_oot) if X_oot is not None else None
        
        # 4. Feature Selection
        print("Running Feature Selection...")
        if final_features:
             print(f"Using provided final features ({len(final_features)})...")
             self.selected_features = final_features
             # Filter iv_df for report consistency if needed, but not strictly required for flow
        else:
            # IV Filter
            features_iv = self.selector.filter_by_iv(self.iv_df, threshold=iv_threshold)
            print(f"Features after IV filter: {len(features_iv)}")
            
            # Stepwise
            if selection_method == 'stepwise':
                self.selected_features = self.selector.stepwise_selection(
                    X_train_woe[features_iv], y_train, method='forward'
                )
            else:
                self.selected_features = features_iv
            
        print(f"Final selected features: {len(self.selected_features)}")
        print(self.selected_features)
        
        # 5. Model Training
        print("Training Model...")
        self.trainer.train(X_train_woe[self.selected_features], y_train)
        
        # 6. Evaluation
        print("Evaluating Model...")
        self.evaluate_performance(X_train_woe, y_train, X_test_woe, y_test, X_oot_woe, y_oot)
        
        # 7. Scorecard Generation
        print("Generating Scorecard...")
        self.sc_transformer = ScorecardTransformer(pdo=pdo, base_score=base_score, base_odds=base_odds)
        self.scorecard_df = self.sc_transformer.fit(self.trainer.model, self.binning, self.selected_features)
        
        return self

    def evaluate_performance(self, X_train_woe, y_train, X_test_woe, y_test, X_oot_woe=None, y_oot=None):
        """
        Evaluate model on Train/Test/OOT.
        """
        # Train metrics
        pred_train = self.trainer.predict_proba(X_train_woe[self.selected_features])
        metrics_train = self.evaluator.evaluate(y_train, pred_train)
        
        # Test metrics
        pred_test = self.trainer.predict_proba(X_test_woe[self.selected_features])
        metrics_test = self.evaluator.evaluate(y_test, pred_test)
        
        # PSI (Train vs Test)
        psi_test = self.evaluator.calculate_psi(pred_train, pred_test)
        
        self.model_metrics = {
            'train': metrics_train,
            'test': metrics_test,
            'psi_train_test': psi_test
        }
        
        if X_oot_woe is not None:
            pred_oot = self.trainer.predict_proba(X_oot_woe[self.selected_features])
            metrics_oot = self.evaluator.evaluate(y_oot, pred_oot)
            psi_oot = self.evaluator.calculate_psi(pred_train, pred_oot)
            self.model_metrics['oot'] = metrics_oot
            self.model_metrics['psi_train_oot'] = psi_oot

    def update_model_weights(self, coef_dict, intercept=None):
        """
        Manually update model weights and re-evaluate.
        coef_dict: {feature: weight}
        """
        if self.trainer.model is None:
            print("Model not trained yet.")
            return
            
        # Get estimator
        estimator = self.trainer.model.best_estimator_ if hasattr(self.trainer.model, 'best_estimator_') else self.trainer.model
        
        # Update coefficients
        # Assuming coef_dict contains all selected features
        new_coefs = []
        for feat in self.selected_features:
            if feat in coef_dict:
                new_coefs.append(coef_dict[feat])
            else:
                # Keep original? Or error?
                # Let's try to find index
                idx = self.selected_features.index(feat)
                new_coefs.append(estimator.coef_[0][idx])
                
        estimator.coef_ = np.array([new_coefs])
        
        if intercept is not None:
            estimator.intercept_ = np.array([intercept])
            
        print("Model weights updated.")
        # Note: Should re-run evaluation and scorecard generation ideally
        # But user needs to provide data for evaluation. 
        # So usually this is done before final save.
        
    def adjust_rules(self, new_rules, X_train=None, y_train=None):
        """
        Manually adjust binning rules and visualize.
        """
        if self.binning is None:
            print("Binning process not initialized. Run fit() first or provide custom rules.")
            return
            
        print("Adjusting rules...")
        # Update specific rules
        for col, rule in new_rules.items():
            if col in self.binning.rules:
                # Assuming user provides 'cutoffs' or 'categories' directly
                if 'cutoffs' in rule:
                    self.binning.rules[col]['cutoffs'] = rule['cutoffs']
                if 'categories' in rule:
                    self.binning.rules[col]['categories'] = rule['categories']
            else:
                print(f"Warning: {col} not found in existing rules.")
                
        # Recalculate stats if data provided (highly recommended)
        if X_train is not None and y_train is not None:
            self.binning.recalculate_stats(X_train, y_train)
            self.iv_df = self.binning.iv_df
            
    def visualize_binning(self, feature_name, filename=None, show=True):
        """
        Visualize binning for a feature.
        """
        if self.binning is None or feature_name not in self.binning.rules:
            print(f"No binning info for {feature_name}")
            return
            
        stats = self.binning.rules[feature_name]['stats']
        plot_binning(stats, feature_name, show=show, filename=filename)

    def save_model(self, filepath):
        """Save the entire AutoScore object."""
        joblib.dump(self, filepath)
        
    def load_model(self, filepath):
        """Load AutoScore object."""
        return joblib.load(filepath)
        
    def predict_proba(self, X):
        X_woe = self.binning.transform(X)
        return self.trainer.predict_proba(X_woe[self.selected_features])
        
    def predict_score(self, X):
        X_woe = self.binning.transform(X)
        return self.sc_transformer.transform(X_woe, self.selected_features, self.trainer.model)
        
    def create_report(self, filename='scorecard_report.xlsx'):
        """
        Export all results to Excel.
        """
        # Prepare DataFrames
        metrics_list = [
            {'Metric': 'Train AUC', 'Value': self.model_metrics['train']['auc']},
            {'Metric': 'Train KS', 'Value': self.model_metrics['train']['ks']},
            {'Metric': 'Test AUC', 'Value': self.model_metrics['test']['auc']},
            {'Metric': 'Test KS', 'Value': self.model_metrics['test']['ks']},
            {'Metric': 'PSI (Train vs Test)', 'Value': self.model_metrics['psi_train_test']}
        ]
        
        if 'oot' in self.model_metrics:
            metrics_list.extend([
                {'Metric': 'OOT AUC', 'Value': self.model_metrics['oot']['auc']},
                {'Metric': 'OOT KS', 'Value': self.model_metrics['oot']['ks']},
                {'Metric': 'PSI (Train vs OOT)', 'Value': self.model_metrics['psi_train_oot']}
            ])
            
        metrics_df = pd.DataFrame(metrics_list)
        
        # Model Coefficients
        model = self.trainer.model.best_estimator_ if hasattr(self.trainer.model, 'best_estimator_') else self.trainer.model
        coef_df = pd.DataFrame({
            'Feature': ['Intercept'] + self.selected_features,
            'Coefficient': [model.intercept_[0]] + list(model.coef_[0])
        })
        
        data_dict = {
            'Summary': self.data_summary,
            'IV Analysis': self.iv_df,
            'Selected Features': pd.DataFrame(self.selected_features, columns=['Feature']),
            'Model Coefficients': coef_df,
            'Model Performance': metrics_df,
            'Scorecard': self.scorecard_df
        }
        
        return self.reporter.generate(filename, data_dict)
