from __future__ import absolute_import, division, print_function
import cctbx.array_family.flex # import dependency
import boost_adaptbx.boost.python as bp
from six.moves import zip
from six.moves import range
ext = bp.import_ext("mmtbx_ncs_ext")
from scitbx.array_family import flex
from cctbx import sgtbx
from libtbx import adopt_init_args
from scitbx import lbfgsb
import math
import scitbx.math
from scitbx.math import matrix
import sys
from scitbx.math import superpose
import mmtbx.alignment
from libtbx.test_utils import approx_equal
from boost_adaptbx import graph
from boost_adaptbx.graph import connected_component_algorithm
import iotbx.pdb

class groups(object):

  def __init__(self,
               pdb_hierarchy,
               crystal_symmetry,
               angular_difference_threshold_deg=5.,
               sequence_identity_threshold=90.,
               quiet=False):
    h = pdb_hierarchy
    superposition_threshold = 2*sequence_identity_threshold - 100.
    n_atoms_all = h.atoms_size()
    s_str = "altloc ' ' and (protein or nucleotide)"
    h = h.select(h.atom_selection_cache().selection(s_str))
    h1 = iotbx.pdb.hierarchy.root()
    h1.append_model(h.models()[0].detached_copy())
    unit_cell = crystal_symmetry.unit_cell()
    result = {}
    if not quiet:
      print("Find groups of chains related by translational NCS")
    # double loop over chains to find matching pairs related by pure translation
    for c1 in h1.chains():
      c1.parent().remove_chain(c1)
      nchains = len(h1.models()[0].chains())
      if([c1.is_protein(), c1.is_na()].count(True)==0): continue
      r1 = list(c1.residues())
      c1_seq = "".join(c1.as_sequence())
      sc_1_tmp = c1.atoms().extract_xyz()
      h1_p1 = h1.expand_to_p1(crystal_symmetry=crystal_symmetry)
      for (ii,c2) in enumerate(h1_p1.chains()):
        orig_c2 = h1.models()[0].chains()[ii%nchains]
        r2 = list(c2.residues())
        c2_seq = "".join(c2.as_sequence())
        sites_cart_1, sites_cart_2 = None,None
        sc_2_tmp = c2.atoms().extract_xyz()
        # chains are identical
        if(c1_seq==c2_seq and sc_1_tmp.size()==sc_2_tmp.size()):
          sites_cart_1 = sc_1_tmp
          sites_cart_2 = sc_2_tmp
          p_identity = 100.
        # chains are not identical, do alignment
        else:
          align_obj = mmtbx.alignment.align(seq_a = c1_seq, seq_b = c2_seq)
          alignment = align_obj.extract_alignment()
          matches = alignment.matches()
          equal = matches.count("|")
          total = len(alignment.a) - alignment.a.count("-")
          p_identity = 100.*equal/max(1,total)
          if(p_identity>superposition_threshold):
            sites_cart_1 = flex.vec3_double()
            sites_cart_2 = flex.vec3_double()
            for i1, i2, match in zip(alignment.i_seqs_a, alignment.i_seqs_b,
                                     matches):
              if(i1 is not None and i2 is not None and match=="|"):
                r1i, r2i = r1[i1], r2[i2]
                assert r1i.resname==r2i.resname, [r1i.resname,r2i.resname,i1,i2]
                for a1 in r1i.atoms():
                  for a2 in r2i.atoms():
                    if(a1.name == a2.name):
                      sites_cart_1.append(a1.xyz)
                      sites_cart_2.append(a2.xyz)
                      break
        # superpose two sequence-aligned chains
        if([sites_cart_1,sites_cart_2].count(None)==0):
          lsq_fit_obj = superpose.least_squares_fit(
            reference_sites = sites_cart_1,
            other_sites     = sites_cart_2)
          angle = lsq_fit_obj.r.rotation_angle()
          t_frac = unit_cell.fractionalize((sites_cart_1-sites_cart_2).mean())
          t_frac = [math.modf(t)[0] for t in t_frac] # put into [-1,1]
          radius = flex.sum(flex.sqrt((sites_cart_1-
            sites_cart_1.mean()).dot()))/sites_cart_1.size()*4./3.
          fracscat = min(c1.atoms_size(),c2.atoms_size())/n_atoms_all
          result.setdefault( frozenset([c1,orig_c2]), [] ).append( [p_identity,[lsq_fit_obj.r, t_frac, angle, radius, fracscat]] )
        else:
          result.setdefault( frozenset([c1,orig_c2]), [] ).append( [p_identity,None] )
    # Build graph
    g = graph.adjacency_list()
    vertex_handle = {}
    for key in result:
      seqid = result[key][0][0]
      sup = min( result[key],key=lambda s:0 if s[1] is None else s[1][2])[1]
      result[key] = [seqid,sup]
      if ((seqid > sequence_identity_threshold) and (sup[2] < angular_difference_threshold_deg)):
        (c1,c2) = key
        if (c1 not in vertex_handle):
          vertex_handle[c1] = g.add_vertex(label=c1)
        if (c2 not in vertex_handle):
          vertex_handle[c2] = g.add_vertex(label=c2)
        g.add_edge(vertex1=vertex_handle[c1],vertex2=vertex_handle[c2])
    # Do connected component analysis and compose final tNCS pairs object
    components = connected_component_algorithm.connected_components(g)
    import itertools
    self.ncs_pairs = []
    self.tncsresults = [ 0, "", [], 0.0 ]
    for (i,group) in enumerate(components):
      chains = [g.vertex_label(vertex=v) for v in group]
      fracscats = []
      radii = []
      for pair in itertools.combinations(chains,2):
        sup = result[frozenset(pair)][1]
        fracscats.append(sup[-1])
        radii.append(sup[-2])
      fs = sum(fracscats)/len(fracscats)
      self.tncsresults[3] = fs # store fracscat in array
      rad = sum(radii)/len(radii)
      #import code, traceback; code.interact(local=locals(), banner="".join( traceback.format_stack(limit=10) ) )
      maxorder = 1
      vectors = []
      previous_id = next(itertools.combinations(chains,2))[0].id
      for pair in itertools.combinations(chains,2):
        sup = result[frozenset(pair)][1]
        ncs_pair = ext.pair(
          r = sup[0],
          t = sup[1],
          radius = rad,
          radius_estimate = rad,
          fracscat = fs,
          rho_mn = flex.double(), # rho_mn undefined, needs to be set later
          id = i)
        self.ncs_pairs.append(ncs_pair)
        # show tNCS pairs in group
        fmt="group %d chains %s <> %s angle: %4.2f trans.vect.: (%s) fracscat: %5.3f"
        t = ",".join([("%6.3f"%t_).strip() for t_ in sup[1]]).strip()
        if not quiet:
          print(fmt%(i, pair[0].id, pair[1].id, sup[2], t, fs))
        if pair[0].id == previous_id:
          maxorder += 1
          orthoxyz = unit_cell.orthogonalize( sup[1] )
          vectors.append(( sup[1], orthoxyz, sup[2] ))
        else:
          previous_id = pair[0].id
          maxorder = 1
          vectors = []
        if maxorder > self.tncsresults[0]:
          self.tncsresults[0] = maxorder
          self.tncsresults[1] = previous_id
          self.tncsresults[2] = vectors
    if not quiet:
      print("Largest TNCS order, peptide chain, fracvector, orthvector, angle, fracscat = ", \
       str(self.tncsresults))
    #import code, traceback; code.interact(local=locals(), banner="".join( traceback.format_stack(limit=10) ) )



