#!/opt/python 3.3/bin/python3
# -*- coding: utf-8 -*-

import tkinter as tk
import os
import sys
import distutils.dir_util
import subprocess
import pprint
#import messagebox

WDprocessing_index = 0
Dataprocessing_index = 0
EnvironSet_index = 0
fields = ["Object ID", "T1 MRI Data Path", "T1 Data Type (nii or mgz)", "dWMRI Data Path", "dWMRI Data Path second", "dwi Data Type (nii or mgz)", "fMRI(Bold) Data Path", \
           "fMRI Data Type (nii or mgz)", "Working Directory", "Number of ROI", "Template Path", "Template Name", "Template in MNI Space?"]

curWD = os.getcwd()
EnvironVariablesDict = {"SGE_ROOT":"/usr/local/ge2011.11/",\
                        "FREESURFER_HOME":"/opt/freesurfer",\
                        "FSLDIR":"/opt/fsl",\
                        "scripts_path":"Scripts",\
                        "MRTrixDIR":"/opt/mrtrix2/bin",\
                        "MRTrix3DIR":"/home/jwang/local/MRtrix3bin",\
                        "FSLSUBALREADYRUN":"true"}

PathList = ["/home/jwang/codes/fsl/bin:/opt/afni:/usr/local/ge2011.11/bin/linux-x64:/sbin"]

SourceList = ["/opt/source/octave.sh", \
              "/opt/source/matlab.sh",\
              "/opt/source/afni_fsl.sh",\
              "/opt/source/freesurfer.sh",\
              "/opt/freesurfer/SetUpFreeSurfer.sh",\
              "/opt/fsl/etc/fslconf/fsl.sh"]

fields_check =[ "ROImask to Template(non MNI) space",\
                "T1 Prerocessing with FreeSurfer",\
                "DWI Prerocessing with FreeSurfer", \
                "DWI Prerocessing with MrTrix2", \
                "fMRI Preprocessing", \
                 "Generate mask",\
                 "Tracking", \
                 "Compute SC matrix", \
                 "Aggregate SC matrix", \
                 "Convert to TVB format and clean up results"]



