from __future__ import absolute_import, division, print_function
import os
from scitbx.array_family import flex
from cctbx.array_family import flex as cctbx_flex
from cctbx import sgtbx, miller
from libtbx import easy_mp, Auto
from scipy import sparse
import numpy as np
from orderedset import OrderedSet
import copy

# Specialization, run only a subset of cosym steps and include plot
from dials.algorithms.symmetry.cosym import CosymAnalysis as BaseClass
from dials.util.multi_dataset_handling import select_datasets_on_identifiers

from dials.algorithms.symmetry.cosym.target import Target

from cctbx.sgtbx.literal_description import literal_description

class CosymAnalysis(BaseClass):

  def __init__(self, *args, **kwargs):
    self.do_plot = kwargs.pop('do_plot', False)
    self.i_plot = kwargs.pop('i_plot', None)
    self.plot_fname = kwargs.pop('plot_fname', None)
    self.plot_format = kwargs.pop('plot_format', None)
    self.output_dir = kwargs.pop('output_dir', None)
    self.cb_op = kwargs.pop('cb_op', None)
    super(CosymAnalysis, self).__init__(*args, **kwargs)

  def run(self):
    super(CosymAnalysis, self).run()
    if self.do_plot: self.plot_after_cluster_analysis()
    # we have the xy embedded coords at this point.

  def plot_after_optimize(self):
          print ("optimized coordinates", self.coords.shape)
          xx = []
          yy = []
          for item in range(self.coords.shape[0]):
            xx.append(self.coords[(item,0)])
            yy.append(self.coords[(item,1)])
          from matplotlib import pyplot as plt
          plt.plot(xx,yy,"r.")
          # denminator of 12 is specific to the use case of P6 (# symops in the metric superlattice)
          plt.plot(xx[::len(xx)//12],yy[::len(yy)//12],"b.")
          plt.plot(xx[:1],yy[:1],"g.")
          plt.axes().set_aspect("equal")
          circle = plt.Circle((0,0),1,fill=False,edgecolor="b")
          ax = plt.gca()
          ax.add_artist(circle)
          plt.show()

  def plot_after_cluster_analysis(self):
          xx = flex.double()
          yy = flex.double()
          for item in range(self.coords.shape[0]):
            xx.append(self.coords[(item,0)])
            yy.append(self.coords[(item,1)])
          from matplotlib import pyplot as plt
          plt.plot(xx, yy, 'r.')
          plt.plot(xx[0:1], yy[0:1], 'k.')
          plt.plot([0,0],[-0.01,0.01],"k-")
          plt.plot([-0.01,0.01],[0,0],"k-")
          ax = plt.gca()
          ax.set_aspect("equal")
          circle = plt.Circle((0,0),1,fill=False,edgecolor="b")
          ax.add_artist(circle)
          if self.plot_fname is None:
            plt.show()
          else:
            plot_path = os.path.join(self.output_dir, self.plot_fname)
            plot_fname = "{}_{}.{}".format(
                plot_path, self.i_plot, self.plot_format)
            plt.savefig(plot_fname)

  def _intialise_target(self):
      if self.params.dimensions is Auto:
          dimensions = None
      else:
          dimensions = self.params.dimensions
      if self.params.lattice_group is not None:
          self.lattice_group = (
              self.params.lattice_group.group()
              .build_derived_patterson_group()
              .info()
              .primitive_setting()
              .group()
          )
      if self.params.twin_axis:
        twin_axis = [tuple(x) for x in self.params.twin_axis]
        twin_rotation = self.params.twin_rotation
      else:
        twin_axis = None
        twin_rotation = None
      self.target = TargetWithCustomSymops(
          self.intensities,
          self.dataset_ids.as_numpy_array(),
          min_pairs=self.params.min_pairs,
          lattice_group=self.lattice_group,
          dimensions=dimensions,
          weights=self.params.weights,
          twin_axes=twin_axis,
          twin_angles=twin_rotation,
          cb_op=self.cb_op
      )



from dials.command_line.cosym import logger
from dials.command_line.cosym import cosym as dials_cl_cosym_wrapper
from dials.util.exclude_images import get_selection_for_valid_image_ranges
from dials.command_line.symmetry import (
    apply_change_of_basis_ops,
    change_of_basis_ops_to_minimum_cell,
    eliminate_sys_absent,
)
from dials.util.filter_reflections import filtered_arrays_from_experiments_reflections

class dials_cl_cosym_subclass (dials_cl_cosym_wrapper):
    def __init__(self, experiments, reflections, uuid_cache_in, params=None,
            do_plot=False, i_plot=None, output_dir=None):
        super(dials_cl_cosym_wrapper, self).__init__(
            events=["run_cosym", "performed_unit_cell_clustering"]
        )
        if params is None:
            params = phil_scope.extract()
        self.params = params

        self._reflections = []
        for refl, expt in zip(reflections, experiments):
            sel = get_selection_for_valid_image_ranges(refl, expt)
            self._reflections.append(refl.select(sel))

        self._experiments, self._reflections = self._filter_min_reflections(
            experiments, self._reflections, uuid_cache_in
        )
        self.ids_to_identifiers_map = {}
        for table in self._reflections:
            self.ids_to_identifiers_map.update(table.experiment_identifiers())
        self.identifiers_to_ids_map = {
            value: key for key, value in self.ids_to_identifiers_map.items()
        }

        if len(self._experiments) > 1:
            # perform unit cell clustering
            identifiers = self._unit_cell_clustering(self._experiments)
            if len(identifiers) < len(self._experiments):
                logger.info(
                    "Selecting subset of %i datasets for cosym analysis: %s"
                    % (len(identifiers), str(identifiers))
                )
                self._experiments, self._reflections = select_datasets_on_identifiers(
                    self._experiments, self._reflections, use_datasets=identifiers
                )
                self.uuid_cache = [self.uuid_cache[int(id)] for id in identifiers]

        # Map experiments and reflections to minimum cell
        cb_ops = change_of_basis_ops_to_minimum_cell(
            self._experiments,
            params.lattice_symmetry_max_delta,
            params.relative_length_tolerance,
            params.absolute_angle_tolerance,
        )
        in_cb_ops = len(cb_ops)
        exclude = [
            expt.identifier
            for expt, cb_op in zip(self._experiments, cb_ops)
            if not cb_op
        ]
        if len(exclude):
            logger.info(
                "Rejecting {} datasets from cosym analysis "\
                "(couldn't determine consistent cb_op to minimum cell):\n"\
                "{}".format(len(exclude), exclude)
            )
            self._experiments, self._reflections = select_datasets_on_identifiers(
                self._experiments, self._reflections, exclude_datasets=exclude
            )
            cb_ops = list(filter(None, cb_ops))

        ex_cb_ops = len(cb_ops)

        #Normally we expect that all the cb_ops are the same (applicable for PSI with P63)
        assertion_dict = {}
        for cb_op in cb_ops:
          key_ = cb_op.as_hkl()
          assertion_dict[key_] = assertion_dict.get(key_, 0)
          assertion_dict[key_] += 1
        if len(assertion_dict) != 1:
          # unexpected, there is normally only 1 cb operator to minimum cell
          from libtbx.mpi4py import MPI
          mpi_rank = MPI.COMM_WORLD.Get_rank()
          mpi_size = MPI.COMM_WORLD.Get_size()
          print ("RANK %02d, # experiments %d, after exclusion %d, unexpectedly there are %d unique cb_ops: %s"%(
            mpi_rank, in_cb_ops, ex_cb_ops, len(assertion_dict),
            ", ".join(["%s:%d"%(key,assertion_dict[key]) for key in assertion_dict])))
          # revisit with different use cases later

          # In fact we need all cb_ops to match because the user might supply
          # a custom reindexing operator and we need to consistently tranform
          # it from the conventional basis into the minimum basis. Therefore,
          # force them all to match, but make sure user is aware.
          if not params.single_cb_op_to_minimum:
            raise RuntimeError('There are >1 different cb_ops to minimum and \
cosym.single_cb_op_to_minimum is not True')
          else:
            best_cb_op_str = max(assertion_dict, key=assertion_dict.get)
            best_cb_op = None
            for cb_op in cb_ops:
              if cb_op.as_hkl() == best_cb_op_str:
                best_cb_op = cb_op
                break
            assert best_cb_op is not None
            cb_ops = [best_cb_op] * len(cb_ops)

        self.cb_op_to_minimum = cb_ops



        # Eliminate reflections that are systematically absent due to centring
        # of the lattice, otherwise they would lead to non-integer miller indices
        # when reindexing to a primitive setting
        self._reflections = eliminate_sys_absent(self._experiments, self._reflections)

        self._experiments, self._reflections = apply_change_of_basis_ops(
            self._experiments, self._reflections, cb_ops
        )

        # transform models into miller arrays
        datasets = filtered_arrays_from_experiments_reflections(
            self.experiments,
            self.reflections,
            outlier_rejection_after_filter=False,
            partiality_threshold=params.partiality_threshold,
        )

        datasets = [
            ma.as_anomalous_array().merge_equivalents().array() for ma in datasets
        ]

        # opportunity here to subclass as defined above, instead of the dials-implemented version
        self.cosym_analysis = CosymAnalysis(
            datasets,
            self.params,
            do_plot=do_plot,
            i_plot=i_plot,
            plot_fname=self.params.plot.filename,
            plot_format=self.params.plot.format,
            output_dir=output_dir,
            cb_op=cb_ops[0]
            )

    def _filter_min_reflections(self, experiments, reflections, uuid_cache_in):
        identifiers = []
        self.uuid_cache = []
        for expt, refl, uuid in zip(experiments, reflections, uuid_cache_in):
            if len(refl) >= self.params.min_reflections:
                identifiers.append(expt.identifier)
                self.uuid_cache.append(uuid)

        return select_datasets_on_identifiers(
            experiments, reflections, use_datasets=identifiers
        )


class TargetWithFastRij(Target):
  def __init__(self, *args, **kwargs):

    # nproc is an init arg that was removed from class
    # dials.algorithms.symmetry.cosym.target.Target in dials commit 1cd5afe4
    self._nproc = kwargs.pop('nproc', 1)

    # if test_data_path is provided, we are constructing this for a unit test
    test_data_path = kwargs.pop('test_data_path', None)
    if test_data_path is None:
      super(TargetWithFastRij, self).__init__(*args, **kwargs)
      return
    else:
      # This is only for unit testing
      import pickle
      import numpy as np
      self._nproc = 1
      with open(test_data_path, 'rb') as f:
        self._lattices = np.array(pickle.load(f))
        self.sym_ops = pickle.load(f)
        self._weights = pickle.load(f)
        self._data = pickle.load(f)
        self._patterson_group = pickle.load(f)
        self._min_pairs = 3 # minimum number of mutual miller indices for a match

        # truncate the input data to save time
        self._lattices = self._lattices[:10]
        i_last = self._lattices[-1]
        self._data = self._data[:i_last]

  def _lattice_lower_upper_index(self, lattice_id):
       lower_index = int(self._lattices[lattice_id])
       upper_index = None
       if lattice_id < len(self._lattices) - 1:
           upper_index = int(self._lattices[lattice_id + 1])
       else:
           assert lattice_id == len(self._lattices) - 1
       return lower_index, upper_index

  def compute_gradients(self, x):
    grad = super(TargetWithFastRij, self).compute_gradients(x)
    return grad.A1

  def _compute_rij_wij(self, use_cache=True, use_super=False):

    if use_super:
      # for testing
      return super()._compute_rij_wij(use_cache=use_cache)

    def _compute_one_row(args):
      """
      Call the compute_one_row method of CC, which is a compute_rij_wij_detail
      instance. Args is a tuple that we unpack because easy_mp.parallel_map
      can only pass a single argument.
      """
      CC, i, NN = args
      rij_row, rij_col, rij_data, wij_row, wij_col, wij_data = [
          list(x) for x in CC.compute_one_row(self._lattices.size, i)
      ]
      rij = sparse.coo_matrix((rij_data, (rij_row, rij_col)), shape=(NN, NN))
      wij = sparse.coo_matrix((wij_data, (wij_row, wij_col)), shape=(NN, NN))
      return rij, wij

    n_lattices = self._lattices.size
    n_sym_ops = len(self.sym_ops)
    NN = n_lattices * n_sym_ops
    lower_i = flex.int()
    upper_i = flex.int()
    for lidx in range(self._lattices.size):
      LL,UU = self._lattice_lower_upper_index(lidx)
      lower_i.append(int(LL))
      if UU is None:  UU = self._data.data().size()
      upper_i.append(int(UU))
    indices = {}
    space_group_type = self._data.space_group().type()
    from xfel.merging import compute_rij_wij_detail
    CC = compute_rij_wij_detail(
        lower_i,
        upper_i,
        self._data.data(),
        self._min_pairs)
    for sym_op in self.sym_ops:
        cb_op = sgtbx.change_of_basis_op(sym_op)
        indices_reindexed = cb_op.apply(self._data.indices())
        miller.map_to_asu(space_group_type, False, indices_reindexed)
        indices[cb_op.as_xyz()] = indices_reindexed
        CC.set_indices(cb_op, indices_reindexed)

    assert self._nproc==1
    results = map(
        _compute_one_row,
        [(CC, i, NN) for i in range(n_lattices)]
    )

    rij_matrix = None
    wij_matrix = None
    for (rij, wij) in results:
        if rij_matrix is None:
            rij_matrix = rij
            wij_matrix = wij
        else:
            rij_matrix += rij
            wij_matrix += wij
    self.rij_matrix = rij_matrix.todense().astype(np.float64)
    self.wij_matrix = wij_matrix.todense().astype(np.float64)
    return self.rij_matrix, self.wij_matrix

class TargetWithCustomSymops(TargetWithFastRij):
  def __init__(
      self,
      intensities,
      lattice_ids,
      weights=None,
      min_pairs=3,
      lattice_group=None,
      dimensions=None,
      nproc=None,
      twin_axes=None,
      twin_angles=None,
      cb_op=None,
  ):
    '''
    A couple extra init arguments permit testing user-defined reindexing ops.
    twin_axes is a list of tuples, e.g. [(0,1,0)] means the twin axis is b.
    twin_angles is a corresponding list to define the rotations; 2 is a twofold
        rotation etc.
    cb_op is the previously determined transformation from the input cells to
        the minimum cell. The data have already been transformed by this
        operator, so we transform the twin operators before testing them.
    '''

    if nproc is not None:
      warnings.warn("nproc is deprecated", DeprecationWarning)
    self._nproc = 1

    if weights is not None:
      assert weights in ("count", "standard_error")
    self._weights = weights
    self._min_pairs = min_pairs

    data = intensities.customized_copy(anomalous_flag=False)
    cb_op_to_primitive = data.change_of_basis_op_to_primitive_setting()
    data = data.change_basis(cb_op_to_primitive).map_to_asu()

    # Convert to uint64 avoids crashes on Windows when later constructing
    # flex.size_t (https://github.com/cctbx/cctbx_project/issues/591)
    order = lattice_ids.argsort().astype(np.uint64)
    sorted_data = data.data().select(flex.size_t(order))
    sorted_indices = data.indices().select(flex.size_t(order))
    self._lattice_ids = lattice_ids[order]
    self._data = data.customized_copy(indices=sorted_indices, data=sorted_data)
    assert isinstance(self._data.indices(), type(cctbx_flex.miller_index()))
    assert isinstance(self._data.data(), type(cctbx_flex.double()))

    # construct a lookup for the separate lattices
    self._lattices = np.array(
        [
            np.where(self._lattice_ids == i)[0][0]
            for i in np.unique(self._lattice_ids)
        ]
    )

    self.sym_ops = OrderedSet(["x,y,z"])
    self._lattice_group = lattice_group
    auto_sym_ops = self._generate_twin_operators()
    if twin_axes is not None:
      assert len(twin_axes) == len(twin_angles)
      lds = [literal_description(cb_op.apply(op)) for op in auto_sym_ops]
      ld_tuples = [(ld.r_info.ev(), ld.r_info.type()) for ld in lds]
      i_symops_to_keep = []
      for i, (axis, angle) in enumerate(ld_tuples):
        if axis in twin_axes and angle in twin_angles:
          i_symops_to_keep.append(i)
      assert len(i_symops_to_keep) == len(twin_axes)
      sym_ops = [auto_sym_ops[i] for i in i_symops_to_keep]
    else:
      sym_ops = auto_sym_ops
    self.sym_ops.update(op.as_xyz() for op in sym_ops)
    if dimensions is None:
      dimensions = max(2, len(self.sym_ops))
    self.set_dimensions(dimensions)

    self._lattice_group = copy.deepcopy(self._data.space_group())
    for sym_op in self.sym_ops:
      self._lattice_group.expand_smx(sym_op)
    self._patterson_group = self._lattice_group.build_derived_patterson_group()

    logger.debug(
        "Lattice group: %s (%i symops)",
        self._lattice_group.info().symbol_and_number(),
        len(self._lattice_group),
    )
    logger.debug(
        "Patterson group: %s", self._patterson_group.info().symbol_and_number()
    )

    self.rij_matrix, self.wij_matrix = self._compute_rij_wij()