def initialize_rho_mn(ncs_pairs, d_spacings_data, binner, rms=0.5):
  """
  Initialize rho_mn
    rhoMN = exp(-(2*pi^2/3)*(rms/d)^2, and rms=0.4-0.8 is probably a good guess.
  """
  n_bins = binner.n_bins_used()
  rho_mn_initial = flex.double(n_bins, 0)
  cntr=0
  for i_bin in binner.range_used():
    sel_bin = binner.selection(i_bin)
    if(sel_bin.count(True)>0):
      arg = (2*math.pi**2/3)*(rms/flex.mean(d_spacings_data.select(sel_bin)))**2
      rho_mn_initial[cntr] = math.exp(-1*arg)
    cntr+=1
  for p in ncs_pairs:
    p.set_rhoMN(rho_mn_initial)

def lbfgs_run(target_evaluator, use_bounds, lower_bound, upper_bound,
       max_iterations=None):
  minimizer = lbfgsb.minimizer(
    n   = target_evaluator.n,
    #factr=1.e+1, XXX Affects speed significantly
    l   = lower_bound, # lower bound
    u   = upper_bound, # upper bound
    nbd = flex.int(target_evaluator.n, use_bounds)) # flag to apply both bounds
  minimizer.error = None
  try:
    icall = 0
    while 1:
      icall += 1
      x, f, g = target_evaluator()
      #print "x,f:", ",".join(["%6.3f"%x_ for x_ in x]), f, icall
      have_request = minimizer.process(x, f, g)
      if(have_request):
        requests_f_and_g = minimizer.requests_f_and_g()
        continue
      assert not minimizer.requests_f_and_g()
      if(minimizer.is_terminated()): break
      # XXX temp fix for python3 failure (run_lbfgs in tncs)
      # Failure is that max_iterations is None
      # temp fix is to break if max_iterations is None or icall>max_iterations
      #if(icall>max_iterations): break
      if ((max_iterations is None) or (icall>max_iterations)): break
  except RuntimeError as e:
    minimizer.error = str(e)
  minimizer.n_calls = icall
  return minimizer