script_dict = {

"title": "#!/bin/bash\n\
echo \"*************************************************************\"\n \
echo \"***               TVB empirical data pipeline             ***\" \n \
echo \"*************************************************************\"\n\
#unset SGE_ROOT \n",

"ROImask to Template(non MNI) space": "echo \"************ ROImask to Template Space *******\"\n\
qsub -V -q bigmem_16.q -l h_vmem=16G -cwd -N Roimask2template_${subID} ${scripts_path}/ROIMask2Template.sh \n",

"T1 Prerocessing with FreeSurfer":"echo \"************ T1 preprocessing with FreeSurfer *******\"\n\
qsub -V -q parallel.q -pe mcore 8 -cwd -N reconAll_${subID} ${scripts_path}/FreeSurfer_recon-all.sh \n",


"DWI Prerocessing with FreeSurfer":"echo \"********** DWI preprocessing with FreeSurfer ********\"\n\
qsub -hold_jid reconAll_${subID} -V -q bigmem_16.q -l h_vmem=16G -cwd -N dtRecon_${subID} ${scripts_path}/FreeSurfer_dt-recon.sh\n",


"DWI Prerocessing with MrTrix2":"echo \"********** DWI preprocessing with MRtrix2 ***********\"\n\
qsub -hold_jid dtRecon_${subID} -V -q bigmem_16.q -l h_vmem=16G -cwd -N mrPre_${subID} ${scripts_path}/MRtrix_prepro.sh\n",

"fMRI Preprocessing": "echo \"**************** fMRI preprocessing *****************\"\n\
qsub -hold_jid reconAll_${subID},dtRecon_${subID}  -V -q all.q -cwd  -N fmri_1_${subID} ${scripts_path}/fmriFC1.sh\n\
qsub -hold_jid fmri_1_${subID} -V -q all.q -cwd -N fmri_2_${subID} ${scripts_path}/fmriFC2.sh\n\
qsub -hold_jid fmri_2_${subID},feat5_stop -V -q all.q -cwd -N fmri_3_${subID} ${scripts_path}/fmriFC3.sh\n",

"Generate mask": "echo \"****************** Generate mask ********************\"\n\
qsub -hold_jid mrPre_${subID} -V -q bigmem_16.q -l h_vmem=16G -cwd -N mask_${subID} ${scripts_path}/mask.sh\n",

"Tracking": "echo \"********************* Tracking **********************\"\n\
qsub -hold_jid mask_${subID} -V -q parallel.q -cwd -N tracking_1_${subID} ${scripts_path}/runTracking1.sh\n\
qsub -hold_jid mask_${subID} -V -q parallel.q -cwd -N tracking_2_${subID} ${scripts_path}/runTracking2.sh\n\
qsub -hold_jid mask_${subID} -V -q parallel.q -cwd -N tracking_3_${subID} ${scripts_path}/runTracking3.sh\n\
qsub -hold_jid mask_${subID} -V -q parallel.q -cwd -N tracking_4_${subID} ${scripts_path}/runTracking4.sh\n\
qsub -hold_jid mask_${subID} -V -q parallel.q -cwd -N tracking_5_${subID} ${scripts_path}/runTracking5.sh\n",

"Compute SC matrix":"echo \"************** Compute SC matrix  *******************\"\n\
qsub -hold_jid tracking_1_${subID},tracking_2_${subID},tracking_3_${subID},tracking_4_${subID},tracking_5_${subID}  -V -q parallel.q -cwd -N SC_${subID}_1 ${scripts_path}/computeSC1.sh\n\
qsub -hold_jid tracking_1_${subID},tracking_2_${subID},tracking_3_${subID},tracking_4_${subID},tracking_5_${subID}  -V -q parallel.q -cwd -N SC_${subID}_2 ${scripts_path}/computeSC2.sh\n\
qsub -hold_jid tracking_1_${subID},tracking_2_${subID},tracking_3_${subID},tracking_4_${subID},tracking_5_${subID}  -V -q parallel.q -cwd -N SC_${subID}_3 ${scripts_path}/computeSC3.sh\n\
qsub -hold_jid tracking_1_${subID},tracking_2_${subID},tracking_3_${subID},tracking_4_${subID},tracking_5_${subID}  -V -q parallel.q -cwd -N SC_${subID}_4 ${scripts_path}/computeSC4.sh\n\
qsub -hold_jid tracking_1_${subID},tracking_2_${subID},tracking_3_${subID},tracking_4_${subID},tracking_5_${subID}  -V -q parallel.q -cwd -N SC_${subID}_5 ${scripts_path}/computeSC5.sh\n\
qsub -hold_jid tracking_1_${subID},tracking_2_${subID},tracking_3_${subID},tracking_4_${subID},tracking_5_${subID}  -V -q parallel.q -cwd -N SC_${subID}_6 ${scripts_path}/computeSC6.sh\n",

"Aggregate SC matrix":"echo \"************** Aggregate SC matrix  *****************\"\n\
qsub -hold_jid SC_${subID}_1,SC_${subID}_2,SC_${subID}_3,SC_${subID}_4,SC_${subID}_5,SC_${subID}_6 -V -q bigmem_16.q -l h_vmem=16G -cwd -N aggSC_${subID} ${scripts_path}/aggregateSC.sh\n",

"Convert to TVB format and clean up results":"echo \"**** Convert to TVB format and clean up results  ****\"\n\
qsub -hold_jid aggSC_${subID},tracking_5_${subID} -V -q bigmem_16.q -l h_vmem=16G -cwd -N convert_${subID} ${scripts_path}/convert2TVB.sh\n"}

"""
files_prepare_scripts = ["#!/bin/bash", \
                             "cd /", \
                             "if [ ! -d "+entries["Working Directory"].get()+"]; then", \
                             "mkdir "+entries["Working Directory"].get(), \
                             "fi", \
                             "cd "+entries["Working Directory"].get() ]"""


"""
files_prepare_scripts = ["#!/bin/bash", \
                             "cd /", \
                             "if [ ! -d "+entries["Working Directory"].get()+"]; then", \
                             "mkdir "+entries["Working Directory"].get(), \
                             "fi", \
                             "cd "+entries["Working Directory"].get() ]"""

def fetch(entries):
    for key in entries:

        text = entries[key].get()
        print('%s: "%s"' % (key, entries[key]))


def makeform(root, fields):
    entries = {}
    for field in fields:
        row = tk.Frame(root)
        lab = tk.Label(row, text = field, anchor = 'w')
        ent = tk.Entry(row)
        row.pack(side = tk.TOP, fill = tk.X, padx = 5, pady = 5)
        lab.pack(side = tk.LEFT)
        ent.pack(side = tk.RIGHT, expand = tk.YES, fill = tk.X)
        entries[field] = ent
    return entries


