from phaser import sequtil

from scitbx import suffixtree
from scitbx.suffixtree import single as singletree

import operator
# Data classes
class Alignment(object):
  """
  A full alignment
  """
  
  def __init__(self, pairs):
    
    self.pairs = pairs
    
    
  def identities(self):
    
    return sum( 1 for ( l, r ) in self.pairs if l == r )
  
  
  def __add__(self, other):
    
    return self.__class__( pairs = self.pairs + other.pairs )
  
  
  def __iadd__(self, other):
    
    self.pairs.extend( other.pairs )
    return self
  
  
  def __len__(self):
    
    return len( self.pairs )
  
  
  @classmethod
  def empty(cls, sequence, start, stop):
    
    length = len( sequence )
    return cls(
      pairs = [
        ( sequence[i] if 0 <= i < length else None, None )
        for i in range( start, stop )
        ]
      )
  

class Fragment(object):
  """
  Unique fragment despite identical sequences
  """
  
  def __init__(self, index, sequence):
    
    self.index = index
    self.sequence = sequence
    
    
  def __len__(self):
    
    return len( self.sequence )
  
  
  def __repr__(self):
    
    return "Fragment(%s, %s)" % ( self.index, self.sequence )
  
  
  def __cmp__(self, other):
    
    return cmp( self.index, other.index )
  
  
  @staticmethod
  def consecutivity(left, right):
  
    return left.index + 1 == right.index
    
  
class HSSMatch(object):
  """
  Represents an HSS
  """
  
  def __init__(self, fragment, length, start):
    
    self.fragment = fragment
    self.length = length
    self.start = start
  
  
  @property
  def stop(self):
    
    return self.start + len( self.fragment )
  
  
  def overlap_region_alignment(self, sequence):
    
    pairs = []
    length = len( sequence )
      
    for ( frag_i, seq_i ) in enumerate( range( self.start, self.stop ) ):
      pairs.append(
        (
          sequence[ seq_i ] if 0 <= seq_i < length  else None,
          self.fragment.sequence[ frag_i ],
          )
        )
    return Alignment( pairs = pairs )
    
    
  def alignment(self, sequence):
      
    return (
      Alignment.empty( sequence = sequence, start = 0, stop = self.start )
      + self.overlap_region_alignment( sequence = sequence )
      + Alignment.empty( sequence = sequence, start = self.stop, stop = len( sequence ) )
      )
  
  
  def compatible_with(self, other):
    
    if self.fragment.index < other.fragment.index:
      return self.stop <= other.start
    
    elif other.fragment.index < self.fragment.index:
      return other.stop <= self.start
    
    return False
  
  
  def remap(self, shift):
    
    self.start += shift
    
    
  def __repr__(self):
    
    return "HSSMatch(fragment = %s, length = %s, start=%s)" % (
      self.fragment,
      self.length,
      self.start,
      )


