from __future__ import division
import phaser
from cctbx import miller
from cctbx.array_family import flex
import os
import sys

class data_adaptor(object):

  def __init__(self, f_obs, r_free_flags, verbose=False, safety_checks=True):
    self.verbose = verbose
    assert f_obs.anomalous_flag()
    if (safety_checks):
      assert f_obs.is_real_array()
      assert f_obs.map_to_asu().indices().all_eq(f_obs.indices())
      assert f_obs.sys_absent_flags().data().count(True) == 0
      assert f_obs.data().all_gt(0)
      assert r_free_flags.is_bool_array()
      assert r_free_flags.indices().all_eq(f_obs.indices())
    matches = miller.match_bijvoet_mates(
      f_obs.space_group_info().type(), f_obs.indices())
    self.sel_pp = matches.pairs_hemisphere_selection("+")
    self.sel_pm = matches.pairs_hemisphere_selection("-")
    self.sel_sp = matches.singles("+")
    self.sel_sm = matches.singles("-")
    if (safety_checks):
      n_conflicting = (   r_free_flags.data().select(self.sel_pp)
                       != r_free_flags.data().select(self.sel_pm)).count(True)
      if (n_conflicting != 0):
        raise RuntimeError(
          "Conflicting R-free flags for %d Bijvoet pair%s."
            % plural_s(n_conflicting))
    self.indices = flex.miller_index()
    self.indices.reserve(
      self.sel_pp.size() + self.sel_sp.size() + self.sel_sm.size())
    self.indices.extend( f_obs.indices().select(self.sel_pp))
    self.indices.extend( f_obs.indices().select(self.sel_sp))
    self.indices.extend(-f_obs.indices().select(self.sel_sm))
    self.f_obs_p, self.f_obs_m = self.split_plus_minus(data=f_obs.data())
    if (f_obs.sigmas() is None):
      self.sigmas_p, self.sigmas_m = flex.double(self.f_obs_p.size(), 1), \
                                     flex.double(self.f_obs_m.size(), 1)
    else:
      self.sigmas_p, self.sigmas_m = self.split_plus_minus(
        data=f_obs.sigmas())
    self.r_free_flags_p, self.r_free_flags_m  = self.split_plus_minus(
      data=r_free_flags.data())
    self.r_free_flags_pm = ~(self.r_free_flags_p | self.r_free_flags_m)
    n_indices = self.indices.size()
    n_pairs = self.sel_pp.size()
    i_singles_plus_minus_boundary = n_pairs + self.sel_sp.size()
    self.present_p = flex.bool()
    self.present_p.reserve(n_indices)
    self.present_p.resize(i_singles_plus_minus_boundary, True)
    self.present_p.resize(n_indices, False)
    self.present_m = flex.bool()
    self.present_m.reserve(n_indices)
    self.present_m.resize(n_pairs, True)
    self.present_m.resize(i_singles_plus_minus_boundary, False)
    self.present_m.resize(n_indices, True)
    self.miller_indices_work = f_obs.indices().select(~r_free_flags.data())
    self.miller_indices_test = f_obs.indices().select( r_free_flags.data())
    self.output = phaser.Output()
    #self.output.setVerbose( b = True, e = True )

  def split_plus_minus(self, data):
    # XXX use reserve
    p = data.select(self.sel_pp)
    m = data.select(self.sel_pm)
    p.extend(data.select(self.sel_sp))
    m.resize(p.size())
    m.extend(data.select(self.sel_sm))
    p.resize(m.size())
    return p, m

  def refine_sad_input_tuple(self, xray_structure, pdb_hierarchy=None):
    inputs = [
      xray_structure.space_group(),
      xray_structure.unit_cell(),
      self.indices,
      self.r_free_flags_pm,
      self.f_obs_p,
      self.f_obs_m,
      self.sigmas_p,
      self.sigmas_m,
      self.present_p,
      self.present_m,
      self.output]
    if (pdb_hierarchy is not None) :
      inputs.append(
        pdb_hierarchy.as_pdb_string(crystal_symmetry=xray_structure))
    return tuple(inputs)

  def target(self,
        xray_structure,
        previous_overall_scaleK=None,
        previous_overall_scaleU=None,
        previous_variances=None,
        pdb_hierarchy=None,
        log=None):
    return target(
      data_adaptor=self,
      xray_structure=xray_structure,
      pdb_hierarchy=pdb_hierarchy,
      previous_overall_scaleK=previous_overall_scaleK,
      previous_overall_scaleU=previous_overall_scaleU,
      previous_variances=previous_variances,
      out=log)

  # Utility method for running substructure completion from other scripts
  # which use CCTBX-centric classes (e.g. mmtbx.f_model).
  # Needs some refactoring.
  def set_ep_auto_data (self, xray_structure, pdb_hierarchy) :
    """
    """
    input = phaser.InputEP_AUTO()
    input.setSPAC_NAME(str(xray_structure.space_group_info()))
    input.setCELL6(xray_structure.unit_cell().parameters())
    input.setCRYS_MILLER(self.indices)
    input.addCRYS_ANOM_DATA("ep",
      "sad",
      self.f_obs_p,
      self.sigmas_p,
      self.present_p,
      self.f_obs_m,
      self.sigmas_m,
      self.present_m)
    # FIXME
    input.setPART_IOTBX(pdb_hierarchy)
    #tmp_file = "tmp_%d.pdb" % os.getpid()
    #tmp = open(tmp_file, "w")
    #tmp.write(pdb_hierarchy.as_pdb_string())
    #tmp.close()
    #input.setPART_PDB(tmp_file)
    #
    input.setPART_VARI("ID")
    input.setPART_DEVI(1.)
    n_res = n_bases = 0
    for chain in pdb_hierarchy.models()[0].chains() :
      main_conf = chain.conformers()[0]
      if (main_conf.is_protein()) :
        n_res += len(main_conf.residues())
      elif (main_conf.is_na()) :
        n_bases += len(main_conf.residues())
    if (n_res > 0) :
      input.addCOMP_PROT_NRES_NUM(n_res, 1)
    if (n_bases > 0) :
      input.addCOMP_NUCL_NRES_NUM(n_bases, 1)
    return input

