import argparse
import json
import gdown
import os
import matplotlib.pyplot as plt
import numpy as np
import streamlit as st

from pathlib import Path

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import tensorflow as tf
from tensorflow.keras.models import load_model


@st.cache(show_spinner=False)
def download_weights(model_choice):
    """
    Downloads model weights for deployment
    """

    # Create directory
    save_dest = Path("models")
    save_dest.mkdir(exist_ok=True)

    # Download weights for the chosen model
    if model_choice == "DenseNet (baseline)":
        url = "https://drive.google.com/uc?id=10-TWkCW_BAZLpGXkxPqXFV8lg-jnWNJD"
        output = "models/densenet.h5"

        if not Path(output).exists():
            with st.spinner("Model weights were not found, downloading them. This may take a while."):
                gdown.download(url, output, quiet=False)

    elif model_choice == "VGG16 (baseline)":
        url = "https://drive.google.com/uc?id=1UaNIHQ-HYeN5v6egV9kAdwU0Nb4CfLBF"
        output = "models/vgg16.h5"

        if not Path(output).exists():
            with st.spinner("Model weights were not found, downloading them. This may take a while."):
                gdown.download(url, output, quiet=False)

    elif model_choice == "DenseNet (best)":
        url = "https://drive.google.com/uc?id=1JUvuzyGQpScHyq2q25yhG962g3PMJ1eu"
        output = "models/densenet_best.h5"

        if not Path(output).exists():
            with st.spinner("Model weights were not found, downloading them. This may take a while."):
                gdown.download(url, output, quiet=False)

    elif model_choice == "VGG16 (best)":
        url = "https://drive.google.com/uc?id=19iu-Qhaofczgl6iMt6DSB_OHDBs9ggsr"
        output = "models/vgg16_best.h5"

        if not Path(output).exists():
            with st.spinner("Model weights were not found, downloading them. This may take a while."):
                gdown.download(url, output, quiet=False)
    else:
        raise ValueError("Unknown model: {}".format(model_choice))


def preprocess_image(image_file):
    """Preprocess image"""

    x, _ = process_path(image_file)
    x = np.expand_dims(x, axis=0)

    return x


def app_dof_predict(model_choice, image_file):

    # Download weights for the chosen model
    download_weights(model_choice)
    image = preprocess_image(image_file)
    prediction = {}

    if model_choice == "DenseNet (baseline)":
        model = load_model("models/densenet.h5", compile=False)
    elif model_choice == "VGG16 (baseline)":
        model = load_model("models/vgg16.h5", compile=False)
    elif model_choice == "DenseNet (best)":
        model = load_model("models/densenet_best.h5", compile=False)
    elif model_choice == "VGG16 (best)":
        model = load_model("models/vgg16_best.h5", compile=False)
    preds = model.predict(image)

    prediction = {
        "class": int(np.argmax(preds)),
        "probability": float(preds[0][np.argmax(preds)]),
    }

    return prediction


def decode_img(img):
    """Decode image and resize"""
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, [200, 300])

    return img


def process_path(file_path):
    """Process input path"""
    img = tf.io.read_file(file_path)
    img = decode_img(img)

    return img, file_path


def plot_results(infer_images, inference_predicted_class, inference_predictions, class_names=["bokeh", "no bokeh"]):
    """Plot four images with predicted class and probabilities"""
    plt.figure(figsize=(40, 60))

    for i, (infer_img, _) in enumerate(infer_images.take(10)):
        ax = plt.subplot(2, 5, i + 1)
        plt.imshow(infer_img.numpy() / 255)

        # Find the predicted class from predictions

        m = "Predicted: {}, {:.2f}%".format(class_names[inference_predicted_class[i]], inference_predictions[i] * 100)
        plt.title(m)
        plt.axis("off")
    plt.show()


def dof_predict(infer_images, model_path):

    trained_model = load_model(model_path, compile=False)

    inference_predicted_class = []
    inference_predictions = []
    results = {}
    for infer_img, img_name in infer_images:
        print(img_name)
        preds = trained_model.predict(tf.expand_dims(infer_img, axis=0))
        inference_predicted_class.append(np.argmax(preds))
        print(preds)
        inference_predictions.append(preds[0][np.argmax(preds)])

        results[str(img_name.numpy().decode("utf8").split("/")[-1])] = {
            "class": int(np.argmax(preds)),
            "prob": float(preds[0][np.argmax(preds)]),
        }

    plot_results(infer_images, inference_predicted_class, inference_predictions)

    return results


def save_results(results):
    """Save results to json"""
    json.dump(results, open("results.json", "w"))


def main(test_dir, model_path):

    # get the count of image files in the train directory
    inference_ds = tf.data.Dataset.list_files(test_dir + "/*", shuffle=False)

    infer_images = inference_ds.map(process_path)

    # inference
    results = dof_predict(infer_images, model_path)

    # save results
    save_results(results)


if __name__ == "__main__":

    # Initiate the parser
    parser = argparse.ArgumentParser()

    parser.add_argument("-data", action="store", help="Dataset path")
    parser.add_argument("-model", action="store", help="Model path")
    arguments = parser.parse_args()

    dataset = arguments.data
    model = arguments.model
    main(dataset, model)