import os
import tempfile
import pydicom
import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr

def dicom_to_nifti_and_display(dicom_file):
    """
    Converte um arquivo DICOM para NIfTI e exibe as imagens antes e depois da conversão.
    
    :param dicom_file: Arquivo DICOM carregado pelo usuário.
    :return: Tupla contendo a imagem DICOM original e a imagem NIfTI convertida.
    """
    try:
        # Ler o arquivo DICOM
        dicom_data = pydicom.dcmread(dicom_file.name)
        dicom_image = dicom_data.pixel_array

        # Extrair o slice central da imagem DICOM
        if len(dicom_image.shape) == 3:
            largest_axis = np.argmax(dicom_image.shape)
            slice_index = dicom_image.shape[largest_axis] // 2
            if largest_axis == 0:
                dicom_slice = dicom_image[slice_index, :, :]
            elif largest_axis == 1:
                dicom_slice = dicom_image[:, slice_index, :]
            else:
                dicom_slice = dicom_image[:, :, slice_index]
        else:
            dicom_slice = dicom_image  # Se for 2D, use a imagem diretamente

        # Criar um objeto NIfTI
        nifti_image = nib.Nifti1Image(dicom_image, affine=np.eye(4))

        # Salvar temporariamente o arquivo NIfTI
        with tempfile.NamedTemporaryFile(suffix=".nii", delete=False) as temp_nifti:
            nib.save(nifti_image, temp_nifti.name)
            nifti_path = temp_nifti.name

        # Carregar o arquivo NIfTI para exibição
        nifti_loaded = nib.load(nifti_path).get_fdata()

        # Extrair o slice central da imagem NIfTI
        if len(nifti_loaded.shape) == 3:
            largest_axis = np.argmax(nifti_loaded.shape)
            slice_index = nifti_loaded.shape[largest_axis] // 2
            if largest_axis == 0:
                nifti_slice = nifti_loaded[slice_index, :, :]
            elif largest_axis == 1:
                nifti_slice = nifti_loaded[:, slice_index, :]
            else:
                nifti_slice = nifti_loaded[:, :, slice_index]
        else:
            nifti_slice = nifti_loaded  # Se for 2D, use a imagem diretamente

        # Exibir as imagens usando Matplotlib
        fig, axes = plt.subplots(1, 2, figsize=(12, 6), constrained_layout=True)
        
        # Imagem DICOM (slice central)
        axes[0].imshow(dicom_slice, cmap="gray", aspect="auto")
        axes[0].set_title("Imagem DICOM Original", fontsize=14, color="blue")
        axes[0].axis("off")

        # Imagem NIfTI (slice central)
        axes[1].imshow(nifti_slice, cmap="gray", aspect="auto")
        axes[1].set_title("Imagem NIfTI Convertida", fontsize=14, color="green")
        axes[1].axis("off")

        # Salvar a figura em um buffer temporário
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_image:
            plt.savefig(temp_image.name, bbox_inches="tight", dpi=300)  # Alta resolução
            plt.close(fig)
            return temp_image.name, nifti_path

    except Exception as e:
        raise gr.Error(f"Erro ao processar o arquivo: {e}")

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Conversor DICOM para NIfTI com Visualização")
    gr.Markdown("Faça upload de um arquivo DICOM, visualize a imagem original, converta para NIfTI e veja o resultado.")

    with gr.Row():
        input_dicom = gr.File(label="Upload de Arquivo DICOM")
        output_images = gr.Image(label="Visualização das Imagens", interactive=False)
        output_nifti = gr.File(label="Download do Arquivo NIfTI")

    convert_button = gr.Button("Converter e Visualizar")

    def process_and_display(dicom_file):
        image_path, nifti_path = dicom_to_nifti_and_display(dicom_file)
        return image_path, nifti_path

    convert_button.click(
        process_and_display,
        inputs=input_dicom,
        outputs=[output_images, output_nifti]
    )

# Executar a interface
demo.launch()