from __future__ import absolute_import, division, print_function
from scitbx.lstbx import normal_eqns_solving
from cctbx import geometry_restraints, adp_restraints, sgtbx, adptbx
from cctbx.array_family import flex
from cctbx.xray import parameter_map
from smtbx.refinement import restraints
from smtbx.refinement.restraints.adp_restraints import\
     adp_similarity_restraints, isotropic_adp_restraints,\
     rigid_bond_restraints, fixed_u_eq_adp_restraints,\
     adp_u_eq_similarity_restraints, adp_volume_similarity_restraints

from cctbx.adp_restraints import adp_restraint_params
import smtbx.utils
import smtbx.development
from smtbx.refinement import constraints, least_squares
from libtbx.test_utils import approx_equal
from libtbx.utils import wall_clock_time
import libtbx
from scitbx import matrix
from six.moves import range

geom = geometry_restraints
adp = adp_restraints

rows_per_restraint = {
  geom.bond_similarity_proxy: 6,
  adp.adp_similarity_proxy: 6,
  adp.isotropic_adp_proxy: 6,
  }

class restraints_test_case:

  def __init__(self):
    self.xray_structure = smtbx.development.sucrose()
    self.tolerance = 1e-4
    for sc in self.xray_structure.scatterers():
      sc.flags.set_grad_site(True)
      if sc.flags.use_u_aniso(): sc.flags.set_grad_u_aniso(True)
      if sc.flags.use_u_iso(): sc.flags.set_grad_u_iso(True)

    self.param_map = parameter_map(self.xray_structure.scatterers())
    assert self.proxies.size() > 0

  def run(self):
    self.exercise_ls_restraints()

  def exercise_ls_restraints(self):
    xs = self.xray_structure.deep_copy_scatterers()
    linearised_eqns = self.manager.build_linearised_eqns(xs, xs.parameter_map())
    design_matrix = linearised_eqns.design_matrix.as_dense_matrix()
    fd_design = flex.double()
    for proxy in self.proxies:
      grads = self.fd_grads(proxy)
      for i, grad in enumerate(grads):
        fd_design.extend(grad)
    assert approx_equal(design_matrix, fd_design, self.tolerance)

class geometry_restraints_test_case(restraints_test_case):

  def exercise_ls_restraints(self):
    match = restraints_test_case.exercise_ls_restraints(self)

  def fd_grads(self, proxy):
    grads = flex.double(self.param_map.n_parameters)
    eps = 1e-8
    uc = self.xray_structure.unit_cell()
    sites_cart = self.xray_structure.sites_cart().deep_copy()
    for i in range(self.param_map.n_scatterers):
      grad_site_cart = [0,0,0]
      for j in range(3):
        h = [0,0,0]
        h[j] = eps
        h = matrix.col(h)
        sites_cart[i] = matrix.col(sites_cart[i])+h
        r = self.restraint_t(uc, sites_cart, proxy)
        d1 = r.delta
        sites_cart[i] = matrix.col(sites_cart[i])-2*h
        r = self.restraint_t(uc, sites_cart, proxy)
        d2 = r.delta
        d_delta = (d1-d2)/(2*eps)
        grad_site_cart[j] = d_delta
      grad_site_frac = uc.fractionalize_gradient(grad_site_cart)
      for j in range(3):
        grads[self.param_map[i].site+j] = grad_site_frac[j]
    return [grads]

class bond_restraint_test_case(geometry_restraints_test_case):
  manager = restraints.manager(
    bond_proxies = geometry_restraints.shared_bond_simple_proxy([
      geom.bond_simple_proxy((0,30), 1.42, 1),
      geom.bond_simple_proxy((1,21), 1.42, 1)
    ]))
  proxies = manager.bond_proxies
  restraint_t = geom.bond

