#!/usr/bin/python3

import os
import json
import matplotlib.pyplot as plt
import nibabel as nib
from make_binary_masks_from_fsdir import make_binary_masks_from_fsdir

# Print disk space
os.system('df -h')

################### Define some plotting functions ###########################

def make_plot(base_image, overlay, title, filename, x, y, z, apect_ratio_vector, output_folder):
    
    # This function simply gets two MRI images as inputs and overlays them 
    # using different colors for each image. Used as a diagnostic image.
        
    fig, (ax1, ax2, ax3) = plt.subplots(1,3)
    fig.suptitle(title, fontsize=20)

    epi_img = nib.load(base_image)
    epi_img_data = epi_img.get_fdata()
    ax1.imshow(epi_img_data[x,:,:], cmap="gray", aspect = apect_ratio_vector[0])
    ax2.imshow(epi_img_data[:,y,:], cmap="gray", aspect = apect_ratio_vector[1])
    ax3.imshow(epi_img_data[:,:,z], cmap="gray", aspect = apect_ratio_vector[2])
    ax1.axis('off')
    ax2.axis('off')
    ax3.axis('off')  
    
    if overlay != 'NA':
        epi_img = nib.load(overlay)
        epi_img_data = epi_img.get_fdata()
        ax1.imshow(epi_img_data[x,:,:], cmap="hot", alpha=0.4, aspect = apect_ratio_vector[3])
        ax2.imshow(epi_img_data[:,y,:], cmap="hot", alpha=0.4, aspect = apect_ratio_vector[4])
        ax3.imshow(epi_img_data[:,:,z], cmap="hot", alpha=0.4, aspect = apect_ratio_vector[5])
        ax1.axis('off')
        ax2.axis('off')
        ax3.axis('off')

    plt.savefig(os.path.join(output_folder, filename))    

###################### Define the Paths and parameters ########################

# Set paths and create some folders
flywheel_base ='/flywheel/v0/'
first_output_dir = '/firstOutput'
os.system('mkdir %s'%first_output_dir)
final_output_dir = os.path.join(flywheel_base, 'output')
manifest_file = os.path.join(flywheel_base, 'manifest.json')
config_path = os.path.join(flywheel_base, 'config.json')
fsl_environment_path = 'FSLDIR=/usr/lib/fsl/5.0;. /etc/fsl/5.0/fsl.sh;PATH=${FSLDIR}:${PATH};export FSLDIR PATH;/usr/lib/fsl/5.0/'
freesurfer_environment_path = 'export FREESURFER_HOME=/freesurfer;/bin/bash -c \'source $FREESURFER_HOME/FreeSurferEnv.sh &>/dev/null\';export SUBJECTS_DIR=%s;/freesurfer/bin/'
mcr_path = '/usr/local/MATLAB/MATLAB_Runtime/v97' 
compiled_white_matter_function = '/opt/compiled_localWM/run_remove_localWM_FwVersion.sh'
compiled_psd_function = '/opt/compiled_psdFunc/run_calcOneSidedPSD_FwVersion.sh'
fsl_identity = '/usr/share/fsl/5.0/etc/flirtsch/ident.mat'

# Unzip the fmriprep zip
fmriprepOutputArchive_folder = os.path.join(flywheel_base, 'input/fmriprepOutputArchive')
fmriprepOutputArchive = os.path.join(fmriprepOutputArchive_folder, os.listdir(fmriprepOutputArchive_folder)[0]) 
unzipped_folder = '/prep'
os.system('mkdir %s' % unzipped_folder)
os.system('unzip -q %s -d %s' % (fmriprepOutputArchive, unzipped_folder))
mainprep_folder = os.path.join(unzipped_folder, os.listdir(unzipped_folder)[0], 'fmriprep')
if ~os.path.exists(mainprep_folder):
    mainprep_folder = os.path.join(unzipped_folder, os.listdir(unzipped_folder)[0])

images_in_T1 = []
for i in os.listdir(mainprep_folder):
    if 'sub' in i and 'html' not in i:
        prep_subject = os.path.join(mainprep_folder, i)
        for ii in os.listdir(prep_subject):
            if 'ses' in ii:
                prep_session = os.path.join(prep_subject, ii)
                prep_func_dir = os.path.join(prep_session, 'func')

                for image in os.listdir(prep_func_dir):
                    if 'space-T1w_desc-preproc_bold.nii.gz' in image:
                        found_image = os.path.join(prep_func_dir, image)
                        images_in_T1.append(found_image)


# Unzip the recon-all archive
reconAllArchive_folder = os.path.join(flywheel_base, 'input/reconAllArchive')
reconAllArchive = os.path.join(reconAllArchive_folder, os.listdir(reconAllArchive_folder)[0]) 
unzipped_folder = '/reconall'
os.system('mkdir %s' % unzipped_folder)
os.system('unzip -q %s -d %s' % (reconAllArchive, unzipped_folder))
mainrecon_folder = os.path.join(unzipped_folder, os.listdir(unzipped_folder)[0])

