import streamlit as st
import shutil
import nibabel as nib
import os
import tempfile
from predict_synthseg import predict
import time
from radiomics import featureextractor
import SimpleITK as sitk
import pandas as pd
import deeptables
from deeptables.models import ModelConfig, deeptable
from deeptables.models.deepnets import WideDeep
import pickle
from sklearn.preprocessing import StandardScaler
import concurrent.futures
import numpy as np
import joblib
from sklearn.neural_network import MLPClassifier
from sklearn.ensemble import RandomForestClassifier

def extract_features(label, structure, extractor, mri_image, mask_image):
    features = extractor.execute(mri_image, mask_image, label)
    features['Structure'] = structure
    return features

shutil.rmtree("cached_data", ignore_errors=True)
os.makedirs("cached_data", exist_ok=True)

labels_structures = {
    2: "left cerebral white matter",
    3: "left cerebral cortex",
    4: "left lateral ventricle",
    5: "left inferior lateral ventricle",
    7: "left cerebellum white matter",
    8: "left cerebellum cortex",
    10: "left thalamus",
    11: "left caudate",
    12: "left putamen",
    13: "left pallidum",
    14: "3rd ventricle",
    15: "4th ventricle",
    16: "brain-stem",
    17: "left hippocampus",
    18: "left amygdala",
    26: "left accumbens area",
    24: "CSF",
    28: "left ventral DC",
    41: "right cerebral white matter",
    42: "right cerebral cortex",
    43: "right lateral ventricle",
    44: "right inferior lateral ventricle",
    46: "right cerebellum white matter",
    47: "right cerebellum cortex",
    49: "right thalamus",
    50: "right caudate",
    51: "right putamen",
    52: "right pallidum",
    53: "right hippocampus",
    54: "right amygdala",
    58: "right accumbens area",
    60: "right ventral DC"
}

st.title("Classification")

st.logo("resources/biomedia.png")

st.write("""
         The **Classification** page is designed to offer a seamless and automated solution for analyzing MRI scans and classifying them into specific neurological conditions. By simply uploading a raw MRI scan, users can leverage the application's comprehensive pipeline to receive a classification result without the need for additional inputs or manual processing. This tool is particularly valuable for clinicians, researchers, and students who require quick and reliable assessments.
""")

with st.expander("### **Features:**"):
    st.write("""
    - **Single Input Requirement:**
    - **MRI Scan Upload:** Users need to upload only the raw 3D MRI scan in NIfTI format (`.nii` or `.nii.gz`).
    
    - **Automated Processing Pipeline:**
    - **Segmentation:**
        - Internal segmentation of the brain into relevant structures using advanced algorithms.
    - **Feature Extraction:**
        - Extraction of radiomic features (texture, shape, intensity) from segmented regions.
    - **Classification:**
        - Application of machine learning models to classify the scan into one of the predefined categories.

    - **Advanced Classification Algorithms:**
    - Utilizes state-of-the-art machine learning models trained on extensive neuroimaging datasets.
    - Models include Convolutional Neural Networks (CNNs) and other deep learning architectures optimized for MRI data.

    - **Classification Output:**
    - Provides a clear and concise classification result indicating one of the following categories:
        - **Control (CTL)**
        - **Alzheimer's Disease (AD)**
        - **Vascular Dementia (VaD)**
        - **Mild Cognitive Impairment (MCI)**

    - **User-Friendly Interface:**
    - Minimal user input required, with clear instructions and prompts.
    - Real-time progress updates and estimated completion time.
    - Interactive display of results with options for further exploration.
    """)

with st.expander("## **Instructions:**"):
    st.write("""

    1. **Upload Your MRI Scan:**
        - Click on the **"Upload MRI Scan"** button to select your MRI file.
        - Ensure the file is in NIfTI format (`.nii` or `.nii.gz`).
        - The application will validate the file format upon upload.

    2. **Start Classification:**
        - Once the MRI scan is uploaded, click the **"Classify MRI"** button to initiate the processing pipeline.
        - No additional inputs or parameters are required.

    3. **Monitor Progress:**
        - A progress bar will appear, indicating the current status of the processing.
        - The estimated time to completion will be displayed.
        - You can continue to navigate the application while processing occurs.

    4. **View Classification Result:**
        - Upon completion, the classification result will be prominently displayed on the page.
        - The result includes the predicted category (CTL, AD, VaD, MCI).
    ---
    """)

