# Distance measures
def distance_square(left, right):
  
  ( x1, y1, z1 ) = left.xyz
  ( x2, y2, z2 ) = right.xyz
  
  return ( ( x1 - x2 ) ** 2 + ( y1 - y2 ) ** 2 + ( z1 - z2 ) ** 2 )


def resseq_distance(left, right):
  """
  Calculates the distance in residue sequence number
  """
  
  return abs( left.resseq_as_int() - right.resseq_as_int() )


# Distance calculation strategies
class DistanceCalculateStrategy(object):
  """
  Calculates distance every time
  """
  
  def __init__(self, method):
    
    self.method = method
    
    
  def __call__(self, left, right):
  
    return self.method( left = left, right = right )


class DistanceCacheStrategy(object):
  """
  Caches the distances already calculated
  """
  
  def __init__(self, method):
    
    self.d_sq_for = {}
    self.method = method
    
    
  def __call__(self, left, right):
    
    key = frozenset( ( left, right ) )
    
    if key not in self.d_sq_for:
      self.d_sq_for[ key ] = self.method( left = left, right = right )
    
    return self.d_sq_for[ key ]
    

# Find minimum distance between two sets
class MinDistanceFromSet(object):
  """
  Minimum distance square from a set of atoms
  """
  
  def __init__(self, others, calculator):
    
    self.others = others
    self.calculator = calculator
    
    
  def __call__(self, entity):
    
    return min( self.calculator( left = entity, right = o ) for o in self.others )
  

class ClosestToSet(object):
  """
  Finds the entity closest to a set of other entities
  """
  
  def __init__(self, selection, others, calculator):
    
    calc = MinDistanceFromSet( others = others, calculator = calculator )
    ( self.entity, self.distance ) = min(
      ( ( a, calc( a ) ) for a in selection ),
      key = lambda p: p[1],
      )


# Hierarchy caching
def parent_info(atom):
  
  residue_group = atom.parent().parent()
  chain = residue_group.parent()
  root = chain.parent().parent()
  
  return ( residue_group, chain, root )
        

class AtomHierarchy(object):
  """
  A fast way to find the children of an entity
  """
  
  def __init__(self):
    
    self.in_residue_group = {}
    self.in_chain = {}
    self.in_root = {}
    
    
  @property
  def residue_groups(self):
    
    return self.in_residue_group
  
  
  @property
  def chains(self):
    
    return self.in_chain
  
  
  @property
  def roots(self):
    
    return self.in_root
  
  
  @property
  def atoms(self):
    
    import operator
    return reduce( operator.or_, self.in_root.values(), set() )
    
    
  def add(self, atom):
    
    ( rg, chain, root ) = parent_info( atom = atom )
    
    self.in_residue_group.setdefault( rg, set() ).add( atom )
    self.in_chain.setdefault( chain, set() ).add( atom )
    self.in_root.setdefault( root, set() ).add( atom )
    
    
  def add_multiple(self, atoms):
    
    for atom in atoms:
      ( rg, chain, root ) = parent_info( atom = atom )
      
      self.in_residue_group.setdefault( rg, set() ).add( atom )
      self.in_chain.setdefault( chain, set() ).add( atom )
      self.in_root.setdefault( root, set() ).add( atom )
    
    
  def remove(self, atom):
    
    ( rg, chain, root ) = parent_info( atom = atom )
    
    self.in_residue_group[ rg ].remove( atom )
    self.in_chain[ chain ].remove( atom )
    self.in_root[ root ].remove( atom )
    
    if not self.in_residue_group[ rg ]:
      del self.in_residue_group[ rg ]
      
    if not self.in_chain[ chain ]:
      del self.in_chain[ chain ]
      
    if not self.in_root[ root ]:
      del self.in_root[ root ]
      
      
  def remove_multiple(self, atoms):
    
    rgs = set()
    chains = set()
    roots = set()
    
    for atom in atoms:
      ( rg, chain, root ) = parent_info( atom = atom )
      self.in_residue_group[ rg ].remove( atom )
      self.in_chain[ chain ].remove( atom )
      self.in_root[ root ].remove( atom )
      rgs.add( rg )
      chains.add( chain )
      roots.add( root )
      
    for rg in rgs:
      if not self.in_residue_group[ rg ]:
        del self.in_residue_group[ rg ]
        
    for chain in chains:
      if not self.in_chain[ chain ]:
        del self.in_chain[ chain ]
        
    for root in roots:
      if not self.in_root[ root ]:
        del self.in_root[ root ]
    

