import argparse
import sys
import os
import time

sys.path.append('./')

from Code.adipose_pipeline import run_adipose_pipeline
from Code.utilities.misc import locate_file,locate_dir
import pandas as pd
import numpy as np

import warnings

if not sys.warnoptions:
    warnings.simplefilter("ignore")


def check_paths(save_folder,subject_id,flags):

    save_path=os.path.join(flags['output_path'],save_folder)

    if not os.path.isdir(save_path):
        os.mkdir(save_path)

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

    final_path = os.path.join(save_path,subject_id)
    return final_path


class Transcript(object):

    def __init__(self, filename):
        self.terminal = sys.stdout
        self.logfile = open(filename, "a")

    def write(self, message):
        self.terminal.write(message)
        self.logfile.write(message)

    def flush(self):
        # this flush method is needed for python 3 compatibility.
        # this handles the flush command by doing nothing.
        # you might want to specify some extra behavior here.
        pass



def option_parse():


    parser = argparse.ArgumentParser(
        description='Adipose Pipeline to segment the  abdominal adipose tissue into VAT and SAT. '
                    'Each subject should have a independent folder with the water and fat images. '
                    'Input images have to be nifti files and should be named consistently in all subjects. '
                    'The Output path is define by the user ; all the outputs from the pipeline will be store under $output_path/$subject_id. '
                    'The predicted segmentation mask is save under ($AAT_pred) and all the statistics under $ATT_stats',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-f", "--file", help="csv file containing the subjects to process, the csv file should be order as follow : $subject_id,$subject_path", required=False,default='participants.csv')
    parser.add_argument("-outp", "--output_folder",
                        help="Main folder where the variables and control images are going to be store", required=False, default='')

    parser.add_argument("-fat", "--fat_image", type=str, help="Name of the fat image", required=False,
                        default='FatImaging_F.nii.gz')
    parser.add_argument("-water", "--water_image", type=str, help="Name of the water image", required=False,
                        default='FatImaging_W.nii.gz')

    parser.add_argument('-No_QC',"--control_images",action='store_true',help='Plot subjects prediction for visual quality control',required=False,default=False)

    parser.add_argument('-loc', "--run_localization", action='store_true',
                        help='run abdominal region localization model ', required=False, default=False)

    parser.add_argument('-axial', "--axial", action='store_true',
                        help='run only axial model ', required=False, default=False)

    parser.add_argument('-order', "--order", type=int,
                        help='interpolation order (0=nearest,1=linear(default),2=quadratic,3=cubic) ', required=False, default=1)

    parser.add_argument('-comp', "--compartments", type=int,
                        help='Number of equal compartments to run the analysis, by default the whole region(wb) is calculated', required=False, default=0)

    parser.add_argument('-AAT', "--increase_threshold", type=float,
                        help='Warning flag for an increase in AAT over threhold between consecutive scans', required=False, default=0.4)

    parser.add_argument('-ratio', "--sat_to_vat_threshold", type=float,
                        help='Warning flag for a high vat to sat ratio', required=False, default=2.0)

    parser.add_argument('-stats', "--run_stats", action='store_true',
                        help='run only stats , segmentation map required ', required=False, default=False)

    parser.add_argument('-gpu_id', "--gpu_id", type=int,
                        help='if using gpu, please give the gpu device name', required=False, default=0)




    args = parser.parse_args()

    FLAGS = {}
    FLAGS['multiviewModel'] = '/tool/Adipose_Seg_Models/Segmentation/'
    FLAGS['singleViewModels'] = '/tool/Adipose_Seg_Models/Segmentation/'
    FLAGS['localizationModels'] = '/tool/Adipose_Seg_Models/Localization/'
    FLAGS['input_path']='/tool/Data'
    FLAGS['output_path']='/tool/Output'
    FLAGS['imgSize'] = [256, 224, 72]
    FLAGS['spacing'] = [float(1.9531), float(1.9531),float(5.0)]
    FLAGS['base_ornt'] = np.array([[0, -1], [1, 1], [2, 1]])
    #FLAGS['compartments']=0
    #control_images = True


    return args,FLAGS


def run_fatsegnet(args,FLAGS):

    # load file
    participant_file=locate_file('*'+args.file,FLAGS['input_path'])
    if participant_file:
        print('Loading participant from file : %s'%participant_file[0])
        df =pd.read_csv(participant_file[0],header=None)
        if df.empty:
            print('Participant file empty ')
        else:
            file_list=df.values
            for sub in file_list:
                id=sub[0]
                path = locate_dir('*'+str(id)+'*',FLAGS['input_path'])
                if path:
                    if os.path.isdir(path[0]):

                        start = time.time()

                        save_path = check_paths(save_folder=args.output_folder, subject_id=str(id),flags=FLAGS)

                        sys.stdout= Transcript(filename=save_path + '/temp.log')

                        run_adipose_pipeline(args=args, flags=FLAGS, save_path=save_path,data_path=path[0],id=str(id))

                        end = time.time() - start

                        print("Total time for computation of segmentation is %0.4f seconds."%end)

                        sys.stdout.logfile.close()
                        sys.stdout = sys.stdout.terminal
                    else:
                        print ('Directory %s not found'%path)
                else :
                    print('Directory name %s not found' % id)
            print('\n')
            print('Thank you for using FatSegNet')
            print('If you find it useful and use it for a publication, please cite: ')
            print('\n')
            print('Estrada S, Lu R, Conjeti S, et al.'
                  'FatSegNet: A fully automated deep learning pipeline for adipose tissue segmentation on abdominal dixon MRI.'
                  'Magn Reson Med. 2019;00:1-13. https:// doi.org/10.1002/mrm.28022')
    else:
        print('No partipant file found, please provide one the input data folder')



if __name__=='__main__':


    args,FLAGS= option_parse()

    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID";
    # The GPU id to use, usually either "0" or "1";
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id);

    run_fatsegnet(args,FLAGS)

    sys.exit(0)

