import os
import ants
import monai
import torch
import shutil
import numpy as np
import pandas as pd
import altair as alt
import nibabel as nib
import streamlit as st
from random import randint
from itertools import chain
import antspynet as antspynet
from torch.utils.data import DataLoader
from monai.transforms import Compose, LoadImaged
from monai.networks.nets.efficientnet import EfficientNetBN

import dicom2nifti


st.set_option('deprecation.showPyplotGlobalUse', False)
np.random.seed(0)
torch.manual_seed(0)

template = ants.image_read('MNI152_T1_1mm_brain.nii.gz')


def pre_process(image):
    with st.spinner('Reading the image...'):
        y = ants.image_read(image)
    with st.spinner('Bias field correction ongoing...'):
        y = ants.utils.n4_bias_field_correction(y)
    with st.spinner('Denoising the image...'):
        yn = y + np.random.randn(*y.shape).astype('float32')*5
        y = ants.denoise_image(yn, ants.get_mask(y))
    with st.spinner('brain_extraction fn. running...'):
        x = antspynet.utilities.brain_extraction(
            y, modality='t1', antsxnet_cache_directory=None, verbose=True)
        y = y*x
    with st.spinner('Registering from template..'):
        y1 = ants.registration(fixed=template, moving=y,
                               type_of_transform='AffineFast')
    with st.spinner('Applying transforms...'):
        y = ants.apply_transforms(
            fixed=template, moving=y, transformlist=y1['fwdtransforms'])
        st.success('Successfully Preprocessed the Image !')
    return y


col1, col2, col3 = st.columns(3)

with col1:
    st.write(' ')

with col2:
    st.image("unilogo.png")

with col3:
    st.write(' ')

st.markdown("<h1 style='text-align: center; color: grey;'>Deep-AD:  Deep Learning Model for Early Detection of Alzheimer’s</h1>", unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: grey;'>Developed by: Deevyankar Agarwal</h6>",
            unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: grey;'> Part Time Ph.D. Student, UVA, Spain</h6>",
            unsafe_allow_html=True)
st.write('**Description**:  Users can upload T1-W MRIs either in NifTI or DICOM format. After preprocessing (N4 bias field correction, noise removal, brain extraction, and registration in the MNI-152 template), the model will classify MRI scans into one of three groups.')

st.markdown('- AD : Alzheimer’s')
st.markdown('- CN : Cognitively Normal')
st.markdown('- SMCI : stable MCI')

st.write('This Application is based on ensemble learning. The output of multiclassification task AD vs. sMCI vs. CN will be validated further by  binary classification models AD vs. CN and sMCI vs. AD implemented by end-to-end learning and 3D transfer learning, respectively. It will provide an extra layer of verification to make robust decisions.')
st.markdown('''<br>''', unsafe_allow_html=True)

element1 = st.write("""
         # MRI Classification :brain:
         """
                    )

if 'key' not in st.session_state:
    st.session_state.key = str( randint(1000, 100000000))
file_upload = st.file_uploader("Upload the MRI scan (either a single NIfTI file or a folder containing multiple DICOM files)", type=[
                               "nii", "gz", "dcm"], accept_multiple_files=True, key=st.session_state.key)
st.set_option('deprecation.showfileUploaderEncoding', False)


if file_upload == []:
    st.text("No file uploaded !")

st.text('Note : Please clear existing files before uploading new files')
if st.button('Clear Uploaded File(s)', help='Please clear existing files before uploading new files') and 'key' in st.session_state.keys():
    st.session_state.pop('key')
    st.experimental_rerun()

st.write("⚠️ [**Feedback form**](https://forms.gle/xuScGN6Cmf69bsUE9) ⚠️ ")