# Unresolved clashes
class UnresolvedClashes(object):
  """
  Clashes not resolved and atoms still in clashes
  """
  
  def __init__(self, clashes):
    
    self.for_atom = {}
    
    for clash in clashes:
      self.for_atom.setdefault( clash.left.atom, set() ).add( clash )
      self.for_atom.setdefault( clash.right.atom, set() ).add( clash )
      
      
  @property
  def atoms(self):
    
    return self.for_atom
  
  
  @property
  def clashes(self):
    
    result = set()
    result.update( *self.for_atom.values() )
  
    return result
      
      
  def remove(self, atom):
    
    changed = set()
    
    for clash in list( self.for_atom[ atom ] ):
      self.for_atom[ clash.left.atom ].discard( clash )
      self.for_atom[ clash.right.atom ].discard( clash )
      changed.add( clash.left.atom )
      changed.add( clash.right.atom )
      
    affecteds = set()
    
    for a in changed:
      if not self.for_atom[ a ]:
        del self.for_atom[ a ]
        affecteds.add( a )
        
    return affecteds
  
  
  def remove_multiple(self, atoms):
    
    changeds = set()
    
    for atom in atoms:
      for clash in list( self.for_atom[ atom ] ):
        self.for_atom[ clash.left.atom ].discard( clash )
        self.for_atom[ clash.right.atom ].discard( clash )
        changeds.add( clash.left.atom )
        changeds.add( clash.right.atom )
      
    affecteds = set()
    
    for a in changeds:
      if not self.for_atom[ a ]:
        del self.for_atom[ a ]
        affecteds.add( a )
        
    return affecteds
      
      
  def clashes_for(self, atom):
    
    return self.for_atom[ atom ]
  
  
  def atoms_clashing_with(self, atom):
    
    clashings = set()
    
    for clash in self.clashes_for( atom = atom ):
      clashings.add( clash.left.atom )
      clashings.add( clash.right.atom )
      
    clashings.remove( atom )
    return clashings
  

# Clash resolution algorithm and caching details
class SequenceAndSpatialDistanceSelection(object):
  """
  Contains the state of the process
  """
  
  def __init__(self, unresolved, spatial_distance, resseq_distance):
    
    self.unresolved = unresolved
    self.hierarchy = AtomHierarchy()
    
    for atom in self.unresolved.atoms:
      self.hierarchy.add( atom )
      
    self.spatial_distance = spatial_distance
    self.resseq_distance = resseq_distance

    
  @property
  def clashes(self):
    
    return self.unresolved.clashes
  
  
  @property
  def atoms(self):
    
    return self.unresolved.atoms
    
    
  def remove(self, atom):
    
    affecteds = self.unresolved.remove( atom = atom )
    
    for a in affecteds:
      self.hierarchy.remove( atom = a )
      
      
  def remove_multiple(self, atoms):
    
    affecteds = self.unresolved.remove_multiple( atoms = atoms )
    self.hierarchy.remove_multiple( atoms = affecteds )
    
    
  def clash_count_for(self, atom):
    
    return len( self.unresolved.clashes_for( atom = atom ) )
  
  
# Algorithm selectors
def select_first_atom_based_on_clash_number(remaining):
  
  return max( remaining.atoms, key = remaining.clash_count_for )