class angle_restraint_test_case(geometry_restraints_test_case):
  manager = restraints.manager(
    angle_proxies = geometry_restraints.shared_angle_proxy([
      geom.angle_proxy((30, 0, 19), 115, 1),
      geom.angle_proxy((21, 1, 2), 110, 1)
    ]))
  proxies = manager.angle_proxies
  restraint_t = geom.angle

class dihedral_restraint_test_case(geometry_restraints_test_case):
  manager = restraints.manager(
    dihedral_proxies = geometry_restraints.shared_dihedral_proxy([
      geom.dihedral_proxy((21, 19, 24, 26), 180, 1),
      geom.dihedral_proxy((5, 26, 28, 7), 60, 1)
    ]))
  proxies = manager.dihedral_proxies
  restraint_t = geom.dihedral

class adp_restraints_test_case(restraints_test_case):

  def __init__(self):
    restraints_test_case.__init__(self)

  def fd_grads(self, proxy):
    dynamic_restraint_proxy_classes = (
      adp.adp_u_eq_similarity_proxy,
      adp.adp_volume_similarity_proxy,
    )
    if isinstance(proxy, (dynamic_restraint_proxy_classes)):
      n_restraints = len(proxy.i_seqs)
    else:
      n_restraints = rows_per_restraint.get(proxy.__class__, 1)
    grads = [flex.double(self.param_map.n_parameters) for i in range(n_restraints)]
    eps = 1e-8
    uc = self.xray_structure.unit_cell()
    xs = self.xray_structure
    u_cart = xs.scatterers().extract_u_cart(uc).deep_copy()
    u_star = xs.scatterers().extract_u_star().deep_copy()
    u_iso = xs.scatterers().extract_u_iso().deep_copy()
    single_delta_classes = (
      adp.fixed_u_eq_adp,
    )
    for n in range(n_restraints):
      for i in range(self.param_map.n_scatterers):
        use_u_aniso = self.param_map[i].u_aniso > -1
        use_u_iso = self.param_map[i].u_iso > -1
        for j in range(6):
          if use_u_aniso:
            h = [0,0,0,0,0,0]
            h[j] = eps
            h = matrix.sym(sym_mat3=h)
            u_star[i]=list((matrix.sym(sym_mat3=u_star[i]) + h).as_sym_mat3())
            r = self.restraint(proxy, u_cart=flex.sym_mat3_double([
              adptbx.u_star_as_u_cart(uc, u) for u in u_star]))
            if isinstance(r, adp.rigid_bond):
              d1 = r.delta_z()
            elif isinstance(r, single_delta_classes):
              d1 = r.delta()
            else:
              d1 = r.deltas()[n]
            u_star[i]=list((matrix.sym(sym_mat3=u_star[i]) - 2*h).as_sym_mat3())
            r = self.restraint(proxy, u_cart=flex.sym_mat3_double([
              adptbx.u_star_as_u_cart(uc, u) for u in u_star]))
            if isinstance(r, adp.rigid_bond):
              d2 = r.delta_z()
            elif isinstance(r, single_delta_classes):
              d2 = r.delta()
            else:
              d2 = r.deltas()[n]
          elif use_u_iso:
            u_iso[i] += eps
            r = self.restraint(proxy, u_iso=u_iso)
            if isinstance(r, adp.rigid_bond):
              d1 = r.delta_z()
            elif isinstance(r, single_delta_classes):
              d1 = r.delta()
            else:
              d1 = r.deltas()[n]
            u_iso[i] -= 2*eps
            r = self.restraint(proxy, u_iso=u_iso)
            if isinstance(r, adp.rigid_bond):
              d2 = r.delta_z()
            elif isinstance(r, single_delta_classes):
              d2 = r.delta()
            else:
              d2 = r.deltas()[n]
          d_delta = (d1-d2)/(2*eps)
          if not isinstance(r, adp.rigid_bond) and j > 2:
            d_delta *= 2 # off diagonals count twice
          if use_u_aniso:
            grads[n][self.param_map[i].u_aniso+j] = d_delta
          elif use_u_iso:
            grads[n][self.param_map[i].u_iso] = d_delta
            break
    return grads

