import streamlit as st
import pydicom
import matplotlib.pyplot as plt
import zipfile
import os

# Streamlit app title
st.title("DICOM Image Viewer")

# Upload a ZIP file containing DICOMs
uploaded_zip = st.file_uploader("Upload a ZIP file containing DICOMs", type=["zip"])

# Initialize variables for DICOM data and image list
dicom_images = []
num_slices = 0
error_message = ""

if uploaded_zip is not None:
    # Unzip the uploaded ZIP file and extract DICOMs
    with zipfile.ZipFile(uploaded_zip, "r") as zip_ref:
        zip_ref.extractall("temp_zip_folder")

    # Get a list of DICOM files
    dicom_files = [f for f in os.listdir("temp_zip_folder") if f.lower().endswith(".dcm")]
    num_slices = len(dicom_files)

    # Initialize a list to store successfully processed DICOMs
    processed_dicoms = []

    # Process DICOM files and skip problematic ones
    for dicom_file in dicom_files:
        try:
            dicom_data = pydicom.dcmread(f"temp_zip_folder/{dicom_file}")
            dicom_images.append(dicom_data.pixel_array)
            processed_dicoms.append(dicom_data)
        except Exception as e:
            error_message += f"Skipping problematic DICOM ({dicom_file}): {str(e)}\n"

    # User selects the slice using a slider
    selected_slice = st.slider("Select a DICOM slice", 0, len(processed_dicoms) - 1)

    if selected_slice < len(processed_dicoms):
        # Display the selected DICOM image
        plt.imshow(dicom_images[selected_slice], cmap=plt.cm.bone)
        plt.axis("off")
        plt.title(f"DICOM Image - Slice {selected_slice + 1}/{len(processed_dicoms)}")
        plt.tight_layout()

        # Show the image in the Streamlit app
        st.pyplot(plt)
    else:
        error_message += "No DICOM slices available for the selected slice number."

    # Remove the temporary folder and its contents
    for dicom_file in dicom_files:
        os.remove(f"temp_zip_folder/{dicom_file}")
    os.rmdir("temp_zip_folder")

st.write("Upload a ZIP file containing DICOMs to view and scroll through the slices.")
st.write(error_message)