def makeform_check(root, fields):
    checks = {}
    for field in fields_check:
        row = tk.Frame(root)
        lab = tk.Label(row, text = "", anchor = 'w')
        var = tk.IntVar()
        ent = tk.Checkbutton(row,text = field, variable = var)
        row.pack(side = tk.TOP, fill = tk.X, padx = 5, pady = 5)
        lab.pack(side = tk.LEFT)
        ent.pack(side = tk.LEFT)
        checks[field] = var
    return checks




def WDprocessing(entries):
    global WDprocessing_index
    ans_cf = input("copy scripts? y/n")
    if WDprocessing_index == 1:
        ans_wd = input("changing working directory? y/n")
        ans_sub = input("changing working subject? y/n")

    Path_Check_Correct = True
    # check if the inputs from users are empty
    if WDprocessing_index == 0 or (WDprocessing_index == 1 and ans_wd.lower() in ["y", "yes"]):

        for key in ["T1 MRI Data Path", "dWMRI Data Path", "dWMRI Data Path second","fMRI(Bold) Data Path", "Working Directory", "Template Path"]:
            if entries[key].get() != ' ' and not os.path.exists(entries[key].get()):
                if key == "Working Directory":
                    print("please input valid"+" "+ key)
                    Path_Check_Correct  = False

    if Path_Check_Correct and ans_cf in ["y", "yes"] and (WDprocessing_index == 0 or (WDprocessing_index == 1 and ans_wd.lower() in ["y", "yes"])):
        #current_path = os.getcwd()
        distutils.dir_util.copy_tree(curWD+"/TVB", entries["Working Directory"].get())

    if Path_Check_Correct and (WDprocessing_index == 0 or \
            (WDprocessing_index == 1 and (ans_sub.lower() in ["y", "yes"] or ans_wd.lower() in ["y", "yes"]))):

        subjects_path = entries["Working Directory"].get()+"/subjects"
        subjectid_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()
        rawdata_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA"
        dti_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI"
        dti_path1 = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI1"
        dti_path2 = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI2"
        T1_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/MPRAGE"
        Bold_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/BOLD-EPI"


        if not os.path.exists(subjects_path):
            os.makedirs(subjects_path)
        if not os.path.exists(subjectid_path):
            os.makedirs(subjectid_path)
        if not os.path.exists(rawdata_path):
            os.makedirs(rawdata_path)
        if not os.path.exists(dti_path) and len(entries["dWMRI Data Path"].get()) > 0:
            os.makedirs(dti_path)
        if not os.path.exists(dti_path1) and len(entries["dWMRI Data Path second"].get()) >  0:
            os.makedirs(dti_path1)
        if not os.path.exists(dti_path2) and len(entries["dWMRI Data Path second"].get()) >  0:
            os.makedirs(dti_path2)
        if not os.path.exists(T1_path):
            os.makedirs(T1_path)
        if not os.path.exists(Bold_path):
            os.makedirs(Bold_path)
        WDprocessing_index = 1

def Dataprocessing(entries):
    global Dataprocessing_index
    dti_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI"
    T1_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/MPRAGE"
    dti_path1 = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI1"
    dti_path2 = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI2"
    Bold_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/BOLD-EPI"
    if WDprocessing_index == 0:
        ans_wd_re = input("need run working directory process? y/n")

    if WDprocessing_index == 1 or (WDprocessing_index == 0 and ans_wd_re.lower() in ["n", "no"]):
        if len(entries["dWMRI Data Path"].get()) > 0 and len(entries["dWMRI Data Path second"].get()) == 0:
            distutils.dir_util.copy_tree(entries["dWMRI Data Path"].get(), dti_path)
        if len(entries["dWMRI Data Path second"].get()) >  0:
            distutils.dir_util.copy_tree(entries["dWMRI Data Path"].get(), dti_path1)
            distutils.dir_util.copy_tree(entries["dWMRI Data Path second"].get(), dti_path2)

        if len(entries["fMRI(Bold) Data Path"].get())>0:
            distutils.dir_util.copy_tree(entries["fMRI(Bold) Data Path"].get(), Bold_path)
        if len(entries["T1 MRI Data Path"].get()) > 0:
            distutils.dir_util.copy_tree(entries["T1 MRI Data Path"].get(), T1_path)
        Dataprocessing_index = 1