class isotropic_adp_test_case(adp_restraints_test_case):
  proxies = isotropic_adp_restraints(
    xray_structure=smtbx.development.sucrose()).proxies
  # no need to test all of them every time
  proxies = adp.shared_isotropic_adp_proxy(
    flex.select(proxies, flags=flex.random_bool(proxies.size(), 0.5)))
  manager = restraints.manager(isotropic_adp_proxies=proxies)

  def restraint(self, proxy, u_iso=None, u_cart=None):
    if u_cart is None:
      u_cart=self.xray_structure.scatterers().extract_u_cart(
        self.xray_structure.unit_cell())
    return adp.isotropic_adp(
      adp_restraint_params(u_cart=u_cart),
      proxy)

class fixed_u_eq_adp_test_case(adp_restraints_test_case):
  proxies = fixed_u_eq_adp_restraints(
    xray_structure=smtbx.development.sucrose(),
    u_eq_ideal=0.025).proxies
  # no need to test all of them every time
  proxies = adp.shared_fixed_u_eq_adp_proxy(
    flex.select(proxies, flags=flex.random_bool(proxies.size(), 0.5)))
  manager = restraints.manager(fixed_u_eq_adp_proxies=proxies)

  def restraint(self, proxy, u_iso=None, u_cart=None):
    if u_cart is None:
      u_cart=self.xray_structure.scatterers().extract_u_cart(
        self.xray_structure.unit_cell())
    return adp.fixed_u_eq_adp(
      adp_restraint_params(u_cart=u_cart),
      proxy)

class adp_similarity_test_case(adp_restraints_test_case):
  proxies = adp_similarity_restraints(
    xray_structure=smtbx.development.sucrose()).proxies
  # no need to test all of them every time
  proxies = adp.shared_adp_similarity_proxy(
    flex.select(proxies, flags=flex.random_bool(proxies.size(), 0.5)))
  manager = restraints.manager(adp_similarity_proxies=proxies)

  def restraint(self, proxy, u_iso=None, u_cart=None):
    if u_cart is None:
      u_cart=self.xray_structure.scatterers().extract_u_cart(
        self.xray_structure.unit_cell())
    if u_iso is None:
      u_iso=self.xray_structure.scatterers().extract_u_iso()
    use_u_aniso=self.xray_structure.use_u_aniso()
    return adp.adp_similarity(
      adp_restraint_params(u_cart=u_cart, u_iso=u_iso, use_u_aniso=use_u_aniso),
      proxy)

class adp_u_eq_similarity_test_case(adp_restraints_test_case):
  proxies = adp_u_eq_similarity_restraints(
    xray_structure=smtbx.development.sucrose()).proxies
  # no need to test all of them every time
  #proxies = adp.shared_adp_u_eq_similarity_proxy(
    #flex.select(proxies, flags=flex.random_bool(proxies.size(), 0.5)))
  manager = restraints.manager(adp_u_eq_similarity_proxies=proxies)

  def restraint(self, proxy, u_iso=None, u_cart=None):
    if u_cart is None:
      u_cart=self.xray_structure.scatterers().extract_u_cart(
        self.xray_structure.unit_cell())
    if u_iso is None:
      u_iso=self.xray_structure.scatterers().extract_u_iso()
    use_u_aniso=self.xray_structure.use_u_aniso()
    return adp.adp_u_eq_similarity(
      adp_restraint_params(u_cart=u_cart, u_iso=u_iso, use_u_aniso=use_u_aniso),
      proxy)