class UnalignedRegion(object):
  """
  A region of sequence and fragments corresponding to it
  """
  
  def __init__(
    self,
    left,
    right,
    fragments,
    applicable_sequence,
    is_inside,
    simple_fill,
    remap,
    left_edge,
    central,
    right_edge,
    alignable_fragments,
    ):
    
    self.left = left
    self.right = right
    self.fragments = fragments
    self._applicable_sequence = applicable_sequence
    self._is_inside = is_inside
    self._simple_fill = simple_fill
    self._remap = remap
    self._left_edge = left_edge
    self._central = central
    self._right_edge = right_edge
    self._alignable_fragments = alignable_fragments
    
    
  def applicable_sequence(self, full):
    
    return self._applicable_sequence( owner = self, full = full )
  
  
  def is_inside(self, position):
    
    return self._is_inside( owner = self, position = position )
  
  
  def compatible_with(self, match):
    
    return ( self.is_inside( position = match.start )
      and self.is_inside( position = match.stop ) )
  
  
  def simple_fill(self):
    
    return self._simple_fill( owner = self )
  
  
  def remap(self, match):
    
    return self._remap( owner = self, match = match)
  
  
  def left_edge(self, right, fragments):
    
    return self._left_edge( owner = self, right = right, fragments = fragments )
  
  
  def central(self, left, right, fragments):
    
    return self._central(
      owner = self,
      left = left,
      right = right,
      fragments = fragments,
      )
  
  
  def right_edge(self, left, fragments):
    
    return self._right_edge( owner = self, left = left, fragments = fragments )
  
  
  def alignable_fragments(self, full):
    
    return self._alignable_fragments( owner = self, full = full )
    
    
  @classmethod
  def LeftEdge(cls, right, fragments):
    
    return cls(
      left = None,
      right = right,
      fragments = fragments,
      applicable_sequence = cls.left_applicable_sequence,
      is_inside = cls.left_is_inside,
      simple_fill = cls.left_simple_fill,
      remap = cls.no_remap,
      left_edge = cls.leftmost_segment,
      central = cls.central_segment,
      right_edge = cls.inner_right_segment,
      alignable_fragments = cls.left_alignable_fragments,
      )
  
  
  @classmethod
  def Central(cls, left, right, fragments):
    
    if not sum( [ len( f ) for f in fragments ] ) <= ( right - left ):
      raise ValueError, "Fragments do not fit into space provided"
    
    return cls(
      left = left,
      right = right,
      fragments = fragments,
      applicable_sequence = cls.central_applicable_sequence,
      is_inside = cls.central_is_inside,
      simple_fill = cls.central_simple_fill,
      remap = cls.do_remap,
      left_edge = cls.inner_left_segment,
      central = cls.central_segment,
      right_edge = cls.inner_right_segment,
      alignable_fragments = cls.central_alignable_fragments,
      )
  
  
  @classmethod
  def RightEdge(cls, left, fragments):
    
    return cls(
      left = left,
      right = None,
      fragments = fragments,
      applicable_sequence = cls.right_applicable_sequence,
      is_inside = cls.right_is_inside,
      simple_fill = cls.right_simple_fill,
      remap = cls.do_remap,
      left_edge = cls.inner_left_segment,
      central = cls.central_segment,
      right_edge = cls.rightmost_segment,
      alignable_fragments = cls.right_alignable_fragments,
      )
    
    
  @classmethod
  def Open(cls, length, fragments):
    
    len_frags = sum( [ len( f ) for f in fragments ] )
    left = ( length - len_frags ) // 2 
    return cls(
      left = left,
      right = left + length,
      fragments = fragments,
      applicable_sequence = cls.open_applicable_sequence,
      is_inside = cls.open_is_inside,
      simple_fill = cls.right_simple_fill,
      remap = cls.no_remap,
      left_edge = cls.leftmost_segment,
      central = cls.central_segment,
      right_edge = cls.rightmost_segment,
      alignable_fragments = cls.open_alignable_fragments,
      )
    
  # Sequence cutting
  @staticmethod
  def left_applicable_sequence(owner, full):
    
    return full[ : owner.right ]
  
  @staticmethod
  def central_applicable_sequence(owner, full):
    
    return full[ owner.left : owner.right ]
  
  @staticmethod
  def right_applicable_sequence(owner, full):
    
    return full[ owner.left : ]
  
  @staticmethod
  def open_applicable_sequence(owner, full):
    
    return full
  
  # Validity checking
  @staticmethod
  def left_is_inside(owner, position):
    
    return position <= owner.right
  
  @staticmethod
  def central_is_inside(owner, position):
    
    return owner.left <= position <= owner.right
  
  @staticmethod
  def right_is_inside(owner, position):
    
    return owner.left <= position
  
  @staticmethod
  def open_is_inside(owner, position):
    
    return True
  
  # Sequence remapping
  @staticmethod
  def no_remap(owner, match):
    
    pass
  
  @staticmethod
  def do_remap(owner, match):
    
    match.remap( shift = owner.left )
  
  # Fill methods
  @staticmethod
  def reverse_simple_fill(fragments, border):
  
    seed = AlignmentSeed()
    
    for f in reversed( fragments ):
      border -= len( f )
      m = HSSMatch( fragment = f, length = 0, start = border )
      seed.append( match = m )
      
    return seed
  
  @staticmethod
  def forward_simple_fill(fragments, border):
  
    seed = AlignmentSeed()
    
    for f in fragments:
      m = HSSMatch( fragment = f, length = 0, start = border )
      seed.append( match = m )
      border += len( f ) 
      
    return seed
  
  @staticmethod
  def left_simple_fill(owner):
    
    return UnalignedRegion.reverse_simple_fill(
      fragments = owner.fragments,
      border = owner.right,
      )
    
  @staticmethod
  def right_simple_fill(owner):
    
    return UnalignedRegion.forward_simple_fill(
      fragments = owner.fragments,
      border = owner.left,
      )

  @staticmethod
  def central_simple_fill(owner):
    
    half = len( owner.fragments ) // 2 + 1
    seed1 = UnalignedRegion.forward_simple_fill(
      fragments = owner.fragments[ : half ],
      border = owner.left,
      )
    seed2 = UnalignedRegion.reverse_simple_fill(
      fragments = owner.fragments[ half : ],
      border = owner.right,
      )
    seed1.merge( seed = seed2 )
    return seed1
  
  # Division management
  @staticmethod
  def leftmost_segment(owner, right, fragments):
    
    if not owner.is_inside( position = right ):
      raise ValueError, "Overflow on right edge"
    
    return UnalignedRegion.LeftEdge( right = right, fragments = fragments )
  
  @staticmethod
  def inner_left_segment(owner, right, fragments):
    
    if not owner.is_inside( position = right ):
      raise ValueError, "Overflow on right edge"
    
    return UnalignedRegion.Central(
      left = owner.left,
      right = right,
      fragments = fragments,
      )
  
  @staticmethod
  def central_segment(owner, left, right, fragments):
    
    if not owner.is_inside( position = left ):
      raise ValueError, "Overflow on left edge"
    
    if not owner.is_inside( position = right ):
      raise ValueError, "Overflow on right edge"
    
    return UnalignedRegion.Central(
      left = left,
      right = right,
      fragments = fragments,
      )
    
  @staticmethod
  def inner_right_segment(owner, left, fragments):
    
    if not owner.is_inside( position = left ):
      raise ValueError, "Overflow on left edge"
    
    return UnalignedRegion.Central(
      left = left,
      right = owner.right,
      fragments = fragments,
      )
  
  @staticmethod
  def rightmost_segment(owner, left, fragments):
    
    if not owner.is_inside( position = left ):
      raise ValueError, "Overflow on left edge"
    
    return UnalignedRegion.RightEdge( left = left, fragments = fragments )
  
  # Fragment partitioning
  @staticmethod
  def left_alignable_fragments(owner, full):
    
    if owner.right == 0:
      return []
    
    length = 0
    
    for ( index, fragment ) in enumerate( reversed( owner.fragments ), start = 1 ):
      length += len( fragment )
      
      if owner.right < length:
        return owner.fragments[ -index : ]
      
    return owner.fragments
  
  @staticmethod
  def central_alignable_fragments(owner, full):
    
    return owner.fragments
  
  @staticmethod
  def right_alignable_fragments(owner, full):
    
    seqlen = len( full )
    
    if seqlen <= owner.left:
      return []
    
    length = owner.left
    
    for ( index, fragment ) in enumerate( owner.fragments, start = 1 ):
      length += len( fragment )
      
      if seqlen < length:
        return owner.fragments[ : index ]
      
    return owner.fragments
  
  @staticmethod
  def open_alignable_fragments(owner, full):
    
    return owner.fragments