class minimizer(object):

  def __init__(self,
               potential,
               use_bounds,
               lower_bound,
               upper_bound,
               initial_values):
    adopt_init_args(self, locals())
    self.x = initial_values
    self.n = self.x.size()

  def run(self):
    self.minimizer = lbfgs_run(
      target_evaluator=self,
      use_bounds=self.use_bounds,
      lower_bound = self.lower_bound,
      upper_bound = self.upper_bound)
    self()
    return self

  def __call__(self):
    self.potential.update(x = self.x)
    self.f = self.potential.target()
    self.g = self.potential.gradient()
    return self.x, self.f, self.g

class potential(object):

  def __init__(self, f_obs, ncs_pairs, reflections_per_bin):
    adopt_init_args(self, locals())
    # Create bins
    f_obs.setup_binner(reflections_per_bin = reflections_per_bin)
    self.binner = f_obs.binner()
    n_bins = self.binner.n_bins_used()
    self.n_bins = n_bins
    self.SigmaN = None
    self.update_SigmaN()
    #
    self.rbin = flex.int(f_obs.data().size(), -1)
    for i_bin in self.binner.range_used():
      for i_seq in self.binner.array_indices(i_bin):
        self.rbin[i_seq] = i_bin-1 # i_bin starts with 1, not 0 !
    assert flex.min(self.rbin)==0
    assert flex.max(self.rbin)==n_bins-1
    # Extract symmetry matrices
    self.sym_matrices = []
    for m_as_string in f_obs.space_group().smx():
      o = sgtbx.rt_mx(symbol=str(m_as_string), t_den=f_obs.space_group().t_den())
      m_as_double = o.r().as_double()
      self.sym_matrices.append(m_as_double)
    self.gradient_evaluator = None
    self.target_and_grads = ext.tncs_eps_factor_refinery(
        tncs_pairs               = self.ncs_pairs,
        f_obs                    = self.f_obs.data(),
        sigma_f_obs              = self.f_obs.sigmas(),
        rbin                     = self.rbin,
        SigmaN                   = self.SigmaN,
        space_group              = self.f_obs.space_group(),
        miller_indices           = self.f_obs.indices(),
        fractionalization_matrix = self.f_obs.unit_cell().fractionalization_matrix(),
        sym_matrices             = self.sym_matrices)

    self.update()

  def update(self, x=None):
    if(self.gradient_evaluator=="rhoMN"):
      size = len(self.ncs_pairs)
      for i, ncs_pair in enumerate(self.ncs_pairs):
        ncs_pair.set_rhoMN(x[i*self.n_bins:(i+1)*self.n_bins])
      self.target_and_grads.update_pairs(self.ncs_pairs)
    elif(self.gradient_evaluator=="radius"):
      for ncs_pair, x_ in zip(self.ncs_pairs, x):
        ncs_pair.set_radius(x_)
      self.target_and_grads = ext.tncs_eps_factor_refinery(
        tncs_pairs               = self.ncs_pairs,
        f_obs                    = self.f_obs.data(),
        sigma_f_obs              = self.f_obs.sigmas(),
        rbin                     = self.rbin,
        SigmaN                   = self.SigmaN,
        space_group              = self.f_obs.space_group(),
        miller_indices           = self.f_obs.indices(),
        fractionalization_matrix = self.f_obs.unit_cell().fractionalization_matrix(),
        sym_matrices             = self.sym_matrices)
      self.target_and_grads.set_compute_gradients_radius()

  def update_SigmaN(self):
    if(self.SigmaN is None):
      eps = self.f_obs.epsilons().data().as_double()
    else:
      eps = self.target_and_grads.tncs_epsfac()
    self.SigmaN = flex.double(self.f_obs.data().size(), 0)
    for i_bin in self.binner.range_used():
      bin_sel = self.f_obs.binner().selection(i_bin)
      f_obs_bin = self.f_obs.select(bin_sel)
      f_obs_bin_data = f_obs_bin.data()
      f_obs_bin_data_size = f_obs_bin_data.size()
      if(f_obs_bin_data_size>0):
        eps_bin = eps.select(bin_sel)
        sn = flex.sum(f_obs_bin_data*f_obs_bin_data/eps_bin)/f_obs_bin_data_size
        self.SigmaN = self.SigmaN.set_selected(bin_sel, sn)
    assert self.SigmaN.all_gt(0)

  def set_refine_radius(self):
    self.gradient_evaluator = "radius"
    self.target_and_grads.set_compute_gradients_radius()
    return self

  def set_refine_rhoMN(self):
    self.gradient_evaluator = "rhoMN"
    self.target_and_grads.set_compute_gradients_rho_mn()
    return self

  def target(self):
    return self.target_and_grads.target()

  def gradient(self):
    if(self.gradient_evaluator=="rhoMN"):
      return self.target_and_grads.gradient_rhoMN()
    elif(self.gradient_evaluator=="radius"):
      return self.target_and_grads.gradient_radius()
    else: assert 0

