from sklearn.model_selection import GridSearchCV
from sklearn.feature_selection import RFECV, RFE

def do_clf(clf, grid_params, data, scaled=True):
    X_train, X_test = data.get_X_train_test(scaled)
    y_train, y_test = data.get_y_train_test()
    id_train, id_test = data.get_id_train_test()

    info = {}
    info['name'] = clf.__class__.__name__
    print(clf.__class__.__name__)

    # Do GridSearch.
    est = GridSearchCV(clf, grid_params)
    # Fit on training data.
    est.fit(X_train, y_train)

    # Save training score.
    info['train_score'] = est.score(X_train, y_train)
    # Save best parameters.
    info['best_params'] = est.best_params_
    try:
        # Save coefficient (if it has some).
        info['coefficients'] = est.best_estimator_.coef_.tolist()
        if info['name'] == 'SVR':
            info['coefficients'] = est.best_estimator_.coef_[0].tolist()
    except (AttributeError, ValueError):
        pass
    try:
        # Save feature importances (random forest).
        info['feature_importances'] = est.best_estimator_.feature_importances_.tolist()
    except AttributeError:
        pass
    # Save test score.
    info['test_score'] = est.score(X_test, y_test)

    # Save predictions.
    info['train_predicted'] = list(zip(id_train, est.predict(X_train)))
    info['test_predicted'] = list(zip(id_test, est.predict(X_test)))
    return info

def do_rfecv_clf(clf, grid_params, data, scaled=True):
    X_train, X_test = data.get_X_train_test(scaled)
    y_train, y_test = data.get_y_train_test()
    id_train, id_test = data.get_id_train_test()

    info = {}
    info['name'] = clf.__class__.__name__
    print(clf.__class__.__name__)

    # Do GridSearch on RFECV.
    selector = RFECV(clf, step=1)
    est = GridSearchCV(selector, grid_params)
    # Fit on training data.
    est.fit(X_train, y_train)

    # Save training score.
    info['train_score'] = est.score(X_train, y_train)
    # Save best parameters.
    info['best_params'] = est.best_params_
    # Save support vector.
    info['support'] = est.best_estimator_.support_.tolist()
    try:
        # Save coefficient (if it has some).
        info['coefficients'] = est.best_estimator_.estimator_.coef_.tolist()
        if info['name'] == 'SVR':
            info['coefficients'] = est.best_estimator_.estimator_.coef_[0].tolist()
    except (AttributeError, ValueError):
        pass
    try:
        # Save feature importances (random forest).
        info['feature_importances'] = est.best_estimator_.estimator_.feature_importances_.tolist()
    except AttributeError:
        pass
    # Save test score.
    info['test_score'] = est.score(X_test, y_test)

    # Save predictions.
    info['train_predicted'] = list(zip(id_train, est.predict(X_train)))
    info['test_predicted'] = list(zip(id_test, est.predict(X_test)))
    
    return info

def do_rfe_clf(clf, grid_params, feature_cnt, data, scaled=True):
    X_train, X_test = data.get_X_train_test(scaled)
    y_train, y_test = data.get_y_train_test()
    id_train, id_test = data.get_id_train_test()

    info = {}
    info['name'] = clf.__class__.__name__
    print(clf.__class__.__name__)

    # Do GridSearch on RFE.
    selector = RFE(clf, step=1, n_features_to_select=feature_cnt)
    est = GridSearchCV(selector, grid_params)
    # Fit on training data.
    est.fit(X_train, y_train)

    # Save training score.
    info['train_score'] = est.score(X_train, y_train)
    # Save best parameters.
    info['best_params'] = est.best_params_
    # Save support vector.
    info['support'] = est.best_estimator_.support_.tolist()
    try:
        # Save coefficient (if it has some).
        info['coefficients'] = est.best_estimator_.estimator_.coef_.tolist()
        if info['name'] == 'SVR':
            info['coefficients'] = est.best_estimator_.estimator_.coef_[0].tolist()
    except (AttributeError, ValueError):
        pass
    try:
        # Save feature importances (random forest).
        info['feature_importances'] = est.best_estimator_.estimator_.feature_importances_.tolist()
    except AttributeError:
        pass
    # Save test score.
    info['test_score'] = est.score(X_test, y_test)

    # Save predictions.
    info['train_predicted'] = list(zip(id_train, est.predict(X_train)))
    info['test_predicted'] = list(zip(id_test, est.predict(X_test)))
    
    return info