def select_atom_from_same_residue_group_maximal_proximity(rgs, remaining, moribund):
  
  return [ min(
    (
      ClosestToSet(
        selection = remaining.hierarchy.in_residue_group[ rg ],
        others = moribund.in_residue_group[ rg ],
        calculator = remaining.spatial_distance,
        )
      for rg in rgs
      ),
    key = lambda s: s.distance
    ).entity ]
  
  
def select_all_atoms_from_same_residue_group_or_one_with_maximal_proximity(rgs, remaining, moribund):
  
  seed = min(
    (
      ClosestToSet(
        selection = remaining.hierarchy.in_residue_group[ rg ],
        others = moribund.in_residue_group[ rg ],
        calculator = remaining.spatial_distance,
        )
      for rg in rgs
      ),
    key = lambda s: s.distance
    ).entity
    
  return select_passed_or_all_atoms_in_residue_group_if_no_self_clashes(
      atom = seed,
      remaining = remaining,
      )
    
    
class select_atom_from_same_residue_group_in_bond_or_closest(object):
  
  def __init__(self, bond_distance = 1.8):
    
    self.bond_distance_sq = bond_distance ** 2
    
  
  def __call__(self, rgs, remaining, moribund):
  
    closest = ( None, None )
    
    for rg in rgs:
      calc = MinDistanceFromSet(
        others = moribund.in_residue_group[ rg ],
        calculator = remaining.spatial_distance,
        )
      
      for atom in remaining.hierarchy.in_residue_group[ rg ]:
        min_dist_square = calc( entity = atom )
        
        if min_dist_square < self.bond_distance_sq:
          return [ atom ]
        
        if closest[1] < -min_dist_square:
          closest = ( atom, -min_dist_square )
          
    return [ closest[0] ]
    
    
class select_atom_from_nearby_residue_group_or_closest_within_a_chain(object):
  
  def __init__(self, max_sequence_distance):
    
    self.max_sequence_distance = max_sequence_distance
    
  
  def __call__(self, chains, remaining, moribund):
  
    closest = min(
      (
        ClosestToSet(
          selection = [ rg for rg in remaining.hierarchy.residue_groups
            if rg.parent() == chain ],
          others = [ rg for rg in moribund.residue_groups
            if rg.parent() == chain ],
          calculator = remaining.resseq_distance,
          )
        for chain in chains
        ),
      key = lambda s: s.distance,
      )
    
    if closest.distance <= self.max_sequence_distance:
      return [ max(
        remaining.hierarchy.in_residue_group[ closest.entity ],
        key = remaining.clash_count_for,
        ) ]
      
    else:
      return [ min(
        (
          ClosestToSet(
            selection = remaining.hierarchy.in_chain[ chain ],
            others = moribund.in_chain[ chain ],
            calculator = remaining.spatial_distance,
            )
          for chain in chains
          ),
        key = lambda s: s.distance
        ).entity ]
      
      
def has_clashes_within_residue_group(atom, unresolved):
  
  return any(
    c.left.residue_group == c.right.residue_group
    for c in unresolved.clashes_for( atom = atom )
    )
  
  
def has_clashes_within_chain(atom, unresolved):
  
  return any(
    c.left.chain == c.right.chain
    for c in unresolved.clashes_for( atom = atom )
    )
      
      
def select_passed_or_all_atoms_in_residue_group_if_no_self_clashes(
  atom,
  remaining,
  ):
  
  candidates = remaining.hierarchy.in_residue_group[ atom.parent().parent() ]
  
  if all(
    not has_clashes_within_residue_group(
      atom = a,
      unresolved = remaining.unresolved,
      )
    for a in candidates
    ):
    return candidates.copy()
  
  else:
    return [ atom ]
      
      
