import json
import numpy as np
import torch
from fastai.vision.all import *
import shap
import gradio as gr

# Load Fastai model
learner = load_learner('model.pkl')
# Extract PyTorch model
pytorch_model = learner.model.eval()

# Get class names
url = "https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json"
with open(shap.datasets.cache(url)) as file:
    class_names = [v[1] for v in json.load(file).values()]

def predict(img):
    img = np.array(img)
    img_tensor = torch.tensor(img.transpose((2, 0, 1))).float().unsqueeze(0)
    with torch.no_grad():
        output = pytorch_model(img_tensor)
    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    sorted_probs, indices = torch.sort(probabilities, descending=True)
    sorted_labels = [class_names[i] for i in indices]
    top_labels = sorted_labels[:10]
    top_probs = sorted_probs[:10].tolist()
    return dict(zip(top_labels, top_probs))

def f(img):
    img = np.array(img)
    img_tensor = torch.tensor(img.transpose((2, 0, 1))).float().unsqueeze(0)
    with torch.no_grad():
        output = pytorch_model(img_tensor)
    return output

# Interpretation function
def interpretation_function(img):
    masker = shap.maskers.Image("inpaint_telea", [224, 224, 3])
    explainer = shap.PartitionExplainer(f, masker)
    pred = f(img).argmax()
    shap_values = explainer(np.expand_dims(img, 0), max_evals=10)
    scores = shap_values.values[0][:, :, :, pred]
    scores = scores.mean(axis=-1)
    max_val, min_val = np.max(scores), np.min(scores)
    scores = (scores - min_val) / (max_val - min_val)
    return {"original": gr.processing_utils.encode_array_to_base64(img),
            "interpretation": scores.tolist()}

# Gradio UI
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="Input Image", shape=(224, 224))
            with gr.Row():
                classify = gr.Button("Classify")
                interpret = gr.Button("Interpret")
        with gr.Column():
            label = gr.Label(label="Predicted Class")
        with gr.Column():
            interpretation = gr.components.Interpretation(input_img)
    classify.click(predict, input_img, label)
    interpret.click(interpretation_function, input_img, interpretation)

demo.launch()