#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ----------------------------------------------------------------------------------------------------------------------
# Author: Lalith Kumar Shiyam Sundar
#         Sebastian Gutschmayer
# Institution: Medical University of Vienna
# Research Group: Quantitative Imaging and Medical Physics (QIMP) Team
# Date: 09.02.2023
# Version: 2.0.0
#
# Description:
# This module handles image conversion for the moosez.
#
# Usage:
# The functions in this module can be imported and used in other modules within the moosez to perform image conversion.
#
# ----------------------------------------------------------------------------------------------------------------------

import os

import SimpleITK
import dicom2nifti
import pydicom
from rich.progress import Progress


def read_dicom_folder(folder_path: str) -> SimpleITK.Image:
    """
    Reads a folder of DICOM files and returns the image

    :param folder_path: str, Directory to get DICOM files from
    """
    reader = SimpleITK.ImageSeriesReader()
    dicom_names = reader.GetGDCMSeriesFileNames(folder_path)
    reader.SetFileNames(dicom_names)

    dicom_image = reader.Execute()
    return dicom_image


def non_nifti_to_nifti(input_path: str, output_directory: str = None) -> None:
    """
    Converts any image format known to ITK to NIFTI

    :param input_path: str, Directory OR filename to convert to nii.gz
    :param output_directory: str, optional output directory to write the image to.
    """
    if os.path.isdir(input_path):
        dicom_info = create_dicom_lookup(input_path)
        nifti_dir = dcm2niix(input_path)
        rename_nifti_files(nifti_dir, dicom_info)
        return
    elif os.path.isfile(input_path):
        if input_path.endswith('.nii.gz') or input_path.endswith('.nii'):
            return
        output_image = SimpleITK.ReadImage(input_path)
        output_image_basename = f"{os.path.splitext(os.path.basename(input_path))[0]}.nii"
    else:
        return

    if output_directory is None:
        output_directory = os.path.dirname(input_path)
    output_image_path = os.path.join(output_directory, output_image_basename)
    SimpleITK.WriteImage(output_image, output_image_path)


def standardize_to_nifti(parent_dir: str):
    """
    Converts all images in a parent directory to NIFTI
    """
    # go through the subdirectories
    subjects = os.listdir(parent_dir)
    # get only the directories
    subjects = [subject for subject in subjects if os.path.isdir(os.path.join(parent_dir, subject))]

    with Progress() as progress:
        task = progress.add_task("[white] Processing subjects...", total=len(subjects))
        for subject in subjects:
            subject_path = os.path.join(parent_dir, subject)
            if os.path.isdir(subject_path):
                images = os.listdir(subject_path)
                for image in images:
                    if os.path.isdir(os.path.join(subject_path, image)):
                        image_path = os.path.join(subject_path, image)
                        non_nifti_to_nifti(image_path)
                    elif os.path.isfile(os.path.join(subject_path, image)):
                        image_path = os.path.join(subject_path, image)
                        non_nifti_to_nifti(image_path)
            else:
                continue
            progress.update(task, advance=1, description=f"[white] Processing {subject}...")


def dcm2niix(input_path: str) -> str:
    """
    Converts DICOM images into Nifti images using dcm2niix
    :param input_path: Path to the folder with the dicom files to convert
    :return: str, Path to the folder with the converted nifti files
    """
    output_dir = os.path.dirname(input_path)
    dicom2nifti.convert_directory(input_path, output_dir, compression=False, reorient=True)
    return output_dir


def create_dicom_lookup(dicom_dir):
    """Create a lookup dictionary from DICOM files.

    Parameters:
    dicom_dir (str): The directory where DICOM files are stored.

    Returns:
    dict: A dictionary where the key is the anticipated filename that dicom2nifti will produce and
          the value is the modality of the DICOM series.
    """

    # a dictionary to store information from the DICOM files
    dicom_info = {}

    # loop over the DICOM files
    for filename in os.listdir(dicom_dir):
        if filename.endswith('.dcm') or filename.endswith('.ima') or filename.endswith(''):
            # read the DICOM file
            ds = pydicom.dcmread(os.path.join(dicom_dir, filename))

            # extract the necessary information
            series_number = ds.SeriesNumber
            series_description = ds.SeriesDescription if 'SeriesDescription' in ds else None
            sequence_name = ds.SequenceName if 'SequenceName' in ds else None
            protocol_name = ds.ProtocolName if 'ProtocolName' in ds else None
            modality = ds.Modality

            # anticipate the filename dicom2nifti will produce and store the modality tag with it
            if series_description is not None:
                anticipated_filename = f"{series_number}_{series_description}.nii"
            elif sequence_name is not None:
                anticipated_filename = f"{series_number}_{sequence_name}.nii"
            elif protocol_name is not None:
                anticipated_filename = f"{series_number}_{protocol_name}.nii"
            else:
                anticipated_filename = f"{series_number}.nii"

            dicom_info[anticipated_filename] = modality

    return dicom_info


def rename_nifti_files(nifti_dir, dicom_info):
    """Rename NIfTI files based on a lookup dictionary.

    Parameters:
    nifti_dir (str): The directory where NIfTI files are stored.
    dicom_info (dict): A dictionary where the key is the anticipated filename that dicom2nifti will produce and
                       the value is the modality of the DICOM series.
    """

    # loop over the NIfTI files
    for filename in os.listdir(nifti_dir):
        if filename.endswith('.nii'):
            # get the corresponding DICOM information
            modality = dicom_info.get(filename, '')

            # create the new filename
            new_filename = f"{modality}_{filename}"

            # rename the file
            os.rename(os.path.join(nifti_dir, filename), os.path.join(nifti_dir, new_filename))
