import streamlit as st
from idc_index import index
from pathlib import Path
import pydicom
import pandas as pd
# import pyarrow as pa
# import pyarrow.parquet as pq
from tempfile import TemporaryDirectory
import os
from pathlib import Path
# import polars
import pydicom.datadict as dd
import shutil
import papermill as pm 
import subprocess 


# Main Streamlit app code
st.title("DICOM Classification Demo")
st.write("Select IDC data to download, extract images and metadata, and perform inference using three pre-trained models")

# Fetch IDC index
client = index.IDCClient()
index_df = client.index

# Option to choose IDC data
st.subheader("Choose IDC Data to Process")
collection_ids = index_df["collection_id"].unique()
selected_collection_id = st.selectbox("Select Collection ID", collection_ids)

# Filter dataframe based on selected collection_id
df_filtered_by_collection = index_df[index_df["collection_id"] == selected_collection_id]

patients = df_filtered_by_collection["PatientID"].unique()
selected_patient_id = st.selectbox("Select Patient ID", patients)

# Filter dataframe based on selected patient_id
df_filtered_by_patient = df_filtered_by_collection[df_filtered_by_collection["PatientID"] == selected_patient_id]

modalities = df_filtered_by_patient["Modality"].unique()
selected_modality = st.selectbox("Select Modality", modalities)

# Filter dataframe based on selected modality
df_filtered_by_modality = df_filtered_by_patient[df_filtered_by_patient["Modality"] == selected_modality]

studies = df_filtered_by_modality["StudyInstanceUID"].unique()
selected_study = st.selectbox("Select Study", studies)

# Filter dataframe based on selected study
df_filtered_by_study = df_filtered_by_modality[df_filtered_by_modality["StudyInstanceUID"] == selected_study]

series = df_filtered_by_study["SeriesInstanceUID"].unique()
selected_series = st.selectbox("Select Series", series)

print('selected_series: ' + str(selected_series))

# Button to run the notebook - which loads the pretrained models and runs inference 
if st.button("Run inference"):

    # Code to run when the button is pressed
    st.write("Button pressed! Running inference")

    if os.path.exists("DICOMScanClassification_user_demo.ipynb"):
        os.remove("DICOMScanClassification_user_demo.ipynb")

    subprocess.run(["wget", "https://raw.githubusercontent.com/deepakri201/DICOMScanClassification_pw41/main/DICOMScanClassification_user_demo.ipynb"])

    pm.execute_notebook(
            "DICOMScanClassification_user_demo.ipynb",
            'output.ipynb',
            parameters = dict(SeriesInstanceUID=selected_series)
    )

    with open('output.ipynb', "rb") as f:
        st.download_button(
            label="Download the output notebook file",
            data=f,
            file_name="output.ipynb",
            mime="application/json"
        )