# Copyright 2019 Population Health Sciences and Image Analysis, German Center for Neurodegenerative Diseases(DZNE)
#
#    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 __future__ import division
import sys
sys.path.append('./')
sys.path.append('../')

import os

import nibabel as nib
import pandas as pd
from Code.utilities.misc import locate_file
from Code.utilities.visualization_misc import multiview_plotting
from Code.utilities.metrics import calculate_statistics_v2
from Code.utilities.models import run_adipose_localization,run_adipose_segmentation
import numpy as np
from keras import backend as K
from Code.utilities.image_processing import largets_connected_componets,find_labels
from Code.utilities.conform import conform



def clean_segmentations(label_map):

    new_label_map=np.copy(label_map)

    new_label_map= largets_connected_componets(new_label_map)

    return new_label_map


def extreme_AAT_increase_flag(predict_array,threshold=0.3):

    extreme_increase_flag = False
    for slice in range(1,(predict_array.shape[0]-1)) :
        previous_sat =np.sum(predict_array[slice-1,:,:] == 1)
        previous_vat =np.sum(predict_array[slice-1,:,:] == 2)
        current_sat=np.sum(predict_array[slice,:,:] == 1)
        current_vat =np.sum(predict_array[slice,:,:] == 2)
        following_sat=np.sum(predict_array[slice+1,:,:] == 1)
        following_vat=np.sum(predict_array[slice+1,:,:] == 2)

        sat_threshold=current_sat*threshold
        vat_threshold = current_vat * threshold

        if np.abs(current_sat-previous_sat) > sat_threshold or np.abs(current_sat-following_sat) > sat_threshold:
            extreme_increase_flag= 'SAT increase over the threshold'
        elif np.abs(current_vat-previous_vat) > vat_threshold or np.abs(current_vat-following_vat) > vat_threshold:
            extreme_increase_flag = 'VAT increase over the threshold'

    return  extreme_increase_flag

def stats_variable_initialization(nb_comparments,weighted=True):

    # initialize Stats Variables
    variable_columns = []

    volume_variable_columns = ['VOL_cm3', 'SAT_VOL_cm3', 'VAT_VOL_cm3', 'AAT_VOL_cm3',
                               'VAT_VOL_TO_SAT_VOL', 'VAT_VOL_TO_AAT_VOL', 'SAT_VOL_TO_AAT_VOL']

    w_volume_variable_columns= ['W_VOL_cm3','WSAT_VOL_cm3', 'WVAT_VOL_cm3',
                               'WAAT_VOL_cm3', 'WVAT_VOL_TO_WSAT_VOL', 'WVAT_VOL_TO_WAAT_VOL', 'WSAT_VOL_TO_WAAT_VOL']

    area_variable_columns = ['HEIGHT_cm', 'AVG_AREA_cm2', 'AVG_PERIMETER_cm']

    base_variable_len={}
    base_variable_len['Area']=len(area_variable_columns)
    base_variable_len['Volume']=len(volume_variable_columns)
    base_variable_len['W_Volume']=len(w_volume_variable_columns)

    roi_areas = ['wb']
    if nb_comparments != 0:
        # From Feet to Head
        for i in range(int(nb_comparments), 0, -1):
            roi_areas.append('Q' + str(i))

    for roi in roi_areas:
        for area_id in area_variable_columns:
            variable_columns.append(roi + '_' + area_id)
        for vol_id in volume_variable_columns:
            variable_columns.append(roi + '_' + vol_id)

        if weighted:
            for w_vol_id in w_volume_variable_columns:
                variable_columns.append(roi + '_' + w_vol_id)


    variable_columns.insert(0, 'imageid')
    variable_columns.insert(1, '#_Slices')
    variable_columns.insert(2,'FLAGS')

    return variable_columns,base_variable_len

def check_image_contrast(water_array,fat_array):

    slice = fat_array.shape[0] // 2

    water_slice=water_array[slice,20:-20,20:-20]
    fat_slice=fat_array[slice,20:-20,20:-20]

    intensity_max=np.max([np.max(water_slice),np.max(fat_slice)])

    water_slice=water_slice/intensity_max
    fat_slice=fat_slice/intensity_max

    new_fat=np.zeros((fat_slice.shape[0],fat_slice.shape[1]))

    new_fat[fat_slice >= (0.10 * np.max(fat_slice))] = 2
    new_fat[fat_slice >= (0.30*np.max(fat_slice))] = 1

    border_idx=np.where(new_fat == 2)

    point_index=np.arange(0,len(border_idx[0]),10)
    point_y=border_idx[0][point_index]
    point_x=border_idx[1][point_index]

    fat_count=0
    no_fat_count=0
    for j in range(len(point_x)):

        value = fat_slice[point_y[j],point_x[j]] -water_slice[point_y[j],point_x[j]]

        if value < 0 :
            fat_count += 1
        else:
            no_fat_count += 1

    if no_fat_count > fat_count or ((no_fat_count/fat_count) > 0.75):
        FLAG='Check image contrast'
    else:
        FLAG = False

    return FLAG

