from __future__ import absolute_import, division, print_function
from libtbx.utils import null_out
from scitbx.array_family import flex
'''
Mask utilities
'''

class create_mask_around_atoms(object):

  '''
    Class to create a map_manager object containing a mask around atoms,
    smooth or adjust density outside, and apply to a supplied map_manager

  '''

  def __init__(self,
      model = None,
      xray_structure = None,
      mask_atoms_atom_radius = None,
      invert_mask = None,
      n_real = None,
      map_manager = None,
      wrapping = None):

    '''
     Create a mask (map object) with values of 1 near atoms in xray_structure

     Parameters are:
       model:   required source of information about where atoms are
       xray_structure:   alternate source of information about where atoms are
       mask_atoms_atom_radius:  radius around atoms to mask
       invert_mask:  outside is 1 and inside is 0
       n_real:  dimensions of map to create, e.g., existing map_data.all()
       map_manager: alternate source of information for n_real. origin (0, 0, 0)
         and source of wrapping
       wrapping:  Whether map wraps around unit cell boundaries, where unit
         cell is that defined by xray_structure or model
    '''

    assert (model is not None) or (xray_structure is not None)
    assert mask_atoms_atom_radius is not None
    assert (n_real is not None) or (map_manager is not None)
    assert (map_manager is not None) or isinstance(wrapping, bool)

    if not n_real:
      assert map_manager.map_data().origin() == (0, 0, 0)
      n_real = map_manager.map_data().all()


    if model:
      xray_structure = model.get_xray_structure()
      self._crystal_symmetry = model.crystal_symmetry()
    else:
      self._crystal_symmetry = xray_structure.crystal_symmetry()

    if map_manager and wrapping is None:
      wrapping = map_manager.wrapping()

    assert wrapping is not None

    if (not wrapping) or (self._crystal_symmetry.space_group_number==1):
      # usual
      sites_frac = xray_structure.sites_frac()
    else:  # wrapping and cs: need to expand
      sites_frac = xray_structure.expand_to_p1().sites_frac()

    import boost_adaptbx.boost.python as bp
    cctbx_maptbx_ext = bp.import_ext("cctbx_maptbx_ext")
    radii = flex.double(
      sites_frac.size(), mask_atoms_atom_radius)
    if invert_mask:
      mask_value_inside_molecule = 0
      mask_value_outside_molecule = 1
    else: # usual
      mask_value_inside_molecule = 1
      mask_value_outside_molecule = 0
    self._mask = cctbx_maptbx_ext.mask(
      sites_frac                  = sites_frac,
      unit_cell                   = xray_structure.unit_cell(),
      n_real                      = n_real,
      mask_value_inside_molecule  = mask_value_inside_molecule,
      mask_value_outside_molecule = mask_value_outside_molecule,
      radii                       = radii,
      wrapping                    = wrapping)
    # Set up map_manager with this mask
    if map_manager:
      self._map_manager = map_manager.customized_copy(map_data = self._mask)
    else:
      from iotbx.map_manager import map_manager as map_man
      self._map_manager = map_man(map_data = self._mask,
        unit_cell_crystal_symmetry = self.crystal_symmetry(),
        unit_cell_grid = self._mask.all(),
        wrapping = wrapping)

    self._map_manager.set_is_mask(True)

    # Initialize soft mask
    self._is_soft_mask = False
    self._is_soft_mask_around_edges = False

  def mask(self):
    return self._mask

  def map_manager(self):
    return self._map_manager

  def is_soft_mask_around_edges(self):
    return self._is_soft_mask_around_edges

  def is_soft_mask(self):
    return self._is_soft_mask

  def solvent_content(self):
    if hasattr(self, '_solvent_content'):
      return self._solvent_content

  def crystal_symmetry(self):
    return self._crystal_symmetry

  def soft_mask(self, soft_mask_radius = None):
    '''
    Make the mask a soft mask
    Parameter:
      soft_mask_radius:   defines distance over which mask is smoothed
    This has to be done in P1 as the mask might not follow map symmetry
      (for example a mask around edges does not)
    '''
    assert soft_mask_radius is not None
    assert not self.is_soft_mask()  # do not do it twice

    from cctbx.maptbx.segment_and_split_map import smooth_mask_data

    if self._crystal_symmetry.space_group_number() != 1:
      from cctbx import crystal
      crystal_symmetry = crystal.symmetry(
          self._crystal_symmetry.unit_cell().parameters(), 1)
    else:
      crystal_symmetry = self._crystal_symmetry

    self._mask = smooth_mask_data(mask_data = self._mask,
      crystal_symmetry = crystal_symmetry,
       rad_smooth = soft_mask_radius)

    # Note that mask still might not follow map symmetry if it did not before

    self._map_manager = self._map_manager.customized_copy(map_data = self._mask)
    self._map_manager.set_is_mask(True)
    self._is_soft_mask_around_edges = True

  def apply_mask_to_other_map_manager(self, other_map_manager = None,
     set_outside_to_mean_inside = False,
     set_mean_to_zero = False):

    '''
    Apply this mask to a map_manager object containing map_data.
    Creates new map_manager

    Parameters:
       other_map_manager: map_manager to be masked in place
       set_outside_to_mean_inside:  if True,
          set the value outside mask to the mean inside the mask
       set_mean_to_zero:  Adjust overall mean of map to zero after masking
    '''

    assert other_map_manager.is_similar(self.map_manager())
    assert other_map_manager.map_data().origin() == (0, 0, 0)

    from cctbx.maptbx.segment_and_split_map import apply_mask_to_map
    new_map_data = apply_mask_to_map(mask_data = self._mask,
      smoothed_mask_data = self._mask,
      set_outside_to_mean_inside = set_outside_to_mean_inside,
      set_mean_to_zero = set_mean_to_zero,
      map_data = other_map_manager.map_data(),
      out = null_out())

    return self.map_manager().customized_copy(map_data = new_map_data)