class adp_volume_similarity_test_case(adp_restraints_test_case):
  proxies = adp_volume_similarity_restraints(
    xray_structure=smtbx.development.sucrose()).proxies
  manager = restraints.manager(adp_volume_similarity_proxies=proxies)
  def __init__(self):
    adp_restraints_test_case.__init__(self)
    # eigen values and eigen vectors are dependent after all...
    # may need to make smaller
    self.tolerance = 0.2
  def restraint(self, proxy, u_iso=None, u_cart=None):
    if u_cart is None:
      u_cart=self.xray_structure.scatterers().extract_u_cart(
        self.xray_structure.unit_cell())
    if u_iso is None:
      u_iso=self.xray_structure.scatterers().extract_u_iso()
    use_u_aniso=self.xray_structure.use_u_aniso()
    return adp.adp_volume_similarity(
      adp_restraint_params(u_cart=u_cart, u_iso=u_iso, use_u_aniso=use_u_aniso),
      proxy)

class rigid_bond_test_case(adp_restraints_test_case):
  proxies = rigid_bond_restraints(
    xray_structure=smtbx.development.sucrose()).proxies
  # no need to test all of them every time
  proxies = adp.shared_rigid_bond_proxy(
    flex.select(proxies, flags=flex.random_bool(proxies.size(), 0.3)))
  manager = restraints.manager(rigid_bond_proxies=proxies)

  def restraint(self, proxy, u_iso=None, u_cart=None):
    if u_cart is None:
      u_cart = self.xray_structure.scatterers().extract_u_cart(
        self.xray_structure.unit_cell())
    sites_cart = self.xray_structure.sites_cart()
    return adp.rigid_bond(
      adp_restraint_params(sites_cart=sites_cart, u_cart=u_cart),
      proxy)

def exercise_restrained_refinement(options):
  import random
  random.seed(1)
  flex.set_random_seed(1)
  xs0 = smtbx.development.random_xray_structure(
    sgtbx.space_group_info('P1'),
    n_scatterers=options.n_scatterers,
    elements="random")
  for sc in xs0.scatterers():
    sc.flags.set_grad_site(True)
  sc0 = xs0.scatterers()
  uc = xs0.unit_cell()

  mi = xs0.build_miller_set(anomalous_flag=False, d_min=options.resolution)
  fo_sq = mi.structure_factors_from_scatterers(
    xs0, algorithm="direct").f_calc().norm()
  fo_sq = fo_sq.customized_copy(sigmas=flex.double(fo_sq.size(), 1))

  i, j, k, l = random.sample(range(options.n_scatterers), 4)
  bond_proxies = geometry_restraints.shared_bond_simple_proxy()
  w = 1e9
  d_ij = uc.distance(sc0[i].site, sc0[j].site)*0.8
  bond_proxies.append(geom.bond_simple_proxy(
    i_seqs=(i, j),
    distance_ideal=d_ij,
    weight=w))
  d_jk = uc.distance(sc0[j].site, sc0[k].site)*0.85
  bond_proxies.append(geom.bond_simple_proxy(
    i_seqs=(j, k),
    distance_ideal=d_jk,
    weight=w))
  d_ki = min(uc.distance(sc0[k].site, sc0[i].site)*0.9, (d_ij + d_jk)*0.8)
  bond_proxies.append(geom.bond_simple_proxy(
    i_seqs=(k, i),
    distance_ideal=d_ki,
    weight=w))
  d_jl = uc.distance(sc0[j].site, sc0[l].site)*0.9
  bond_proxies.append(geom.bond_simple_proxy(
    i_seqs=(j, l),
    distance_ideal=d_jl,
    weight=w))
  d_lk = min(uc.distance(sc0[l].site, sc0[k].site)*0.8, 0.75*(d_jk + d_jl))
  bond_proxies.append(geom.bond_simple_proxy(
    i_seqs=(l, k),
    distance_ideal=d_lk,
    weight=w))
  restraints_manager = restraints.manager(bond_proxies=bond_proxies)

  xs1 = xs0.deep_copy_scatterers()
  xs1.shake_sites_in_place(rms_difference=0.1)

  def ls_problem():
    xs = xs1.deep_copy_scatterers()
    reparametrisation = constraints.reparametrisation(
      structure=xs,
      constraints=[],
      connectivity_table=smtbx.utils.connectivity_table(xs),
      temperature=20)
    return least_squares.crystallographic_ls(
      fo_sq.as_xray_observations(),
      reparametrisation=reparametrisation,
      restraints_manager=restraints_manager)

  gradient_threshold, step_threshold = 1e-6, 1e-6
  eps = 5e-3

  ls = ls_problem()
  t = wall_clock_time()
  cycles = normal_eqns_solving.naive_iterations(
    ls,
    gradient_threshold=gradient_threshold,
    step_threshold=step_threshold,
    track_all=True)
  if options.verbose:
    print("%i %s steps in %.6f s" % (cycles.n_iterations, cycles, t.elapsed()))
  sc = ls.xray_structure.scatterers()
  for p in bond_proxies:
    d = uc.distance(*[ sc[i_pair].site for i_pair in p.i_seqs ])
    assert approx_equal(d, p.distance_ideal, eps)

  ls = ls_problem()
  t = wall_clock_time()
  cycles = normal_eqns_solving.levenberg_marquardt_iterations(
    ls,
    gradient_threshold=gradient_threshold,
    step_threshold=step_threshold,
    tau=1e-3,
    track_all=True)
  if options.verbose:
    print("%i %s steps in %.6f s" % (cycles.n_iterations, cycles, t.elapsed()))
  sc = ls.xray_structure.scatterers()
  sc = ls.xray_structure.scatterers()
  for p in bond_proxies:
    d = uc.distance(*[ sc[i].site for i in p.i_seqs ])
    assert approx_equal(d, p.distance_ideal, eps)

