import lightgbm as lgb
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler, StandardScaler, PolynomialFeatures
from sklearn.feature_selection import SelectKBest, mutual_info_regression, f_regression, SelectFromModel
from sklearn.decomposition import PCA
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.neural_network import MLPRegressor
from sklearn.pipeline import Pipeline
from sklearn.linear_model import LinearRegression, Ridge, RidgeCV
from sklearn.svm import SVR

from config.conmmon_config import RANDOM_STATE


def feature_manipulation_method(manip_method: str):
    if manip_method == 'polynomial':
        return PolynomialFeatures(degree=2, interaction_only=True, include_bias=False, order='C')
    else:
        raise ValueError(f"Invalid feature manipulation method: {manip_method}")


def normalize(normalization_method: str):
    if normalization_method == 'standard':
        return StandardScaler()
    elif normalization_method == 'minmax':
        return MinMaxScaler()
    else:
        raise ValueError(f"Invalid normalization method: {normalization_method}")
        

def feature_selection(feature_selection_method: str, top_features: int):
    if feature_selection_method == 'mutual_info_regression':
        return SelectKBest(mutual_info_regression, k=top_features)
    elif feature_selection_method == 'f_regression':
        return SelectKBest(f_regression, k=top_features)
    elif feature_selection_method == 'pca':
        return PCA(n_components=top_features, random_state=RANDOM_STATE)
    elif feature_selection_method == 'model_rf':
        rf = RandomForestRegressor(n_estimators=50, random_state=RANDOM_STATE)
        # selector = SelectFromModel(rf, max_features=top_features, threshold=-np.inf)  # threshold='median',
        selector = SelectFromModel(rf, threshold='median')
        return selector
    else:
        raise ValueError(f"Invalid feature selection method: {feature_selection_method}")
    

def get_model(model_name: str):
    if model_name.startswith('random_forest'):
        return RandomForestRegressor(random_state=RANDOM_STATE)
    elif model_name.startswith('gbm'):
        return GradientBoostingRegressor(random_state=RANDOM_STATE)
    elif model_name.startswith('lightgbm'):
        return lgb.LGBMRegressor(random_state=RANDOM_STATE, deterministic=True, objective='regression')
    elif model_name.startswith('mlp'):
        return MLPRegressor(random_state=RANDOM_STATE)
    elif model_name.startswith('LR'):
        return LinearRegression(fit_intercept=False)
    elif model_name.startswith('Ridge'):
        return Ridge(random_state=RANDOM_STATE)
    elif model_name.startswith('RidgeCV'):
        return RidgeCV()
    elif model_name.startswith('svm'):
        return SVR()
    else:
        raise ValueError(f"Invalid model name: {model_name}")


def get_model_params(model_name: str, para_dict: dict):
    if model_name.startswith('lightgbm'):
        model = lgb.LGBMRegressor(
            num_leaves=para_dict['model__num_leaves'],
            max_depth=para_dict['model__max_depth'],
            learning_rate=para_dict['model__learning_rate'],
            n_estimators=para_dict['model__n_estimators'],
            min_child_samples=para_dict['model__min_child_samples'],
            min_child_weight=para_dict['model__min_child_weight'],
            reg_alpha=para_dict['model__lambda_l1'],
            reg_lambda=para_dict['model__lambda_l2'],
            colsample_bytree=para_dict['model__feature_fraction'],
            subsample=para_dict['model__bagging_fraction'],
            random_state=RANDOM_STATE,
            deterministic=True,
            objective='regression')
        return model
    elif model_name.startswith('mlp'):
        model = MLPRegressor(
            hidden_layer_sizes=para_dict['model__hidden_layer_sizes'],
            learning_rate=para_dict['model__learning_rate'],
            max_iter=para_dict['model__max_iter'],
            solver=para_dict['model__solver'],
            alpha=para_dict['model__alpha'],
            activation=para_dict['model__activation'],
            random_state=RANDOM_STATE)
        return model
    elif model_name.startswith('gbm'):
        model = GradientBoostingRegressor(
            n_estimators=para_dict['model__n_estimators'],
            learning_rate=para_dict['model__learning_rate'],
            max_depth=para_dict['model__max_depth'],
            max_leaf_nodes=para_dict['model__max_leaf_nodes'],
            subsample=para_dict['model__subsample'],
            min_samples_split=para_dict['model__min_samples_split'],
            min_samples_leaf=para_dict['model__min_samples_leaf'],
            max_features=para_dict['model__max_features'],
            loss=para_dict['model__loss'],
            alpha=para_dict['model__alpha'],# only useful for loss=huber
            random_state=RANDOM_STATE
        )
        return model
    elif model_name.startswith('random_forest'):
        model = RandomForestRegressor(
            n_estimators=para_dict['model__n_estimators'],
            max_depth=para_dict['model__max_depth'],
            min_samples_split=para_dict['model__min_samples_split'],
            min_samples_leaf=para_dict['model__min_samples_leaf'],
            max_features=para_dict['model__max_features'],
            criterion=para_dict['model__criterion'],
            random_state=RANDOM_STATE
        )
        return model
    else:
        raise ValueError(f"Invalid model name: {model_name}")


def get_lgb_features_from_pipeline(pipeline: Pipeline):
    all_steps = pipeline.steps
    for name, model in all_steps:
        if name == 'scaler':
            return model.feature_names_in_
        else:
            print('failed to get feature names from pipeline')
            return None


def get_lgb_feature_importance_from_pipeline(pipeline: Pipeline, feature_names=None):
    all_steps = pipeline.steps
    if feature_names is None:
        feature_names = get_lgb_features_from_pipeline(pipeline)
    for name, model in all_steps:
        if name == 'selector':
            feature_scores = model.scores_
            break
    select_k = model.k
    ordered_index = np.argsort(feature_scores)
    select_index = ordered_index[::-1][:select_k]
    select_features = [feature_names[i] for i in sorted(select_index)]

    for name, model in all_steps:
        if name == 'model':
            f_values = model.feature_importances_
            break
    f_df = pd.DataFrame({'features': select_features, 'importance': f_values})
    f_df = f_df.sort_values(by='importance', ascending=False)
    f_df.reset_index(drop=True, inplace=True)
    return f_df