#!/usr/bin/env python
# XDSCorrectHelpers.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.
#
# 18th October 2006
#
# Helpers for XDS when running the correct step - this will -
#
#  - check that all input files are present and correct
#  - run xds to do integration, with help from the input parameters
#    and a generic xds writer
#  - parse the output from CORRECT.LP

import os
import re
import sys
import copy

if not os.environ.has_key('XIA2CORE_ROOT'):
  raise RuntimeError, 'XIA2CORE_ROOT not defined'
if not os.environ.has_key('XIA2_ROOT'):
  raise RuntimeError, 'XIA2_ROOT not defined'

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

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

# helper methods/functions - these can be used externally for the purposes
# of testing...

def _resolution_estimate(ordered_pair_list, cutoff):
  '''Come up with a linearly interpolated estimate of resolution at
  cutoff cutoff from input data [(resolution, i_sigma)].'''

  x = []
  y = []

  for o in ordered_pair_list:
    x.append(o[0])
    y.append(o[1])

  if max(y) < cutoff:
    # there is no point where this exceeds the resolution
    # cutoff
    return -1.0

  # this means that there is a place where the resolution cutof
  # can be reached - get there by working backwards

  x.reverse()
  y.reverse()

  if y[0] >= cutoff:
    # this exceeds the resolution limit requested
    return x[0]

  j = 0
  while y[j] < cutoff:
    j += 1

  resolution = x[j] + (cutoff - y[j]) * (x[j - 1] - x[j]) / \
               (y[j - 1] - y[j])

  return resolution

def _parse_correct_lp(filename):
  '''Parse the contents of the CORRECT.LP file pointed to by filename.'''

  if not os.path.split(filename)[-1] == 'CORRECT.LP':
    raise RuntimeError, 'input filename not CORRECT.LP'

  file_contents = open(filename, 'r').readlines()

  postrefinement_stats = { }

  # default values - this may not be refined if the multiplicity
  # is low...
  postrefinement_stats['sdcorrection'] = (1.0, 0.0)

  for i in range(len(file_contents)):
    if 'OF SPOT    POSITION (PIXELS)' in file_contents[i]:
      rmsd_pixel = float(file_contents[i].split()[-1])
      postrefinement_stats['rmsd_pixel'] = rmsd_pixel

    if 'OF SPINDLE POSITION (DEGREES)' in file_contents[i]:
      rmsd_phi = float(file_contents[i].split()[-1])
      postrefinement_stats['rmsd_phi'] = rmsd_phi

    # want to convert this to mm in some standard setting!
    if 'DETECTOR COORDINATES (PIXELS) OF DIRECT BEAM' in file_contents[i]:
      beam = map(float, file_contents[i].split()[-2:])
      postrefinement_stats['beam'] = beam

    if 'CRYSTAL TO DETECTOR DISTANCE (mm)' in file_contents[i]:
      distance = float(file_contents[i].split()[-1])
      postrefinement_stats['distance'] = distance

    if 'UNIT CELL PARAMETERS' in file_contents[i]:
      cell = map(float, file_contents[i].split()[-6:])
      postrefinement_stats['cell'] = cell

    if 'E.S.D. OF CELL PARAMETERS' in file_contents[i]:
      # bug # 3132 - check that the last token is not
      # "-1.0E+00-1.0E+00-1.0E+00-1.0E+00-1.0E+00-1.0E+00" -
      # if it is it means that the refinement didn't
      # happen (for some reason...)

      if '-1.0E+00-1.0E+00-1.0E+00' in file_contents[i]:
        cell_esd = [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0]
      else:
        cell_esd = map(float, file_contents[i].split()[-6:])
      postrefinement_stats['cell_esd'] = cell_esd

    if 'REFLECTIONS ACCEPTED' in file_contents[i]:
      postrefinement_stats['n_ref'] = int(file_contents[i].split()[0])

    # look for I/sigma (resolution) information...
    if 'RESOLUTION RANGE  I/Sigma  Chi^2  R-FACTOR  R-FACTOR' in \
       file_contents[i]:
      resolution_info = []
      j = i + 3
      while not '-----' in file_contents[j]:
        try:
          l = file_contents[j].split()
          resolution_info.append((float(l[1]),float(l[2])))
        except ValueError, e:
          l = file_contents[j].split()
          m = re.match(r'(\d+\.\d{2})(\d+\.\d+)', l[2])
          resolution_info.append((float(l[1]),float(m.group(1))))
        j += 1

      # bug # 2409 - this seems a little harsh set as 1.0 so
      # set this to 0.75 - even then 0.5 may be better..
      resolution_old = _resolution_estimate(resolution_info, 0.5)
      postrefinement_stats['resolution_estimate_old'] = resolution_old

      # also recover the highest resolution limit of the data
      j += 1
      postrefinement_stats['highest_resolution'] = float(
          file_contents[j].split()[1])

    if 'a          b              INPUT DATA SET' in file_contents[i]:
      sdcorrection = map(float, file_contents[i + 1].split()[:2])

      postrefinement_stats['sdcorrection'] = tuple(sdcorrection)

    if 'CORRELATION  NPAIR  Rmeas  COMPARED  ESD' in file_contents[i]:
      j = i + 2
      while file_contents[j].strip():
        if '*' in file_contents[j]:
          postrefinement_stats['reindex_op'] = map(
              int, file_contents[j].split()[-12:])
        j += 1

  return postrefinement_stats

if __name__ == '__main__':
  correct_lp = os.path.join(os.environ['XIA2_ROOT'], 'Wrappers', 'XDS',
                            'Doc', 'CORRECT.LP')
  if len(sys.argv) > 1:
    correct_lp = sys.argv[1]
  print _parse_correct_lp(correct_lp)