from libtbx.utils import time_log
timer_init = time_log(label="ml_sad_init")
timer_set_f_calc = time_log(label="ml_sad_set_f_calc")
timer_refine_variance_terms = time_log(label="ml_sad_refine_variance_terms")
timer_functional = time_log(label="ml_sad_functional")
timer_gradients_raw = time_log(label="ml_sad_gradients_raw")
timer_gradients = time_log(label="ml_sad_gradients (raw excl.)")
timer_curvatures_raw = time_log(label="ml_sad_curvatures_raw")
timer_curvatures = time_log(label="ml_sad_curvatures (raw excl.)")

class target(object):

  def __init__(self,
        data_adaptor,
        xray_structure,
        pdb_hierarchy=None,
        previous_overall_scaleK=None,
        previous_overall_scaleU=None,
        previous_variances=None,
        out=None):
    if (out is None) :
      out = sys.stdout
    self.out = out
    timer_init.start()
    self.data_adaptor = data_adaptor
    self.refine_sad_instance = phaser.RefineSAD(
      *data_adaptor.refine_sad_input_tuple(xray_structure=xray_structure,
         pdb_hierarchy=pdb_hierarchy))
    if (previous_overall_scaleK is not None and previous_overall_scaleK > 0):
      self.refine_sad_instance.set_scaleK(previous_overall_scaleK)
    if (previous_overall_scaleU is not None):
      self.refine_sad_instance.set_scaleU(previous_overall_scaleU)
    if (previous_variances is None):
      self.have_variances = False
    else:
      self.refine_sad_instance.set_variance_array(previous_variances)
      self.have_variances = True
    print >> self.out, timer_init.log()

  def reject_outliers(self):
    self.refine_sad_instance.calcIntegrationPoints()
    self.refine_sad_instance.calcOutliers()

  def set_f_calc(self, f_calc) :
    timer_set_f_calc.start()
    self.f_calc = f_calc
    p, m = self.data_adaptor.split_plus_minus(f_calc.data())
    self.refine_sad_instance.setFC(p, flex.conj(m))
    if (not self.have_variances):
      self.refine_sad_instance.initVariancesFromFC()
      self.refine_variance_terms()
      self.have_variances = True
    print >> self.out, timer_set_f_calc.log()

  def refine_variance_terms(self):
    timer_refine_variance_terms.start()
    assert self.f_calc is not None
    protocol = phaser.ProtocolSAD()
    minimizer = phaser.Minimizer()
    # refine SD (plus overall k/B) against the test data
    protocol.setRefineSigmaA(
      fix_sa=True, fix_sb=True, fix_sp=True, fix_sd=False,
      fix_k=False, fix_b=False, fix_pk=False, fix_pb=False)
    #protocol.setNCYC(ncyc=10)
    self.refine_sad_instance.useTestSet()
    minimizer.run(
      self.refine_sad_instance,
      protocol,
      self.data_adaptor.output)
    # refine SP (plus overall k/B) against the working data
    protocol.setRefineSigmaA(
      fix_sa=True, fix_sb=True, fix_sp=False, fix_sd=False,
      fix_k=False, fix_b=False, fix_pk=False, fix_pb=False)
    #protocol.setNCYC(ncyc=10)
    self.refine_sad_instance.useWorkingSet()
    minimizer.run(
      self.refine_sad_instance,
      protocol,
      self.data_adaptor.output)
    va = self.refine_sad_instance.get_variance_array()
    print >> self.out, "LOOK: phaser ScaleK, ScaleU:", \
      self.refine_sad_instance.get_refined_scaleK(), \
      self.refine_sad_instance.get_refined_scaleU()
    print >> self.out, timer_refine_variance_terms.log()

  def functional(self, use_working_set):
    timer_functional.start()
    assert self.f_calc is not None
    if (use_working_set):
      self.refine_sad_instance.useWorkingSet()
    else:
      self.refine_sad_instance.useTestSet()
    self.use_working_set = use_working_set
    result = self.refine_sad_instance.targetFn()
    print >> self.out, timer_functional.log(), result
    return result

  def gradients_raw(self):
    timer_gradients_raw.start()
    assert self.f_calc is not None
    result = self.refine_sad_instance.dL_by_dFC()
    print >> self.out, timer_gradients_raw.log()
    return result

  def curvatures_raw(self):
    timer_curvatures_raw.start()
    assert self.f_calc is not None
    result = self.refine_sad_instance.d2L_by_dFC2()
    print >> self.out, timer_curvatures_raw.log()
    return result

  def llg_map_coeffs (self) :
    fllg = flex.double()
    phllg = flex.double()
    self.refine_variance_terms() #AJM
    llg_coeffs = self.refine_sad_instance.calcGradCoefs(
      "AX", False, fllg, phllg)
    indices = self.refine_sad_instance.miller
    assert (indices.size() == llg_coeffs.size())
    llg_coeffs_array = miller.set(
      crystal_symmetry=self.f_calc,
      indices=indices,
      anomalous_flag=False).array(data=llg_coeffs)
    return llg_coeffs_array

  def derivatives(self, curvs):
    grads_raw = self.gradients_raw()
    timer_gradients.start()
    if (self.use_working_set):
      result_indices = self.data_adaptor.miller_indices_work
    else:
      result_indices = self.data_adaptor.miller_indices_test
    matches = miller.match_indices(result_indices, grads_raw.miller)
    self.number_of_outliers = matches.singles(0)
    assert matches.singles(1).size() == 0
    def raw_as_miller(raw_data, data_type):
      result_data = data_type(result_indices.size())
      result_data.set_selected(
        matches.pairs().column(0),
        raw_data.select(matches.pairs().column(1)))
      return self.f_calc.customized_copy(
        indices=result_indices,
        data=result_data)
    da_db = raw_as_miller(
      raw_data=grads_raw.dL_by_dFC, data_type=flex.complex_double)
    print >> self.out, timer_gradients.log()
    if (not curvs): return da_db
    curvs_raw = self.curvatures_raw()
    timer_curvatures.start()
    assert curvs_raw.miller.all_eq(grads_raw.miller)
    daa_dbb_dab = raw_as_miller(
      raw_data=curvs_raw.d2L_by_dFC2, data_type=flex.vec3_double)
    print >> self.out, timer_curvatures.log()
    return da_db, daa_dbb_dab

  def gradients(self):
    return self.derivatives(curvs=False)