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

# Set up the Streamlit app title
st.title("DICOM Viewer for CBCT Scans")

# Directory paths where DICOM files are stored
single_image_directory = "00002067/"
multiple_images_directory = "00007FA6/"

# Function to get the slice location or image position for sorting
def get_slice_location(dicom_file):
    ds = pydicom.dcmread(dicom_file)
    # Try to get the 'Slice Location'; if not available, use 'Image Position Patient'
    if 'SliceLocation' in ds:
        return float(ds.SliceLocation)
    elif 'ImagePositionPatient' in ds:
        # 'ImagePositionPatient' is a list; use the first element (z-coordinate)
        return float(ds.ImagePositionPatient[2])
    else:
        return 0  # Default to 0 if neither tag is available

# Display the single DICOM image from the first folder
st.header("Single DICOM Image from 00002067")

# Check if the directory exists
if os.path.exists(single_image_directory):
    # Get all files in the single_image_directory (no extension filtering)
    single_image_files = [os.path.join(single_image_directory, f) for f in os.listdir(single_image_directory) if os.path.isfile(os.path.join(single_image_directory, f))]

    if single_image_files:
        # Read and display the first DICOM image found
        ds = pydicom.dcmread(single_image_files[0])
        single_image = ds.pixel_array

        # Display the image
        fig, ax = plt.subplots()
        ax.imshow(single_image, cmap='gray')
        ax.axis('off')
        st.pyplot(fig)
    else:
        st.write("No DICOM files found in the directory 00002067.")
else:
    st.write("Directory 00002067 does not exist.")

# Display the multiple DICOM images from the second folder using a slider
st.header("Browse Multiple DICOM Images with slider")

# Check if the directory exists
if os.path.exists(multiple_images_directory):
    # Get all files in the multiple_images_directory (no extension filtering)
    multiple_image_files = [os.path.join(multiple_images_directory, f) for f in os.listdir(multiple_images_directory) if os.path.isfile(os.path.join(multiple_images_directory, f))]

    if multiple_image_files:
        # Sort DICOM files by slice location or image position
        multiple_image_files.sort(key=lambda x: get_slice_location(x))
        
        # Read all DICOM images (sorted)
        dicom_images = [pydicom.dcmread(file).pixel_array for file in multiple_image_files]

        # Display the selected DICOM image dynamically
        st.header(f"Displaying Images from folder 00007FA6")

        # Interactive slider to select image index (below the image)
        image_index = st.slider("Select Image Index", 0, len(dicom_images) - 1, 0)

        # Display the image based on the slider position
        fig, ax = plt.subplots()
        ax.imshow(dicom_images[image_index], cmap='gray')
        ax.axis('off')
        st.pyplot(fig)

        # Automatic image display section
        st.header("Automatic Image Display - view animated images (click start)")
        placeholder = st.empty()  # Placeholder for the images

        # Initialize session state for animation control
        if 'animation_running' not in st.session_state:
            st.session_state.animation_running = False
        if 'current_index' not in st.session_state:
            st.session_state.current_index = 0

        # Button to start the animation
        if st.button("Start Animation"):
            st.session_state.animation_running = True
        
        # Button to pause the animation
        if st.button("Pause Animation"):
            st.session_state.animation_running = False
        
        # Button to resume the animation
        if st.button("Resume Animation"):
            st.session_state.animation_running = True

        # Run the animation loop if running
        if st.session_state.animation_running:
            while st.session_state.animation_running and st.session_state.current_index < len(dicom_images):
                with placeholder:
                    # Display the current image
                    fig, ax = plt.subplots()
                    ax.imshow(dicom_images[st.session_state.current_index], cmap='gray')
                    ax.axis('off')
                    st.pyplot(fig)
                time.sleep(0.5)  # Display each image for 0.5 seconds
                st.session_state.current_index += 1
                
                # Reset to loop from the start if the end is reached
                if st.session_state.current_index >= len(dicom_images):
                    st.session_state.current_index = 0

    else:
        st.write("No DICOM files found in the directory 00007FA6.")
else:
    st.write("Directory 00007FA6 does not exist.")