#!/usr/bin/env python
"""PAR/REC to NIfTI converter
"""
from __future__ import division, print_function, absolute_import

from optparse import OptionParser, Option
import numpy as np
import sys
import os
import gzip
import nibabel
import nibabel.parrec as pr
from nibabel.parrec import one_line
from nibabel.mriutils import calculate_dwell_time, MRIError
import nibabel.nifti1 as nifti1
from nibabel.filename_parser import splitext_addext
from nibabel.volumeutils import array_to_file, fname_ext_ul_case


def get_opt_parser():
    # use module docstring for help output
    p = OptionParser(
        usage="%s [OPTIONS] <PAR files>\n\n" % sys.argv[0] + __doc__,
        version="%prog " + nibabel.__version__)
    p.add_option(
        Option("-v", "--verbose", action="store_true", dest="verbose",
               default=False,
               help="""Make some noise."""))
    p.add_option(
        Option("-o", "--output-dir", action="store", type="string",
               dest="outdir", default=None,
               help=one_line("""Destination directory for NIfTI files.
                             Default: current directory.""")))
    p.add_option(
        Option("-c", "--compressed", action="store_true",
               dest="compressed", default=False,
               help="Whether to write compressed NIfTI files or not."))
    p.add_option(
        Option("-p", "--permit-truncated", action="store_true",
               dest="permit_truncated", default=False,
               help=one_line(
                   """Permit conversion of truncated recordings. Support for
                   this is experimental, and results *must* be checked
                   afterward for validity.""")))
    p.add_option(
        Option("-b", "--bvs", action="store_true", dest="bvs", default=False,
               help=one_line(
                   """Output bvals/bvecs files in addition to NIFTI
                   image.""")))
    p.add_option(
        Option("-d", "--dwell-time", action="store_true", default=False,
               dest="dwell_time",
               help=one_line(
                   """Calculate the scan dwell time. If supplied, the magnetic
                   field strength should also be supplied using
                   --field-strength (default 3). The field strength must be
                   supplied because it is not encoded in the PAR/REC
                   format.""")))
    p.add_option(
        Option("--field-strength", action="store", type="float",
               dest="field_strength",
               help=one_line(
                   """The magnetic field strength of the recording, only needed
                   for --dwell-time. The field strength must be supplied
                   because it is not encoded in the PAR/REC format.""")))
    p.add_option(
        Option("--origin", action="store", dest="origin", default="scanner",
               help=one_line(
                   """Reference point of the q-form transformation of the NIfTI
                   image. If 'scanner' the (0,0,0) coordinates will refer to
                   the scanner's iso center. If 'fov', this coordinate will be
                   the center of the recorded volume (field of view). Default:
                   'scanner'.""")))
    p.add_option(
        Option("--minmax", action="store", nargs=2, dest="minmax",
               help=one_line(
                   """Mininum and maximum settings to be stored in the NIfTI
                   header. If any of them is set to 'parse', the scaled data is
                   scanned for the actual minimum and maximum.  To bypass this
                   potentially slow and memory intensive step (the data has to
                   be scaled and fully loaded into memory), fixed values can be
                   provided as space-separated pair, e.g. '5.4 120.4'. It is
                   possible to set a fixed minimum as scan for the actual
                   maximum (and vice versa). Default: 'parse parse'.""")))
    p.set_defaults(minmax=('parse', 'parse'))
    p.add_option(
        Option("--store-header", action="store_true", dest="store_header",
               default=False,
               help=one_line(
                   """If set, all information from the PAR header is stored in
                   an extension ofthe NIfTI file header.  Default: off""")))
    p.add_option(
        Option("--scaling", action="store", dest="scaling", default='dv',
               help=one_line(
                   """Choose data scaling setting. The PAR header defines two
                   different data scaling settings: 'dv' (values displayed on
                   console) and 'fp' (floating point values). Either one can be
                   chosen, or scaling can be disabled completely ('off').  Note
                   that neither method will actually scale the data, but just
                   store the corresponding settings in the NIfTI header, unless
                   non-uniform scaling is used, in which case the data is
                   stored in the file in scaled form. Default: 'dv'""")))
    p.add_option(
        Option("--overwrite", action="store_true", dest="overwrite",
               default=False,
               help=one_line("""Overwrite file if it exists. Default:
                             False""")))
    return p


def verbose(msg, indent=0):
    if verbose.switch:
        print("%s%s" % (' ' * indent, msg))


def error(msg, exit_code):
    sys.stderr.write(msg + '\n')
    sys.exit(exit_code)


