#!/usr/bin/env python
# xia2setup.py
#   Copyright (C) 2006 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# xia2setup.py - an application to generate the .xinfo file for data
# reduction from a directory full of images, optionally with scan and
# sequence files which will be used to add matadata.
#
# 18th December 2006
#


import os
import sys
import exceptions
import time
import traceback

if not os.environ.has_key('XIA2_ROOT'):
  raise RuntimeError, 'XIA2_ROOT not defined'
if not 'XIA2CORE_ROOT' in os.environ:
  os.environ['XIA2CORE_ROOT'] = os.path.join(os.environ['XIA2_ROOT'], 'core')

if not os.environ['XIA2_ROOT'] in sys.path:
  sys.path.append(os.path.join(os.environ['XIA2_ROOT']))

from Schema.Sweep import SweepFactory
from Experts.FindImages import image2template_directory
from Handlers.CommandLine import CommandLine
from Handlers.Flags import Flags
from Handlers.Phil import PhilIndex
from Wrappers.CCP4.Chooch import Chooch
from Modules.LabelitBeamCentre import compute_beam_centre
from Handlers.Streams import streams_off

known_image_extensions = ['img', 'mccd', 'mar2300', 'mar1200', 'mar1600',
                          'mar3450', 'osc', 'cbf', 'mar2000']

xds_file_names = ['ABS', 'ABSORP', 'BKGINIT', 'BKGPIX', 'BLANK', 'DECAY',
                  'X-CORRECTIONS', 'Y-CORRECTIONS', 'MODPIX', 'FRAME',
                  'GX-CORRECTIONS', 'GY-CORRECTIONS', 'DX-CORRECTIONS',
                  'DY-CORRECTIONS', 'GAIN']

known_sweeps = { }

known_scan_extensions = ['scan']

known_sequence_extensions = ['seq']

latest_sequence = None

latest_chooch = None

target_template = None

def is_scan_name(file):
  global known_scan_extensions

  if os.path.isfile(file):
    if file.split('.')[-1] in known_scan_extensions:
      return True

  return False

def is_sequence_name(file):
  global known_sequence_extensions

  if os.path.isfile(file):
    if file.split('.')[-1] in known_sequence_extensions:
      return True

  return False

def is_image_name(filename):

  global known_image_extensions
  from Wrappers.XDS.XDSFiles import XDSFiles

  if os.path.isfile(filename):

    if os.path.split(filename)[-1] in XDSFiles:
      return False

    # also XDS scaling files from previous xia2 job - hard coded, messy :o(
    for xds_file in 'ABSORP', 'DECAY', 'MODPIX':
      if os.path.join('scale', xds_file) in filename:
        return False

    exten = filename.split('.')[-1]
    if exten in known_image_extensions:
      return True

    # check for files like foo_bar.0001 - c/f FIXME for 25/SEP/06
    end = filename.split('.')[-1]
    try:
      j = int(end)
      return True
    except:
      pass

  return False

def is_xds_file(f):
  filename = os.path.split(f)[1]

  xds_files = ['ABS', 'ABSORP', 'BKGINIT', 'BKGPIX', 'BLANK', 'DECAY',
               'DX-CORRECTIONS', 'DY-CORRECTIONS', 'FRAME', 'GAIN',
               'GX-CORRECTIONS', 'GY-CORRECTIONS', 'MODPIX',
               'X-CORRECTIONS', 'Y-CORRECTIONS']

  return (filename.split('.')[0].split('_') in xds_files)

def get_template(f):

  global target_template

  if not is_image_name(f):
    return

  if is_xds_file(f):
    return

  # in here, check the permissions on the file...

  if not os.access(f, os.R_OK):
    from Handlers.Streams import Debug
    Debug.write('No read permission for %s' % f)

  try:
    template, directory = image2template_directory(f)

    if target_template:
      if template != target_template:
        return

  except exceptions.Exception, e:
    from Handlers.Streams import Debug
    Debug.write('Exception: %s (%s)' % (str(e), f))
    # traceback.print_exc(file = sys.stdout)

  return os.path.join(directory, template)


