import os
import warnings
warnings.filterwarnings('ignore')
import numpy as np
import pandas as pd
import witwidget
from witwidget.notebook.visualization import (
    WitWidget, 
    WitConfigBuilder,
)
pd.options.display.max_columns = 50

"""
Explanation: Inclusive ML - Understanding Bias
Learning Objectives

Invoke the What-if Tool against a deployed Model
Explore attributes of the dataset
Examine aspects of bias in model results
Evaluate how the What-if Tool provides suggestions to remediate bias

Introduction
This notebook shows use of the What-If Tool inside of a Jupyter notebook.  The What-If Tool, among many other things, allows you to explore the impacts of Fairness in model design and deployment.
The notebook invokes a previously deployed XGBoost classifier model on the UCI census dataset which predicts whether a person earns more than $50K based on their census information.
You will then visualize the results of the trained classifier on test data using the What-If Tool.  
First, you will import various libaries and settings that are required to complete the lab.
End of explanation
"""


PROJECT = "<YOUR PROJECT>"

BUCKET = "gs://{project}".format(project=PROJECT)
MODEL = 'xgboost_model'
VERSION = 'v1'
MODEL_DIR = os.path.join(BUCKET, MODEL)


os.environ['PROJECT'] = PROJECT
os.environ['BUCKET'] = BUCKET
os.environ['MODEL'] = MODEL
os.environ['VERSION'] = VERSION
os.environ['MODEL_DIR'] = MODEL_DIR

"""
Explanation: In the next cell, replace &lt;YOUR PROJECT&gt; (inside the double-quotes) with your GCP project id (for example qwiklabs-gcp-fe367d9e174dfbd3):
End of explanation
"""


%%bash

gcloud config set project $PROJECT

gsutil mb $BUCKET
gsutil cp gs://cloud-training-demos/mlfairness/model.bst $MODEL_DIR/model.bst

gcloud ai-platform models list | grep $MODEL || gcloud ai-platform models create $MODEL

gcloud ai-platform versions list --model $MODEL | grep $VERSION ||
gcloud ai-platform versions create $VERSION \
  --model=$MODEL \
  --framework='XGBOOST' \
  --runtime-version=1.14 \
  --origin=$MODEL_DIR \
  --python-version=3.5 \
  --project=$PROJECT

"""
Explanation: Set up the notebook environment
First you must perform a few environment and project configuration steps.  
These steps may take 8 to 10 minutes, please wait until you see the following response before proceeding:
"Creating version (this might take a few minutes)......done."
End of explanation
"""


%%bash

gsutil cp gs://cloud-training-demos/mlfairness/income.pkl .
gsutil cp gs://cloud-training-demos/mlfairness/x_test.npy .
gsutil cp gs://cloud-training-demos/mlfairness/y_test.npy .  

features = pd.read_pickle('income.pkl')
x_test = np.load('x_test.npy')
y_test = np.load('y_test.npy')

"""
Explanation: Finally, download the data and arrays needed to use the What-if Tool.
End of explanation
"""


features.head()

"""
Explanation: Now take a quick look at the data.  The ML model type used for this analysis is XGBoost.  XGBoost is a machine learning framework that uses decision trees and gradient boosting to build predictive models. It works by ensembling multiple decision trees together based on the score associated with different leaf nodes in a tree. 
XGBoost requires all values to be numeric so the orginial dataset was slightly modified.  The biggest change made was to assign a numeric value to Sex.  The originial dataset only had the values "Female" and "Male" for Sex.  The decision was made to assign the value "1" to Female and "2" to Male. As part of the data prepartion effort the Pandas function "get_dummies" was used to convert the remaining domain values into numerical equivalent.  For instance the "Education" column was turned into several sub-columns named after the value in the column.  For instance the "Education_HS-grad" has a value of "1" for when that was the orginial categorical value and a value of "0" for other cateogries.
End of explanation
"""


# Combine the features and labels into one array for the What-if Tool

num_wit_examples = 2000

test_examples = np.hstack((
    x_test[:num_wit_examples],
    y_test[:num_wit_examples].reshape(-1, 1)
))

"""
Explanation: To connect the What-if Tool to an AI Platform model, you need to pass it a subset of your test examples.  The commannd below will create a Numpy array of 2000 from our test examples.
End of explanation
"""


# TODO 1
FEATURE_NAMES = features.columns.tolist() + ['income_prediction']


def adjust(pred):
    return [1 - pred, pred]


config_builder = (
    WitConfigBuilder(test_examples.tolist(), FEATURE_NAMES)
    .set_ai_platform_model(PROJECT, MODEL, VERSION, adjust_prediction=adjust)
    .set_target_feature('income_prediction')
    .set_label_vocab(['low', 'high'])
)

WitWidget(config_builder, height=800)

"""
Explanation: Instantiating the What-if Tool is as simple as creating a WitConfigBuilder object and passing it the AI Platform model desired to be analyzed.
The optional "adjust_prediction" parameter is used because the What-if Tool expects a list of scores for each class in our model (in this case 2). Since the model only returns a single value from 0 to 1, it must be transformed to the correct format in this function.   Lastly, the name 'income_prediction' is used as the ground truth label.
It may take 1 to 2 minutes for the What-if Tool to load and render the visualization palette, please be patient.
End of explanation
"""