class AlignmentSeed(object):
  """
  A set of HSS matches for an alignment
  """
  
  def __init__(self, hssmatches = []):
    
    self.match_for = {}
    
    for m in hssmatches:
      self.append( match = m )
    
  
  def _append_unchecked(self, match):
    
    self.match_for[ match.fragment ] = match
    
    
  def compatible_with(self, match):
    
    return all( match.compatible_with( other = m ) for m in self.matches )
    
    
  def append(self, match):
    
    if not self.compatible_with( match = match):
      raise ValueError, "%s incompatible with seed" % match
    
    self._append_unchecked( match = match )
    
    
  def remove(self, match):
    
    if match not in self.matchset():
      raise ValueError, "Unknown match %s" % match
    
    del self.match_for[ match.fragment ]
    
    
  def merge(self, seed):
    
    if not all( self.compatible_with( match = m ) for m in seed.matches ):
      raise ValueError, "Incompatible matches during seed merge"
    
    for match in seed.matches:
      self._append_unchecked( match = match )
      
      
  def overlaps_with(self, start, stop):
    
    return sum(
      ( max( min( stop, m.stop ), start ) - min( max( start, m.start ), stop ) )
      for m in self.matches
      )
    
  
  def remap(self, shift):
    
    for match in self.matches:
      match.remap( shift = shift )
    
    
  @property
  def matches(self):
    
    return self.match_for.values()
  
  
  def matchset(self):
    
    return frozenset( self.matches )
  
  
  @property
  def fragments(self):
    
    return self.match_for.keys()
    
    