def save_datablock(filename):
  from Schema import imageset_cache
  from dxtbx.datablock import DataBlock
  from dxtbx.serialize import dump

  datablock = DataBlock([])
  for imagesets in imageset_cache.values():
    for imageset in imagesets.values():
      datablock.append(imageset)

  dump.datablock(datablock, filename)


def load_datablock(filename):
  from Schema import imageset_cache
  from dxtbx.serialize import load
  from libtbx.containers import OrderedDict

  datablocks = load.datablock(filename, check_format=False)

  for datablock in datablocks:
    for imageset in datablock.extract_imagesets():
      template = imageset.get_template()
      if template not in imageset_cache:
        imageset_cache[template] = OrderedDict()
      imageset_cache[template][imageset.get_scan().get_image_range()[0]] = imageset


def parse_sequence(sequence_file):
  sequence = ''

  for record in open(sequence_file).readlines():
    if record[0].upper() in \
       'ABCDEFGHIJKLMNOPQRSTUVWXYZ ':
      sequence += record.strip().upper()

  global latest_sequence
  latest_sequence = sequence
  return

def visit(root, directory, files):
  files.sort()

  templates = set()

  for f in files:

    full_path = os.path.join(directory, f)
    if is_image_name(full_path):
      template = get_template(full_path)
      if template is not None:
        templates.add(template)

    if is_scan_name(full_path):
      global latest_chooch
      try:
        latest_chooch = Chooch()
        if CommandLine.get_atom_name():
          latest_chooch.set_atom(CommandLine.get_atom_name())
        latest_chooch.set_scan(full_path)
        latest_chooch.scan()
      except:
        latest_chooch = None

    if is_sequence_name(full_path):
      parse_sequence(full_path)

  return templates