class select_atoms_in_nearby_residue_group_or_closest_within_a_chain(object):
  
  def __init__(self, max_sequence_distance):
    
    self.max_sequence_distance = max_sequence_distance
    
    
  def __call__(self, chains, remaining, moribund):
    
    closest = min(
      (
        ClosestToSet(
          selection = [ rg for rg in remaining.hierarchy.residue_groups
            if rg.parent() == chain ],
          others = [ rg for rg in moribund.residue_groups
            if rg.parent() == chain ],
          calculator = remaining.resseq_distance,
          )
        for chain in chains
        ),
      key = lambda s: s.distance,
      )
    
    if closest.distance <= self.max_sequence_distance:
      
      seed = max(
        remaining.hierarchy.in_residue_group[ closest.entity ],
        key = remaining.clash_count_for,
        )
      
    else:
      seed = min(
        (
          ClosestToSet(
            selection = remaining.hierarchy.in_chain[ chain ],
            others = moribund.in_chain[ chain ],
            calculator = remaining.spatial_distance,
            )
          for chain in chains
          ),
        key = lambda s: s.distance
        ).entity
        
    return select_passed_or_all_atoms_in_residue_group_if_no_self_clashes(
      atom = seed,
      remaining = remaining,
      )
      

def select_atom_from_with_highest_number_of_clashes(roots, remaining, moribund):
  
  return [ max(
    [
      max(
        [ ( a, remaining.clash_count_for( atom = a ) )
          for a in remaining.hierarchy.in_root[ root ] ],
        key = lambda p: p[1],
        )
      for root in roots
      ],
    key = lambda p: p[1]
    )[0] ]
  
  
def select_atom_from_root_with_highest_number_of_clashes(root, remaining, moribund):
  
  return max(
    remaining.hierarchy.in_root[ root ],
    key = remaining.clash_count_for,
    )
  
  
class ClashResolution(object):
  """
  Contains a clash list and a set of atoms that need to be deleted to resolve
  these. Acts mostly as a milestone
  """
  
  def __init__(self, clashes, algorithm, preassigned = []):
    
    self.clashes = clashes
    self.moribund = algorithm.resolve(
      clashes = self.clashes,
      preassigned = preassigned,
      )
    
    
  @property
  def atoms(self):
    
    atoms = set()
    
    for c in self.clashes:
      atoms.add( c.left.atom )
      atoms.add( c.right.atom )
      
    return atoms
  
    
class SequenceAndSpatialDistanceAlgorithm(object):
  """
  Selects atoms for removal.
  
  If there are atoms from the same model (i.e. clash must be with symmetry
  equivalent), the following algorithm is used (in order of priority):
   
  1. select an atom from the same residue_group (with maximal spatial proximity)
  2. select a residue_group that is from the same chain and close in sequence
     (up to N residues); select the atom with the largest number of clashes.
     If no such residue_group exists, then based on spatial proximity.
  3. select an atom with the largest number of clashes
  
  If there are no atoms from the same model(s), select an atom with the largest
  number of clashes.
  """
  
  def __init__(
    self,
    first_atom_selection = select_first_atom_based_on_clash_number,
    from_same_residue_group = select_atom_from_same_residue_group_maximal_proximity,
    from_same_chain = select_atoms_in_nearby_residue_group_or_closest_within_a_chain(
      max_sequence_distance = 4,
      ),
    from_same_root = select_atom_from_with_highest_number_of_clashes,
    spatial_distance_strategy = DistanceCalculateStrategy,
    resseq_distance_strategy = DistanceCalculateStrategy,
    ):
    
    self.first_atom_selection = first_atom_selection
    self.from_same_residue_group = from_same_residue_group
    self.from_same_chain = from_same_chain
    self.from_same_root = from_same_root
    
    self.spatial_distance_strategy = spatial_distance_strategy
    self.resseq_distance_strategy = resseq_distance_strategy
  
  
  def consecutive_atom_selection(self, moribund, remaining):
    
    # Same residue_group
    common = set( moribund.residue_groups ).intersection(
      remaining.hierarchy.residue_groups
      )
    
    if common:
      return self.from_same_residue_group(
        rgs = common,
        remaining = remaining,
        moribund = moribund,
        )
        
    # Same chain
    common = set( moribund.chains ).intersection(
      remaining.hierarchy.chains
      )
    
    if common:
      return self.from_same_chain(
        chains = common,
        remaining = remaining,
        moribund = moribund,
        )
          
    # Same root
    common = set( moribund.roots ).intersection(
      remaining.hierarchy.roots
      )
    
    if common:
      return self.from_same_root(
        roots = common,
        remaining = remaining,
        moribund = moribund,
        )
        
    # New model
    return self.first_atom_selection( remaining = remaining )
  
  
  def resolve(self, clashes, preassigned = []):
    
    # no clashes
    if not clashes:
      return set()
    
    # initialize data structures
    remaining = SequenceAndSpatialDistanceSelection(
      unresolved = UnresolvedClashes( clashes = clashes ),
      spatial_distance = self.spatial_distance_strategy( method = distance_square ),
      resseq_distance = self.resseq_distance_strategy( method = resseq_distance ),
      )
    moribund = AtomHierarchy()
    
    # remove requested atoms or select new one
    if preassigned:
      moribund.add_multiple( atoms = preassigned )
      remaining.remove_multiple( atoms = preassigned )
    
    else:
      first = self.first_atom_selection( remaining = remaining )
      moribund.add( first )
      remaining.remove( first )
    
    # Add atoms until no clashes
    while remaining.clashes:
      selecteds = self.consecutive_atom_selection(
        moribund = moribund,
        remaining = remaining,
        )
      
      remaining.remove_multiple( atoms = selecteds )
      moribund.add_multiple( atoms = selecteds )
      
    return moribund.atoms
  
  
