from multiprocessing import set_start_method
import numpy as np
from umcglib.utils import apply_parallel,print_stats_np
from os import listdir
import matplotlib.pyplot as plt

# mypath = f'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/untarred'
# patient_id = 'pat00123'
# seg = np.load(f'{mypath}/{patient_id}/{patient_id}_seg.npy').astype(int)
# print(np.shape(seg))
# print_stats_np(seg,'seg_sdgsd')

def load_seg(file,mypath):          
     patient_id = file
     print(patient_id)
     seg = np.load(f'{mypath}/{patient_id}/{patient_id}_seg.npy').astype(int)
     background = np.sum(seg == 0)
     femoral_cartilage = np.sum(seg == 1)
     tibial_cartilage = np.sum(seg == 2)
     patellar_cartilage = np.sum(seg == 3)
     femur = np.sum(seg == 4)
     tibia = np.sum(seg == 5)
     patella = np.sum(seg == 6)
     output = [background,femoral_cartilage,tibial_cartilage,patellar_cartilage,femur,tibia,patella]
     # print('background',background)
     return output

def calculate_hist(segs):
     background = np.sum(segs == 0)
     femoral_cartilage = np.sum(segs == 1)
     tibial_cartilage = np.sum(segs == 2)
     patellar_cartilage = np.sum(segs == 3)
     femur = np.sum(segs == 4)
     tibia = np.sum(segs == 5)
     patella = np.sum(segs == 6)
     output = [background,femoral_cartilage,tibial_cartilage,patellar_cartilage,femur,tibia,patella]
     # print('background',background)
     return output

if __name__ == '__main__':
     path = f'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/untarred'
     files = [f for f in listdir(path)]
     print('files length',len(files))
     # For multiprocessing
     set_start_method("spawn")
     print('spawn done')
     # segs_list = apply_parallel(files, load_seg, 4, mypath = path)
     segs_list = apply_parallel(files, load_seg, 12, mypath = path)
     print('done loading 50')
     # segs_list_100 = apply_parallel(files[50:100], load_seg, 4, mypath = path)
     # print('done loading 100')
     # segs_list_150 = apply_parallel(files[100:150], load_seg, 4, mypath = path)
     # print('done loading 150')
     # segs_list_200 = apply_parallel(files[150:200], load_seg, 4, mypath = path)
     # print('done loading 200')
     # segs_list_250 = apply_parallel(files[200:250], load_seg, 4, mypath = path)
     # print('done loading 250')
     # segs_list_300 = apply_parallel(files[250:300], load_seg, 4, mypath = path)
     print('done loading 300')
     # segs_list.extend(segs_list_100)
     # segs_list.extend(segs_list_150)
     # segs_list.extend(segs_list_200)
     # segs_list.extend(segs_list_250)
     # segs_list.extend(segs_list_300)
     print('segs_list is done')

     output = np.stack(segs_list, axis=0)
     # print('load done',np.shape(segs_n))

     # output = apply_parallel(segs_list,calculate_hist,4)
     print('hist calc done',np.shape(output))
     # print(np.stack(output)[:,0])

     dict = {
          'background': np.stack(output)[:,0],
          'femoral_cartilage': np.stack(output)[:,1],
          'tibial_cartilage': np.stack(output)[:,2],
          'patellar_cartilage': np.stack(output)[:,3],
          'femur': np.stack(output)[:,4],
          'tibia': np.stack(output)[:,5],
          'patella': np.stack(output)[:,6]
     }

     np.save('../dict.npy',dict)

     for keys in dict:
          plt.figure()
          data = dict[f'{keys}']
          plt.hist(data)
          plt.title(f"historgram of {keys}")
          plt.savefig(f"../{keys}.png", dpi=300)
          print('done  ',keys)