#!/usr/bin/env python
#
# import.py
#
#  Copyright (C) 2013 Diamond Light Source
#
#  Author: James Parkhurst
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package.
from __future__ import division
from dxtbx.datablock import DataBlockFactory, DataBlockDumper

help_message = '''

This program is used to import image data files into a format that can be used
within dials. The program looks at the metadata for each image along with the
filenames to determine the relationship between sets of images. Once all the
images have been analysed, a datablock object is written to file which specifies
the relationship between files. For example if two sets of images which belong
to two rotation scans have been given, two image sweeps will be saved. Images to
be processed are specified as command line arguments. Sometimes, there is a
maximum number of arguments that can be given on the command line and the number
of files may exceed this. In this case image filenames can be input on stdin
delmited by a new line using the -i option (see below for examples).

Examples::

  dials.import image_*.cbf

  dials.import image_1_*.cbf image_2_*.cbf

  find . -name "image_*.cbf" | dials.import

  dials.import <<eof
  image_1.cbf
  image_2.cbf
  eof

'''

# Create the phil parameters
from libtbx.phil import parse
phil_scope = parse('''

  output = datablock.json
    .type = str
    .help = "The output JSON or pickle file"

  compact = False
    .type = bool
    .help = "For JSON output use compact representation"

  input {
    template = None
      .type = str
      .help = "The image sweep template"
      .multiple = True
    reference_geometry = None
      .type = path
      .help = "Experimental geometry from this datablock.json or "
              "experiments.json will override the geometry from the "
              "image headers."
  }

  mosflm_beam_centre = None
    .type = floats(size=2)
    .help = "Override the beam centre from the image headers, following "
            "the mosflm convention."

''')


class Script(object):
  ''' Class to parse the command line options. '''

  def __init__(self):
    ''' Set the expected options. '''
    from dials.util.options import OptionParser
    import libtbx.load_env

    # Create the option parser
    usage = "usage: %s [options] /path/to/image/files" % libtbx.env.dispatcher_name
    self.parser = OptionParser(
      usage=usage,
      sort_options=True,
      phil=phil_scope,
      read_datablocks_from_images=True,
      epilog=help_message)

  def run(self):
    ''' Parse the options. '''
    from dxtbx.datablock import DataBlockTemplateImporter
    from dials.util.options import flatten_datablocks

    # Parse the command line arguments
    params, options = self.parser.parse_args(show_diff_phil=True)
    datablocks = flatten_datablocks(params.input.datablock)

    reference_detector = None
    reference_beam = None

    if params.input.reference_geometry is not None:
      from dxtbx.serialize import load
      try:
        experiments = load.experiment_list(
          params.input.reference_geometry, check_format=False)
        assert len(experiments.detectors()) == 1
        assert len(experiments.beams()) == 1
        reference_detector = experiments.detectors()[0]
        reference_beam = experiments.beams()[0]
      except Exception, e:
        datablock = load.datablock(params.input.reference_geometry)
        assert len(datablock) == 1
        imageset = datablock[0].extract_imagesets()[0]
        reference_detector = imageset.get_detector()
        reference_beam = imageset.get_beam()

    # Check we have some filenames
    if len(datablocks) == 0:

      # Check if a template has been set and print help if not, otherwise try to
      # import the images based on the template input
      if len(params.input.template) == 0:
        self.parser.print_help()
        exit(0)
      else:
        importer = DataBlockTemplateImporter(
          params.input.template,
          options.verbose)
        datablocks = importer.datablocks

    # Loop through the data blocks
    for i, datablock in enumerate(datablocks):

      # Extract any sweeps
      sweeps = datablock.extract_sweeps()

      # Extract any stills
      stills = datablock.extract_stills()
      if not stills:
        num_stills = 0
      else:
        num_stills = len(stills)

      if reference_beam is not None and reference_detector is not None:
        for sweep in sweeps:
          assert reference_detector.is_similar_to(sweep.get_detector())
          sweep.set_beam(reference_beam)
          sweep.set_detector(reference_detector)

        for still in stills:
          assert reference_detector.is_similar_to(still.get_detector())
          still.set_beam(reference_beam)
          still.set_detector(reference_detector)

      elif params.mosflm_beam_centre is not None:
        from dxtbx.model.detector_helpers import set_mosflm_beam_centre
        for sweep in sweeps:
          set_mosflm_beam_centre(
            sweep.get_detector(), sweep.get_beam(), params.mosflm_beam_centre)

        for still in stills:
          set_mosflm_beam_centre(
            still.get_detector(), still.get_beam(), params.mosflm_beam_centre)

      # Print some data block info
      print "-" * 80
      print "DataBlock %d" % i
      print "  format: %s" % str(datablock.format_class())
      print "  num images: %d" % datablock.num_images()
      print "  num sweeps: %d" % len(sweeps)
      print "  num stills: %d" % num_stills

      # Loop through all the sweeps
      if options.verbose > 1:
        for j, sweep in enumerate(sweeps):
          print ""
          print "Sweep %d" % j
          print "  length %d" % len(sweep)
          print sweep.get_beam()
          print sweep.get_goniometer()
          print sweep.get_detector()
          print sweep.get_scan()

    # Write the datablock to a JSON or pickle file
    if params.output:
      print "-" * 80
      print 'Writing datablocks to %s' % params.output
      dump = DataBlockDumper(datablocks)
      dump.as_file(params.output, compact=params.compact)

if __name__ == '__main__':
  from dials.util import halraiser
  try:
    script = Script()
    script.run()
  except Exception as e:
    halraiser(e)

@@ -1053,16 +1053,17 @@
 tdin%0Adel
+i
 mited by
@@ -1119,16 +1119,199 @@
 amples).
+%0AAlternatively a template can be specified using the template= parameter where%0Athe consecutive digits representing the image numbers in the filenames are%0Areplaced with '#' characters.
 %0A%0AExampl
@@ -1388,16 +1388,58 @@
 _*.cbf%0A%0A
+  dials.import template=image_1_####.cbf%0A%0A
   find .

fix typo in doc; add template= example