class PartialAlignment(object):
  """
  A partial alignment from hss matches
  """
  
  def __init__(self, sequence, seed, unaligneds):
    
    self.sequence = sequence
    self.seed = seed
    self.unaligned_in = {}
    
    for uar in unaligneds:
      for frag in uar.fragments:
        self.unaligned_in[ frag ] = uar 
    
  @property
  def unaligneds(self):
    
    return set( self.unaligned_in.values() )
  
  
  @property
  def unmatcheds(self):
    
    return sorted( self.unaligned_in )
    
    
  def merge(self, seed):
    
    for match in seed.matches:
      if match.fragment not in self.unaligned_in:
        raise ValueError, "%s unknown" % match.fragment
      
      if not self.seed.compatible_with( match = match ):
        raise ValueError, "%s incompatible with seed" % match
      
    for match in seed.matches:
      uar = self.unaligned_in[ match.fragment ]
      uar.fragments.remove( match.fragment )
      del self.unaligned_in[ match.fragment ]
      self.seed.append( match = match )
  
  
  def identities(self):
    
    return sum(
      m.overlap_region_alignment( sequence = self.sequence ).identities()
      for m in self.seed.matches
      )
    
    
  def overlaps(self):
    
    return self.seed.overlaps_with( start = 0, stop = len( self.sequence ) )
  
  
  def alignment_start(self):
    
    matches = self.seed.matches
    
    if not matches:
      return 0
    
    first = min( matches, key = lambda m: m.fragment )
    return min ( first.start, 0 )
  
  
  def alignment_stop(self):
    
    length = len( self.sequence )
    matches = self.seed.matches
    
    if not matches:
      return length
    
    last = max( matches, key = lambda m: m.fragment )
    return max ( last.stop, length )
    
    
  def alignment(self):
    
    matches = sorted( self.seed.matches, key = lambda m: m.fragment )
    
    if not matches:
      return Alignment( pairs = [ ( c, None ) for c in self.sequence ] )
    
    previous = matches[0]
    alignment = Alignment.empty(
      sequence = self.sequence,
      start = 0,
      stop = previous.start,
      )
    alignment += previous.overlap_region_alignment( sequence = self.sequence ) 
    
    for following in matches[1:]:
      alignment += Alignment.empty(
        sequence = self.sequence,
        start = previous.stop,
        stop = following.start,
        )
      alignment += following.overlap_region_alignment( sequence = self.sequence )
      previous = following
      
    alignment += Alignment.empty(
      sequence = self.sequence,
      start = previous.stop,
      stop = len( self.sequence ),
      )
    
    return alignment
  
# Fragment search
class HSSFind(object):
  """
  Prepares a sequence for quick HSS calculations
  """
  
  def __init__(self, sequence):
    
    self.tree = singletree.tree()
    builder = singletree.ukkonen( tree = self.tree )
    
    for char in sequence:
      builder.append( glyph = char )
      
    builder.append( glyph = object() )
    builder.detach()
    self.leaf_labels_under = suffixtree.calculate_leaf_indices( root = self.tree.root )
    
    
  def overlap_with(self, sequence, min_length = 1):
    
    msi = singletree.matching_statistics(
      tree = self.tree,
      iterable = iter( sequence ),
      )
    
    for ( start2, ( length, ( edge, pos ) ) ) in enumerate( msi ):
      if length < min_length:
        continue
      
      for start1 in self.leaf_labels_under[ edge ]:
        yield ( start1 - start2, length )
        
        
  def frameshift_with(self, sequence, min_length = 1):
    
    length_for = {}
    
    for ( i, l ) in self.overlap_with( sequence = sequence, min_length = min_length ):
      length_for[ i ] = max( l, length_for.get( i, 0 ) )
      
    return length_for
  
  
  def matches_with(self, fragment, min_length = 1):
    
    frameshift_with = self.frameshift_with(
      sequence = fragment.sequence,
      min_length = min_length,
      )
    
    return [
      HSSMatch( fragment = fragment, length = length, start = start )
      for ( start, length ) in frameshift_with.items()
      ]
  

