import os
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from mpl_toolkits.mplot3d import Axes3D
from skimage import measure
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from lungmask import LMInferer
import SimpleITK as sitk
import gradio as gr

# --- Lung Segmentation Functions ---

def process_dcm_file(file_path, inferer):
    """Loads a DCM file, performs lung segmentation, and returns the mask."""
    input_image = sitk.ReadImage(file_path)
    segmentation = inferer.apply(input_image)
    newseg = segmentation.reshape(512, 512)  # Assuming 512x512 images
    return newseg

def segment_lungs_from_dicom(dcm_folder):
    """Segments lungs from DICOM files in a folder and returns a 3D volume."""
    if not os.path.exists(dcm_folder) or not os.path.isdir(dcm_folder):
        raise ValueError("Invalid DICOM folder path.")

    inferer = LMInferer()
    segmentation_masks = []
    for filename in os.listdir(dcm_folder):
        if filename.endswith(".dcm"):
            file_path = os.path.join(dcm_folder, filename)
            mask = process_dcm_file(file_path, inferer)
            segmentation_masks.append(mask)
    volume = np.stack(segmentation_masks, axis=0)
    return volume

# --- 3D Visualization Function ---

def plot_3d_lungs(lungs_volume, threshold=0.5):
    """Creates an interactive 3D plot of segmented lungs using Plotly (upright)."""
    verts, faces, normals, values = measure.marching_cubes(lungs_volume.transpose(2, 1, 0), threshold)

    # Apply rotation to make lungs upright
    # Assuming you want to rotate 90 degrees counter-clockwise around the X-axis
    rotation_angle_degrees = -90
    rotation_angle_radians = np.radians(rotation_angle_degrees)
    rotation_matrix = np.array([[1, 0, 0],
                               [0, np.cos(rotation_angle_radians), -np.sin(rotation_angle_radians)],
                               [0, np.sin(rotation_angle_radians), np.cos(rotation_angle_radians)]])
    rotated_verts = np.dot(verts, rotation_matrix)

    x, y, z = zip(*rotated_verts)  # Use rotated vertices
    i, j, k = zip(*faces)

    mesh = go.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, opacity=0.7, color='lightblue')
    fig = go.Figure(data=[mesh])
    fig.update_layout(scene_aspectmode='data')  # Maintain aspect ratio
    return fig

# --- Gradio Interface ---

def process_and_visualize(selected_folder):
    if selected_folder not in ["tumor", "lung", "tumor2"]:
        return "Invalid folder selection."  # Handle invalid input

    volume = segment_lungs_from_dicom(selected_folder)
    visualization = plot_3d_lungs(volume)
    return visualization

inputs = gr.Dropdown(choices=["tumor", "lung", "tumor2"], label="Select DICOM Folder")
output = gr.Plot(label="3D Segmented Lungs")

iface = gr.Interface(
    fn=process_and_visualize,
    inputs=inputs,
    outputs=output,
    title="3D Lung Segmentation Visualization",
    description="Visualize segmented lungs from DICOM images.",
)

iface.launch()