if len(file_upload) == 1:

    for file in file_upload:
        file.name = file.name
    with open(file.name, "wb") as f:
        f.write(file.getbuffer())

    saved_path = f"{file.name}"

    display_image = ants.image_read(saved_path)
    element2 = st.pyplot(ants.plot(display_image))

    processed_image = pre_process(saved_path)
    a = processed_image.to_nibabel()
    saved_preprocessed_path = 'input_image'
    nib.save(a, saved_preprocessed_path)
    element3 = st.text("Preprocessed Image")
    element4 = st.pyplot(ants.plot(f"{saved_preprocessed_path}.nii", cmap="seismic"))

    transformsv = Compose(
        [
            LoadImaged(keys=["img"])
        ]
    )

    test_files = [{"img": f"{saved_preprocessed_path}.nii", "label": "NA"}]
    test_ds = monai.data.Dataset(data=test_files, transform=transformsv)
    test_loader = DataLoader(test_ds, batch_size=1,
                             pin_memory=torch.cuda.is_available())

    for test_data in test_loader:
        test_images, test_labels = test_data["img"], test_data["label"]

        with st.spinner('Performing Inference...'):
            model = EfficientNetBN(
                "efficientnet-b0", spatial_dims=3, in_channels=1, num_classes=3)
            model.load_state_dict(torch.load(
                'MCEBNfold3.pth', map_location='cpu'))
            model.eval()
            prediction = model(test_images.unsqueeze(1))
            pred = prediction.argmax(dim=1).item()
            class_names = ["SMCI", "AD", "CN"]
            predicted_label = class_names[pred]

        graph_input = list(chain.from_iterable(prediction.tolist()))
        "Plot depicting Class Probabilities"
        source = pd.DataFrame({
            'Model output': graph_input,
            'class': ["SMCI", "AD", "CN"]
        })

        bar_chart = alt.Chart(source).mark_bar().encode(
            y='Model output:Q',
            x='class:O',
        )

        element5 = st.altair_chart(bar_chart, use_container_width=True)

        element6 = st.write(
            f"The MRI Scan belong to the class **{predicted_label}**")


        if pred == 0:
            with st.spinner('Please wait...verifying the model output with another model'):
                model_verify = monai.networks.nets.DenseNet264(spatial_dims=3, in_channels=1, out_channels=2)
                model_verify.load_state_dict(torch.load(
                    'DENSENET264ADvsCNbest_metric_model_classification3d_dict.pth', map_location='cpu'))
                model_verify.eval()
                prediction_verify = model_verify(test_images.unsqueeze(1))
                pred_verify = prediction_verify.argmax(dim=1).item()
                class_names_verify = ["CN", "AD"]
                predicted_label_verify = class_names_verify[pred_verify]

                if pred_verify == 0:
                
                    if predicted_label_verify == predicted_label:
                        st.write(
                            f"Succesfully Verified the result, both models classified the scan as **{predicted_label_verify}**")
                    else:
                        st.write(
                            f"Verifying gave a different result ! **First model predicted as {predicted_label}, other predicted {predicted_label_verify}**")
                
                if pred_verify == 1 : 

                        model_verify = EfficientNetBN(
                            "efficientnet-b0", spatial_dims=3, in_channels=1, num_classes=2)
                        model_verify.load_state_dict(torch.load(
                            'EBNfold3.pth', map_location='cpu'))
                        model_verify.eval()
                        prediction_verify = model_verify(test_images.unsqueeze(1))
                        pred_verify = prediction_verify.argmax(dim=1).item()
                        class_names_verify = ["SMCI", "AD"]
                        predicted_label_verify = class_names_verify[pred_verify]

                        if predicted_label_verify == predicted_label:
                            st.write(
                                f"Succesfully Verified the result, both models classified the scan as **{predicted_label_verify}**")
                        else:
                            st.write(
                                f"Verifying gave a different result ! **First model predicted as {predicted_label}, other predicted {predicted_label_verify}**")

                                              
                            
                    

        if pred == 1:
            with st.spinner('Please wait...verifying the model output with another model'):
                model_verify = EfficientNetBN(
                    "efficientnet-b0", spatial_dims=3, in_channels=1, num_classes=2)
                model_verify.load_state_dict(torch.load(
                    'EBNfold3.pth', map_location='cpu'))
                model_verify.eval()
                prediction_verify = model_verify(test_images.unsqueeze(1))
                pred_verify = prediction_verify.argmax(dim=1).item()
                class_names_verify = ["SMCI", "AD"]
                predicted_label_verify = class_names_verify[pred_verify]

                if predicted_label_verify == predicted_label:
                    st.write(
                        f"Succesfully Verified the result, both models classified the scan as **{predicted_label_verify}**")
                else:
                    st.write(
                        f"Verifying gave a different result ! **First model predicted as {predicted_label}, other predicted {predicted_label_verify}**")

                    

        if pred == 2:
            with st.spinner('Please wait...verifying the model output with another model'):
                model_verify = EfficientNetBN(
                    "efficientnet-b0", spatial_dims=3, in_channels=1, num_classes=2)
                model_verify.load_state_dict(torch.load(
                    'ENB0ADvsCNbest_metric_model_classification3d_dict.pth', map_location='cpu'))
                model_verify.eval()
                prediction_verify = model_verify(test_images.unsqueeze(1))
                pred_verify = prediction_verify.argmax(dim=1).item()
                class_names_verify = ["CN", "AD"]
                predicted_label_verify = class_names_verify[pred_verify]

                if predicted_label_verify == predicted_label:
                    st.write(
                        f"Succesfully Verified the result, both models classified the scan as **{predicted_label_verify}**")
                else:
                    st.write(
                        f"Verifying gave a different result ! **First model predicted as {predicted_label}, other predicted {predicted_label_verify}**")

                     


        graph_input_1 = list(chain.from_iterable(prediction_verify.tolist()))

        "Plot depicting verifying model outputs"
        source_1 = pd.DataFrame({
            'Model output': graph_input_1,
            'class': class_names_verify
        })

        bar_chart_1 = alt.Chart(source_1).mark_bar().encode(
            y='Model output:Q',
            x='class:O',
        )

        st.altair_chart(bar_chart_1, use_container_width=True)


