

import dash

from catboost import CatBoostClassifier, CatBoostRegressor
from xgboost import XGBClassifier, XGBRegressor
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor

from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer
from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names
from explainerdashboard.dashboards import ExplainerDashboard


def get_classification_explainer(xgboost=False, include_y=True):
    X_train, y_train, X_test, y_test = titanic_survive()
    if xgboost:
        model = XGBClassifier().fit(X_train, y_train)
    else:
        model = RandomForestClassifier(n_estimators=50, max_depth=10).fit(X_train, y_train)
    if include_y:
        explainer = ClassifierExplainer(
                            model, X_test, y_test, 
                            cats=['Sex', 'Deck', 'Embarked'],
                            labels=['Not survived', 'Survived'])
    else:
        explainer = ClassifierExplainer(
                            model, X_test, 
                            cats=['Sex', 'Deck', 'Embarked'],
                            labels=['Not survived', 'Survived'])

    explainer.calculate_properties()
    return explainer


def get_regression_explainer(xgboost=False, include_y=True):
    X_train, y_train, X_test, y_test = titanic_fare()
    train_names, test_names = titanic_names()
    if xgboost:
        model = XGBRegressor().fit(X_train, y_train)
    else:
        model = RandomForestRegressor(n_estimators=50, max_depth=10).fit(X_train, y_train)

    if include_y:
        reg_explainer = RegressionExplainer(model, X_test, y_test, 
                                        cats=['Sex', 'Deck', 'Embarked'], 
                                        idxs=test_names, 
                                        units="$")
    else:
        reg_explainer = RegressionExplainer(model, X_test, 
                                        cats=['Sex', 'Deck', 'Embarked'], 
                                        idxs=test_names, 
                                        units="$")

    reg_explainer.calculate_properties()
    return reg_explainer

def get_multiclass_explainer(xgboost=False, include_y=True):
    X_train, y_train, X_test, y_test = titanic_embarked()
    train_names, test_names = titanic_names()
    if xgboost:
        model = XGBClassifier().fit(X_train, y_train)
    else:
        model = RandomForestClassifier(n_estimators=50, max_depth=10).fit(X_train, y_train)

    if include_y:
        if xgboost:
            multi_explainer = ClassifierExplainer(model, X_test, y_test,
                                            model_output='logodds',
                                            cats=['Sex', 'Deck'], 
                                            labels=['Queenstown', 'Southampton', 'Cherbourg'])
        else:
            multi_explainer = ClassifierExplainer(model, X_test, y_test,
                                            cats=['Sex', 'Deck'], 
                                            labels=['Queenstown', 'Southampton', 'Cherbourg'])
    else:
        if xgboost:
            multi_explainer = ClassifierExplainer(model, X_test, 
                                            model_output='logodds',
                                            cats=['Sex', 'Deck'], 
                                            labels=['Queenstown', 'Southampton', 'Cherbourg'])
        else:
            multi_explainer = ClassifierExplainer(model, X_test, 
                                            cats=['Sex', 'Deck'], 
                                            labels=['Queenstown', 'Southampton', 'Cherbourg'])

    multi_explainer.calculate_properties()
    return multi_explainer


def get_catboost_classifier():
    X_train, y_train, X_test, y_test = titanic_survive()
    train_names, test_names = titanic_names()

    model = CatBoostClassifier(iterations=100, verbose=0).fit(X_train, y_train)
    explainer = ClassifierExplainer(
                        model, X_test, y_test, 
                        cats=[{'Gender': ['Sex_female', 'Sex_male', 'Sex_nan']}, 
                                            'Deck', 'Embarked'],
                        labels=['Not survived', 'Survived'],
                        idxs=test_names)

    X_cats, y_cats = explainer.X_merged, explainer.y.astype("int")
    model = CatBoostClassifier(iterations=5, verbose=0).fit(X_cats, y_cats, cat_features=[5, 6, 7])
    explainer = ClassifierExplainer(model, X_cats, y_cats, idxs=X_test.index)
    explainer.calculate_properties(include_interactions=False)
    return explainer


def get_catboost_regressor():
    X_train, y_train, X_test, y_test = titanic_fare()

    model = CatBoostRegressor(iterations=5, verbose=0).fit(X_train, y_train)
    explainer = RegressionExplainer(model, X_test, y_test, 
                                    cats=["Sex", 'Deck', 'Embarked'])
    X_cats, y_cats = explainer.X_merged, explainer.y
    model = CatBoostRegressor(iterations=5, verbose=0).fit(X_cats, y_cats, cat_features=[5, 6, 7])
    explainer = RegressionExplainer(model, X_cats, y_cats, idxs=X_test.index)
    explainer.calculate_properties(include_interactions=False)
    return explainer


def test_classification_dashboard(dash_duo):
    explainer = get_classification_explainer()
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_regression_dashboard(dash_duo):
    explainer = get_regression_explainer()
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=20)
    assert dash_duo.get_logs() == [], "browser console should contain no error"