def check_flags(predicted_array,water_array,fat_array,ratio_vat_sat,threshold=0.30,sat_to_vat_threshold=2.0):
    FLAG = check_image_contrast(water_array,fat_array)

    if FLAG == False:

        FLAG=extreme_AAT_increase_flag(predicted_array,threshold=threshold)

        if ratio_vat_sat > sat_to_vat_threshold:
            FLAG = 'High VAT to SAT ratio'

    return FLAG





def run_adipose_pipeline(args,flags,save_path='/',data_path='/',id='Test'):

    output_stats = 'AAT_stats.tsv'
    output_pred_fat = 'AAT_pred.nii.gz'
    output_pred = 'ALL_pred.nii.gz'
    qc_images = []



    print('-' * 30)
    print('Loading Subject')
    print(id)
    sub = id


    fat_file = locate_file('*'+str(args.fat_image), data_path)
    water_file = locate_file('*'+str(args.water_image), data_path)



    # Check fat
    if fat_file:
        print('-' * 30)
        print('Loading Fat Image')
        print(fat_file[0])
        #Load Fat Images
        fat_img = nib.load(fat_file[0])
        ishape = fat_img.shape

        #Check if  data from example_data_folder was loaded : Only contains the value -9999
        if len(np.unique(fat_img.get_data())) > 2:
            if len(ishape) > 3 and ishape[3] != 1:
                print('ERROR: Multiple input frames (' + format(fat_img.shape[3]) + ') not supported!')
            else:
                fat_img = conform(fat_img, flags=flags, order=args.order, save_path=save_path, mod='fat',
                                  axial=args.axial)
                fat_array = fat_img.get_data()
                fat_array = np.swapaxes(fat_array, 0, 2)
                fat_zooms = fat_img.header.get_zooms()


            print('-' * 30)
            print('Loading Water Image')
            #Check water image
            if not water_file:
                weighted=False
                print('No water image found, weighted volumes would not be calculated')
                water_array=np.zeros(fat_array.shape)
            else:
                print(water_file[0])
                weighted=True
                water_img = nib.load(water_file[0])
                ishape = fat_img.shape
                if len(ishape) > 3 and ishape[3] != 1:
                    print('ERROR: Multiple input frames (' + format(water_img.shape[3]) + ') not supported!')
                    weighted = False
                    print('No water image found, weighted volumes would not be calculated')
                    water_array = np.zeros(fat_array.shape)
                else:
                    water_img = conform(water_img, flags=flags, order=args.order, save_path=save_path, mod='water',
                                       axial=args.axial)
                    water_array = water_img.get_data()
                    water_array = np.swapaxes(water_array, 0, 2)


            variable_columns, base_variable_len = stats_variable_initialization(args.compartments,weighted)
            ratio_position = variable_columns.index('wb_VAT_VOL_TO_SAT_VOL')

            pixel_matrix = np.zeros((1, len(variable_columns)), dtype=object)
            row_px = 0

            img_spacing=np.copy(fat_zooms)

            if not args.run_stats:

                if args.run_localization:
                    high_idx,low_idx=run_adipose_localization(fat_array,flags)
                    K.clear_session()
                else:
                    high_idx=fat_array.shape[0]
                    low_idx= 0

                print('the index values are %d, %d' % (low_idx, high_idx))

                # Image Segmentation
                pred_array=run_adipose_segmentation(fat_array,flags,args)
                K.clear_session()

            else:
                pred_file = locate_file('*AAT_pred.nii.gz', data_path)

                if pred_file :
                    pred_img = nib.load(pred_file[0])
                    pred_array = pred_img.get_data()
                    pred_array = np.swapaxes(pred_array, 0, 2)
                    pred_zooms = pred_img.header.get_zooms()
                    img_spacing = np.copy(pred_zooms)
                    # img_spacing[0] = pred_zooms[2]
                    # img_spacing[2] = pred_zooms[0]


                    high_idx, low_idx = find_labels(pred_array)
                else :
                    print('Subject has no prediction map, a ATT_pred.nii.gz file is required to run the stats option')
                    print('-' * 30)
                    print('ERROR: Subject doesnt have a AAT_pred.nii.gz')

            print('-' * 30)
            print('Calculating Stats')

            pred_array[0:low_idx,:,:]=0
            pred_array[high_idx:,:,:]=0

            pred_array [low_idx:high_idx, :, :] = clean_segmentations(pred_array[low_idx:high_idx, :, :])

            pixel_matrix[row_px:row_px + 1, 0] = sub


            pixel_matrix[row_px:row_px + 1, 3:] = calculate_statistics_v2(pred_array[low_idx:high_idx, :, :],
                                                                          water_array[low_idx:high_idx, :, :],
                                                                          fat_array[low_idx:high_idx, :, :],
                                                                          low_idx, high_idx, variable_columns[3:],
                                                                          base_variable_len, img_spacing,
                                                                          args.compartments, weighted=weighted)


            pixel_matrix[row_px:row_px + 1, 1] = int(high_idx-low_idx)
            pixel_matrix[row_px:row_px + 1, 2] = check_flags(pred_array[low_idx:high_idx, :, :],water_array=water_array,fat_array=fat_array,
                                                             ratio_vat_sat=pixel_matrix[row_px, ratio_position],
                                                             threshold=args.increase_threshold,sat_to_vat_threshold=args.sat_to_vat_threshold)

            df = pd.DataFrame(pixel_matrix[row_px:row_px+1, :], columns=variable_columns)

            if not os.path.isdir(os.path.join(save_path, 'Segmentations')):
                os.mkdir(os.path.join(save_path, 'Segmentations'))

            seg_path=os.path.join(save_path, 'Segmentations')

            df.to_csv(seg_path+'/'+output_stats, sep='\t', index=False)
            df.to_json(seg_path+ '/AAT_variables_summary.json', orient='records')

            row_px += 1

            # Modified images for display
            disp_fat = np.flipud(fat_array[:])
            disp_fat = np.fliplr(disp_fat[:])
            disp_pred=np.copy(pred_array)
            disp_pred = np.flipud(disp_pred)
            disp_pred = np.fliplr(disp_pred)

            #only display SAT and VAT
            disp_pred[disp_pred>=3]=0

            idx = (np.where(disp_pred > 0))
            low_idx = np.min(idx[0])
            high_idx = np.max(idx[0])

            interval = (high_idx - low_idx) // 4

            # Control images of the segmentation
            if not args.control_images:
                if not os.path.isdir(os.path.join(save_path, 'QC')):
                    os.mkdir(os.path.join(save_path, 'QC'))
                for i in range(4):
                    control_point = [0, int(np.ceil(disp_fat.shape[1] / 2)), int(np.ceil(disp_fat.shape[2] / 2))]
                    control_point[0] = int(np.ceil(np.random.uniform(high_idx - interval * i, high_idx - interval * ((i + 1)))))
                    multiview_plotting(disp_fat, disp_pred, control_point, save_path+'/QC/QC_%s.png' % i,
                                       classes=5, alpha=0.5, nbviews=3)

            print('-' * 30)
            print('Saving Segmentation')
            # Save prediction
            pred_array=np.swapaxes(pred_array,2,0)
            pred_img = nib.Nifti1Image(pred_array, fat_img.affine, fat_img.header)
            nib.save(pred_img, seg_path+'/'+output_pred)

            pred_array[pred_array>=3]=0
            pred_img = nib.Nifti1Image(pred_array, fat_img.affine, fat_img.header)
            nib.save(pred_img, seg_path+'/'+output_pred_fat)

            print('-' * 30)

            print('Finish Subject %s' % sub)

            print('-' * 30)


        else :
            print('ERROR: Input image empty \n'
                  'Note : Volumes from the example_data_folder are empty \n'
                  'The example_data_folder is only a ilustrative example on how volumes have to be organized for FatSegNet to work.')
            print('Please provided your own dixon MR scans')

    else:
        print('')
        print('-' * 30)
        print('ERROR: Subject doesnt have a Fat Image named %s,\n'
              'Please verified that the name provided to the -fat argument matches the one in the participants folder (default : FatImaging_F.nii.gz )'%str(args.fat_image))