if len(file_upload) > 1:

    print(len(file_upload))

    if os.path.exists('tmp') == True:
        shutil.rmtree('tmp')
    os.makedirs('tmp')

    for file in file_upload:
        file.name = file.name
        with open(file.name, "wb") as f:
            f.write(file.getbuffer())
        shutil.copy(file.name, 'tmp')
    print(len(file_upload))

    display_image = st.empty()
    # display_image = ants.core.ants_image_io.dicom_read('tmp')
    saved_path = 'uploaded_image'
    display_image = dicom2nifti.dicom_series_to_nifti('tmp', saved_path, reorient_nifti=True)
    # nib.save(display_image, saved_path)
    display_image = ants.image_read(f"{saved_path}.nii")
    element2 = st.pyplot(ants.plot(display_image))

    # b = display_image.to_nibabel()
    # saved_path = 'uploaded_image'
    # nib.save(b, saved_path)

    processed_image = pre_process(f"{saved_path}.nii")
    a = processed_image.to_nibabel()
    saved_preprocessed_path = 'input_image'
    nib.save(a, saved_preprocessed_path)
    element3 = st.text("Preprocessed Image")
    element4 = st.pyplot(ants.plot(f"{saved_preprocessed_path}.nii", cmap="seismic"))

    transformsv = Compose(
        [
            LoadImaged(keys=["img"])
        ]
    )

    test_files = [{"img": f"{saved_preprocessed_path}.nii", "label": 1}]
    test_ds = monai.data.Dataset(data=test_files, transform=transformsv)
    test_loader = DataLoader(test_ds, batch_size=1,
                             pin_memory=torch.cuda.is_available())

    for test_data in test_loader:
        test_images, test_labels = test_data["img"], test_data["label"]
        with st.spinner('Performing Inference...'):
            model = EfficientNetBN(
                "efficientnet-b0", spatial_dims=3, in_channels=1, num_classes=3)
            model.load_state_dict(torch.load(
                'MCEBNfold3.pth', map_location='cpu'))
            model.eval()
            prediction = model(test_images.unsqueeze(1))
            pred = prediction.argmax(dim=1).item() 
            class_names = ["SMCI", "AD", "CN"]
            predicted_label = class_names[pred]

        graph_input = list(chain.from_iterable(prediction.tolist()))
        "Plot depicting Class Probabilities"
        source = pd.DataFrame({
            'Model output': graph_input,
            'class': ["SMCI", "AD", "CN"]
        })

        bar_chart = alt.Chart(source).mark_bar().encode(
            y='Model output:Q',
            x='class:O',
        )

        element5 = st.altair_chart(bar_chart, use_container_width=True)

        element6 = st.write(
            f"The MRI Scan belong to the class **{predicted_label}**")



        if pred == 0:
            with st.spinner('Please wait...verifying the model output with another model'):
                model_verify = monai.networks.nets.DenseNet264(spatial_dims=3, in_channels=1, out_channels=2)
                model_verify.load_state_dict(torch.load(
                    'DENSENET264ADvsCNbest_metric_model_classification3d_dict.pth', map_location='cpu'))
                model_verify.eval()
                prediction_verify = model_verify(test_images.unsqueeze(1))
                pred_verify = prediction_verify.argmax(dim=1).item()
                class_names_verify = ["CN", "AD"]
                predicted_label_verify = class_names_verify[pred_verify]

                if pred_verify == 0:
                
                    if predicted_label_verify == predicted_label:
                        st.write(
                            f"Succesfully Verified the result, both models classified the scan as **{predicted_label_verify}**")
                    else:
                        st.write(
                            f"Verifying gave a different result ! **First model predicted as {predicted_label}, other predicted {predicted_label_verify}**")
                
                if pred_verify == 1 : 

                        model_verify = EfficientNetBN(
                            "efficientnet-b0", spatial_dims=3, in_channels=1, num_classes=2)
                        model_verify.load_state_dict(torch.load(
                            'EBNfold3.pth', map_location='cpu'))
                        model_verify.eval()
                        prediction_verify = model_verify(test_images.unsqueeze(1))
                        pred_verify = prediction_verify.argmax(dim=1).item()
                        class_names_verify = ["SMCI", "AD"]
                        predicted_label_verify = class_names_verify[pred_verify]

                        if predicted_label_verify == predicted_label:
                            st.write(
                                f"Succesfully Verified the result, both models classified the scan as **{predicted_label_verify}**")
                        else:
                            st.write(
                                f"Verifying gave a different result ! **First model predicted as {predicted_label}, other predicted {predicted_label_verify}**")

                                                                           


        if pred == 1:
            with st.spinner('Please wait...verifying the model output with another model'):
                model_verify = EfficientNetBN(
                    "efficientnet-b0", spatial_dims=3, in_channels=1, num_classes=2)
                model_verify.load_state_dict(torch.load(
                    'EBNfold3.pth', map_location='cpu'))
                model_verify.eval()
                prediction_verify = model_verify(test_images.unsqueeze(1))
                pred_verify = prediction_verify.argmax(dim=1).item()
                class_names_verify = ["SMCI", "AD"]
                predicted_label_verify = class_names_verify[pred_verify]

                if predicted_label_verify == predicted_label:
                    st.write(
                        f"Succesfully Verified the result, both models classified the scan as **{predicted_label_verify}**")
                else:
                    st.write(
                        f"Verifying gave a different result ! **First model predicted as {predicted_label}, other predicted {predicted_label_verify}**")

                    

        if pred == 2:
            with st.spinner('Please wait...verifying the model output with another model'):
                model_verify = monai.networks.nets.DenseNet264(spatial_dims=3, in_channels=1, out_channels=2)
                model_verify.load_state_dict(torch.load(
                    'F3DENSENET264ADvsCNbest_metric_model_classification3d_dict.pth', map_location='cpu'))
                model_verify.eval()
                prediction_verify = model_verify(test_images.unsqueeze(1))
                pred_verify = prediction_verify.argmax(dim=1).item()
                class_names_verify = ["CN", "AD"]
                predicted_label_verify = class_names_verify[pred_verify]

                if predicted_label_verify == predicted_label:
                    st.write(
                        f"Succesfully Verified the result, both models classified the scan as **{predicted_label_verify}**")
                else:
                    st.write(
                        f"Verifying gave a different result ! **First model predicted as {predicted_label}, other predicted {predicted_label_verify}**")

                      



        graph_input_1 = list(chain.from_iterable(prediction_verify.tolist()))

        "Plot depicting verifying model outputs"
        source_1 = pd.DataFrame({
            'Model output': graph_input_1,
            'class': class_names_verify
        })

        bar_chart_1 = alt.Chart(source_1).mark_bar().encode(
            y='Model output:Q',
            x='class:O',
        )

        st.altair_chart(bar_chart_1, use_container_width=True)
        