def test_simple_classification_dashboard(dash_duo):
    explainer = get_classification_explainer()
    db = ExplainerDashboard(explainer, title="testing", responsive=False, simple=True)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("#simple-classifier-composite-title", "testing", timeout=20)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_simple_regression_dashboard(dash_duo):
    explainer = get_regression_explainer()
    db = ExplainerDashboard(explainer, title="testing", responsive=False, simple=True)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("#simple-regression-composite-title", "testing", timeout=20)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_multiclass_dashboard(dash_duo):
    explainer = get_multiclass_explainer()
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_xgboost_classification_dashboard(dash_duo):
    explainer = get_classification_explainer(xgboost=True)
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_xgboost_regression_dashboard(dash_duo):
    explainer = get_regression_explainer(xgboost=True)
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_xgboost_multiclass_dashboard(dash_duo):
    explainer = get_multiclass_explainer(xgboost=True)
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_classification_dashboard_no_y(dash_duo):
    explainer = get_classification_explainer(include_y=False)
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_regression_dashboard_no_y(dash_duo):
    explainer = get_regression_explainer(include_y=False)
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_multiclass_dashboard_no_y(dash_duo):
    explainer = get_multiclass_explainer(include_y=False)
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_catboost_classification_dashboard(dash_duo):
    explainer = get_catboost_classifier()
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"


def test_cat_boost_regression_dashboard(dash_duo):
    explainer = get_catboost_regressor()
    db = ExplainerDashboard(explainer, title="testing", responsive=False)
    dash_duo.start_server(db.app)
    dash_duo.wait_for_text_to_equal("h1", "testing", timeout=30)
    assert dash_duo.get_logs() == [], "browser console should contain no error"
[(15, 'explainerdashboard.datasets.titanic_survive', 'titanic_survive', False, 'from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names\n'), (36, 'explainerdashboard.datasets.titanic_fare', 'titanic_fare', False, 'from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names\n'), (37, 'explainerdashboard.datasets.titanic_names', 'titanic_names', False, 'from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names\n'), (58, 'explainerdashboard.datasets.titanic_embarked', 'titanic_embarked', False, 'from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names\n'), (59, 'explainerdashboard.datasets.titanic_names', 'titanic_names', False, 'from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names\n'), (91, 'explainerdashboard.datasets.titanic_survive', 'titanic_survive', False, 'from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names\n'), (92, 'explainerdashboard.datasets.titanic_names', 'titanic_names', False, 'from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names\n'), (95, 'explainerdashboard.explainers.ClassifierExplainer', 'ClassifierExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (104, 'explainerdashboard.explainers.ClassifierExplainer', 'ClassifierExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (110, 'explainerdashboard.datasets.titanic_fare', 'titanic_fare', False, 'from explainerdashboard.datasets import titanic_survive, titanic_fare, titanic_embarked, titanic_names\n'), (113, 'explainerdashboard.explainers.RegressionExplainer', 'RegressionExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (117, 'explainerdashboard.explainers.RegressionExplainer', 'RegressionExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (124, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (132, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (139, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (147, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (155, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (163, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (171, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (179, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (187, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (195, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (203, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (211, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (219, 'explainerdashboard.dashboards.ExplainerDashboard', 'ExplainerDashboard', False, 'from explainerdashboard.dashboards import ExplainerDashboard\n'), (21, 'explainerdashboard.explainers.ClassifierExplainer', 'ClassifierExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (26, 'explainerdashboard.explainers.ClassifierExplainer', 'ClassifierExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (44, 'explainerdashboard.explainers.RegressionExplainer', 'RegressionExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (49, 'explainerdashboard.explainers.RegressionExplainer', 'RegressionExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (67, 'explainerdashboard.explainers.ClassifierExplainer', 'ClassifierExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (72, 'explainerdashboard.explainers.ClassifierExplainer', 'ClassifierExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (77, 'explainerdashboard.explainers.ClassifierExplainer', 'ClassifierExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (82, 'explainerdashboard.explainers.ClassifierExplainer', 'ClassifierExplainer', False, 'from explainerdashboard.explainers import ClassifierExplainer, RegressionExplainer\n'), (94, 'catboost.CatBoostClassifier', 'CatBoostClassifier', False, 'from catboost import CatBoostClassifier, CatBoostRegressor\n'), (103, 'catboost.CatBoostClassifier', 'CatBoostClassifier', False, 'from catboost import CatBoostClassifier, CatBoostRegressor\n'), (112, 'catboost.CatBoostRegressor', 'CatBoostRegressor', False, 'from catboost import CatBoostClassifier, CatBoostRegressor\n'), (116, 'catboost.CatBoostRegressor', 'CatBoostRegressor', False, 'from catboost import CatBoostClassifier, CatBoostRegressor\n'), (17, 'xgboost.XGBClassifier', 'XGBClassifier', False, 'from xgboost import XGBClassifier, XGBRegressor\n'), (19, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', False, 'from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor\n'), (39, 'xgboost.XGBRegressor', 'XGBRegressor', False, 'from xgboost import XGBClassifier, XGBRegressor\n'), (41, 'sklearn.ensemble.RandomForestRegressor', 'RandomForestRegressor', False, 'from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor\n'), (61, 'xgboost.XGBClassifier', 'XGBClassifier', False, 'from xgboost import XGBClassifier, XGBRegressor\n'), (63, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', False, 'from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor\n')]