#!/usr/bin/env python
# coding: utf-8

# # Assess predictions on image classification fridge data with a pytorch model

# This notebook demonstrates the use of the `responsibleai` API to assess an image classification pytorch model trained on the fridge dataset. It walks through the API calls necessary to create a widget with model analysis insights, then guides a visual analysis of the model.

# ## Launch Responsible AI Toolbox

# The following section examines the code necessary to create datasets and a model. It then generates insights using the `responsibleai` API that can be visually analyzed.

# ### Load Model and Data

# In[ ]:


import os
import sys
from zipfile import ZipFile
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from responsibleai_vision.common.constants import ImageColumns
import json
from fastai.learner import load_learner
from raiutils.common.retries import retry_function

try:
    from urllib import urlretrieve
except ImportError:
    from urllib.request import urlretrieve

EPOCHS = 10
LEARNING_RATE = 1e-4
IM_SIZE = 300
BATCH_SIZE = 16
FRIDGE_MODEL_NAME = 'fridge_model'
FRIDGE_MODEL_WINDOWS_NAME = 'fridge_model_windows'
WIN = 'win'

def load_fridge_dataset():
    # create data folder if it doesnt exist.
    os.makedirs("data", exist_ok=True)

    # download data
    download_url = ("https://publictestdatasets.blob.core.windows.net/" +
                    "computervision/fridgeObjects.zip")
    data_file = "./data/fridgeObjects.zip"
    urlretrieve(download_url, filename=data_file)

    # extract files
    with ZipFile(data_file, "r") as zip:
        print("extracting files...")
        zip.extractall(path="./data")
        print("done")
    # delete zip file
    os.remove(data_file)
    # get all file names into a pandas dataframe with the labels
    data = pd.DataFrame(columns=[ImageColumns.IMAGE.value,
                                 ImageColumns.LABEL.value])
    
    feats = []
    for folder in os.listdir("./data/fridgeObjects"):
        for file in os.listdir("./data/fridgeObjects/" + folder):
            image_path = "./data/fridgeObjects/" + folder + "/" + file
            feats.append({ImageColumns.IMAGE.value: image_path,
                     ImageColumns.LABEL.value: folder})
    
    # get all file names into a pandas dataframe with the labels
    data = pd.DataFrame(feats, columns=[ImageColumns.IMAGE.value,
                                        ImageColumns.LABEL.value])

    return data


class FetchModel(object):
    def __init__(self):
        pass

    def fetch(self):
        if sys.platform.startswith(WIN):
            model_name = FRIDGE_MODEL_WINDOWS_NAME
        else:
            model_name = FRIDGE_MODEL_NAME
        url = ('https://publictestdatasets.blob.core.windows.net/models/' +
               model_name)
        urlretrieve(url, FRIDGE_MODEL_NAME)


def retrieve_or_train_fridge_model(df, force_train=False):
    if force_train:
        model = train_fastai_image_classifier(df)
        # Save model to disk
        model.export(FRIDGE_MODEL_NAME)
    else:
        fetcher = FetchModel()
        action_name = "Dataset download"
        err_msg = "Failed to download dataset"
        max_retries = 4
        retry_delay = 60
        retry_function(fetcher.fetch, action_name, err_msg,
                       max_retries=max_retries,
                       retry_delay=retry_delay)
        model = load_learner(FRIDGE_MODEL_NAME)
    return model


# ### Load Pretrained Faster RCNN model

# In[ ]:


data = load_fridge_dataset()
model = retrieve_or_train_fridge_model(data)

train_data = data
test_data = data
class_names = data[ImageColumns.LABEL.value].unique()


# ### Create Model and Data Insights

# In[ ]:


from raiwidgets import ResponsibleAIDashboard
from responsibleai_vision import ModelTask, RAIVisionInsights


# To use Responsible AI Toolbox, initialize a RAIInsights object upon which different components can be loaded.
# 
# RAIInsights accepts the model, the full dataset, the test dataset, the target feature string and the task type string as its arguments.

# In[ ]:


rai_insights = RAIVisionInsights(model, test_data.sample(10, random_state=42),
                                 "label",
                                 task_type=ModelTask.IMAGE_CLASSIFICATION,
                                 classes=class_names)

rai_insights


# Add the components of the toolbox that are focused on model assessment.

# In[ ]:


# Interpretability
rai_insights.explainer.add()
# Error Analysis
rai_insights.error_analysis.add()


# Once all the desired components have been loaded, compute insights on the test set.

# In[ ]:


rai_insights.compute()


# Finally, visualize and explore the model insights. Use the resulting widget or follow the link to view this in a new tab.

# In[ ]:


ResponsibleAIDashboard(rai_insights)