def EnvironSet(entries):
    global EnvironSet_index
    dti_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI"
    T1_path = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/MPRAGE"
    dti_path1 = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI1"
    dti_path2 = entries["Working Directory"].get()+"/subjects"+"/"+entries["Object ID"].get()+"/RAWDATA/DTI2"
    
    if WDprocessing_index*Dataprocessing_index == 0:
        ans_env_conf = input("are you sure you have done working directory and data process? y/n")
    variables_Check_Correct = True
    # check if the inputs from users are empty
 
    for key in ["Object ID", "T1 Data Type (nii or mgz)","dwi Data Type (nii or mgz)", \
           "fMRI Data Type (nii or mgz)", "Number of ROI", "Template Name", "Template in MNI Space?"]:
        if len(entries[key].get()) == 0:
            print("please input valid"+" "+ key)
            variables_Check_Correct  = False
        elif key in ["T1 Data Type (nii or mgz)","dwi Data Type (nii or mgz)","fMRI Data Type (nii or mgz)"] and \
               entries[key].get().lower() not in ["nii","mgz"]:
            print("please input valid"+" "+ key+" nii or mgz")
            variables_Check_Correct  = False
        elif key in ["Number of ROI"] and entries[key].get().lower() not in ["96","68"]:
            print("please input valid"+" "+ key+" 96 or 68")
            variables_Check_Correct  = False
        elif key in ["Template in MNI Space?"] and entries[key].get().lower() not in ["yes","no"]:
            print("please input valid"+" "+ key+" yes or no")
            variables_Check_Correct  = False


    if variables_Check_Correct and WDprocessing_index == 1 & Dataprocessing_index == 1 or \
            (variables_Check_Correct and WDprocessing_index*Dataprocessing_index == 0 and ans_env_conf.lower()in ["y", "yes"]):
        #os.environ["PATH"] += os.pathsep + os.pathsep.join(PathList)
        #os.environ["LD_LIBRARY_PATH"]+=os.pathsep+os.pathsep.join(["/opt/lib.exported"])
        for key in EnvironVariablesDict:
            os.environ[key] = EnvironVariablesDict[key]
        os.environ["subID"] = entries["Object ID"].get()
        os.environ["numROI"] = entries["Number of ROI"].get()
        os.environ["T1_dtype"] = entries["T1 Data Type (nii or mgz)"].get().lower()
        os.environ["dwi_dtype"] = entries["dwi Data Type (nii or mgz)"].get().lower()
        os.environ["fmri_dtype"] = entries["fMRI Data Type (nii or mgz)"].get().lower()
        os.environ["templateMNI"] = entries["Template in MNI Space?"].get().lower()
        os.environ["templatePath"] = entries["Template Path"].get()
        os.environ["templateFilename"] = entries["Template Name"].get()
 
        if not os.path.exists(dti_path1) and not os.path.exists(dti_path):
            os.environ["numDwi"] = '0'
        if os.path.exists(dti_path) and not os.path.exists(dti_path1):
            os.environ["numDwi"] = '1'
        
        if os.path.exists(dti_path1):
            os.environ["numDwi"] = '2'

        for source in SourceList:
            command = ['bash', '-c', 'source '+source]
            proc = subprocess.Popen(command, stdout=subprocess.PIPE)

            for line in proc.stdout:
                (key, _, value) = line.partition("=")
                os.environ[key]=value

            proc.communicate()

        pprint.pprint(dict(os.environ))

        EnvironSet_index = 1


def sh_file_generation(entries, checks):
    open(entries["Working Directory"].get()+"/TVB_pipe.sh", "w")
    f = open(entries["Working Directory"].get()+"/TVB_pipe.sh", "r+")
    os.chmod(entries["Working Directory"].get()+"/TVB_pipe.sh", 0o755)
    f.write(script_dict["title"])
    for field in fields_check:
        if checks[field].get() ==1:
            f.write(script_dict[field])
    f.close()
    os.chdir(entries["Working Directory"].get())
    #qsub_call = "qsub -V -q bigmem_16.q -l h_vmem=16G -cwd %s"
    #qsub_call = "./"
    #subprocess.call("unset SGE_ROOT", shell=True)
    subprocess.call("./TVB_pipe.sh", shell=True)


"""
        open("files_prepare_scripts", "a")
        
        f = open("files_prepare_scripts", "r+")
        for script in files_prepare_scripts:
            f.write(script+"\n")

        f.close()"""




# class Checkbar(tk.Frame):
# 	def __init__(self, parent=None, picks=[], side = tk.LEFT, anchor = tk.W):
# 		tk.Frame.__init__(self, parent)
# 		self.vars = []
# 		for pick in picks:
# 			var = tk.IntVar()
# 			self.row = tk.Frame(self)
# 			chk = tk.Checkbutton(self.row,text = pick, variable = var)
# 			chk.pack(side = side, anchor = anchor, expand = tk.YES)
# 			self.vars.append(var)