def exercise_add_equation():
  linearised_eqns = restraints.linearised_eqns_of_restraint(10, 10)
  delta = 0.5
  grads = flex.double((0,0,1,0,0,2,0,0,-1, 0))
  w = 10
  linearised_eqns.add_equation(delta, grads, w)
  assert linearised_eqns.n_restraints() == 1
  linearised_eqns.add_equation(delta, grads, w)
  linearised_eqns.add_equation(delta, grads, w)
  assert linearised_eqns.n_restraints() == 3
  from scitbx import sparse
  assert approx_equal(
    linearised_eqns.design_matrix.as_dense_matrix(),
    sparse.matrix(rows=10, columns=10,
                  elements_by_columns=[ { 0: 0, 1: 0, 2: 0 },
                                        { 0: 0, 1: 0, 2: 0 },
                                        { 0: 1, 1: 1, 2: 1 },
                                        { 0: 0, 1: 0, 2: 0 },
                                        { 0: 0, 1: 0, 2: 0 },
                                        { 0: 2, 1: 2, 2: 2 },
                                        { 0: 0, 1: 0, 2: 0 },
                                        { 0: 0, 1: 0, 2: 0 },
                                        { 0: -1, 1: -1, 2: -1 },
                                        { 0: 0, 1: 0, 2: 0 }, ]).as_dense_matrix())


def exercise_ls_restraints(options):
  exercise_add_equation()
  exercise_restrained_refinement(options)
  bond_restraint_test_case().run()
  angle_restraint_test_case().run()
  dihedral_restraint_test_case().run()

  isotropic_adp_test_case().run()
  adp_similarity_test_case().run()
  rigid_bond_test_case().run()
  fixed_u_eq_adp_test_case().run()
  adp_u_eq_similarity_test_case().run()
  adp_volume_similarity_test_case().run()

def run():
  libtbx.utils.show_times_at_exit()
  import sys
  from libtbx.option_parser import option_parser
  command_line = (option_parser()
    .option(None, "--verbose",
            action="store_true")
    .option(None, "--scatterers",
            dest='n_scatterers',
            type="int",
            default=5)
    .option(None, "--resolution",
            type="float",
            default=0.2)
  ).process(args=sys.argv[1:])
  exercise_ls_restraints(command_line.options)

if __name__ == '__main__':
  run()