#!/usr/bin/env python
# AnalyseMyIntensities.py
#   Copyright (C) 2007 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# 16th June 2007
#
# A tool to use for the analysis and gathering of scaled intensity data
# from a single macromolecular crystal. This will be both a module (for
# use in xia2) and an application in it's own right, AMI.
#
# Example usage:
#
# ami hklin1 PEAK.HKL hklin2 INFL.HKL hklin3 LREM.HKL HKLOUT merged.mtz << eof
# drename file 1 pname demo xname only dname peak
# drename file 2 pname demo xname only dname infl
# drename file 3 pname demo xname only dname lrem
# solvent 0.53
# symm P43212
# reindex h,k,l
# cell 55.67 55.67 108.92 90.0 90.0 90.0
# anomalous on
# eof
#
# should also allow for a HKLREF.

import os
import sys
import math
import copy

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

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

from Wrappers.CCP4.CCP4Factory import CCP4Factory

from Modules.Scalepack2Mtz import Scalepack2Mtz
from Modules.Mtz2Scalepack import Mtz2Scalepack
from Modules.XDS2Mtz import XDS2Mtz

from lib.bits import is_mtz_file, is_xds_file, is_scalepack_file
from lib.NMolLib import compute_nmol, compute_solvent

from Handlers.Streams import Chatter, Debug
from Handlers.Files import FileHandler