# 	def state(self):
# 		return map((lambda var: var.get()), self.vars)





def quit_process():
    os.chdir(curWD)
    top.quit



if __name__ == '__main__':
    top = tk.Tk()
    ents = makeform(top, fields)
    #top.bind('<Return>', (lambda event, e=ents: fetch(e)))
    row = tk.Frame(top)
    title_check = tk.Label(row, text = "Choose the processes:")
    row.pack(side = tk.TOP, fill = tk.X, padx = 5, pady = 5)
    title_check.pack(side = tk.LEFT)
    chks = makeform_check(top,fields_check)
    b1 = tk.Button(top,text = 'Quit', command = top.quit)
    b1.pack(side =tk.LEFT, padx = 5, pady = 5)
    b2 = tk.Button(top,text = 'Run Working Directory Process', command = (lambda e=ents: WDprocessing(e)))
    b2.pack(side =tk.LEFT, padx = 5, pady = 5)

    b4 = tk.Button(top,text = 'Run Data Copy Process', command = (lambda e=ents: Dataprocessing(e)))
    b4.pack(side =tk.LEFT, padx = 5, pady = 5)

    b3 = tk.Button(top,text = 'Environ Set', command = (lambda e=ents: EnvironSet(e)))
    b3.pack(side =tk.LEFT, padx = 5, pady = 5)

    b5 = tk.Button(top,text = 'Sh Script', command = (lambda e1=ents, e2=chks: sh_file_generation(e1,e2)))
    b5.pack(side =tk.LEFT, padx = 5, pady = 5)


    top.mainloop()


#     CheckVar1 = tk.IntVar()
# CheckVar2 = tk.IntVar()
# CheckVar3 = tk.IntVar()
# CheckVar4 = tk.IntVar()
# CheckVar5 = tk.IntVar()
# CheckVar6 = tk.IntVar()
# CheckVar7 = tk.IntVar()
# CheckVar8 = tk.IntVar()
# CheckVar9 = tk.IntVar()

# tk.Label(top, text = "Object ID").grid(row = 0)
# tk.Label(top, text = "T1 MRI Data Path").grid(row = 1)
# tk.Label(top, text = "dWMRI Data Path").grid(row = 2)
# tk.Label(top, text = "fMRI(Bolt) Data Path").grid(row = 3)

# tk.Label(top, text = "Working Directory").grid(row = 4)





# e1 = tk.Entry(top)
# e2 = tk.Entry(top)
# e3 = tk.Entry(top)
# e4 = tk.Entry(top)
# e5 = tk.Entry(top)

# e1.grid(row = 0, column = 1)
# e2.grid(row = 1, column = 1)
# e3.grid(row = 2, column = 1)
# e4.grid(row = 3, column = 1)
# e5.grid(row = 4, column = 1)

# tk.Label(top, text = "Choose the processes:").grid(row = 6, sticky = tk.W)

# k=7;

# ObjectId = e1.get()
# T1Path = e2.get()
# DtiPath = e3.get()
# BoltPath = e4.get（）
# WorkPath = e5.get ()


# tk.Checkbutton(top, text = "T1 Prerocessing with FreeSurfer", variable = CheckVar1).grid(row = 0+k, sticky = tk.W)

# tk.Checkbutton(top, text = "DWI Prerocessing with FreeSurfer", variable = CheckVar2).grid(row = 1+k, sticky = tk.W)

# tk.Checkbutton(top, text = "DWI Prerocessing with MrTrix2", variable = CheckVar3).grid(row = 2+k, sticky = tk.W)
# tk.Checkbutton(top, text = "fMRI Preprocessing ", variable = CheckVar4).grid(row = 3+k, sticky = tk.W)
# tk.Checkbutton(top, text = "Generate mask ", variable = CheckVar5).grid(row = 4+k, sticky = tk.W)
# tk.Checkbutton(top, text = "Tracking ", variable = CheckVar6).grid(row = 5+k, sticky = tk.W)
# tk.Checkbutton(top, text = "Compute SC matrix ", variable = CheckVar7).grid(row = 6+k, sticky = tk.W)
# tk.Checkbutton(top, text = "Aggregate SC matrix", variable = CheckVar8).grid(row = 7+k, sticky = tk.W)
# tk.Checkbutton(top, text = "Convert to TVB format and clean up results ", variable = CheckVar9).grid(row = 8+k, sticky = tk.W)






#top.mainloop()