def finite_differences_grad_radius(ncs_pairs, f_obs, reflections_per_bin,
      tolerance):
  reflections_per_bin = min(f_obs.data().size(), reflections_per_bin)
  f_obs.setup_binner(reflections_per_bin = reflections_per_bin)
  binner = f_obs.binner()
  n_bins = binner.n_bins_used()
  #
  radii = flex.double()
  for ncs_pair in ncs_pairs:
    radii.append(ncs_pair.radius)
  #
  pot = potential(f_obs = f_obs, ncs_pairs = ncs_pairs,
      reflections_per_bin = reflections_per_bin)
  pot = pot.set_refine_radius()
  t = pot.target()
  g_exact = pot.gradient()
  #print "Exact:", list(g_exact)
  #
  eps = 1.e-4
  #
  g_fd = []
  for i, rad in enumerate(radii):
    radii_p = radii.deep_copy()
    radii_m = radii.deep_copy()
    radii_p[i] = radii[i]+eps
    radii_m[i] = radii[i]-eps
    #
    pot.update(x = flex.double(radii_p))
    t1 = pot.target()
    #
    pot.update(x = flex.double(radii_m))
    t2 = pot.target()
    #
    g_fd_ = (t1-t2)/(2*eps)
    g_fd.append(g_fd_)
  #print "Finite diff.:",g_fd
  relative_error = flex.double()
  for g1,g2 in zip(g_exact, g_fd):
    #print "exact: %10.6f fd: %10.6f"%(g1,g2)
    relative_error.append( abs((g1-g2)/(g1+g2))*2.*100. )
  mmm = relative_error.min_max_mean().as_tuple()
  print("min/max/mean of |(g_exact-g_fd)/(g_exact+g_fd)|*100.*2:",\
    "%6.4f %6.4f %6.4f"%mmm)
  # assert approx_equal(mmm, [0,0,0], tolerance) # reinstate after proper constraints

def finite_differences_rho_mn(ncs_pairs, f_obs, reflections_per_bin,
      tolerance):
  reflections_per_bin = min(f_obs.data().size(), reflections_per_bin)
  f_obs.setup_binner(reflections_per_bin = reflections_per_bin)
  binner = f_obs.binner()
  n_bins = binner.n_bins_used()
  #
  pot = potential(f_obs = f_obs, ncs_pairs = ncs_pairs,
      reflections_per_bin = reflections_per_bin)
  pot = pot.set_refine_rhoMN()
  t = pot.target()
  g_exact = pot.gradient()
  #
  rho_mn = flex.double()
  for p in ncs_pairs:
    rho_mn.extend(p.rho_mn)
  #
  eps = 1.e-6
  #
  g_fd = []
  for i, rho_mn_i in enumerate(rho_mn):
    rho_mn_p = rho_mn.deep_copy()
    rho_mn_p[i] = rho_mn_i + eps
    rho_mn_m = rho_mn.deep_copy()
    rho_mn_m[i] = rho_mn_i - eps
    #
    pot.update(x = rho_mn_p)
    t1 = pot.target()
    #
    pot.update(x = rho_mn_m)
    t2 = pot.target()
    #
    g_fd_ = (t1-t2)/(2*eps)
    g_fd.append(g_fd_)
  relative_error = flex.double()
  for g1,g2 in zip(g_exact, g_fd):
    #print "exact: %10.6f fd: %10.6f"%(g1,g2)
    relative_error.append( abs((g1-g2)/(g1+g2))*2.*100. )
  mmm = relative_error.min_max_mean().as_tuple()
  print("min/max/mean of |(g_exact-g_fd)/(g_exact+g_fd)|*100.*2:",\
    "%6.4f %6.4f %6.4f"%mmm)
  # assert approx_equal(mmm, [0,0,0], tolerance) # reinstate after proper constraints

