#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author: oesteban
# @Date:   2016-03-16 11:28:27
# @Last Modified by:   oesteban
# @Last Modified time: 2018-03-12 11:48:39

"""
Batch export freesurfer results to animated gifs

"""

import os
import os.path as op
import subprocess as sp
from shutil import rmtree
from argparse import ArgumentParser
from argparse import RawTextHelpFormatter
from tempfile import mkdtemp
from errno import EEXIST
import numpy as np
import nibabel as nb
from skimage import exposure


def main():
    """Entry point"""
    parser = ArgumentParser(description='Batch export freesurfer results to animated gifs',
                            formatter_class=RawTextHelpFormatter)
    g_input = parser.add_argument_group('Inputs')
    g_input.add_argument('-s', '--subject-id', action='store')
    g_input.add_argument('-t', '--temp-dir', action='store')
    g_input.add_argument('--keep-temp', action='store_true', default=False)
    g_input.add_argument('--zoom', action='store_true', default=False)
    g_input.add_argument('--hist-eq', action='store_true', default=False)
    g_input.add_argument('--use-xvfb', action='store_true', default=False)

    g_outputs = parser.add_argument_group('Outputs')
    g_outputs.add_argument('-o', '--output-dir', action='store', default='fs2gif')

    opts = parser.parse_args()

    if opts.temp_dir is None:
        tmpdir = mkdtemp()
    else:
        tmpdir = op.abspath(opts.temp_dir)
        try:
            os.makedirs(tmpdir)
        except OSError as exc:
            if exc.errno != EEXIST:
                raise exc

    out_dir = op.abspath(opts.output_dir)
    try:
        os.makedirs(out_dir)
    except OSError as exc:
        if exc.errno != EEXIST:
            raise exc

    subjects_dir = os.getenv('SUBJECTS_DIR', op.abspath('subjects'))
    subject_list = [opts.subject_id]
    if opts.subject_id is None:
        subject_list = [op.basename(name) for name in os.listdir(subjects_dir)
                        if op.isdir(os.path.join(subjects_dir, name))]
    environ = os.environ.copy()
    environ['SUBJECTS_DIR'] = subjects_dir
    if opts.use_xvfb:
        environ['doublebufferflag'] = 1

    # tcl_file = pkgr.resource_filename('mriqc', 'data/fsexport.tcl')
    tcl_contents = """
SetOrientation 0
SetCursor 0 128 128 128
SetDisplayFlag 3 0
SetDisplayFlag 22 1
set i 0
"""

    for subid in subject_list:
        sub_path = op.join(subjects_dir, subid)
        tmp_sub = op.join(tmpdir, subid)
        try:
            os.makedirs(tmp_sub)
        except OSError as exc:
            if exc.errno != EEXIST:
                raise exc

        niifile = op.join(tmp_sub, '%s.nii.gz') % subid
        ref_file = op.join(sub_path, 'mri', 'T1.mgz')
        sp.call(['mri_convert', op.join(sub_path, 'mri', 'norm.mgz'), niifile],
                cwd=tmp_sub)
        data = nb.load(niifile).get_data()
        data[data > 0] = 1

        # Compute brain bounding box
        indexes = np.argwhere(data)
        bbox_min = indexes.min(0)
        bbox_max = indexes.max(0) + 1
        center = np.average([bbox_min, bbox_max], axis=0)

        if opts.hist_eq:
            modnii = op.join(tmp_sub, '%s.nii.gz' % subid)
            ref_file = op.join(tmp_sub, '%s.mgz' % subid)
            img = nb.load(niifile)
            data = exposure.equalize_adapthist(img.get_data(), clip_limit=0.03)
            nb.Nifti1Image(data, img.affine, img.header).to_filename(modnii)
            sp.call(['mri_convert', modnii, ref_file], cwd=tmp_sub)

        if not opts.zoom:
            # Export tiffs for left hemisphere
            tcl_file = op.join(tmp_sub, '%s.tcl' % subid)
            with open(tcl_file, 'w') as tclfp:
                tclfp.write(tcl_contents)
                tclfp.write('for { set slice %d } { $slice < %d } { incr slice } {' % (
                    bbox_min[2], bbox_max[2]))
                tclfp.write('    SetSlice $slice\n')
                tclfp.write('    RedrawScreen\n')
                tclfp.write('    SaveTIFF [format "%s/%s-' % (tmp_sub, subid) + '%03d.tif" $i]\n')
                tclfp.write('    incr i\n')
                tclfp.write('}\n')
                tclfp.write('QuitMedit\n')
            cmd = ['tkmedit', subid, 'T1.mgz', 'lh.pial', '-aux-surface', 'rh.pial', '-tcl',
                   tcl_file]
            if opts.use_xvfb:
                cmd = _xvfb_run() + cmd

            print('Running tkmedit: %s' % ' '.join(cmd))
            sp.call(cmd, env=environ)
            # Convert to animated gif
            print('Stacking coronal slices')
            sp.call(['convert', '-delay', '10', '-loop', '0', '%s/%s-*.tif' % (tmp_sub, subid),
                     '%s/%s.gif' % (out_dir, subid)])

        else:
            # Export tiffs for left hemisphere
            tcl_file = op.join(tmp_sub, 'lh-%s.tcl' % subid)
            with open(tcl_file, 'w') as tclfp:
                tclfp.write(tcl_contents)
                tclfp.write('SetZoomLevel 2')
                tclfp.write('for { set slice %d } { $slice < %d } { incr slice } {' % (
                    bbox_min[2], bbox_max[2]))
                tclfp.write('    SetZoomCenter %d %d $slice\n' % (center[0] + 30, center[1] - 10))
                tclfp.write('    SetSlice $slice\n')
                tclfp.write('    RedrawScreen\n')
                tclfp.write('    SaveTIFF [format "{}/{}-lh-%03d.tif" $i]\n'.format(
                    tmp_sub, subid))
                tclfp.write('    incr i\n')
                tclfp.write('}\n')
                tclfp.write('QuitMedit\n')
            cmd = ['tkmedit', subid, 'norm.mgz', 'lh.white', '-tcl', tcl_file]
            if opts.use_xvfb:
                cmd = _xvfb_run() + cmd

            print('Running tkmedit: %s' % ' '.join(cmd))
            sp.call(cmd, env=environ)
            # Convert to animated gif
            print('Stacking coronal slices')

            # Export tiffs for right hemisphere
            tcl_file = op.join(tmp_sub, 'rh-%s.tcl' % subid)
            with open(tcl_file, 'w') as tclfp:
                tclfp.write(tcl_contents)
                tclfp.write('SetZoomLevel 2')
                tclfp.write('for { set slice %d } { $slice < %d } { incr slice } {' % (
                    bbox_min[2], bbox_max[2]))
                tclfp.write('    SetZoomCenter %d %d $slice\n' % (center[0] - 30, center[1] - 10))
                tclfp.write('    SetSlice $slice\n')
                tclfp.write('    RedrawScreen\n')
                tclfp.write('    SaveTIFF [format "{}/{}-rh-%03d.tif" $slice]\n'.format(
                    tmp_sub, subid))
                tclfp.write('    incr i\n')
                tclfp.write('}\n')
                tclfp.write('QuitMedit\n')
            cmd = ['tkmedit', subid, 'norm.mgz', 'rh.white', '-tcl', tcl_file]
            if opts.use_xvfb:
                cmd = _xvfb_run() + cmd

            print('Running tkmedit: %s' % ' '.join(cmd))
            sp.call(cmd, env=environ)
            # Convert to animated gif
            print('Stacking coronal slices')
            sp.call(['convert', '-delay', '10', '-loop', '0', '%s/%s-lh-*.tif' % (tmp_sub, subid),
                     '%s/%s-lh.gif' % (out_dir, subid)])
            sp.call(['convert', '-delay', '10', '-loop', '0', '%s/%s-rh-*.tif' % (tmp_sub, subid),
                     '%s/%s-rh.gif' % (out_dir, subid)])

        if not opts.keep_temp:
            rmtree(tmp_sub, ignore_errors=True, onerror=_myerror)


def _xvfb_run(wait=5, server_args='-screen 0, 1600x1200x24', logs=None):
    """
    Wrap command with xvfb-run. Copied from:
    https://github.com/VUIIS/seam/blob/1dabd9ca5b1fc7d66ef7d41c34ea8d42d668a484/seam/util.py

    """
    if logs is None:
        logs = op.join(mkdtemp(), 'fs2gif_xvfb')

    return ['xvfb-run',
            '-a',  # automatically get a free server number
            '-f {}.out'.format(logs),
            '-e {}.err'.format(logs),
            '--wait={:d}'.format(wait),
            '--server-args="{}"'.format(server_args)]


def _myerror(msg):
    print('WARNING: Error deleting temporal files: %s' % msg)


if __name__ == '__main__':
    main()
