text: #    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
  #
  #    Licensed under the Apache License, Version 2.0 (the "License");
  #    you may not use this file except in compliance with the License.
  #    You may obtain a copy of the License at
  #
  #        http://www.apache.org/licenses/LICENSE-2.0
  #
  #    Unless required by applicable law or agreed to in writing, software
  #    distributed under the License is distributed on an "AS IS" BASIS,
  #    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #    See the License for the specific language governing permissions and
  #    limitations under the License.
  
  
  from copy import deepcopy
  
  from batchgenerators.utilities.file_and_folder_operations import *
  import shutil
  import SimpleITK as sitk
  from multiprocessing import Pool
  from medpy.metric import dc
  import numpy as np
  import nibabel as nib
  import os
  import sys
  sys.path.append('../..')
  from nnunet.paths import network_training_output_dir
  from scipy.ndimage import label
  
  
  def compute_dice_scores(ref: str, pred: str):
      ref = sitk.GetArrayFromImage(sitk.ReadImage(ref))
      pred = sitk.GetArrayFromImage(sitk.ReadImage(pred))
      kidney_mask_ref = ref > 0
      kidney_mask_pred = pred > 0
      if np.sum(kidney_mask_pred) == 0 and kidney_mask_ref.sum() == 0:
          kidney_dice = np.nan
      else:
          kidney_dice = dc(kidney_mask_pred, kidney_mask_ref)
  
      tumor_mask_ref = ref == 2
      tumor_mask_pred = pred == 2
      if np.sum(tumor_mask_ref) == 0 and tumor_mask_pred.sum() == 0:
          tumor_dice = np.nan
      else:
          tumor_dice = dc(tumor_mask_ref, tumor_mask_pred)
  
      geometric_mean = np.mean((kidney_dice, tumor_dice))
      return kidney_dice, tumor_dice, geometric_mean
  
  
  def evaluate_folder(folder_gt: str, folder_pred: str):
      p = Pool(8)
      niftis = subfiles(folder_gt, suffix=".nii.gz", join=False)
      images_gt = [join(folder_gt, i) for i in niftis]
      images_pred = [join(folder_pred, i) for i in niftis]
      results = p.starmap(compute_dice_scores, zip(images_gt, images_pred))
      p.close()
      p.join()
  
      with open(join(folder_pred, "results.csv"), 'w') as f:
          for i, ni in enumerate(niftis):
              f.write("%s,%0.4f,%0.4f,%0.4f\n" % (ni, *results[i]))
  
  
  def remove_all_but_the_two_largest_conn_comp(img_itk_file: str, file_out: str):
      """
      This was not used. I was just curious because others used this. Turns out this is not necessary for my networks
      """
      img_itk = sitk.ReadImage(img_itk_file)
      img_npy = sitk.GetArrayFromImage(img_itk)
  
      labelmap, num_labels = label((img_npy > 0).astype(int))
  
      if num_labels > 2:
          label_sizes = []
          for i in range(1, num_labels + 1):
              label_sizes.append(np.sum(labelmap == i))
          argsrt = np.argsort(label_sizes)[::-1] # two largest are now argsrt[0] and argsrt[1]
          keep_mask = (labelmap == argsrt[0] + 1) | (labelmap == argsrt[1] + 1)
          img_npy[~keep_mask] = 0
          new = sitk.GetImageFromArray(img_npy)
          new.CopyInformation(img_itk)
          sitk.WriteImage(new, file_out)
          print(os.path.basename(img_itk_file), num_labels, label_sizes)
      else:
          shutil.copy(img_itk_file, file_out)
  
  
  def manual_postprocess(folder_in,
                         folder_out):
      """
      This was not used. I was just curious because others used this. Turns out this is not necessary for my networks
      """
      maybe_mkdir_p(folder_out)
      infiles = subfiles(folder_in, suffix=".nii.gz", join=False)
  
      outfiles = [join(folder_out, i) for i in infiles]
      infiles = [join(folder_in, i) for i in infiles]
  
      p = Pool(8)
      _ = p.starmap_async(remove_all_but_the_two_largest_conn_comp, zip(infiles, outfiles))
      _ = _.get()
      p.close()
      p.join()
  
  
  
  
  def copy_npz_fom_valsets():
      '''
      this is preparation for ensembling
      :return:
      '''
      base = join(network_training_output_dir, "3d_lowres/Task48_KiTS_clean")
      folders = ['nnUNetTrainerNewCandidate23_FabiansPreActResNet__nnUNetPlans',
                 'nnUNetTrainerNewCandidate23_FabiansResNet__nnUNetPlans',
                 'nnUNetTrainerNewCandidate23__nnUNetPlans']
      for f in folders:
          out = join(base, f, 'crossval_npz')
          maybe_mkdir_p(out)
          shutil.copy(join(base, f, 'plans.pkl'), out)
          for fold in range(5):
              cur = join(base, f, 'fold_%d' % fold, 'validation_raw')
              npz_files = subfiles(cur, suffix='.npz', join=False)
              pkl_files = [i[:-3] + 'pkl' for i in npz_files]
              assert all([isfile(join(cur, i)) for i in pkl_files])
              for n in npz_files:
                  corresponding_pkl = n[:-3] + 'pkl'
                  shutil.copy(join(cur, n), out)
                  shutil.copy(join(cur, corresponding_pkl), out)
  
  
  def ensemble(experiments=('nnUNetTrainerNewCandidate23_FabiansPreActResNet__nnUNetPlans',
                 'nnUNetTrainerNewCandidate23_FabiansResNet__nnUNetPlans'), out_dir="/media/fabian/Results/nnUNet/3d_lowres/Task48_KiTS_clean/ensemble_preactres_and_res"):
      from nnunet.inference.ensemble_predictions import merge
      folders = [join(network_training_output_dir, "3d_lowres/Task48_KiTS_clean", i, 'crossval_npz') for i in experiments]
      merge(folders, out_dir, 8)
  
  
  def prepare_submission(fld= "/home/fabian/drives/datasets/results/nnUNet/test_sets/Task48_KiTS_clean/predicted_ens_3d_fullres_3d_cascade_fullres_postprocessed", # '/home/fabian/datasets_fabian/predicted_KiTS_nnUNetTrainerNewCandidate23_FabiansResNet',
                         out='/home/fabian/drives/datasets/results/nnUNet/test_sets/Task48_KiTS_clean/submission'):
      nii = subfiles(fld, join=False, suffix='.nii.gz')
      maybe_mkdir_p(out)
      for n in nii:
          outfname = n.replace('case', 'prediction')
          shutil.copy(join(fld, n), join(out, outfname))
  
  
  def pretent_to_be_nnUNetTrainer(base, folds=(0, 1, 2, 3, 4)):
      """
      changes best checkpoint pickle nnunettrainer class name to nnUNetTrainer
      :param experiments:
      :return:
      """
      for fold in folds:
          cur = join(base, "fold_%d" % fold)
          pkl_file = join(cur, 'model_best.model.pkl')
          a = load_pickle(pkl_file)
          a['name_old'] = deepcopy(a['name'])
          a['name'] = 'nnUNetTrainer'
          save_pickle(a, pkl_file)
  
  
  def reset_trainerName(base, folds=(0, 1, 2, 3, 4)):
      for fold in folds:
          cur = join(base, "fold_%d" % fold)
          pkl_file = join(cur, 'model_best.model.pkl')
          a = load_pickle(pkl_file)
          a['name'] = a['name_old']
          del a['name_old']
          save_pickle(a, pkl_file)
  
  
  def nnUNetTrainer_these(experiments=('nnUNetTrainerNewCandidate23_FabiansPreActResNet__nnUNetPlans',
                 'nnUNetTrainerNewCandidate23_FabiansResNet__nnUNetPlans',
                 'nnUNetTrainerNewCandidate23__nnUNetPlans')):
      """
      changes best checkpoint pickle nnunettrainer class name to nnUNetTrainer
      :param experiments:
      :return:
      """
      base = join(network_training_output_dir, "3d_lowres/Task48_KiTS_clean")
      for exp in experiments:
          cur = join(base, exp)
          pretent_to_be_nnUNetTrainer(cur)
  
  
  def reset_trainerName_these(experiments=('nnUNetTrainerNewCandidate23_FabiansPreActResNet__nnUNetPlans',
                 'nnUNetTrainerNewCandidate23_FabiansResNet__nnUNetPlans',
                 'nnUNetTrainerNewCandidate23__nnUNetPlans')):
      """
      changes best checkpoint pickle nnunettrainer class name to nnUNetTrainer
      :param experiments:
      :return:
      """
      base = join(network_training_output_dir, "3d_lowres/Task48_KiTS_clean")
      for exp in experiments:
          cur = join(base, exp)
          reset_trainerName(cur)
  
  
  if __name__ == "__main__":
      base_image = "/dfsdata2/zhangyao_data/DB/Abdomen1K/Task100_FourAbdOrgans/imagesTs"
      base_label = "/dfsdata2/zhangyao_data/DB/Abdomen1K/Task100_FourAbdOrgans/labelsTs"
      out = "../data/nnUNet_raw_data/Task440_MSDLiverBase"
      cases = np.sort(os.listdir(base_label))
      cases = [x for x in cases if x.split("_")[0] == "pancreas" and int(x.split(".")[0].split("_")[1]) <= 421]
  
      maybe_mkdir_p(out)
      maybe_mkdir_p(join(out, "imagesTr"))
      maybe_mkdir_p(join(out, "imagesTs"))
      maybe_mkdir_p(join(out, "labelsTr"))
  
      for c in cases:
          shutil.copy(join(base_image, c[:-7] + "_0000.nii.gz"), join(out, "imagesTr", c[:-7] + "_0000.nii.gz"))
          # shutil.copy(join(base_label, c), join(out, "labelsTr", c))
          nii_volume = nib.load(join(base_label, c))
          seg_volume = nii_volume.get_fdata()
          seg_volume[seg_volume!=1] = 0
          seg_volume[seg_volume==1] = 1
          nib.save(nib.Nifti1Image(seg_volume, nii_volume.affine), join(out, "labelsTr", c))
          print(join(out, "imagesTr", c[:-7] + "_0000.nii.gz"))
  
      json_dict = {}
      json_dict['name'] = "MSDLiverBase"
      json_dict['description'] = ""
      json_dict['tensorImageSize'] = "4D"
      json_dict['reference'] = ""
      json_dict['licence'] = ""
      json_dict['release'] = "0.0"
      json_dict['modality'] = {
          "0": "CT",
      }
      json_dict['labels'] = {
          "0": "background",
          "1": "Liver",
      }
      json_dict['numTraining'] = len(cases)
      json_dict['numTest'] = 0
      json_dict['training'] = [{'image': "./imagesTr/%s" % i, "label": "./labelsTr/%s" % i} for i in
                               cases]
      json_dict['test'] = []
  
      save_json(json_dict, os.path.join(out, "dataset.json"))
  
  