def print_sweeps(out = sys.stdout):

  global known_sweeps, latest_sequence

  sweeplists = known_sweeps.keys()
  sweeplists.sort()

  # analysis pass

  wavelengths = []

  for sweep in sweeplists:
    sweeps = known_sweeps[sweep]
    # this should sort on exposure epoch ...?
    sweeps.sort()
    for s in sweeps:

      if len(s.get_images()) < Flags.get_min_images():
        continue

      wavelength = s.get_wavelength()

      if not wavelength in wavelengths:
        wavelengths.append(wavelength)

  wavelength_map = { }

  project = CommandLine.get_project_name()
  if not project:
    project = 'AUTOMATIC'

  crystal = CommandLine.get_crystal_name()
  if not crystal:
    crystal = 'DEFAULT'

  out.write('BEGIN PROJECT %s\n' % project)
  out.write('BEGIN CRYSTAL %s\n' % crystal)

  out.write('\n')

  # check to see if a user spacegroup has been assigned - if it has,
  # copy it in...

  settings = PhilIndex.params.xia2.settings

  if settings.space_group is not None:
    out.write(
      'USER_SPACEGROUP %s\n' % settings.space_group.type().lookup_symbol())
    out.write('\n')

  if settings.unit_cell is not None:
    out.write('USER_CELL %.2f %.2f %.2f %.2f %.2f %.2f\n' % \
              settings.unit_cell.parameters())
    out.write('\n')

  if Flags.get_freer_file():
    out.write('FREER_FILE %s\n' % Flags.get_freer_file())
    out.write('\n')

  if latest_sequence:
    out.write('BEGIN AA_SEQUENCE\n')
    out.write('\n')
    for sequence_chunk in [latest_sequence[i:i + 60] \
                           for i in range(0, len(latest_sequence), 60)]:
      out.write('%s\n' % sequence_chunk)
    out.write('\n')
    out.write('END AA_SEQUENCE\n')
    out.write('\n')

  if CommandLine.get_atom_name():
    out.write('BEGIN HA_INFO\n')
    out.write('ATOM %s\n' % CommandLine.get_atom_name().lower())
    if CommandLine.get_atom_name().lower() == 'se' and latest_sequence:
      # assume that this is selenomethionine
      out.write('! If this is SeMet uncomment next line...\n')
      out.write('!NUMBER_PER_MONOMER %d\n' % latest_sequence.count('M'))
      out.write('!NUMBER_TOTAL M\n')
    else:
      out.write('!NUMBER_PER_MONOMER N\n')
      out.write('!NUMBER_TOTAL M\n')
    out.write('END HA_INFO\n')
    out.write('\n')

  for j in range(len(wavelengths)):

    global latest_chooch

    if latest_chooch:
      name = latest_chooch.id_wavelength(wavelengths[j])
      first_name = name
      counter = 1

      while name in [wavelength_map[w] for w in wavelength_map]:
        counter += 1
        name = '%s%d' % (first_name, counter)

      fp, fpp = latest_chooch.get_fp_fpp(wavelengths[j])
    else:
      fp, fpp = 0.0, 0.0
      if len(wavelengths) == 1 and CommandLine.get_atom_name():
        name = 'SAD'
      elif len(wavelengths) == 1:
        name = 'NATIVE'
      else:
        name = 'WAVE%d' % (j + 1)

    wavelength_map[wavelengths[j]] = name

    out.write('BEGIN WAVELENGTH %s\n' % name)

    dmin = Flags.get_resolution_high()
    dmax = Flags.get_resolution_low()

    if dmin and dmax:
      out.write('RESOLUTION %f %f\n' % (dmin, dmax))
    elif dmin:
      out.write('RESOLUTION %f\n' % dmin)

    out.write('WAVELENGTH %f\n' % wavelengths[j])
    if fp != 0.0 and fpp != 0.0:
      out.write('F\' %5.2f\n' % fp)
      out.write('F\'\' %5.2f\n' % fpp)

    out.write('END WAVELENGTH %s\n' % name)
    out.write('\n')

  j = 0
  for sweep in sweeplists:
    sweeps = known_sweeps[sweep]
    # this should sort on exposure epoch ...?
    sweeps.sort()
    for s in sweeps:

      # require at least n images to represent a sweep...

      if len(s.get_images()) < Flags.get_min_images():
        continue

      j += 1
      name = 'SWEEP%d' % j

      out.write('BEGIN SWEEP %s\n' % name)

      if Flags.get_reversephi():
        out.write('REVERSEPHI\n')

      out.write('WAVELENGTH %s\n' % wavelength_map[s.get_wavelength()])

      out.write('DIRECTORY %s\n' % s.get_directory())
      out.write('IMAGE %s\n' % os.path.split(s.imagename(min(
          s.get_images())))[-1])

      if Flags.get_start_end():
        start, end = Flags.get_start_end()

        if start < min(s.get_images()):
          raise RuntimeError, 'requested start %d < %d' % \
                (start, min(s.get_images()))

        if end > max(s.get_images()):
          raise RuntimeError, 'requested end %d > %d' % \
                (end, max(s.get_images()))

        out.write('START_END %d %d\n' % (start, end))
      else:
        out.write('START_END %d %d\n' % (min(s.get_images()),
                                         max(s.get_images())))

      # really don't need to store the epoch in the xinfo file
      # out.write('EPOCH %d\n' % int(s.get_collect()[0]))
      user_beam_centre = settings.beam_centre
      if user_beam_centre is not None:
        out.write('BEAM %6.2f %6.2f\n' % tuple(user_beam_centre))
      elif not settings.trust_beam_centre:
        interactive = Flags.get_interactive()
        Flags.set_interactive(False)
        beam_centre = compute_beam_centre(s)
        if beam_centre:
          out.write('BEAM %6.2f %6.2f\n' % tuple(beam_centre))
        Flags.set_interactive(interactive)
      out.write('END SWEEP %s\n' % name)

      out.write('\n')

  out.write('END CRYSTAL %s\n' % crystal)
  out.write('END PROJECT %s\n' % project)


