#!/usr/bin/env python3

# Copyright (c) 2008-2021 the MRtrix3 contributors.
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#
# Covered Software is provided under this License on an "as is"
# basis, without warranty of any kind, either expressed, implied, or
# statutory, including, without limitation, warranties that the
# Covered Software is free of defects, merchantable, fit for a
# particular purpose or non-infringing.
# See the Mozilla Public License v. 2.0 for more details.
#
# For more details, see http://www.mrtrix.org/.

VENDORS = ['ge', 'siemens']

def usage(cmdline): #pylint: disable=unused-variable
  cmdline.set_author('Ben Jeurissen (ben.jeurissen@uantwerpen.be)')
  cmdline.set_synopsis('Create a vendor-specific gradient scheme')
  cmdline.add_description('This script creates a vendor-specific text file containing diffusion gradient information. Supported vendors are: ' + str(VENDORS))
  cmdline.add_example_usage('To convert the gradient scheme generated by gen_scheme to a Siemens .dvs file','vendorscheme dw_scheme.txt dw_scheme.dvs -vendor siemens')
  cmdline.add_argument('input', help='The input gradient scheme')
  cmdline.add_argument('output', help='The output vendor-specific textfile')
  options = cmdline.add_argument_group('Options for the vendorscheme command')
  options.add_argument('-vendor', help='Specify the vendor for which to output the gradient scheme. Options are: ' + str(VENDORS), default='none')
  options.add_argument('-recipient', help='Specify who/what this file was created for (e.g. specific person, site or project)')

def execute(): #pylint: disable=unused-variable
  from mrtrix3 import MRtrixError #pylint: disable=no-name-in-module, import-outside-toplevel
  from mrtrix3 import app  #pylint: disable=no-name-in-module, import-outside-toplevel
  from numpy import loadtxt, sqrt, reshape, around, unique, sum
  if not app.ARGS.vendor in VENDORS:
    raise MRtrixError('Vendor name must be one of ' + str(VENDORS) + '. Provided: ' + app.ARGS.vendor)
  try:
    input = loadtxt(app.ARGS.input)
  except:
    raise MRtrixError('Could not load gradient scheme from ' + str(app.ARGS.input))
  app.check_output_path(app.ARGS.output)
  output = open(app.ARGS.output, "w")

  norm = sum(input[:,0:3]**2,axis=-1)**(1./2)
  dirs = input[:,0:3]/reshape(norm,(input.shape[0],1))
  dirs = dirs*reshape(sqrt(input[:,3]/input[:,3].max()),(dirs.shape[0],1))
  bvals = unique(around(input[:,3]),return_counts=True)
  numbers = bvals[1]
  bvals = bvals[0]
  output.write('# This is a set of diffusion gradient directions generated by MRtrix\' vendorscheme (www.mrtrix.org)\n')
  if app.ARGS.recipient:
    output.write('# for ' + app.ARGS.recipient + '\n')
  output.write('# It consists of:\n')
  for i in range(0,bvals.shape[0]):
    output.write('# - %i b = %i s/mm^2 images\n' % (numbers[i],bvals[i]))
  if app.ARGS.vendor == 'ge':
    output.write('# To use it, select %i directions and set the b-value to %i s/mm^2\n' % (dirs.shape[0],bvals[-1]))
    output.write('%i\n' % dirs.shape[0])
    for i in range(0,dirs.shape[0]):
      output.write('%0.6f %0.6f %0.6f\n' % (dirs[i,0],dirs[i,1],dirs[i,2]))
  elif app.ARGS.vendor == 'siemens':
    output.write('# To use it, choose Free mode, select %i directions, and set the b-value to %i s/mm^2\n' % (dirs.shape[0],bvals[-1]))
    output.write('[directions=%i]\n' % dirs.shape[0])
    output.write('CoordinateSystem = xyz\n')
    output.write('Normalisation = none\n')
    for i in range(0,dirs.shape[0]):
      output.write('Vector[%i] = (%0.6f,%0.6f,%0.6f)\n' % (i,dirs[i,0],dirs[i,1],dirs[i,2]))
  output.close()

# Execute the script
import mrtrix3
mrtrix3.execute() #pylint: disable=no-member