class HierarchicalAlgorithm(object):
  """
  Selects atoms for removal.
  
  Processing is done according to PDB hierarchy:
  1. First, a suitable atom is selected as the starting point.
  2. Within the same root:
    a. Within the same chain:
      - Discard all atoms that clash with other chains.
      - Go through the remaining residue groups, and resolve clashes
    b. If there are still atoms from the same root, select it, and repeat a.
  3. If there are still unresolved clashes, select a new atom, and repeat 2.
  """
  
  def __init__(
    self,
    first_atom_selection = select_first_atom_based_on_clash_number,
    from_same_residue_group = select_atom_from_same_residue_group_maximal_proximity,
    from_same_chain = select_atoms_in_nearby_residue_group_or_closest_within_a_chain(
      max_sequence_distance = 4,
      ),
    from_same_root = select_atom_from_root_with_highest_number_of_clashes,
    spatial_distance_strategy = DistanceCalculateStrategy,
    resseq_distance_strategy = DistanceCalculateStrategy,
    ):
    
    self.first_atom_selection = first_atom_selection
    self.from_same_residue_group = from_same_residue_group
    self.from_same_chain = from_same_chain
    self.from_same_root = from_same_root
    
    self.spatial_distance_strategy = spatial_distance_strategy
    self.resseq_distance_strategy = resseq_distance_strategy
  
  
  def resolve(self, clashes, preassigned = []):
    """
    A generic algorithm that takes the PDB hierarchy into account for clash resolution
    """
    
    if not clashes:
      return set()
    
    # initialize data structures
    remaining = SequenceAndSpatialDistanceSelection(
      unresolved = UnresolvedClashes( clashes = clashes ),
      spatial_distance = self.spatial_distance_strategy( method = distance_square ),
      resseq_distance = self.resseq_distance_strategy( method = resseq_distance ),
      )
    moribund = AtomHierarchy()
    
    if preassigned:
      moribund.add_multiple( atoms = preassigned )
      remaining.remove_multiple( atoms = preassigned )
      
      for chain in set( a.parent().parent().parent() for a in preassigned ):
        self.resolve_chain(
          chain = chain,
          moribund = moribund,
          remaining = remaining,
          )
    
    while remaining.unresolved.clashes:
      selected = self.first_atom_selection( remaining = remaining )
      moribund.add( atom = selected )
      remaining.remove( atom = selected )
      self.resolve_root_of(
        atom = selected,
        moribund = moribund,
        remaining = remaining,
        )
        
    return moribund.atoms
  
  
  def resolve_root_of(self, atom, moribund, remaining):
        
    chain = atom.parent().parent().parent()
    self.resolve_chain(
      chain = chain,
      moribund = moribund,
      remaining = remaining,
      )
    root = chain.parent().parent()
    
    while root in remaining.hierarchy.roots:
      selected = self.from_same_root(
        root = root,
        remaining = remaining,
        moribund = moribund,
        )
      moribund.add( atom = selected )
      remaining.remove( atom = selected )
      chain = selected.parent().parent().parent()
      self.resolve_chain(
        chain = chain,
        moribund = moribund,
        remaining = remaining,
        )
      assert chain not in remaining.hierarchy.chains
  
  
  def resolve_chain(self, chain, moribund, remaining):
      
    clashing_with_aliens = [ a for a in remaining.hierarchy.in_chain.get( chain, [] )
      if not has_clashes_within_chain( atom = a, unresolved = remaining.unresolved ) ]
    
    remaining.remove_multiple( atoms = clashing_with_aliens )
    moribund.add_multiple( atoms = clashing_with_aliens )
    
    while chain in remaining.hierarchy.chains:
      # Same residue_group
      common = set( moribund.residue_groups ).intersection(
        remaining.hierarchy.residue_groups
        )
      
      if common:
        selecteds = self.from_same_residue_group(
          rgs = common,
          remaining = remaining,
          moribund = moribund,
          )
        remaining.remove_multiple( atoms = selecteds )
        moribund.add_multiple( atoms = selecteds )
        continue
          
      # Same chain
      common = set( moribund.chains ).intersection(
        remaining.hierarchy.chains
        )
      assert common
      
      selecteds = self.from_same_chain(
        chains = common,
        remaining = remaining,
        moribund = moribund,
        )
      remaining.remove_multiple( atoms = selecteds )
      moribund.add_multiple( atoms = selecteds )


