import os
import shutil
import nibabel as nib
import warnings
import torch

from totalsegmentator.python_api import totalsegmentator

from mosamaticdesktop.tasks.task import Task
from mosamaticdesktop.logger import Logger
from mosamaticdesktop.utils import createNameWithTimestamp
from mosamaticdesktop.tasks.totalsegmentatortask.checksegmentation import CheckSegmentation
from mosamaticdesktop.tasks.totalsegmentatortask.sliceselector import SliceSelector
from mosamaticdesktop.utils import currentTimeInSeconds, elapsedSeconds

warnings.filterwarnings('ignore', 'Invalid value for VR UI')

LOGGER = Logger()

ROIS = [
    'vertebrae_S1', 'vertebrae_C1', 'vertebrae_C2', 'vertebrae_C3', 'vertebrae_C4', 'vertebrae_C5', 'vertebrae_C6', 'vertebrae_C7', 
    'vertebrae_L1', 'vertebrae_L2', 'vertebrae_L3', 'vertebrae_L4', 'vertebrae_L5', 'vertebrae_T1', 'vertebrae_T2', 'vertebrae_T3', 
    'vertebrae_T4', 'vertebrae_T5', 'vertebrae_T6', 'vertebrae_T7', 'vertebrae_T8', 'vertebrae_T9', 'vertebrae_T10', 'vertebrae_T11', 
    'vertebrae_T12'
]


