import time 
import os
import numpy as np
import SimpleITK as sitk
from multiprocessing import Pool

def get_image_from_mask_area(image_path, mask_path):
    image = sitk.ReadImage(image_path)
    mask = sitk.ReadImage(mask_path)
    image_np = sitk.GetArrayFromImage(image).astype(float)
    mask_np = sitk.GetArrayFromImage(mask).astype(float)
    image_mask_np  = image_np[mask_np > 0]
    data = cal_avg_std(image_mask_np) # 1-dimension
    return data 

def cal_avg_std(image_mask_np):
    image_mask_np.sort()
    value_min = np.percentile(image_mask_np, 0.1)
    value_max = np.percentile(image_mask_np, 99.9)
    data = np.clip(image_mask_np, value_min, value_max)
    return data
    
def main():
    num_threads = 8
    root_path = '/home/xxx/Kidney_tumor/05_data/mianyang_new'
    case_list = os.listdir(root_path)
    num = len(case_list)
    index = 0
    images_mask_path = []
    for case in case_list:
        index += 1
        print('{} [{}] --- [{}]'.format(case, index, num))
        case_path = os.path.join(root_path, case)
        image_path = os.path.join(case_path, 'image.nii.gz')
        mask_path = os.path.join(case_path, 'mask.nii.gz')
        
        images_mask_path.append((image_path, mask_path))      # num_threads      
    p = Pool(num_threads)
    set = p.starmap(get_image_from_mask_area, images_mask_path)
    set = np.array(set)
    set1 = np.concatenate(set, axis=0)
    p.close()
    p.join()
    print('mean is {}, std is {}'.format(np.mean(set1),np.std(set1)))  
if __name__ == '__main__':
    main()