class compute_eps_factor(object):

  def __init__(self, f_obs, pdb_hierarchy, reflections_per_bin):
    f_obs = f_obs.deep_copy()
    if(not f_obs.sigmas_are_sensible()):
      f_obs.set_sigmas(sigmas = flex.double(f_obs.data().size(), 0.0))
    reflections_per_bin = min(f_obs.data().size(), reflections_per_bin)
    f_obs.setup_binner(reflections_per_bin = reflections_per_bin)
    self.unit_cell = f_obs.unit_cell()
    #
    self.ncs_pairs = groups(
      pdb_hierarchy    = pdb_hierarchy,
      crystal_symmetry = f_obs.crystal_symmetry()).ncs_pairs
    initialize_rho_mn(
      ncs_pairs       = self.ncs_pairs,
      d_spacings_data = f_obs.d_spacings().data(),
      binner          = f_obs.binner())
    self.epsfac = None
    if(len(self.ncs_pairs)>0):
      # Radii
      radii = flex.double()
      rad_lower_bound = flex.double()
      rad_upper_bound = flex.double()
      for ncs_pair in self.ncs_pairs:
        radii.append(ncs_pair.radius)
        rad_lower_bound.append(ncs_pair.radius/3)
        rad_upper_bound.append(ncs_pair.radius*3)
      # Target and gradients evaluator
      pot = potential(f_obs = f_obs, ncs_pairs = self.ncs_pairs,
        reflections_per_bin = reflections_per_bin)
      for it in range(2):
        # refine eps fac
        rho_mn = flex.double()
        for ncs_pair in self.ncs_pairs:
          rho_mn.extend(ncs_pair.rho_mn)
        m = minimizer(
          potential      = pot.set_refine_rhoMN(),
          use_bounds     = 2,
          lower_bound    = flex.double(rho_mn.size(), 0.),
          upper_bound    = flex.double(rho_mn.size(), 1.),
          initial_values = rho_mn).run()
        # refine radius
        radii = flex.double()
        for ncs_pair in self.ncs_pairs:
          radii.append(ncs_pair.radius)
        m = minimizer(
          potential      = pot.set_refine_radius(),
          use_bounds     = 2,
          lower_bound    = rad_lower_bound,
          upper_bound    = rad_upper_bound,
          initial_values = radii).run()
      self.epsfac = pot.target_and_grads.tncs_epsfac()

  def show_summary(self, log=None):
    if(self.epsfac is None): return None
    if(log is None): log = sys.stdout
    for i, ncs_pair in enumerate(self.ncs_pairs):
      print("tNCS pair: %d"%i, file=log)
      print("  Group ID:", ncs_pair.id, file=log)
      angle = matrix.sqr(ncs_pair.r).rotation_angle()
      t = ",".join([("%6.3f"%t_).strip() for t_ in ncs_pair.t]).strip()
      t_cart = ",".join([("%6.3f"%t_).strip()
        for t_ in self.unit_cell.orthogonalize(ncs_pair.t)]).strip()
      r = ",".join([("%8.6f"%r_).strip() for r_ in ncs_pair.r]).strip()
      print("  Translation (fractional): (%s)"%t, file=log)
      print("  Translation (Cartesian):  (%s)"%t_cart, file=log)
      print("  Rotation (deg): %-5.2f"%angle, file=log)
      print("  Rotation matrix: (%s)"%r, file=log)
      print("  Radius: %-6.3f"%ncs_pair.radius, file=log)
      print("  Radius (estimate): %-6.1f"%ncs_pair.radius_estimate, file=log)
      print("  fracscat:", ncs_pair.fracscat, file=log)
    print("tNCS eps factor: min,max,mean: %6.4f %6.4f %6.4f"%\
      self.epsfac.min_max_mean().as_tuple(), file=log)




if __name__ == '__main__':
  xtal = iotbx.pdb.input(file_name= sys.argv[1] )
  hroot = xtal.construct_hierarchy()
  xtalsym = xtal.crystal_symmetry()
  groups(hroot, xtalsym, float(sys.argv[2]))