def get_sweep(args):
  assert len(args) == 1
  directory, template = os.path.split(args[0])

  try:
    sweeplist = SweepFactory(template, directory)

  except exceptions.Exception, e:
    from Handlers.Streams import Debug
    Debug.write('Exception: %s (%s)' % (str(e), args[0]))
    return None

  return sweeplist


def get_sweeps(templates):
  global known_sweeps

  from libtbx import easy_mp
  from Handlers.Phil import PhilIndex
  params = PhilIndex.get_python_object()
  mp_params = params.xia2.settings.multiprocessing
  njob = mp_params.njob

  if njob > 1:
    from Driver.DriverFactory import DriverFactory
    drivertype = DriverFactory.get_driver_type()
    if drivertype == "qsub":
      method = "qsub"
    else:
      method = "multiprocessing"

    args = [(template,) for template in templates]
    results_list = easy_mp.parallel_map(
      get_sweep, args, params=None,
      processes=njob,
      method=method,
      asynchronous=True,
      callback=None,
      preserve_order=True,
      preserve_exception_message=True)

  else:
    results_list = [get_sweep((template,)) for template in templates]

  from Schema import imageset_cache
  from libtbx.containers import OrderedDict

  for template, sweeplist in zip(templates, results_list):
    if sweeplist is not None:
      known_sweeps[template] = sweeplist
      for sweep in sweeplist:
        imageset = sweep.get_imageset()
        if template not in imageset_cache:
          imageset_cache[template] = OrderedDict()
        imageset_cache[template][imageset.get_scan().get_image_range()[0]] = imageset


def rummage(path):
  '''Walk through the directories looking for sweeps.'''
  templates = set()
  for root, dirs, files in os.walk(path, followlinks=True):
    templates.update(visit(os.getcwd(), root, files))

  get_sweeps(templates)

  return

def write_xinfo(filename, path, template = None):

  global target_template

  target_template = template

  crystal = CommandLine.get_crystal_name()

  if not crystal:
    crystal = 'DEFAULT'

  if not os.path.isabs(filename):
    filename = os.path.abspath(filename)

  directory = os.path.join(os.getcwd(), crystal, 'setup')

  try:
    os.makedirs(directory)
  except OSError, e:
    if not 'File exists' in str(e):
      raise e

  # FIXME should I have some exception handling in here...?

  start = os.getcwd()
  os.chdir(directory)

  # if we have given a template and directory on the command line, just
  # look there (i.e. not in the subdirectories)

  if CommandLine.get_template() and CommandLine.get_directory():
    get_sweeps(visit(None, CommandLine.get_directory(),
                     os.listdir(CommandLine.get_directory())))
  else:
    rummage(path)

  save_datablock(os.path.join(start, 'xia2-datablock.json'))

  fout = open(filename, 'w')
  print_sweeps(fout)

  # change back directory c/f bug # 2693 - important for error files...
  os.chdir(start)

  return

if __name__ == '__main__':

  streams_off()

  argv = sys.argv

  # test to see if sys.argv[-2] + path is a valid path - to work around
  # spaced command lines

  path = argv.pop()

  # perhaps move to a new directory...

  crystal = CommandLine.get_crystal_name()

  fout = open(os.path.join(os.getcwd(), 'automatic.xinfo'), 'w')

  if not crystal:
    crystal = 'DEFAULT'

  start = os.path.abspath(os.getcwd())

  directory = os.path.join(os.getcwd(), crystal, 'setup')

  try:
    os.makedirs(directory)
  except OSError, e:
    if not 'File exists' in str(e):
      raise e

  os.chdir(directory)

  while not os.path.exists(path):
    path = '%s %s' % (argv.pop(), path)

  if not os.path.isabs(path):
    path = os.path.abspath(path)

  rummage(path)
  print_sweeps(fout)

  save_datablock(os.path.join(start, 'xia2-datablock.json'))