#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Sat Sep  5 20:34:46 2020

@author: ana
"""


# import general libraries
import os
import numpy as np
import pandas as pd
import scipy
import nibabel as nib
import pydicom
import glob
import warnings
from copy import deepcopy
from scipy.ndimage import find_objects
from scipy.ndimage.morphology import binary_fill_holes
from skimage import measure
from matplotlib.patches import Polygon
from libraries.utils_nii_dicom import set_header_info


###############################################################################
###################################   FUNCTIONS   #############################
###############################################################################

def delete_all(obj):
    for i in vars(obj):
        del obj.i

def overwrite_ct_threshold(ct_image, body = None, artefact = None, contrast = None):
    # Change the HU out of the body to air: -1000    
    if body is not None:
        # Change the HU outside the body to -1000
        ct_image[body==0]=-1000
    if artefact is not None:
        # Change the HU to muscle: 14
        ct_image[artefact==1]=14
    if contrast is not None:
        # Change the HU to water: 0 Houndsfield Unit: CT unit
        ct_image[contrast==1]=0
    # Threshold above 1560HU
    ct_image[ct_image > 1560] = 1560
    return ct_image

def remove_dose_outside_mask(dose, mask = None):
    # Put dose = 0 outside the mask (typically the body)   
    if mask is not None:
        dose[mask==0]=0
    return dose

def get_and_save_tv(roi, nib_header, tv_names, dst_dir):
    tv = np.zeros(nib_header['dim'][1:4])
    roi_names = list(roi.keys())
    if tv_names is not None:
        for c in tv_names:
            if c in roi_names:
                tv[roi[c]>0]=int(c.split('_')[1])/100
                
    # create nii, update header and save
    save_nii_image(tv, nib_header,dst_dir, 'struct_tv')

    return tv

def get_and_save_oars(roi, nib_header, oar_names, dst_dir):
    oars = np.zeros(nib_header['dim'][1:4])
    
    roi_names = list(roi.keys())
    contoursexist = []
    
    if oar_names is not None:
        for i in range(len(oar_names)):
            if oar_names[i] in roi_names:
                if roi[oar_names[i]] is not None:
                    contoursexist.append(1)
                    oars += (2**i)*roi[oar_names[i]]
                else:
                    contoursexist.append(0)        
                    
    # create nii, update header and save 
    save_nii_image(oars, nib_header,dst_dir, 'struct_oar', contours_exist = contoursexist)

        
def get_and_save_sample_probability(tv,dst_dir):
    # get sample probability for patch-based training
    bufftv=np.zeros_like(tv)
    bufftv[tv!=0]=1
    m=bufftv.sum(axis=0).sum(axis=0).sum()/1000
    sample_probability_slc=(bufftv.sum(axis=0).sum(axis=0)+m)/(bufftv.sum(axis=0).sum(axis=0)+m).sum()
    sample_probability_row=(bufftv.sum(axis=1).sum(axis=1)+m)/(bufftv.sum(axis=1).sum(axis=1)+m).sum()
    sample_probability_col=(bufftv.sum(axis=0).sum(axis=-1)+m)/(bufftv.sum(axis=0).sum(axis=-1)+m).sum()
    
    np.savez_compressed(os.path.join(dst_dir,'sample_probability_row.npz'),sample_probability_row)
    np.savez_compressed(os.path.join(dst_dir,'sample_probability_col.npz'),sample_probability_col)
    np.savez_compressed(os.path.join(dst_dir,'sample_probability_slc.npz'),sample_probability_slc)
    

def decompress_struct(struct_nib,struct_list):
    
    struct_data = struct_nib.get_fdata()
    roi = dict.fromkeys(struct_list, None) 
    # get contourexists from header
    contoursexist = list(struct_nib.header.extensions[0].get_content())
    for i in range(len(contoursexist)):
        if contoursexist[i] == 1:
            roi[struct_list[i]] = np.bitwise_and(struct_data.astype(int), 2 ** i).astype(bool)

    return roi

def binary_to_integers(struct,select_list):
        
    roi_names = list(struct.keys())
    # get shape
    for k in roi_names:
        if struct[k] is not None:
            struct_shape = struct[k].shape
    # initialize roi
    roi = np.zeros(struct_shape)     
    
    # initialize label
    label = 0
    for i in range(len(select_list)):
        label = label + 1
        if select_list[i] in roi_names:
            if struct[select_list[i]] is not None:
                # populate roi matrix with selected rois in select_list
                roi[struct[select_list[i]]] = label
        
    return roi

def save_nii_image(nib_img, nib_header,dst_dir, image_name, contours_exist = None):
    
    image_nii = nib.Nifti1Image(nib_img, affine=np.eye(4)) # for Nifti1 header, change for a Nifti2 type of header
    # Update header fields
    if contours_exist is not None:
        image_nii = set_header_info(image_nii, nib_header['pixdim'][1:4], [nib_header['qoffset_x'],nib_header['qoffset_y'],nib_header['qoffset_z']], contours_exist = contours_exist)
    else:
        image_nii = set_header_info(image_nii, nib_header['pixdim'][1:4], [nib_header['qoffset_x'],nib_header['qoffset_y'],nib_header['qoffset_z']])
    # Save  nii 
    nib.save(image_nii, os.path.join(dst_dir,image_name +'.nii.gz'))