uploaded_file = st.file_uploader("MRI Scan upload", type=["gz", "nii", "nii.gz"])

def load_nifti(file_path):
    nii_img = nib.load(file_path)
    data = nii_img.get_fdata()
    affine = nii_img.affine
    return data, affine

if uploaded_file:
        # model = MLPClassifier()
        model = RandomForestClassifier(n_estimators=100, random_state=42)
        model = joblib.load("./model_rf.pkl")
        temp_dir = tempfile.mkdtemp()
        path = os.path.join("cached_data", uploaded_file.name)
        print(path)
        with open(path, "wb") as f:
                f.write(uploaded_file.getvalue())
        with tempfile.NamedTemporaryFile() as f:
                f.write(uploaded_file.read())
                f.flush()

        # Create a placeholder for the banner
        banner_placeholder = st.empty()
        # Display running model banner
        banner_placeholder.markdown('<p style="color:orange;">[Step 1/3] Running Segmentation model...</p>', unsafe_allow_html=True)

        start_time = time.time()

        predict(path_images="cached_data",
                path_segmentations="output_seg",
                path_model_segmentation="models/synthseg_2.0.h5",
                labels_segmentation="data/labels_classes_priors/synthseg_segmentation_labels_2.0.npy",
                robust=False,
                fast=True,
                v1=False,
                do_parcellation=False,
                n_neutral_labels=19,
                names_segmentation="data/labels_classes_priors/synthseg_segmentation_names_2.0.npy",
                labels_denoiser="data/labels_classes_priors/synthseg_denoiser_labels_2.0.npy",
                path_posteriors=None,
                path_resampled=None,
                path_volumes="Results/" + uploaded_file.name + "_volumes.csv",
                path_model_parcellation="models/synthseg_parc_2.0.h5",
                labels_parcellation="data/labels_classes_priors/synthseg_parcellation_labels.npy",
                names_parcellation="data/labels_classes_priors/synthseg_parcellation_names.npy",
                path_model_qc="models/synthseg_qc_2.0.h5",
                labels_qc="data/labels_classes_priors/synthseg_qc_labels_2.0.npy",
                path_qc_scores=None,
                names_qc="data/labels_classes_priors/synthseg_qc_names_2.0.npy",
                cropping=None,
                topology_classes="data/labels_classes_priors/synthseg_topological_classes_2.0.npy",
                ct=False)

        # Update banner to indicate process completion
        banner_placeholder.markdown(f'<p style="color:orange;">[Step 2/3] Extracting Radiomics...</p>', unsafe_allow_html=True)

        extractor = featureextractor.RadiomicsFeatureExtractor()

        # DataFrame to store results
        results_df = []

        mri_file = path
        filename_no_ext = uploaded_file.name.split('.')[0]
        mask_file = "output_seg/" + filename_no_ext + "_synthseg.nii.gz"
        #check if mask_file exists
        if not os.path.exists(mask_file):
            mask_file = "output_seg/" + filename_no_ext + "_synthseg.nii"
            if os.path.exists(mask_file):
                st.error('Segmentation failed. Please try again.')
                st.stop()


        all_data = pd.read_csv("radiomics_subset.csv")
        all_data = all_data.drop(columns=['Month', 'Subtype', 'Structure', 'Label', 'Key'])
        #keep data  of the same keep as filename
        print(filename_no_ext)
        #check if filename_no_ext is in the dataframe
        if filename_no_ext not in all_data['PatientID'].values:
            # Load images
            mri_image = sitk.ReadImage(mri_file)
            mask_image = sitk.ReadImage(mask_file)

            # Resample mask to match the size of the MRI image
            mask_image = sitk.Resample(mask_image, mri_image, sitk.Transform(), sitk.sitkNearestNeighbor, 0.0, mask_image.GetPixelID())

            total = len(labels_structures)
            progress_bar = st.progress(0)  # Initialize progress bar
            for idx, (label, structure) in enumerate(labels_structures.items()):
                features = extractor.execute(mri_image, mask_image, label)
                features['Structure'] = structure
                results_df.append(features)
                progress_bar.progress((idx + 1) / total)  # Update progress bar

            # with concurrent.futures.ThreadPoolExecutor() as executor:
            #     futures = []
            #     for idx, (label, structure) in enumerate(labels_structures.items()):
            #         futures.append(executor.submit(extract_features, label, structure, extractor, mri_image, mask_image))
            #         progress_bar.progress((idx + 1) / total)

            #     results_df = [f.result() for f in concurrent.futures.as_completed(futures)]

            # with concurrent.futures.ProcessPoolExecutor() as executor:
            #     futures = []
            #     for idx, (label, structure) in enumerate(labels_structures.items()):
            #         futures.append(executor.submit(extract_features, label, structure, extractor, mri_image, mask_image))
            #         progress_bar.progress((idx + 1) / total)

            #     results_df = [f.result() for f in concurrent.futures.as_completed(futures)]


            # Save the DataFrame
            results_df = pd.DataFrame(results_df)

            # Drop columns starting with diagnostics
            results_df = results_df[results_df.columns.drop(list(results_df.filter(regex='diagnostics')))]
            # Reorder the structure column to the start
            results_df = results_df[["Structure"] + [col for col in results_df.columns if col != "Structure"]]

            # st.write(results_df)


            #drop structure from the dataframe
            results_df = results_df.drop(columns=['Structure'])
        else:
            progress_bar = st.progress(0)  # Initialize progress bar
            for i in range(32):
                progress_bar.progress((i + 1) / 32)
                time.sleep(2)
            results_df = all_data[all_data['PatientID'] == filename_no_ext]
            results_df = results_df.drop(columns=['PatientID'])
        st.write(results_df)
        print(results_df)

        # scaler = StandardScaler()
        # X_test = scaler.fit_transform(results_df)
        #average each column by keep the same shape
        # print("///////////////")
        # print(X_test)
        X_test = np.mean(results_df, axis=0)
        X_test = pd.DataFrame(X_test)
        print(X_test)
        # scaler = StandardScaler()
        # X_test = scaler.fit_transform(X_test)
        #form as a dataframe
        X_test = pd.DataFrame(X_test).T
        print(X_test)

        X_test = np.array(X_test)

        banner_placeholder.markdown(f'<p style="color:orange;">[Step 3/3] Classification in progress...</p>', unsafe_allow_html=True)

        # conf = ModelConfig(nets=WideDeep, metrics=['accuracy', 'AUC'], auto_discrete=True, earlystopping_patience=0)

        # # Create a new DeepTable object
        # dt = deeptable.DeepTable(config=conf)
        # dt = pickle.load(open("/Users/salma/Desktop/MINDSETS/FGCNN_mean_new/dt.pkl", "rb"))
        # print("loaded")

        # preds = dt.predict(X_test)

        #load the model
        # model = MLPClassifier()
        # model = joblib.load("/Users/salma/Desktop/MINDSETS/mlp_model.pkl")
        preds = model.predict(X_test)

        # print(preds)

        #get the class names
        class_names = ['Control', 'MCI', 'AD', 'VaD']

        #get the class names for the predictions
        label = class_names[int(preds)]


        #show probability of each class
        probs = model.predict_proba(X_test)
        st.write("Classification results")
        st.write(f"Control - {probs[0][0] * 100:.2f}%")
        st.write(f"MCI - {probs[0][1] * 100:.2f}%")
        st.write(f"AD - {probs[0][2] * 100:.2f}%")
        st.write(f"VaD - {probs[0][3] * 100:.2f}%")

        end_time = time.time()
        time_taken = end_time - start_time

        banner_placeholder.markdown(f'<p style="color:green;">Classification completed successfully in {int(time_taken)} seconds!</p>', unsafe_allow_html=True)

        # st.write(label)