"""

=============================================================================================

Project     : Chest X-Ray Pathology Detection and Localization using Deep Learning

Author Name : Rammuni Ravidu Suien Silva

UoW No      : 16267097

IIT No      : 2016134

Module      : Final Year Project 20/21

Supervisor  : Mr Pumudu Fernando

Prototype   : Web Interface - BackEnd [Draft: .v01]

University of Westminster, UK || IIT Sri Lanka

=============================================================================================

"""

import json

import os

from datetime import datetime

import numpy as np

# Flask Imports

from flask import Flask, request, render_template

from flask import send_file

from flask_jsglue import JSGlue

# Tensorflow Keras imports

from tensorflow.keras.models import load_model

# For secure src links

from werkzeug.utils import secure_filename

# System Library import

from lab_cxr_scripts.lab_cxr import CXRPrediction, CXRLocalization

# Model 0, 2 :- xray_labels_set[0] || Model 1 :- xray_labels_set[1]

xray_labels_set = [["Enlarged Cardiomediastinum", "Cardiomegaly", "Lung Lesion", "Lung Opacity", "Edema",

                    "Consolidation", "Pneumonia", "Atelectasis", "Pneumothorax", "Pleural Effusion",

                    "Pleural Other", "Fracture", "Support Devices"],

                   ["Nodule", "Cardiomegaly", "Emphysema", "Fibrosis", "Edema", "Consolidation", "Pneumonia",

                    "Atelectasis", "Pneumothorax", "Pleural Effusion", "Mass", "Infiltration", "Hernia",

                    "Plueral Thickening"]]

# Labels for classification tasks

xray_labels = xray_labels_set[0]

# Dependency pip install pyopenssl

# Flask Configs

app = Flask(__name__)

app.config['MAX_CONTENT_LENGTH'] = 20 * 1024 * 1024  # Request data limited to 20MB

jsglue = JSGlue(app)

# TODO: USER GUIDE

# model load

models = [[

    load_model('models/MIMIC/PAR-64-MODEL-MIMIC-FINAL-2.h5',

               custom_objects={'weighted_loss': CXRPrediction.get_weighted_loss(1, 1)}),

    load_model('models/MIMIC/PAR-128-MODEL-MIMIC-FINAL-2.h5',

               custom_objects={'weighted_loss': CXRPrediction.get_weighted_loss(1, 1)})

], [

    load_model('models/NIH/PAR-64-MODEL-FINAL-NIH-2.h5',

               custom_objects={'weighted_loss': CXRPrediction.get_weighted_loss(1, 1)}),

    load_model('models/NIH/PAR-128-MODEL-FINAL-NIH-2.h5',

               custom_objects={'weighted_loss': CXRPrediction.get_weighted_loss(1, 1)})

]]

model = models[0]

cur_cxr_hash = 'none'

"""

==================================================================================================================

                                            Web request functions

==================================================================================================================

"""

# Web page startup

@app.route('/')

def start_web():

    return render_template("index.html")

# CXR Image upload API

@app.route('/predict/<int:model_id>', methods=['GET', 'POST'])

def upload(model_id):

    if request.method == 'POST':

        print("Model ID", model_id)

        # Selecting Model and labels set

        global model, xray_labels

        model = models[model_id % len(models)]

        xray_labels = xray_labels_set[model_id % len(xray_labels_set)]

        global cur_cxr_hash

        preds = []

        file_count = len(request.files)

        if file_count > 8:

            return

        for file_num in range(file_count):

            # Getting image file from post request through the Web

            cxr_img_file = request.files['file_' + str(file_num)]

            # Generating Hash of the image file

            hashed_filename = CXRPrediction.hash_cxr(cxr_img_file)

            print(hashed_filename)

            cur_cxr_hash = hashed_filename

            # Saving the CXR image to uploads

            cxr_img_path = os.path.dirname(__file__)

            file_path = os.path.join(

                cxr_img_path, 'uploads', secure_filename(hashed_filename))

            cxr_img_file.save(file_path)

            # Detection results calculation

            preds.append(np.array(CXRPrediction.model_predict(file_path, model)[0]).tolist())

        # Final results calculation considering the results of all the uploaded images

        final_preds = np.round(np.multiply(np.mean(preds, axis=0), 100), 2)

        final_preds_max = np.round(np.multiply(np.max(preds, axis=0), 100), 2)

        final_preds_min = np.round(np.multiply(np.min(preds, axis=0), 100), 2)

        print(final_preds)

        # Creating the detection results dictionary/ JSON

        predictions_dict = {}

        for i in range(0, len(xray_labels)):

            det_rate_str = str(final_preds[i]) + "% (" + str(final_preds_max[i]) + "% - " + str(

                final_preds_min[i]) + "%)"

            predictions_dict[xray_labels[i]] = det_rate_str

        # Creating detection result JSON to be sent

        json_predictions = json.dumps(predictions_dict, indent=4)

        result = json_predictions

        return result

    return None

@app.route('/localize')

def localization():  # Localization API

    global cur_cxr_hash

    start = datetime.now()

    filepath = 'localizations/' + cur_cxr_hash.split('.')[0]

    if os.path.exists(filepath):

        file_count = len([name for name in os.listdir(filepath) if os.path.isfile(os.path.join(filepath, name))])

        if not file_count == len(xray_labels):

            # If the localized img is already there no need to re-process

            CXRLocalization.create_cxr_localization_heatmap(cur_cxr_hash, model[len(model) - 1], xray_labels)

    else:

        # Calling Localization Function

        CXRLocalization.create_cxr_localization_heatmap(cur_cxr_hash, model[len(model) - 1], xray_labels)

    print(datetime.now() - start)

    return str(len(xray_labels))  # Returning the localized labels

# Function for sending the localized CXR image

@app.route('/get_cxr_detect_img/<int:pathology_id>')

def get_cxr_detect_img(pathology_id):

    print(pathology_id)

    global cur_cxr_hash

    localized_image_name = xray_labels[pathology_id] + '-localizedHeatmap-' + cur_cxr_hash

    filepath = 'localizations/' + cur_cxr_hash.split('.')[0] + '/'

    return send_file(filepath + localized_image_name, mimetype='image/jpg')

# Function for getting symptoms

@app.route('/get_symptoms')

def get_symptoms():

    return send_file('static/files/Symptoms.json', mimetype='application/json')

print("Server Running...")

if __name__ == '__main__':

    app.run(debug=True)  # Debugging