# -*- coding: utf-8 -*-
"""
The script to demo the feature extraction procedure

Objective: 
1. Read in the image from the resampled nii files (generated by the resize_volume.py)
                          the lung lobe segmentation files 
                          the lesion segmentation files
2. Segment the lesion regions within certain lung lobe
3. Extract radiomic features of the lung lobe lesion region
4. Run the manuscript in a multi-processing way

Input:
1. Resampled CT images 
2. Lung lobe masks ()
3. Lesion masks ()

Output:
Radiomic feature list of each lung lobe of each patient/image
"""

from __future__ import print_function

import numpy as np
import SimpleITK as sitk
from radiomics import featureextractor

import sys
import os
import os.path as osp
import shutil

from multiprocessing import Pool
from multiprocessing import TimeoutError as MP_TimeoutError
from time import sleep

import csv

START = "START"
FINISH = "FINISH"
WARNING = "WARNING"
ERROR = "ERROR"
LOG = "LOG"

# file locations
lobe_mask_loc = 'data_location' # lobe segmentation files
lesion_mask_loc = 'data_location' # lesion segmentation files
image_loc = 'data_location' # CT images
save_root_path = 'data_location' # file save location

# feature extractor settings
# First define the settings
settings = {}
settings['binWidth'] = 20
settings['sigma'] = [0.5, 1.5, 2.5, 3.5, 4.5]

# Instantiate the extractor
extractor = featureextractor.RadiomicsFeatureExtractor(**settings)
extractor.disableAllImageTypes()
extractor.enableImageTypeByName('Original')
extractor.enableImageTypeByName('Square')
extractor.enableImageTypeByName('SquareRoot')
extractor.enableImageTypeByName('Logarithm')
extractor.enableImageTypeByName('Exponential')
extractor.enableImageTypeByName('Wavelet')
extractor.enableImageTypeByName('LoG')

# the lists of features we need
extractor.enableAllFeatures()
extractor.enableFeaturesByName(shape=['Compactness1',
                                      'Compactness2',
                                      'Elongation',
                                      'Flatness',
                                      'LeastAxisLength',
                                      'MajorAxisLength',
                                      'Maximum2DDiameterColumn',
                                      'Maximum2DDiameterRow',
                                      'Maximum2DDiameterSlice',
                                      'Maximum3DDiameter',
                                      'MeshVolume',
                                      'MinorAxisLength',
                                      'SphericalDisproportion',
                                      'Sphericity',
                                      'SurfaceArea',
                                      'SurfaceVolumeRatio',
                                      'VoxelVolume'])

print('Extraction parameters:\n\t', extractor.settings)

# read nii file
def read_nii_file(file_loc):
    reader = sitk.ImageFileReader()
    reader.SetImageIO("NiftiImageIO")
    reader.SetFileName(file_loc)
    image = reader.Execute()
    return image

# read dicom files
def read_dcm_file(file_loc):
    reader = sitk.ImageSeriesReader()
    dicom_names = reader.GetGDCMSeriesFileNames(file_loc)
    reader.SetFileNames(dicom_names)
    pcr_img = reader.Execute()
    return pcr_img


def log_print(pid, comment, logs):
    if type(logs) is str:
        logs = [logs]
    for log in logs:
        print("# JOB %d : --%s-- %s" % (pid, comment, log))
    sys.stdout.flush()


def result_scan(results):
    unfinish = 1
    while unfinish > 0:
        unfinish = 0
        for i, res in enumerate(results):
            try:
                res.get()
            except Exception as e:
                if type(e) == MP_TimeoutError:
                    unfinish += 1
                    continue
                else:
                    print("\n\n\nERROR OCCUR: PID ##%d##, ERRORTYPE: %s\n\n\n" %(i, type(e)))
                    raise e

# feature extraction
def extract_radiomic_features(input_file_loc, mask_lesion_loc, output_file_name, pid, title_flag = 1):
    '''
    # debug only
    input_file_loc = input_file
    mask_lobe_loc = mask_lobe_file
    mask_lesion_loc = mask_lesion_file
    '''
    # read input image
    pcr_img = read_nii_file(input_file_loc)
    # read lesion seg
    mask_lesion = read_nii_file(mask_lesion_loc)
    pcr_img.SetDirection(mask_lesion.GetDirection())
    
    log_print(pid, START, 'Vol:%s' % input_file_loc)
    with open(output_file_name, 'w', newline='') as outfile:
        # extract radiomic features from whole lung lobe lesion region
        result = extractor.execute(pcr_img, mask_lesion)
        keys, values = [],[]
        if title_flag == 1:
            title_flag += 1
            for key, value in result.items():
                keys.append(key)
                values.append(value) 
            csvwriter = csv.writer(outfile)
            csvwriter.writerow(keys)
            csvwriter.writerow(values)
        else:
            for key, value in result.items():
                values.append(value)
            csvwriter = csv.writer(outfile)
            csvwriter.writerow(values)
    log_print(pid, FINISH, 'Vol:%s' % input_file_loc)



_ = None
while _ not in ['y', 'n']:
    _ = input('About to remove dir %s, START? [y/n]' % save_root_path).lower()

if _ == 'n':
    exit()
    

shutil.rmtree(save_root_path, ignore_errors=True)
os.makedirs(save_root_path)
processes = 8
pool = Pool(processes)
results = list()
pid = 0

datasetList = os.listdir(image_loc)
datasetList.sort()

for filename in datasetList:
    file = os.path.splitext(filename)[0]
    
    # image location
    input_file = os.path.join(image_loc, filename)
    # lesion mask location
    mask_lesion_file = os.path.join(lesion_mask_loc, filename)
    
    save_file = osp.join(save_root_path, file+'.csv')
    
    results.append(
        pool.apply_async(
            extract_radiomic_features,
            args=(input_file, mask_lesion_file, save_file, pid)))
    pid += 1

pool.close()
result_scan(results)
pool.join()