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

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

# Directories where DICOM files are stored
dicom_directories = {
    "00002067": "./00002067/",
    "00007FA6": "./00007FA6/"
}

# Sidebar for selecting DICOM folder
st.sidebar.header("Select DICOM Folder")
selected_folder = st.sidebar.selectbox("Choose a folder to view DICOM files", list(dicom_directories.keys()))

# Get the selected directory path
dicom_directory = dicom_directories[selected_folder]

# 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

# Get list of DICOM files in the selected directory
dicom_files = [os.path.join(dicom_directory, f) for f in os.listdir(dicom_directory) if f.endswith('.dcm')]

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

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

    #