class create_mask_with_map_data(create_mask_around_atoms):

  '''
    Class to create a map_manager object containing a supplied mask

  '''

  def __init__(self, map_data,
      map_manager = None):

    '''
     Create a mask (map object) with values supplied in map_data
    '''
    assert isinstance(map_data, flex.double)
    assert map_manager.map_data().all() == map_data.all()
    assert map_data.origin() == (0,0,0)
    assert map_manager.origin_is_zero()

    self._crystal_symmetry = map_manager.crystal_symmetry()

    self._mask = map_data

    # Set up map_manager with this mask
    self._map_manager = map_manager.customized_copy(map_data = self._mask)
    self._map_manager.set_is_mask(True)

    # Initialize soft mask
    self._is_soft_mask = False
    self._is_soft_mask_around_edges = False

class create_mask_around_edges(create_mask_around_atoms):

  '''
    Class to create a map_manager object containing a mask a few pixels in
    around the boundaries of the map

  '''

  def __init__(self,
      boundary_radius = None,
      map_manager = None):

    '''
     Create a mask (map object) with values of 1 except at the map boundaries

     This will not be a soft mask. To make it a soft mask, follow this with
       soft_mask(soft_mask_radius = soft_mask_radius)

     Parameters are:
       boundary_radius:  radius for masking
       map_manager: source of information for n_real and crystal_symmetryi
          . origin (0, 0, 0)
    '''

    assert boundary_radius is not None
    assert (map_manager is not None)

    n_real = map_manager.map_data().all()

    self._crystal_symmetry = map_manager.crystal_symmetry()

    from cctbx.maptbx.segment_and_split_map import get_zero_boundary_map

    self._mask = get_zero_boundary_map(
      map_data = map_manager.map_data(),
      crystal_symmetry = self._crystal_symmetry,
      radius = boundary_radius)
    # Set up map_manager with this mask
    self._map_manager = map_manager.customized_copy(map_data = self._mask)
    self._map_manager.set_is_mask(True)

    # Initialize soft mask
    self._is_soft_mask = False
    self._is_soft_mask_around_edges = False

