# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/09_external_data.ipynb.

# %% auto #0
__all__ = ['MURLs', 'download_ixi_data', 'download_ixi_tiny', 'download_spine_test_data', 'download_example_spine_data',
           'download_medmnist3d_dataset', 'download_example_endometrial_cancer_data']

# %% ../nbs/09_external_data.ipynb #942e1644-0dac-4701-8a1e-5693f80f3ca0
from pathlib import Path
from glob import glob
from numpy import load 
import pandas as pd
from monai.apps import download_url, download_and_extract
from torchio.datasets.ixi import IXITiny
from torchio import ScalarImage
import multiprocessing as mp
from functools import partial

# %% ../nbs/09_external_data.ipynb #6d5600a3-9d70-4714-8cc2-8f4682a7e5a3
class MURLs():
    """A class with external medical dataset URLs."""

    IXI_DATA = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T1.tar'
    IXI_DEMOGRAPHIC_INFORMATION = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI.xls'
    CHENGWEN_CHU_SPINE_DATA = 'https://drive.google.com/uc?id=1rbm9-KKAexpNm2mC9FsSbfnS8VJaF3Kn&confirm=t'
    EXAMPLE_SPINE_DATA = 'https://drive.google.com/uc?id=1K1J2wmWxYUPY0NbylcV4Otv88EqCL4Ys'
    MEDMNIST_DICT = {'OrganMNIST3D': 'https://zenodo.org/record/6496656/files/organmnist3d.npz?download=1',	
                     'NoduleMNIST3D': 'https://zenodo.org/record/6496656/files/nodulemnist3d.npz?download=1',
                     'AdrenalMNIST3D': 'https://zenodo.org/record/6496656/files/adrenalmnist3d.npz?download=1',	
                     'FractureMNIST3D': 'https://zenodo.org/record/6496656/files/fracturemnist3d.npz?download=1',
                     'VesselMNIST3D': 'https://zenodo.org/record/6496656/files/vesselmnist3d.npz?download=1', 
                     'SynapseMNIST3D': 'https://zenodo.org/record/6496656/files/synapsemnist3d.npz?download=1'}
    EXAMPLE_EC_DATA = 'https://drive.google.com/uc?id=1P_J8MfzbfObTkAyWfY-LlEC_IA_NjYn-'
    EXAMPLE_MULTI_EC_DATA = 'https://drive.google.com/uc?id=1yOmbA9-nKfzzsfysgp9OcPw8GAarXPvE'

# %% ../nbs/09_external_data.ipynb #cb4c061a-d7f5-4ba0-9e7b-ee0c6ec480b0
def _process_ixi_xls(xls_path: (str, Path), img_path: Path) -> pd.DataFrame:
    """Private method to process the demographic information for the IXI dataset.

    Args:
        xls_path: File path to the xls file with the demographic information.
        img_path: Folder path to the images.

    Returns:
        A processed dataframe with image path and demographic information.

    Raises:
        ValueError: If xls_path or img_path do not exist.
    """

    print('Preprocessing ' + str(xls_path))

    df = pd.read_excel(xls_path)

    duplicate_subject_ids = df[df.duplicated(['IXI_ID'], keep=False)].IXI_ID.unique()

    for subject_id in duplicate_subject_ids:
        age = df.loc[df.IXI_ID == subject_id].AGE.nunique()
        if age != 1: df = df.loc[df.IXI_ID != subject_id]  # Remove duplicates with two different age values

    df = df.drop_duplicates(subset='IXI_ID', keep='first').reset_index(drop=True)

    df['subject_id'] = ['IXI' + str(subject_id).zfill(3) for subject_id in df.IXI_ID.values]
    df = df.rename(columns={'SEX_ID (1=m, 2=f)': 'gender'})
    df['age_at_scan'] = df.AGE.round(2)
    df = df.replace({'gender': {1: 'M', 2: 'F'}})

    img_list = list(img_path.glob('*.nii.gz'))
    for path in img_list:
        subject_id = path.parts[-1].split('-')[0]
        df.loc[df.subject_id == subject_id, 't1_path'] = str(path)

    df = df.dropna()
    df = df[['t1_path', 'subject_id', 'gender', 'age_at_scan']]
    
    return df