def proc_file(infile, opts):
    # figure out the output filename, and see if it exists
    basefilename = splitext_addext(os.path.basename(infile))[0]
    if opts.outdir is not None:
        # set output path
        basefilename = os.path.join(opts.outdir, basefilename)

    # prep a file
    if opts.compressed:
        verbose('Using gzip compression')
        outfilename = basefilename + '.nii.gz'
    else:
        outfilename = basefilename + '.nii'
    if os.path.isfile(outfilename) and not opts.overwrite:
        raise IOError('Output file "%s" exists, use --overwrite to '
                      'overwrite it' % outfilename)

    # load the PAR header and data
    scaling = None if opts.scaling == 'off' else opts.scaling
    infile = fname_ext_ul_case(infile)
    pr_img = pr.load(infile,
                     permit_truncated=opts.permit_truncated,
                     scaling=scaling)
    pr_hdr = pr_img.header
    raw_data = pr_img.dataobj.get_unscaled()
    affine = pr_hdr.get_affine(origin=opts.origin)
    nimg = nifti1.Nifti1Image(raw_data, affine, pr_hdr)
    nhdr = nimg.header

    if 'parse' in opts.minmax:
        # need to get the scaled data
        verbose('Loading (and scaling) the data to determine value range')
        if opts.scaling == 'off':
            scaled_data = raw_data
        else:
            slope, intercept = pr_hdr.get_data_scaling(method=opts.scaling)
            scaled_data = slope * raw_data + intercept
    if opts.minmax[0] == 'parse':
        nhdr.structarr['cal_min'] = scaled_data.min()
    else:
        nhdr.structarr['cal_min'] = float(opts.minmax[0])
    if opts.minmax[1] == 'parse':
        nhdr.structarr['cal_max'] = scaled_data.max()
    else:
        nhdr.structarr['cal_max'] = float(opts.minmax[1])

    # container for potential NIfTI1 header extensions
    if opts.store_header:
        exts = nifti1.Nifti1Extensions()
        # dump the full PAR header content into an extension
        with open(infile, 'r') as fobj:
            hdr_dump = fobj.read()
            dump_ext = nifti1.Nifti1Extension('comment', hdr_dump)
        exts.append(dump_ext)
    # finalize the header: set proper data offset, pixdims, ...
    nimg.update_header()

    # prep a file
    if opts.compressed:
        verbose('Using gzip compression')
        outfile = gzip.open(outfilename, 'wb')
    else:
        outfile = open(outfilename, 'wb')

    # get original scaling, and decide if we scale in-place or not
    if opts.scaling == 'off':
        slope = np.array([1.])
        intercept = np.array([0.])
    else:
        verbose('Using data scaling "%s"' % opts.scaling)
        slope, intercept = pr_hdr.get_data_scaling(method=opts.scaling)
    verbose('Writing %s' % outfilename)
    if not np.any(np.diff(slope)) and not np.any(np.diff(intercept)):
        # Single scalefactor case
        nhdr.set_slope_inter(slope.ravel()[0], intercept.ravel()[0])
        data_obj = raw_data
    else:
        # Multi scalefactor case
        nhdr.set_slope_inter(1, 0)
        nhdr.set_data_dtype(np.float64)
        data_obj = pr_img.dataobj
    nhdr.write_to(outfile)
    offset = nhdr.get_data_offset()
    array_to_file(data_obj,
                  outfile,
                  out_dtype=nhdr.get_data_dtype(), # for endianness
                  offset=offset)
    # write out bvals/bvecs if requested
    if opts.bvs:
        bvals, bvecs = pr_hdr.get_bvals_bvecs()
        if (bvals, bvecs) == (None, None):
            verbose('No DTI volumes detected, bvals and bvecs not written')
        else:
            verbose('Writing .bvals and .bvecs files')
            with open(basefilename + '.bvals', 'w') as fid:
                # np.savetxt could do this, but it's just a loop anyway
                for val in bvals:
                    fid.write('%s ' % val)
                fid.write('\n')
            with open(basefilename + '.bvecs', 'w') as fid:
                for row in bvecs.T:
                    for val in row:
                        fid.write('%s ' % val)
                    fid.write('\n')

    # write out dwell time if requested
    if opts.dwell_time:
        try:
            dwell_time = calculate_dwell_time(
                pr_hdr.get_water_fat_shift(),
                pr_hdr.get_echo_train_length(),
                opts.field_strength)
        except MRIError:
            verbose('No EPI factors, dwell time not written')
        else:
            verbose('Writing dwell time (%r sec) calculated assuming %sT '
                    'magnet' % (dwell_time, opts.field_strength))
            with open(basefilename + '.dwell_time', 'w') as fid:
                fid.write('%r\n' % dwell_time)
    # done


def main():
    parser = get_opt_parser()
    (opts, infiles) = parser.parse_args()

    verbose.switch = opts.verbose

    if opts.origin not in ['scanner', 'fov']:
        error("Unrecognized value for --origin: '%s'." % opts.origin, 1)
    if opts.dwell_time and opts.field_strength is None:
        error('Need --field-strength for dwell time calculation', 1)

    # store any exceptions
    errs = []
    for infile in infiles:
        verbose('Processing %s' % infile)
        try:
            proc_file(infile, opts)
        except Exception as e:
            errs.append('%s: %s' % (infile, e))

    if len(errs):
        error('Caught %i exceptions. Dump follows:\n\n %s'
              % (len(errs), '\n'.join(errs)), 1)
    else:
        verbose('Done')


if __name__ == '__main__':
    main()