st.markdown('''<br><br>''', unsafe_allow_html=True)
st.markdown('''#### Publications :book:''', unsafe_allow_html=True)

st.markdown("""1. [Transfer Learning for Alzheimer’s Disease through Neuroimaging Biomarkers: A Systematic Review](https://www.mdpi.com/1424-8220/21/21/7259 ) \n
 <small>Q1 Sensors</small> <br><br>

2.  [End-to-End Deep Learning Architectures Using 3D Neuroimaging Biomarkers for Early Alzheimer’s Diagnosis](https://www.mdpi.com/2227-7390/10/15/2575) \n
  <small>Q2 mathematics</small> <br><br>

  3.  [Automated Medical Diagnosis of Alzheimer´s Disease Using an Efficient Net Convolutional Neural Network](https://link.springer.com/article/10.1007/s10916-023-01941-4) \n
  <small>Q1 Springer Nature ,Journal of Medical Systems</small> <br><br>
<br>""", unsafe_allow_html=True)



st.markdown('''#### Contact details :mailbox:''', unsafe_allow_html=True)

st.markdown('''
<b>Group :busts_in_silhouette: &nbsp;</b>: &nbsp; http://www.sigte.tel.uva.es/index.php/en/homepage/  
<small>The eHealth and Telemedicine Group (GTe) of the University of Valladolid is a multidisciplinary international group consisting of telecommunications, informatics and medical doctors from different specialties.</small> \n

<br>

<b>Email :e-mail: &nbsp;</b> : &nbsp; deevynkar@gmail.com''', unsafe_allow_html=True)