#!/usr/bin/python
'''
Script to convert a field map from Toshiba to one that can be used by 
SimMARS and SimG4 (i.e. G4beamline BLFieldMap format).
By default the transformations z --> -x, x --> z, bx --> bz and bz --> -bx are applied

Uses fieldMapTools to read in the file and apply the transformation
'''

import sys
from optparse import OptionParser
import fieldMapTools as ft

def main(argv=None):
  usageString="%prog [options] inputFileName"
  descriptionString="""Converts the format of a field map from Toshiba to one that can be used by SimMARS and SimG4 (i.e. G4beamline BLFieldMap format)
  The default transformation applied is : z --> -x, x --> z, bx --> bz, bz --> -bx
  """
  parser = OptionParser(usage=usageString, description=descriptionString)
  parser.add_option("-d", "--dipole-map", action="store_true",
                            help="Input file is a dipole field map. Do not apply transformation.")
  parser.add_option('-o' ,'--output-file', dest='outputFile',
                            help= 'Specify a filename for the converted fieldmap. Default="output"')
  parser.add_option('-u' ,'--units', dest='units',default="mm",
                            help= 'Specify the units for length used in the input file.  Default="mm"')
  (options,args)=parser.parse_args()

  
  if options.outputFile:
    outputFile=options.outputFile
  else:
    outputFile='output'
  f_out=open(outputFile, 'w')
  
  if len(args)== 1:
    filename=args[0] 
  else:
    print "Error parsing arguments: Only one argument (i.e. the input filename) must be given after options."
    parser.print_help()
    sys.exit(1)
  
  if options.dipole_map:
    transform=False
  else:
    transform=True    
  toshibaMap=ft.readFieldMapFile(filename,transform=transform, units=options.units)
  
  # Extract information for header file and then write header and map to the output file
  f=open(filename, 'r')
  aLine=f.readline().split()
  if (len(aLine) != 4):
    print "Error reading header of file: Expected 4 numbers, read %i numbers" % len(aLine)
    sys.exit(1)
  if options.dipole_map:
    nX=aLine[2];nY=aLine[1];nZ=aLine[0]
  else:
    nX=aLine[0];nY=aLine[1];nZ=aLine[2]
  f.close()
  
  x0=min(toshibaMap[:,0]);y0=min(toshibaMap[:,1]);z0=min(toshibaMap[:,2])

  dX=(max(toshibaMap[:,0])-x0)/(float(nX)-1.0)
  dY=(max(toshibaMap[:,1])-y0)/(float(nY)-1.0)
  dZ=(max(toshibaMap[:,2])-z0)/(float(nZ)-1.0)
  
  f_out.write("#This is the BLFieldMap input file converted from " + filename)
  f_out.write("\ngrid X0=%s Y0=%s Z0=%s nX=%s nY=%s nZ=%s dX=%s dY=%s dZ=%s " % (x0,y0,z0,nX,nY,nZ,dX,dY,dZ))
  
  if options.dipole_map:
    f_out.write("\nextendX flip=Bx ")
    f_out.write("\nextendY flip=Bx,Bz ")
    f_out.write("\nextendZ flip=Bz ")
  else:
    # Assume solenoid field which needs to be extended in Y only.
    #f_out.write("\nextendY flip=By")
    pass
  f_out.write("\ndata")
  
  count=0
  for a in toshibaMap:
    f_out.write("\n %s %s %s %s %s %s " % (a[0],a[1],a[2],a[3],a[4],a[5]))
    count+=1
  print ("Wrote %i grid lines to %s" % (count,outputFile))
  f_out.close()

if __name__ == '__main__':
  main(sys.argv)
  sys.exit(0)