class TotalSegmentatorTask(Task):
    """ This task requires Torch + CUDA. Install using pip:
    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    """
    def __init__(self) -> None:
        super(TotalSegmentatorTask, self).__init__()
        self.addDescriptionParameter(
            name='description',
            description='Extracts anatomical ROIs from list of full CT scans'
        )
        self.addPathParameter(
            name='rootDirectoryPath',
            labelText='Root Directory of CT Scans (each scan in separate subdirectory)'
        )
        self.addOptionGroupParameter(
            name='vertebra',
            labelText='Select Vertebral Level',
            options=ROIS,
        )
        self.addOptionGroupParameter(
            name='device',
            labelText='Select Device (CPU or GPU)',
            options=['GPU', 'CPU'],
            defaultValue='GPU',
        )
        self.addBooleanParameter(
            name='fast',
            labelText='Enable Fast Calculation (for CPU only)',
            defaultValue=False,
        )
        self.addBooleanParameter(
            name='qualityCheck',
            labelText='Perform Quality Check After Segmentation',
            defaultValue=True,
        )
        self.addPathParameter(
            name='outputDirectoryPath',
            labelText='Output Directory'
        )
        self.addTextParameter(
            name='outputDirectoryName',
            labelText='Output Directory Name',
            optional=True,
        )
        self.addBooleanParameter(
            name='overwriteOutputDirectory',
            labelText='Overwrite Output Directory',
            defaultValue=True,
        )

    def execute(self) -> None:

        # Check we have a GPU available
        if torch.cuda.is_available():

            # Get root directory path containing sub-directories for each CT scan
            rootDirectoryPath = self.parameter('rootDirectoryPath').value()
            if rootDirectoryPath:
                outputDirectoryName = self.parameter('outputDirectoryName').value()
                if outputDirectoryName is None:
                    outputDirectoryName = createNameWithTimestamp(name=outputDirectoryName)
                outputDirectoryPath = self.parameter('outputDirectoryPath').value()
                outputDirectoryPath = os.path.join(outputDirectoryPath, outputDirectoryName)
                LOGGER.info(f'Output directory path: {outputDirectoryPath}')

                # Get vertebra
                vertebra = self.parameter('vertebra').value()
                LOGGER.info(f'vertebra = {vertebra}')

                # Quality check yes/no
                qualityCheck = self.parameter('qualityCheck').value()
                LOGGER.info(f'Quality check = {qualityCheck}')

                # Get device and fast options. Disable fast if we're using the GPU
                device = self.parameter('device').value().lower()
                fast = self.parameter('fast').value()
                if device == 'gpu' and fast:
                    fast = False
                LOGGER.info(f'device = {device}, fast = {fast}')

                # Check whether to overwrite previous output directory
                overwriteOutputDirectory = self.parameter('overwriteOutputDirectory').value()
                LOGGER.info(f'Overwrite output fileset: {overwriteOutputDirectory}')
                if overwriteOutputDirectory:
                    if os.path.isdir(outputDirectoryPath):
                        shutil.rmtree(outputDirectoryPath)
                os.makedirs(outputDirectoryPath, exist_ok=True)

                selectedSlicesDirectoryPath = os.path.join(outputDirectoryPath, 'slices')
                if os.path.isdir(selectedSlicesDirectoryPath):
                    shutil.rmtree(selectedSlicesDirectoryPath)
                os.makedirs(selectedSlicesDirectoryPath, exist_ok=False)

                # Each CT scan's files should be in a separate subdirectory inside the root directory
                step = 0
                nrSteps = 0
                for scanDirectoryName in os.listdir(rootDirectoryPath):
                    scanDirectoryPath = os.path.join(rootDirectoryPath, scanDirectoryName)
                    if os.path.isdir(scanDirectoryPath):
                        nrSteps += 1

                startTimeTotal = currentTimeInSeconds()
                for scanDirectoryName in os.listdir(rootDirectoryPath):
                    scanDirectoryPath = os.path.join(rootDirectoryPath, scanDirectoryName)
                    if os.path.isdir(scanDirectoryPath):
                        startTime = currentTimeInSeconds()
                        LOGGER.info(f'Running TotalSegmentator on scan directory {scanDirectoryPath}...')
                        outputScanDirectoryPath = os.path.join(outputDirectoryPath, scanDirectoryName)
                        os.makedirs(outputScanDirectoryPath, exist_ok=True)

                        try:
                            # Run Total Segmentator twice, once to get individual NIFTI files for each ROI. 
                            # And once to store all labels in a single NIFTI. This volume will be used for quality checking
                            totalsegmentator(
                                scanDirectoryPath, outputScanDirectoryPath, fast=fast, device=device)
                            roiFilePath = os.path.join(outputScanDirectoryPath, vertebra + '.nii.gz')

                            ok = True
                            if qualityCheck:
                                totalsegmentator(
                                    scanDirectoryPath, outputScanDirectoryPath, fast=fast, device=device, ml=True)
                                segmentationFilePath = os.path.join(outputDirectoryPath, scanDirectoryName + '.nii') # No .gz extension!
                                segmentation = nib.load(segmentationFilePath)
                                checker = CheckSegmentation(segmentation=segmentation, scanName=scanDirectoryName)
                                ok = checker.execute()

                            if ok:
                                # Get requested ROI and select DICOM slice running through it
                                roi = nib.load(roiFilePath)
                                selector = SliceSelector(roi=roi, volume=segmentation, dicomDirectory=scanDirectoryPath)
                                output_files = selector.execute()
                                if len(output_files) > 0:
                                    LOGGER.info(f'Found median slice {output_files[0]} for {vertebra}')
                                    selectedSlice = os.path.join(selectedSlicesDirectoryPath, scanDirectoryName + '-' + vertebra + '-' + os.path.split(output_files[0])[1])
                                    shutil.copyfile(output_files[0], selectedSlice)
                                    LOGGER.info(f'Elapsed time after one scan: {elapsedSeconds(startTime)} seconds')
                                else:
                                    LOGGER.error(f'Output of slice selector contains more than one file')
                            else:
                                LOGGER.error(f'Error checking segmentation')
                        except Exception as e:
                            LOGGER.error(f'Exception running Total Segmentator on DICOM files in {scanDirectoryPath} ({e})')

                        self.updateProgress(step=step, nrSteps=nrSteps)
                        step += 1

                self.dataManager().createFileSet(fileSetPath=selectedSlicesDirectoryPath)
                LOGGER.info(f'Total elapsed time: {elapsedSeconds(startTimeTotal)} seconds')

                LOGGER.info('Finished')
                LOGGER.info(f'Results can be found in: {outputDirectoryPath}')
            else:
                LOGGER.error(f'PyTorch: no GPU found. This is going to take forever...')