# %% ../nbs/09_external_data.ipynb #6714a68f-1378-46b3-aeff-ef940213ac2f
def download_ixi_data(path: (str, Path) = '../data') -> Path:
    """Download T1 scans and demographic information from the IXI dataset.
    
    Args:
        path: Path to the directory where the data will be stored. Defaults to '../data'.

    Returns:
        The path to the stored CSV file.
    """

    path = Path(path) / 'IXI'
    img_path = path / 'T1_images'

    # Check whether image data already present in img_path:
    is_extracted = False
    try:
        if len(list(img_path.iterdir())) >= 581:  # 581 imgs in the IXI dataset
            is_extracted = True
            print(f"Images already downloaded and extracted to {img_path}")
    except (FileNotFoundError, StopIteration, OSError):
        is_extracted = False

    if not is_extracted:
        download_and_extract(url=MURLs.IXI_DATA, filepath=path / 'IXI-T1.tar', output_dir=img_path)
        (path / 'IXI-T1.tar').unlink()

    download_url(url=MURLs.IXI_DEMOGRAPHIC_INFORMATION, filepath=path / 'IXI.xls')

    processed_df = _process_ixi_xls(xls_path=path / 'IXI.xls', img_path=img_path)
    processed_df.to_csv(path / 'dataset.csv', index=False)

    return path

# %% ../nbs/09_external_data.ipynb #7753da8a-93e8-4bf3-8f78-bb158b4280d0
def download_ixi_tiny(path: (str, Path) = '../data') -> Path:
    """Download the tiny version of the IXI dataset provided by TorchIO.

    Args:
        path: The directory where the data will be 
            stored. If not provided, defaults to '../data'.

    Returns:
        The path to the directory where the data is stored.
    """
    
    path = Path(path) / 'IXITiny'
    
    IXITiny(root=str(path), download=True)
    download_url(url=MURLs.IXI_DEMOGRAPHIC_INFORMATION, filepath=path/'IXI.xls')
    
    processed_df = _process_ixi_xls(xls_path=path/'IXI.xls', img_path=path/'image')
    processed_df['labels'] = processed_df['t1_path'].str.replace('image','label')
    
    processed_df.to_csv(path/'dataset.csv', index=False)
    
    return path

# %% ../nbs/09_external_data.ipynb #b466174a-4b49-4a8f-92c6-1e5e3ca9fc2a
def _create_spine_df(dir: Path) -> pd.DataFrame:
    """Create a pandas DataFrame containing information about spinal images.

    Args:
        dir: Directory path where data (image and segmentation 
            mask files) are stored.

    Returns:
         A DataFrame containing the paths to the image files and their 
            corresponding mask files, the subject IDs, and a flag indicating that 
            these are test data.
    """
    
    img_list = glob(str(dir / 'img/*.nii.gz'))
    mask_list = [str(fn).replace('img', 'seg') for fn in img_list]
    subject_id_list = [fn.split('_')[-1].split('.')[0] for fn in mask_list]
    
    test_data = {
        't2_img_path': img_list,
        't2_mask_path': mask_list,
        'subject_id': subject_id_list,
        'is_test': True,
    }

    return pd.DataFrame(test_data)

# %% ../nbs/09_external_data.ipynb #26256dca-d9df-43f6-b8eb-f36ce2a445dc
def download_spine_test_data(path: (str, Path) = '../data') -> pd.DataFrame:
    """Downloads T2w scans from the study 'Fully Automatic Localization and 
    Segmentation of 3D Vertebral Bodies from CT/MR Images via a Learning-Based 
    Method' by Chu et. al. 

    Args:
        path: Directory where the downloaded data 
            will be stored and extracted. Defaults to '../data'.

    Returns:
        Processed dataframe containing image paths, label paths, and subject IDs.
    """
    
    study = 'chengwen_chu_2015'
    
    download_and_extract(
        url=MURLs.CHENGWEN_CHU_SPINE_DATA, 
        filepath=f'{study}.zip', 
        output_dir=path
    )
    Path(f'{study}.zip').unlink()
    
    return _create_spine_df(Path(path) / study)

# %% ../nbs/09_external_data.ipynb #93b77ec9-a93a-42cc-b707-4e1e75063533
def download_example_spine_data(path: (str, Path) = '../data') -> Path:
    """Downloads example T2w scan and corresponding predicted mask.
    
    Args:
        path: Directory where the downloaded data 
            will be stored and extracted. Defaults to '../data'.

    Returns:
        Path to the directory where the example data has been extracted.
    """
    
    study = 'spine_example_data'
    
    download_and_extract(
        url=MURLs.EXAMPLE_SPINE_DATA, 
        filepath=f'{study}.zip', 
        output_dir=path
    )
    Path(f'{study}.zip').unlink()
    
    return Path(path) / study

# %% ../nbs/09_external_data.ipynb #89f60f76-5e5d-4a40-b49a-617ec2d35731
def _process_medmnist_img(path, idx_arr):
    """Save tensor as NIfTI."""
    
    idx, arr = idx_arr
    img = ScalarImage(tensor=arr[None, :])
    fn = path/f'{idx}_nodule.nii.gz'
    img.save(fn)
    return str(fn)