# Helpers
def reassemble(hssmatches):
  
  for match in hssmatches:
    others = [ m for m in hssmatches if match.compatible_with( other = m ) ]
    
    if not others:
      yield ( match, )
      yield ()
    
    else:
      for compat in reassemble( hssmatches = others ):
        yield ( match, ) + compat
        
      without = [ m for m in hssmatches if m.fragment != match.fragment ]
      
      for compat in reassemble( hssmatches = without ):
        yield compat


def alignment_gap_fill_search(sequence, unaligned, max_hss_count = None):
  
  applicable = unaligned.applicable_sequence( full = sequence )
  hssfind = HSSFind( sequence = applicable )
  alignable = unaligned.alignable_fragments( full = sequence )
  
  hssmatches = reduce(
    operator.add,
    [ hssfind.matches_with( fragment = f, min_length = 1 ) for f in alignable ],
    [],
    )
  
  for m in hssmatches:
    unaligned.remap( match = m )
    
  hssmatches = [ m for m in hssmatches if unaligned.compatible_with( match = m ) ]
  hssmatches.sort( key = lambda m: m.length, reverse = True )
  hssmatches = hssmatches[ : max_hss_count ]
  
  found = set()
  
  for comb in reassemble( hssmatches = hssmatches ):
    if not comb:
      continue
    
    seed = AlignmentSeed( hssmatches = comb )
    
    if seed.matchset() in found:
      continue
    
    found.add( seed.matchset() )
    
    present = set( seed.fragments )
    min_p = min( present )
    max_p = max( present )
    missing = [ f for f in unaligned.fragments if f not in present ]
    
    still_unaligneds = []
    
    try:
      still_unaligneds.append(
        unaligned.left_edge(
          right = seed.match_for[ min_p ].start,
          fragments = [ f for f in missing if f < min_p ],
          )
        )
      
    except ValueError:
      continue
    
    middle = [ f for f in missing if min_p < f < max_p ]
    
    try:
      for s in sequtil.split( middle, consecutivity = Fragment.consecutivity ):
        leftside = max( [ f for f in present if f < s[0] ] )
        rightside = min( [ f for f in present if s[-1] < f ] )
        still_unaligneds.append(
          unaligned.central(
            left = seed.match_for[ leftside ].stop,
            right = seed.match_for[ rightside ].start,
            fragments = s,
            )
          )
      
    except ValueError:
      continue
    
    try:
      still_unaligneds.append(
        unaligned.right_edge(
          left = seed.match_for[ max_p ].stop,
          fragments = [ f for f in missing if max_p < f ],
          )
        )
      
    except ValueError:
      continue
    
    for uar in still_unaligneds:
      seed.merge( seed = uar.simple_fill() )
      
    yield seed
    

# Main functions
def empty_partial_alignment(sequence, fragments):
  
  uar = UnalignedRegion.Open(
    length = len( sequence ),
    fragments = [
      Fragment( index = i, sequence = f ) for ( i, f ) in enumerate( fragments )
      ],
    )
  return PartialAlignment(
    sequence = sequence,
    seed = AlignmentSeed(), 
    unaligneds = [ uar ],
    )