###################Parse Config Here##########################################

with open(config_path) as config_file:
    job_config = json.load(config_file)
config = job_config['config']

radius = config['radius']
TR = config['TR']
saveIntermediate = config['saveIntermediate']
NvoxToErodeWM = config['NvoxToErodeWM']

# ############################# Process images ##################################

# Run the mask maker script
intermediate_files = '/intermediate_files'
os.system('mkdir %s' % intermediate_files)
template_bold = images_in_T1[0]
resampled_T1, brain_mask, white_matter_mask, gray_matter_mask = make_binary_masks_from_fsdir(mainrecon_folder, template_bold, fsl_identity, freesurfer_environment_path, fsl_environment_path, NvoxToErodeWM, intermediate_files)

# Save intermediate files if requested
if saveIntermediate:
    os.system('cd %s; zip -q -r %s *' % (intermediate_files, os.path.join(final_output_dir,'intermediate_files.zip')))

# Plot diagnostics
main_diagnostics_folder = '/main_diagnostics'
os.system('mkdir %s' % main_diagnostics_folder)
images_folder = os.path.join(main_diagnostics_folder, 'images')
os.system('mkdir %s' % images_folder)

raw_loaded = nib.load(brain_mask)
n_i, n_j, n_k = raw_loaded.shape
center_i = (n_i - 1) // 2  # // for integer division
center_j = (n_j - 1) // 2
center_k = (n_k - 1) // 2

make_plot(resampled_T1, brain_mask, 'Brain mask', 'T1.png', center_i, center_j, center_k, [0.7,1,1.3,0.7,1,1.4], images_folder)
make_plot(resampled_T1, white_matter_mask, 'White matter mask', 'whiteMask.png', center_i, center_j, center_k, [0.7,1,1.3,0.7,1,1.4], images_folder)
make_plot(resampled_T1, gray_matter_mask, 'Gray matter mask', 'grayMask.png', center_i, center_j, center_k, [0.7,1,1.3,0.7,1,1.4], images_folder)

# Initiate the html content string
html_content = '''
<h1>Resampled T1 overlaid with brain mask</h1>
<img src="./%s" style="float: left; width: 50%%; margin-right: 1%%; margin-bottom: 0.5em;" alt="Left_lateral">
<p style="clear: both;">
<h1>Resampled T1 overlaid with wm mask</h1>
<img src="./%s" style="float: left; width: 50%%; margin-right: 1%%; margin-bottom: 0.5em;" alt="Left_lateral">
<p style="clear: both;">
<h1>Resampled T1 overlaid with gm mask</h1>
<img src="./%s" style="float: left; width: 50%%; margin-right: 1%%; margin-bottom: 0.5em;" alt="Left_lateral">
<p style="clear: both;">
''' % ('images/T1.png', 'images/whiteMask.png', 'images/grayMask.png')
        
# Loop through images and do the white matter regression
for image_to_process in images_in_T1:
    image_name = os.path.split(image_to_process)[1][:-7]
    specific_file = os.path.join('/tmp', image_name)
    os.system('mkdir %s' % specific_file)
    cmd = '%s %s %s %s %s %s %s' % (compiled_white_matter_function, mcr_path,
                                    image_to_process, brain_mask, white_matter_mask,
                                    first_output_dir, radius)
    os.system(cmd)
    processed_image = os.path.join(first_output_dir, 'noWm_%s.nii.gz' % image_name)
    cmd2 = '%s %s %s %s %s %s %s' % (compiled_psd_function, mcr_path,
                                    image_to_process, processed_image, gray_matter_mask, TR,
                                    specific_file)
    os.system(cmd2)
    final_diagnostic_image = os.path.join(specific_file, os.listdir(specific_file)[0])
    os.system('cp %s %s/%s.png' % (final_diagnostic_image, images_folder, image_name))
    
    image_string = '''
    <h1>Before and after processing frequency-power graph for %s</h1>
    <img src="./%s" style="float: left; width: 50%%; margin-right: 1%%; margin-bottom: 0.5em;" alt="Left_lateral">
    <p style="clear: both;">
    ''' % (image_name, 'images/%s.png' % image_name)
    html_content = html_content + image_string
    
# MAKE HTML HERE
html_file = open('%s/index.html' % main_diagnostics_folder,'w')
html_file.write(html_content)
html_file.close()

os.system('cd %s; zip -q -r %s *' % (main_diagnostics_folder, os.path.join(final_output_dir,'diagnostic_images.html.zip')))
os.system('cd %s; zip -q -r %s *' % (first_output_dir, os.path.join(final_output_dir,'results.zip')))