class AnalyseMyIntensities(object):
  '''A class to use for intensity analysis. This will gather intensities
  (merged or unmerged) from multiple data sets and merge them together
  as well as telling you all about your data.'''

  def __init__(self):
    self._hklin_list = []
    self._project_info = []
    self._hklout = ''
    self._solvent = 0.0
    self._nres = 0
    self._nmol = 0
    self._cell = None
    self._symmetry = None
    self._reindex = None
    self._anomalous = False
    self._rotation_function = False

    self._hklin_stats = { }

    self._working_directory = os.getcwd()

    self._factory = CCP4Factory()

    # working space
    self._resolution = 0.0

    # places to store the merging statistics etc.
    self._merging_statistics = { }
    self._merging_statistics_keys = []

    self._truncate_statistics = { }
    self._truncate_statistics_keys = []
    self._truncate_hklout = []

    self._sfcheck_statistics = { }
    self._sfcheck_statistics_keys = []

    self._huge_log_file = []

    return

  # admin functions

  def set_working_directory(self, working_directory):
    self._working_directory = working_directory
    self._factory.set_working_directory(working_directory)
    return

  def get_working_directory(self):
    return self._working_directory

  def write_log_file(self, filename):
    fout = open(filename, 'w')
    for o in self._huge_log_file:
      fout.write(o)
    fout.close()

  def get_log_file(self):
    return self._huge_log_file

  # input functions

  def add_hklin(self, hklin, project_info = None):
    self._hklin_list.append(hklin)
    if project_info:
      self._project_info.append(project_info)
      pname, xname, dname = project_info
      Chatter.write('Storing %s/%s/%s with %s' % \
                    (pname, xname, dname, hklin))
    else:

      if not is_mtz_file(hklin):
        raise RuntimeError, 'hklin %s not MTZ and no project info' % \
              hklin

      # try to get this from the reflection file
      mtzdump = self._factory.Mtzdump()
      mtzdump.set_hklin(hklin)
      mtzdump.dump()

      datasets = mtzdump.get_datasets()
      if len(datasets) > 1:
        raise RuntimeError, 'more than one dataset in %s' % hklin
      pname, xname, dname = datasets[0].split('/')
      Chatter.write('Found %s in %s' % \
                    (datasets[0], hklin))
      self._project_info.append((pname, xname, dname))

    return

  def set_hklout(self, hklout):
    self._hklout = hklout
    return

  def set_solvent(self, solvent):
    self._solvent = solvent
    return

  def set_nres(self, nres):
    self._nres = nres
    return

  def set_nmol(self, nmol):
    self._nmol = nmol
    return

  def set_cell(self, cell):
    self._cell = cell
    return

  def set_symmetry(self, symmetry):
    self._symmetry = symmetry
    return

  def set_reindex(self, reindex):
    self._reindex = reindex
    return

  def set_anomalous(self, anomalous):
    self._anomalous = anomalous
    return

  def set_rotation_function(self, rotation_function):
    self._rotation_function = rotation_function
    return

  def compute_average_cell(self, hklin_list):

    if len(hklin_list) == 0:
      raise RuntimeError, \
            'no input reflection files to compute cell from'

    cell_a = 0.0
    cell_b = 0.0
    cell_c = 0.0
    cell_alpga = 0.0
    cell_beta = 0.0
    cell_gamma = 0.0
    n_input = 0
    sg = None

    for hklin in hklin_list:
      mtzdump = self._factory.Mtzdump()
      mtzdump.set_hklin(hklin)
      mtzdump.dump()

      resolution = min(mtzdump.get_resolution_range())
      if resolution < self._resolution or self._resolution == 0:
        self._resolution = resolution

      datasets = mtzdump.get_datasets()
      reflections = mtzdump.get_reflections()
      if len(datasets) > 1:
        raise RuntimeError, 'more than one dataset in %s' % hklin
      info = mtzdump.get_dataset_info(datasets[0])

      if not sg:
        sg = info['spacegroup']
      else:
        if sg != info['spacegroup']:
          raise RuntimeError, 'inconsistent spacegroup'

      # check that this u/c is in agreement with the others -
      # allow 10% grace (!)

      if n_input == 0:
        cell_a = info['cell'][0] * reflections
        cell_b = info['cell'][1] * reflections
        cell_c = info['cell'][2] * reflections
        cell_alpha = info['cell'][3] * reflections
        cell_beta = info['cell'][4] * reflections
        cell_gamma = info['cell'][5] * reflections
        n_input += reflections
      else:
        if math.fabs(n_input * info['cell'][0] - cell_a) / \
               cell_a > 0.1:
          raise RuntimeError, 'inconsistent unit cell'
        if math.fabs(n_input * info['cell'][1] - cell_b) / \
               cell_b > 0.1:
          raise RuntimeError, 'inconsistent unit cell'
        if math.fabs(n_input * info['cell'][2] - cell_c) / \
               cell_c > 0.1:
          raise RuntimeError, 'inconsistent unit cell'
        if math.fabs(n_input * info['cell'][3] - cell_alpha) / \
               cell_alpha > 0.1:
          raise RuntimeError, 'inconsistent unit cell'
        if math.fabs(n_input * info['cell'][4] - cell_beta) / \
               cell_beta > 0.1:
          raise RuntimeError, 'inconsistent unit cell'
        if math.fabs(n_input * info['cell'][5] - cell_gamma) / \
               cell_gamma > 0.1:
          raise RuntimeError, 'inconsistent unit cell'

        cell_a += info['cell'][0] * reflections
        cell_b += info['cell'][1] * reflections
        cell_c += info['cell'][2] * reflections
        cell_alpha += info['cell'][3] * reflections
        cell_beta += info['cell'][4] * reflections
        cell_gamma += info['cell'][5] * reflections
        n_input += reflections

    cell_a /= n_input
    cell_b /= n_input
    cell_c /= n_input
    cell_alpha /= n_input
    cell_beta /= n_input
    cell_gamma /= n_input

    average_cell = (cell_a, cell_b, cell_c,
                    cell_alpha, cell_beta, cell_gamma)

    return average_cell, sg

  def _get_solvent(self):
    '''Get the solvent content, either from an mol/asu guess or
    from the solvent content if available.'''

    if self._nmol == 0 and self._nres == 0 and self._solvent == 0:
      raise RuntimeError, 'no solvent content information available'

    if self._solvent:
      return self._solvent

    # else try and compute an estimate

    if self._cell:
      cell = self._cell
    else:
      cell = self._average_cell

    if self._symmetry:
      symmetry = self._symmetry
    else:
      symmetry = self._average_sg

    if not self._nmol:

      self._nmol = compute_nmol(
          cell[0], cell[1], cell[2],
          cell[3], cell[4], cell[3],
          symmetry, self._resolution, self._nres)

      Chatter.write('Estimated %d molecules / ASU' % self._nmol)

    self._solvent = compute_solvent(
        cell[0], cell[1], cell[2],
        cell[3], cell[4], cell[3],
        symmetry, self._nmol, self._nres)

    return self._solvent


  def convert_to_mtz(self):
    '''Convert all HKLIN to MTZ format, storing merging statistics
    if available.'''

    hashes = '################################' + \
             '###############################'

    Chatter.write(hashes)
    Chatter.write('### Gathering data')
    Chatter.write(hashes)

    mtz_in = []

    for j in range(len(self._hklin_list)):
      hklin = self._hklin_list[j]

      if is_mtz_file(hklin):
        mtz_in.append(hklin)

        Chatter.write('Including MTZ file: %s' % hklin)

      elif is_xds_file(hklin):

        Chatter.write('Converting XDS file: %s' % hklin)

        hklout = os.path.join(
            self.get_working_directory(),
            'AMI_HKLIN%d.mtz' % j)

        xds2mtz = XDS2Mtz()
        s = xds2mtz.xds_to_mtz(hklin, hklout, self._anomalous,
                               spacegroup = self._symmetry,
                               cell = self._cell,
                               project_info = self._project_info[j])

        if s:
          k = (j, self._project_info[j])
          self._merging_statistics[j] = s
          self._merging_statistics_keys.append(k)

        mtz_in.append(hklout)

      elif is_scalepack_file(hklin):

        Chatter.write('Converting SCALEPACK file: %s' % hklin)

        hklout = os.path.join(
            self.get_working_directory(),
            'AMI_HKLIN%d.mtz' % j)

        scalepack2mtz = Scalepack2Mtz()
        s = scalepack2mtz.scalepack_to_mtz(hklin, hklout,
                                           self._anomalous,
                                           self._symmetry, self._cell,
                                           self._project_info[j])

        if s:
          k = (j, self._project_info[j])
          self._merging_statistics[j] = s
          self._merging_statistics_keys.append(k)

        mtz_in.append(hklout)

      else:
        raise RuntimeError, 'file %s unrecognised' % hklin

    # next work through this list and apply reindexing operators
    # etc if set... also store resolution limits via a quick
    # mtzdump.

    for hklin in mtz_in:
      mtzdump = self._factory.Mtzdump()
      mtzdump.set_hklin(hklin)
      mtzdump.dump()

      resolution = min(mtzdump.get_resolution_range())

      Chatter.write('Resolution for %s: %.2f' % (hklin, resolution))

      self._hklin_stats[os.path.split(hklin)[-1]] = {
          'resolution':resolution}

    if not self._symmetry and not self._reindex:
      Chatter.write('Determining unit cell')

      self._average_cell, self._average_sg = self.compute_average_cell(
          mtz_in)

      Chatter.write(
          'Determined unit cell: %.2f %.2f %.2f %.2f %.2f %.2f' % \
          self._average_cell)

      self._hklin_list = mtz_in
    else:
      hklin_list = []

      for j in range(len(mtz_in)):

        hklin = mtz_in[j]
        hklout = os.path.join(
            self.get_working_directory(),
            'AMI_HKLIN%d_reindex.mtz' % j)

        Chatter.write('Reindexing %s' % os.path.split(hklin)[-1])

        reindex = self._factory.Reindex()
        reindex.set_hklin(hklin)
        reindex.set_hklout(hklout)
        if self._symmetry:
          reindex.set_spacegroup(self._symmetry)
        if self._reindex:
          reindex.set_operator(self._reindex)
        reindex.reindex()

        hklin_list.append(hklout)

        resolution = self._hklin_stats[os.path.split(hklin)[-1]]
        del(self._hklin_stats[os.path.split(hklin)[-1]])
        self._hklin_stats[os.path.split(hklout)[-1]] = resolution

      # build up the average unit cell here

      Chatter.write('Determining unit cell')

      self._average_cell, self._average_sg = self.compute_average_cell(
          hklin_list)

      Chatter.write(
          'Determined unit cell: %.2f %.2f %.2f %.2f %.2f %.2f' % \
          self._average_cell)
      self._hklin_list = hklin_list

    return

  def analyse_input_hklin(self):
    '''Analyse all converted input reflection files.'''

    hashes = '################################' + \
             '###############################'

    Chatter.write(hashes)
    Chatter.write('### Individual analysis')
    Chatter.write(hashes)

    j = 0
    for hklin in self._hklin_list:

      hklout = os.path.join(
          self.get_working_directory(),
              'TRUNCATE%d.mtz' % j)

      # run truncate

      if not self._project_info[j]:
        Chatter.write('Truncating %s' % os.path.split(hklin)[-1])
      else:
        pname, xname, dname = self._project_info[j]
        Chatter.write('Truncating %s/%s/%s' % (pname, xname, dname))

      resolution = self._hklin_stats[os.path.split(hklin)[-1]][
          'resolution']

      truncate = self._factory.Truncate()
      truncate.set_hklin(hklin)
      truncate.set_hklout(hklout)
      if self._anomalous:
        truncate.set_anomalous(True)
      if resolution > 3.4:
        Chatter.write(
            'Switching off Wilson scaling as data low resolution')
        truncate.set_wilson(False)
      truncate.truncate()

      FileHandler.record_log_file('%s %s %s truncate' % \
                                  self._project_info[j],
                                  truncate.get_log_file())

      # look at the wilson plot fit stats -
      # y = A e ^ - m x
      m, dm, A, da = truncate.get_wilson_fit()
      dmax, dmin = truncate.get_wilson_fit_range()

      # in here I also need to get and store the truncate B
      # factor etc. so that they can be recovered by the calling
      # program

      if dm / m < 0.1:
        Chatter.write(
            'Over range %.2f %.2f get dm / m = %.3f: Wilson OK!' % \
            (dmax, dmin, math.fabs(dm / m)))
      else:
        Chatter.write(
            'Over range %.2f %.2f get dm / m = %.3f: Wilson BAD!' % \
            (dmax, dmin, math.fabs(dm / m)))


      for o in truncate.get_all_output():
        self._huge_log_file.append(o)

      s = truncate.parse_ccp4_loggraph()

      k = (j, self._project_info[j])
      self._truncate_statistics[k] = s
      self._truncate_statistics[k][
          'Wilson B factor'] = truncate.get_b_factor()
      self._truncate_statistics_keys.append(k)

      self._truncate_hklout.append(hklout)

      j += 1

    j = 0

    for hklin in self._truncate_hklout:

      # run sfcheck

      if not self._project_info[j]:
        Chatter.write('Sfchecking %s' % os.path.split(hklin)[-1])
      else:
        pname, xname, dname = self._project_info[j]
        Chatter.write('Sfchecking %s/%s/%s' % (pname, xname, dname))

      sfcheck = self._factory.Sfcheck()
      sfcheck.set_hklin(hklin)
      sfcheck.analyse()

      twinning_score = sfcheck.get_twinning()
      optical_resolution = sfcheck.get_optical_resolution()
      eigenvalues = sfcheck.get_anisotropic_eigenvalues()
      eigenvalues_reliable = \
                           sfcheck.get_anisotropic_eigenvalues_reliable()

      Chatter.write('Optical resolution: %.2f' % optical_resolution)

      if eigenvalues_reliable:
        Chatter.write('Eigenvalue ratios:  %.3f %.3f %.3f' % \
                      eigenvalues)
      else:
        Chatter.write(
            'Eigenvalue ratios:  %.3f %.3f %.3f (unreliable)' % \
            eigenvalues)

      if math.fabs(twinning_score - 1.5) < 0.2:
        Chatter.write('<I^2>/<I>^2 score:  %.2f (twinned)' % \
                      twinning_score)
      elif twinning_score > 2.5:
        Chatter.write('<I^2>/<I>^2 score:  %.2f (nasty)' % \
                      twinning_score)
      else:
        Chatter.write('<I^2>/<I>^2 score:  %.2f (ok)' % \
                      twinning_score)


      for o in sfcheck.get_all_output():
        self._huge_log_file.append(o)

      # then whatever else for the analysis

      if not self._rotation_function:
        j += 1
        continue

      Chatter.write('Analysing self-rotation function')

      hklout = os.path.join(
          self.get_working_directory(),
          'ECALC%d.mtz' % j)

      ecalc = self._factory.Ecalc()
      ecalc.set_hklin(hklin)
      ecalc.set_hklout(hklout)
      ecalc.ecalc()

      for o in ecalc.get_all_output():
        self._huge_log_file.append(o)

      hklin = hklout

      polarrfn = self._factory.Polarrfn()
      polarrfn.set_hklin(hklin)
      polarrfn.set_labin('E', 'SIGE')
      polarrfn.polarrfn()

      for o in polarrfn.get_all_output():
        self._huge_log_file.append(o)

      j += 1

    return

  def merge_analyse(self):
    '''Merge and analyse all of the data sets together, now.'''

    cad_hklin = []

    hashes = '################################' + \
             '###############################'

    Chatter.write(hashes)
    Chatter.write('### Combined analysis')
    Chatter.write(hashes)


    Chatter.write('Merging all reflection files together')

    for j in range(len(self._truncate_hklout)):
      hklin = self._truncate_hklout[j]
      hklout = os.path.join(
          self.get_working_directory(),
          'CAD%d.mtz' % j)

      cad_hklin.append(hklout)

      cad = self._factory.Cad()
      cad.add_hklin(hklin)
      cad.set_hklout(hklout)
      if self._project_info[j]:
        pname, xname, dname = self._project_info[j]
        cad.set_project_info(pname, xname, dname)
        cad.set_new_suffix(dname)

      if self._cell:
        cad.set_new_cell(self._cell)
      else:
        cad.set_new_cell(self._average_cell)

      cad.update()

    # now merge them together...

    cad = self._factory.Cad()
    for hklin in cad_hklin:
      cad.add_hklin(hklin)
    cad.set_hklout(self._hklout)

    cad.merge()

    # then run scaleit to look at the scaling statistics

    Chatter.write('Running scaleit analysis')

    hklout = os.path.join(
        self.get_working_directory(), 'SCALEIT.mtz')

    scaleit = self._factory.Scaleit()

    scaleit.set_hklin(self._hklout)
    scaleit.set_hklout(hklout)
    scaleit.set_anomalous(self._anomalous)

    scaleit.scaleit()

    statistics = scaleit.get_statistics()

    wavelengths = statistics['mapping']
    b_factors = statistics['b_factor']

    derivatives = wavelengths.keys()
    derivatives.sort()

    for j in derivatives[:1]:
      name = b_factors[j]['dname']
      b = b_factors[j]['b']
      r = b_factors[j]['r']

      Chatter.write('%s %5.2f %4.2f (reference)' % (name, b, r))

    for j in derivatives[1:]:
      name = b_factors[j]['dname']
      b = b_factors[j]['b']
      r = b_factors[j]['r']

      Chatter.write('%s %5.2f %4.2f' % (name, b, r))

    for o in scaleit.get_all_output():
      self._huge_log_file.append(o)

    return

  def get_truncate_statistics(self):
    return copy.deepcopy(self._truncate_statistics)


if __name__ == '__main__':
  infl = os.path.join(os.environ['XIA2_ROOT'],
                      'Data', 'Test', 'AMI', 'xds_unmerged',
                      'TS03_INFL_ANOM.hkl')
  lrem = os.path.join(os.environ['XIA2_ROOT'],
                      'Data', 'Test', 'AMI', 'xds_unmerged',
                      'TS03_LREM_ANOM.hkl')

  ami = AnalyseMyIntensities()

  ami.add_hklin(infl, project_info = ('AMI', 'TEST', 'INFL'))
  ami.add_hklin(lrem, project_info = ('AMI', 'TEST', 'LREM'))

  ami.set_anomalous(True)
  ami.set_hklout('out.mtz')

  ami.set_nres(180)

  ami.convert_to_mtz()
  ami.analyse_input_hklin()
  ami.merge_analyse()
  ami._get_solvent()

  ami.write_log_file('ami.log')