class create_mask_around_density(create_mask_around_atoms):

  '''
    Class to create a map_manager object containing a mask around density

  '''

  def __init__(self,
      map_manager = None,
      resolution = None,
      molecular_mass = None,
      sequence = None,
      solvent_content = None):

    '''
     Create a mask (map object) with values of 1 near molecule

     Parameters are:
       map_manager: source of information about density
       resolution : optional resolution of map
       molecular_mass: optional mass (Da) of object in density
       sequence: optional sequence of object in density
       solvent_content : optional solvent_content of map
    '''

    assert (map_manager is not None)

    if not resolution:
      from cctbx.maptbx import d_min_from_map
      resolution = d_min_from_map(
           map_data=map_manager.map_data(),
           unit_cell=map_manager.crystal_symmetry().unit_cell())

    self._crystal_symmetry = map_manager.crystal_symmetry()

    if (molecular_mass or sequence ) and (
          not solvent_content):
      # Try to get a good starting value of solvent_content

      from cctbx.maptbx.segment_and_split_map import get_solvent_fraction
      solvent_content = get_solvent_fraction(
           params = None,
           molecular_mass = molecular_mass,
           sequence = sequence,
           do_not_adjust_dalton_scale = True,
           crystal_symmetry = self._crystal_symmetry,
           out = null_out())

    # Now use automatic procedure to get a mask
    from cctbx.maptbx.segment_and_split_map import \
          get_iterated_solvent_fraction
    self._mask, self._solvent_content = get_iterated_solvent_fraction(
          crystal_symmetry = self._crystal_symmetry,
          fraction_of_max_mask_threshold = 0.05, #
          solvent_content = solvent_content,
          cell_cutoff_for_solvent_from_mask = 1, # Use low-res method always
          use_solvent_content_for_threshold = True,
          mask_resolution = resolution,
          return_mask_and_solvent_fraction = True,
          map = map_manager.map_data(),
          verbose = False,
          out = null_out())

    if self._solvent_content is None:
      from libtbx.utils import Sorry
      raise Sorry("Unable to get solvent content in auto-masking")


    # Set up map_manager with this mask
    self._map_manager = map_manager.customized_copy(map_data = self._mask)
    self._map_manager.set_is_mask(True)

    # Initialize soft mask
    self._is_soft_mask = False
    self._is_soft_mask_around_edges = False

class expand_mask(create_mask_around_atoms):

  '''
    Class to create a map_manager object containing a mask expanded from
     previous mask

  '''

  def __init__(self,
      map_manager = None,
      resolution = None,
      buffer_radius = 5,
      minimum_fraction_of_max = 0.01):


    '''
     Create a mask (map object) with values of 1 expanded by buffer_radius

     Parameters are:
       map_manager: source of previous mask.
       resolution : optional resolution of map
       buffer_radius:  radius to expand in A
       minimum_fraction_of_max: smallest-sized region to expand
         as ratio to biggest
    '''

    assert (map_manager is not None)

    if not resolution:
      from cctbx.maptbx import d_min_from_map
      resolution = d_min_from_map(
           map_data=map_manager.map_data(),
           unit_cell=map_manager.crystal_symmetry().unit_cell())

    self._crystal_symmetry = map_manager.crystal_symmetry()

    from cctbx.maptbx.segment_and_split_map import estimate_expand_size, get_co

    expand_size = estimate_expand_size(
         crystal_symmetry = map_manager.crystal_symmetry(),
         map_data = map_manager.map_data(),
         expand_target = buffer_radius,
         minimum_expand_size = 0,
         out = null_out())

    if expand_size < 1:
      return # do nothing
    co, sorted_by_volume, min_b, max_b = get_co(
       map_data = map_manager.map_data(),
       threshold = 0.5, wrapping = False)
    if len(sorted_by_volume)<2:
      return # do nothing

    minimum_size = sorted_by_volume[1][0] * minimum_fraction_of_max
    s = None
    for v1, i1 in sorted_by_volume[1:]:
      if v1 < minimum_size: break
      bool_region_mask = co.expand_mask(
        id_to_expand = i1, expand_size = expand_size)
      if s is None:
        s = (bool_region_mask == True)
      else:
        s |=  (bool_region_mask == True)
    self._mask = map_manager.deep_copy().map_data()
    self._mask.set_selected(s, 1)
    self._mask.set_selected(~s, 0)

    self._solvent_content = s.count(False)/s.size()


    # Set up map_manager with this mask
    self._map_manager = map_manager.customized_copy(map_data = self._mask)
    self._map_manager.set_is_mask(True)

    # Initialize soft mask
    self._is_soft_mask = False
    self._is_soft_mask_around_edges = False