# %% ../nbs/09_external_data.ipynb #b0199679-707b-445d-8467-3d342246322a
def _df_sort_and_add_columns(df, label_list, is_val):
    """Sort the dataframe based on img_idx and add labels and if it is validation data column."""
    
    df = df.sort_values(by='img_idx').reset_index(drop=True)
    df['labels'], df['is_val'] = label_list, is_val     
    #df = df.replace({"labels": {0:'b', 1:'m'}})
    df = df.drop('img_idx', axis=1)
    
    return df 

# %% ../nbs/09_external_data.ipynb #d82157b8-ab69-4a38-9323-af4b58c6b54a
def _create_nodule_df(pool, output_dir, imgs, labels, is_val=False): 
    """Create dataframe for MedMNIST data."""
    
    img_path_list = pool.map(partial(_process_medmnist_img, output_dir), enumerate(imgs))
    img_idx = [float(Path(fn).parts[-1].split('_')[0]) for fn in img_path_list]
    
    df = pd.DataFrame(list(zip(img_path_list, img_idx)), columns=['img_path','img_idx'])        
    return  _df_sort_and_add_columns(df, labels, is_val)

# %% ../nbs/09_external_data.ipynb #04095111-a38f-4d00-a0fb-2dbfac263eb6
def _dataset_available_locally(path: Path):
    """Checks if the dataset exists in the local directory."""
    
    return path.exists() and any(path.iterdir())

# %% ../nbs/09_external_data.ipynb #a1b456f0-8e14-42c6-a1cd-a332312145ff
def download_medmnist3d_dataset(study: str, path: (str, Path) = '../data',
                                max_workers: int = 1, skip_download: bool = True):
    """Downloads and processes a particular MedMNIST3D dataset.

    Args:
        study: MedMNIST dataset ('OrganMNIST3D', 'NoduleMNIST3D', 
               'AdrenalMNIST3D', 'FractureMNIST3D', 'VesselMNIST3D', 'SynapseMNIST3D')
        path: Directory to store and extract downloaded data. Defaults to '../data'.
        max_workers: Maximum number of worker processes for data processing. 
                     Defaults to 1.
        skip_download: Flag to indicate if the dataset should be re-downloaded if the dataset exists locally. 
                  Default set to False.

    Returns:
        Two pandas DataFrames. The first DataFrame combines training and validation 
        data, and the second DataFrame contains the testing data.
    """
                                    
    path = Path(path) / study                           
    if _dataset_available_locally(path) and skip_download: 
        return pd.read_csv(path / 'train_val.csv'), pd.read_csv(path / 'test.csv')                                     
                                    
    dataset_file_path = path / f'{study}.npz'

    try:
        download_url(url=MURLs.MEDMNIST_DICT[study], filepath=dataset_file_path)
    except KeyError:
        raise ValueError(f"Dataset '{study}' does not exist.")

    data = load(dataset_file_path)
    keys = ['train_images', 'val_images', 'test_images']

    for key in keys:
        (path / key).mkdir(exist_ok=True)

    train_imgs = data[keys[0]]
    val_imgs = data[keys[1]]
    test_imgs = data[keys[2]]

    with mp.Pool(processes=max_workers) as pool:
        train_df = _create_nodule_df(pool, path / keys[0], train_imgs, 
                                     data['train_labels'])
        val_df = _create_nodule_df(pool, path / keys[1], val_imgs, 
                                   data['val_labels'], is_val=True)
        test_df = _create_nodule_df(pool, path / keys[2], test_imgs, 
                                    data['test_labels'])

    train_val_df = pd.concat([train_df, val_df], ignore_index=True)

    dataset_file_path.unlink()

    train_val_df.to_csv(path / 'train_val.csv', index=False)
    test_df.to_csv(path / 'test.csv', index=False)

    return train_val_df, test_df

# %% ../nbs/09_external_data.ipynb #67491c4a-8491-4b10-8f03-87ffa012fee0
def download_example_endometrial_cancer_data(path: (str, Path) = '../data', multi_channel: bool = True):
    """Download example data for endometrial cancer.
    
    Args:
        path (str or Path): Path to save the downloaded data.
        multi_channel (bool): If True, download multi-channel data (Vibe, T2, 
            and ADC). Otherwise, download only Vibe.
    """
    
    download_url = (MURLs.EXAMPLE_MULTI_EC_DATA if multi_channel 
                    else MURLs.EXAMPLE_EC_DATA)
    
    download_and_extract(
        url=download_url, 
        filepath='ec_example_data.zip', 
        output_dir=path)
    
    Path('ec_example_data.zip').unlink()