def partial_alignments(sequence, fragments, min_hss_length, max_hss_count = None):
  
  frags = [
    Fragment( index = i, sequence = f ) for ( i, f ) in enumerate( fragments )
    ]
  hssfind = HSSFind( sequence = sequence )
  
  hssmatches = reduce(
    operator.add,
    [
      hssfind.matches_with( fragment = f, min_length = min_hss_length )
      for f in frags
      ],
    [],
    )
  hssmatches.sort( key = lambda m: m.length, reverse = True )
  hssmatches = hssmatches[ : max_hss_count ]
  
  found = set()
  
  for comb in reassemble( hssmatches = hssmatches ):
    if not comb:
      continue
    
    seed = AlignmentSeed( hssmatches = comb )
    
    if seed.matchset() in found:
      continue
    
    found.add( seed.matchset() )
    
    present = set( seed.fragments )
    min_p = min( present )
    max_p = max( present )
    missing = [ f for f in frags if f not in present ]
    
    unaligneds = []
    unaligneds.append(
      UnalignedRegion.LeftEdge(
        right = seed.match_for[ min_p ].start,
        fragments = [ f for f in missing if f < min_p ],
        )
      )
    middle = [ f for f in missing if min_p < f < max_p ]
    
    try:
      for s in sequtil.split( middle, consecutivity = Fragment.consecutivity ):
        unaligneds.append(
          UnalignedRegion.Central(
            left = seed.match_for[ frags[ s[0].index - 1 ] ].stop,
            right = seed.match_for[ frags[ s[-1].index + 1 ] ].start,
            fragments = s,
            )
          )
      
    except ValueError:
      continue
    
    unaligneds.append(
      UnalignedRegion.RightEdge(
        left = seed.match_for[ max_p ].stop,
        fragments = [ f for f in missing if max_p < f ],
        )
      )
      
    yield PartialAlignment(
      sequence = sequence,
      seed = seed, 
      unaligneds = unaligneds,
      )


def complete_partial_alignment(
  partial,
  max_hss_count = None,
  min_identity_fraction = 0,
  ):
  
  for unaligned in partial.unaligneds:
    seeds = list(
      alignment_gap_fill_search(
        sequence = partial.sequence,
        unaligned = unaligned,
        max_hss_count = max_hss_count,
        )
      )
    
    if seeds:
      best = max(
        seeds,
        key = lambda s: (
          s.overlaps_with( start = 0, stop = len( partial.sequence ) ),
          sum(
            m.overlap_region_alignment( sequence = partial.sequence ).identities()
            for m in s.matches
            ),
          )
        )
      
    else:
      best = unaligned.simple_fill()
      
    for match in best.matches:
      ali = match.overlap_region_alignment( sequence = partial.sequence )
      
      if ali.identities() < len (ali ) * min_identity_fraction:
        best.remove( match = match )
      
    partial.merge( seed = best )
    

class Equivalent(object):
  """
  Establish correspondence between a value and how it should be compared
  """
  
  def __init__(self, value, compare):
    
    self.value = value
    self.compare = compare
    
    
  def __eq__(self, other):
    
    if isinstance( other, self.__class__ ):
      return self.compare == other.compare
    
    else:
      return self.compare == other
    
    
  def __ne__(self, other):
    
    return not ( self == other )
  
  
  def __hash__(self):
    
    return hash( self.compare )
  

def fragments_to_sequence(
  sequence,
  fragments,
  one_letter_for,
  unknown_one_letter = "X",
  min_hss_length = 3,
  max_seed_hss_count = 12,
  max_completion_hss_count = 6,
  ):
  
  processeds = []
  
  for rgs in fragments:
    ags = [ rg.atom_groups() for rg in rgs ]
    compares = [
      one_letter_for.get( ag[0].resname if ag else "UNK", unknown_one_letter )
      for ag in ags
      ]
    seq = [
      Equivalent( value = rg, compare = comp ) for ( rg, comp ) in zip( rgs, compares )
      ]
    processeds.append( seq )
  
  try:
    best = max(
      partial_alignments(
        sequence = sequence,
        fragments = processeds,
        min_hss_length = min_hss_length,
        max_hss_count = max_seed_hss_count,
        ),
      key = lambda p: p.identities(),
      )
    
  except ValueError:
    best = empty_partial_alignment(
      sequence = sequence,
      fragments = processeds,
      )
    
  complete_partial_alignment(
    partial = best,
    max_hss_count = max_completion_hss_count,
    )
  
  return best


def map_to_sequence_frame(sequence, gap_char = "-"):
  
  return [
    Equivalent( value = i, compare = c ) for ( i, c ) in enumerate( sequence )
    if c != gap_char
    ]


def map_back_to_alignment_frame(partial):
  
  if not partial.sequence:
    return []
    
  alignment = partial.alignment()
  sequence = [ p for p in alignment.pairs if p[0] is not None ]
  assert sequence
  last = sequence[-1][0].value
  mapping = [ None ] * ( last + 1 )
  
  for ( seq, frag ) in sequence:
    if frag is None:
      continue
    
    assert seq
    assert seq.value <= last
    mapping[ seq.value ] = frag.value
    
  return mapping