def calculate(collection, min_clash_network = 4):
  
  from boost_adaptbx import graph

  g = graph.adjacency_list(
    graph_type = "undirected",
    vertex_type = "vector",
    edge_type = "set",
    )

  clashes_incident_on = {}

  for clash in set( collection.pairs() ):
    vid = g.add_vertex( label = clash )
    clashes_incident_on.setdefault( clash.left.atom, set() ).add( vid )
    clashes_incident_on.setdefault( clash.right.atom, set() ).add( vid )
      
  for network in clashes_incident_on.values():
    first = network.pop()
    
    for desc in network:
      g.add_edge( first, desc )
    
  from graph import connected_component_algorithm as cca
    
  for component in cca.connected_components( graph = g ):
    
    if len( component ) < min_clash_network:
      continue
    
    yield [ g.vertex_label( vertex = v ) for v in component ]


def alternatives(resolution, algorithm):
  
  clashes = resolution.clashes
  unresolved = UnresolvedClashes( clashes = clashes )
  atoms = resolution.atoms
  accounted = atoms.difference( resolution.moribund )
  
  for c in resolution.clashes:
    if c.left.atom == c.right.atom:
      accounted.add( c.left.atom )
    
  resolutions = []
  
  while atoms != accounted:
    reso = ClashResolution(
      clashes = clashes,
      algorithm = algorithm,
      preassigned = accounted,
      )
    
    # Reinstate atoms that would not clash anyway
    for atom in accounted:
      clashing_atoms = unresolved.atoms_clashing_with( atom = atom )
      
      if reso.moribund.issuperset( clashing_atoms ):
        reso.moribund.remove( atom )
        accounted.add( atom )
      
    for ( index, prev ) in enumerate( resolutions ):
      if ( reso.moribund.issuperset( prev.moribund )
        or reso.moribund.issuperset( resolution.moribund )
        ):
        return resolutions
      
      if prev.moribund.issubset( reso.moribund ):
        resolutions[ index ] = reso
        break
      
    else:
      resolutions.append( reso )
    
    surviving = atoms.difference( reso.moribund )
    
    if accounted.issuperset( surviving ):
      break
    
    accounted.update( surviving )
    
  return resolutions