# -*- coding: utf-8 -*-
"""
Created on Tue Mar  3 10:28:59 2020

@author: lytlemn
"""

import os
import shutil
import json
import tarfile


def dicom_proc(bidspath, info, shifted_date): 
    

    if ".nii" in info.dicom:
        
        tmp_path = bidspath + "tmp/" + str(info.dcm_fname)
    
        #create temporary folder
        os.makedirs(tmp_path)
        #copy the dicom to the temporary folder
        shutil.copy(info.dicom,tmp_path)
        
        orig_file = tmp_path + "/" + info.dcm_fname
    
    else:
        tmp_path = bidspath + "tmp/" + info.dcm_fname
    
        #create temporary folder
        os.makedirs(tmp_path)
        #copy the dicom to the temporary folder
        shutil.copy(info.dicom,tmp_path)
        
        #untar/unzip the dicom
        tar_file = tmp_path + "/" + info.dcm_fname + ".tar.gz"
        tar = tarfile.open(tar_file)
        tar.extractall(path=tmp_path)
        tar.close()
        #execute dcm2niix conversion command
        convert_command = "dcm2niix -b y -o " + tmp_path + " -f tmp_output " + tmp_path + "/"
        os.system(convert_command)
    
        if "fieldmap" in info.task: 
            orig_file = tmp_path + "/tmp_output_e2_ph"
        else:
            orig_file = tmp_path + "/tmp_output"
        #add shifted date to json
        json_file = open((orig_file + ".json"), "r")
        json_data = json.load(json_file)
        shift_data = {"ShiftedAquisitionDate":shifted_date}
        json_data.update(shift_data)
        json_file.close()
    
        json_file = open((orig_file + ".json"), "w")
        json.dump(json_data,json_file,indent = 2)
        json_file.close()
    

    
    #rename nii and json file and move to proper folder
    day = str(info.day_num)
    ser = str(info.series)
    sub = str(info.subject)
    ses = str(info.session)

    if "MPRAGE" in info.task:
        out_folder = bidspath + "sub-" + sub + "/" + ses + "/anat/"
        source_folder = bidspath + "sourcedata/sub-" + sub + "/" + ses + "/anat/"
        nii_name = "sub-" + sub + "_" + ses + "_acq-D" + day + "S" + ser
        filename = nii_name + "_T1w"
    elif "SBRef" in info.task:
        out_folder = bidspath + "sub-" + sub + "/" + ses + "/dwi/"
        source_folder = bidspath + "sourcedata/sub-" + sub + "/" + ses + "/dwi/"
        nii_name = "sub-" + sub + "_" + ses + "_acq-D" + day + "S" + ser
        filename = nii_name + "_sbref"
    elif "DTI" in info.task:
        out_folder = bidspath + "sub-" + sub + "/" + ses + "/dwi/"
        source_folder = bidspath + "sourcedata/sub-" + sub + "/" + ses + "/dwi/"
        nii_name = "sub-" + sub + "_" + ses + "_acq-D" + day + "S" + ser
        filename = nii_name + "_dwi"
    elif "fieldmap" in info.task:
        out_folder = bidspath + "sub-" + sub + "/" + ses + "/fmap/"
        source_folder = bidspath + "sourcedata/sub-" + sub + "/" + ses + "/fmap/"
        nii_name = "sub-" + sub + "_" + ses + "_acq-D" + day + "S" + ser
        filename = nii_name + "_phasediff"
        #add shifted date to json
        json_file = open((orig_file + ".json"), "r")
        json_data = json.load(json_file)
        intend = ses + "/dwi/" + info.intend
        update_data = {"EchoTime1":0.00492,"EchoTime2":0.00738,"IntendedFor":intend}
        json_data.update(update_data)
        json_file.close()
    
        json_file = open((orig_file + ".json"), "w")
        json.dump(json_data,json_file,indent = 2)
        json_file.close()
    
    else:
    #create bids folders if they do not yet exist
        out_folder = bidspath + "sub-" + sub + "/" + ses + "/func/"
        source_folder = bidspath + "sourcedata/sub-" + sub + "/" + ses + "/func/"
        task = info.task[:-1]
        run_num = info.task[-1]
        run = '0' + run_num
        nii_name = "sub-" + sub + "_" + ses + "_task-" + task + "_acq-D" + day + "S" + ser + "_run-" + run
        filename = nii_name + "_bold"
        
    if not os.path.isdir(out_folder):
        os.makedirs(out_folder)
    if not os.path.isdir(source_folder):
        os.makedirs(source_folder)       
    
    os.rename((orig_file + ".nii"),(out_folder + filename + ".nii"))
    
    if os.path.isfile((orig_file + ".json")):
        os.rename((orig_file + ".json"),(out_folder + filename + ".json"))
    
    if os.path.isfile((orig_file + ".bvec")):
        os.rename((orig_file + ".bvec"),(out_folder + filename + ".bvec"))
    
    if os.path.isfile((orig_file + ".bval")):
        os.rename((orig_file + ".bval"),(out_folder + filename + ".bval")) 

    #rename dicom tar
    #move dicom tar to sourcedata 
    if os.path.isfile((tmp_path + "/" + info.dcm_fname + ".tar.gz")):
        os.rename((tmp_path + "/" + info.dcm_fname + ".tar.gz"),(source_folder + filename + ".tar.gz"))
        dcm_files = os.listdir(tmp_path)
        for file in dcm_files:
            if file.endswith(".dcm"):
                os.remove(os.path.join(tmp_path, file))
    
    os.removedirs(tmp_path)
    #delete dicom files and temp folder

    
    return nii_name, source_folder, out_folder