#!/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


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

def get_dict(dict_path):
    """
    get dictionary in NAS/public_info
    :param : dict_path(string) location of the dictionary
    :return: dictionary in pandas format
    """
    if dict_path == 'default':
        dictionary = pd.read_excel(r'/home/ana/NAS_database/public_info/RT_dictionary.xlsx',  engine='openpyxl')
    else:
        dictionary = pd.read_excel(dict_path, engine='openpyxl')
    
    return dictionary

      
def set_header_info(nii_file, voxelsize, image_position_patient, contours_exist = None):
    nii_file.header['pixdim'][1] = voxelsize[0]
    nii_file.header['pixdim'][2] = voxelsize[1]
    nii_file.header['pixdim'][3] = voxelsize[2]
    
    #affine - voxelsize
    nii_file.affine[0][0] = voxelsize[0]
    nii_file.affine[1][1] = voxelsize[1]
    nii_file.affine[2][2] = voxelsize[2]
    #affine - imagecorner
    nii_file.affine[0][3] = image_position_patient[0]
    nii_file.affine[1][3] = image_position_patient[1]
    nii_file.affine[2][3] = image_position_patient[2]
    if contours_exist is not None:
        nii_file.header.extensions.append(nib.nifti1.Nifti1Extension(0, bytearray(contours_exist)))
    return nii_file


def get_struct_and_contoursexist(roi,arrayshape):
    struct = np.zeros((arrayshape[0],arrayshape[1],arrayshape[2]))
    roi_names = list(roi.keys())   
    contoursexist = []
    
    for i in range(len(roi_names)):
        if(not(roi[roi_names[i]] is None)):
            contoursexist.append(1)
            struct += (2**i)*roi[roi_names[i]]
        else:
            contoursexist.append(0)
    return struct, contoursexist


def save_images(dst_dir, voxelsize, image_position_patient, image, image_type, roi=None, dose=None, bdoses=None, prior_knowledge = None, spotmap = None):
    
    # encode in nii  and save at dst_dir
    # IMPORTANT WE NEED TO CONFIRM THE SIGNS OF THE ENTRIES IN THE AFFINE, 
    # ALTHOUGH MAYBE AT THE END THE IMPORTANCE IS HOW WE WILL USE THIS DATA .... 
    # also instead of changing field by field, the pixdim and affine can be encoded
    # using the set_sform method --> info here: https://nipy.org/nibabel/nifti_images.html
    
    # IMAGE (CT, MR ...)
    image_shape = image.shape
    image_nii = nib.Nifti1Image(image, affine=np.eye(4)) # for Nifti1 header, change for a Nifti2 type of header
    # Update header fields
    image_nii = set_header_info(image_nii, voxelsize, image_position_patient)
    # Save  nii 
    nib.save(image_nii, os.path.join(dst_dir,image_type.lower()+'.nii.gz'))
    
    # DOSE
    if dose is not None:
        dose_nii = nib.Nifti1Image(dose, affine=np.eye(4)) # for Nifti1 header, change for a Nifti2 type of header
        # Update header fields
        dose_nii = set_header_info(dose_nii, voxelsize, image_position_patient)
        # Save nii 
        nib.save(dose_nii, os.path.join(dst_dir,'dose.nii.gz'))

    # BDOSES
    if bdoses is not None:
        for b in range(len(bdoses)):
            bdose_nii = nib.Nifti1Image(bdoses[b].Image, affine=np.eye(4)) # for Nifti1 header, change for a Nifti2 type of header
            # Update header fields
            bdose_nii = set_header_info(bdose_nii, voxelsize, image_position_patient)
            # Save nii
            nib.save(bdose_nii, os.path.join(dst_dir,'dose_b{}.nii.gz'.format(b+1)))

    # PRIOR KNOWLEDGE 
    if prior_knowledge is not None:
        for charact, pk in prior_knowledge.items():
            prior_knowledge_nii = nib.Nifti1Image(pk, affine=np.eye(4)) # for Nifti1 header, change for a Nifti2 type of header
            # Update header fields
            prior_knowledge_nii = set_header_info(prior_knowledge_nii, voxelsize, image_position_patient)
            # Save nii 
            nib.save(prior_knowledge_nii, os.path.join(dst_dir,'prior_knowledge_{}.nii.gz'.format(charact.lower())))

    # RTSTRUCT
    if roi is not None:
        struct_compressed, contours_exist = get_struct_and_contoursexist(roi, image_shape)

        struct_nii_compressed = nib.Nifti1Image(struct_compressed, affine=np.eye(4)) # for Nifti1 header, change for a Nifti2 type of header
        # struct_nii_compressed.set_data_dtype('smallest')
        # Update header fields
        struct_nii_compressed = set_header_info(struct_nii_compressed, voxelsize, image_position_patient, contours_exist=contours_exist)
        # Save nii 
        nib.save(struct_nii_compressed, os.path.join(dst_dir,'struct.nii.gz'))