from __future__ import division
import operator

import libtbx.phil
from phaser import rsam
from phaser import output


class SequenceModelCorrespondence(object):
    """
    Keep together sequence and model
    """
    
    def __init__(self, sequence, chain):
        
        self.residue_groups = chain.residue_groups()
        self.rg_index_for = dict( [ ( rg.resseq.strip(), index )
            for ( index, rg ) in enumerate( self.residue_groups ) ] )
        
        from phaser import mmt
        
        self.mmt = mmt.determine( chain = chain )
        self.sequence = sequence
        self.alignment = rsam.align_residues_to_sequence(
            residues = self.residue_groups,
            sequence = self.sequence,
            consecutivity = self.mmt.consecutivity,
            one_letter_for = self.mmt.one_letter_for(),
            unknown = self.mmt.unknown_one_letter,
            )
        
    
    def residue_group_identified_with(self, resseq):
        
        return self.rg_index_for[ resseq ]
    
    
    def residue_group_segment(self, start, stop):
        
        return self.residue_groups[
            self.residue_group_identified_with( resseq = start )
            :
            self.residue_group_identified_with( resseq = stop ) + 1 
            ]
        
        
    def sequence_position_for(self, residue_group):
        
        try:
            index = self.alignment.index( residue_group )
            
        except ValueError:
            index = None
        
        return index
        
        
    def sequence_segment(self, start, stop):
        
        posis = [ self.sequence_position_for( residue_group = rg )
            for rg in self.residue_group_segment( start = start, stop = stop ) ]
        known_posis = [ p for p in posis if p is not None ]
        
        if not known_posis:
            return []
        
        return range( min( known_posis ), max( known_posis ) + 1 )
            

class SequenceRemapInfo(object):
    """
    Contains data to remap ChainSequenceSelections to other sequences
    """            
    
    def __init__(self, seq_from, seq_to, gap):
        
        if len( seq_from ) != len( seq_to ):
            raise ValueError, "Unequal sequence lengths"
        
        n_seq_from = self.number_sequence( sequence = seq_from, gap_char = gap )
        n_seq_to = self.number_sequence( sequence = seq_to, gap_char = gap )
        
        self.other_length = max( n_seq_to ) + 1
        self.other_resseq_for = dict(
            [ ( f, t ) for ( f, t ) in zip( n_seq_from, n_seq_to )
                if f is not None ]
            )
        
    @staticmethod
    def number_sequence(sequence, gap_char):
        
        import itertools
        counter = itertools.count()
        return [ counter.next() if p != gap_char else None for p in sequence ]
    
    
class SegmentedSequence(object):
    """
    A sequence that is in multiple segments
    """
    
    def __init__(self, segments):
        
        assert segments
        self.segments = segments
        self.i_current = 0
        
        
    def current(self):
        
        return self.segments[ self.i_current ]
    
    
    def current_stats(self):
        
        edge = sum( len( self.segments[ i ] ) for i in range( self.i_current ) ) 
        return ( edge, edge + len( self.segments[ self.i_current ] ) )
    
    
    def left(self):
        
        index = self.i_current - 1
        
        if index < 0:
            return []
        
        else:
            return self.segments[ index ]
        
        
    def right(self):
        
        index = self.i_current + 1
        
        if len( self.segments ) <= index:
            return []
        
        else:
            return self.segments[ index ]
        
        
    def merged(self):
        
        return reduce( operator.add, self.segments )
        
    
    def at_left_edge(self):
        
        return self.i_current == 0
    
    
    def at_right_edge(self):
        
        return self.i_current == len( self.segments ) - 1
    
    
    def forward(self):
        
        if not self.at_right_edge():
            self.i_current += 1
        
        
    def backward(self):
        
        if not self.at_left_edge():
            self.i_current -= 1
        
    
    def merge_left(self):
        
        if not self.at_left_edge():
            self.left().extend( self.current() )
            del self.segments[ self.i_current ]
            self.backward()
        
     
    def merge_right(self):
        
        if not self.at_right_edge():
            self.current().extend( self.right() )
            del self.segments[ self.i_current + 1 ]
            
            
    def transform(self, transformation):
        
        current = self.current()
        
        for ( index, value ) in enumerate( current ):
            current[ index ] = transformation( value )

        
class ChainSequenceSelection(object):
    """
    Allows selecting sequence stretches based on residue ids
    """
    
    def __init__(self, count):
        
        self.count = count
        self.iselection = set()
        
        
    @property
    def selection(self):
        
        return [ True if i in self.iselection else False for i in range( self.count ) ]
    
    
    @property
    def chain_length(self):
        
        return self.count
        
        
    def copy(self):
        
        return self.from_iselection( count = self.count, iselection = self.iselection )
                
    
    def insert(self, position):
        
        self.iselection.add( position )
        
        
    def remove(self, position):
        
        self.iselection.remove( position )
    
    
    def add(self, segment):
        
        for index in segment:
            self.insert( position = index )
            
            
    def delete(self, segment):
        
        for index in segment:
            self.remove( position = index )
    
    
    def remap(self, remap_info):
        
        raw = [ remap_info.other_resseq_for[ i ] for i in self.iselection ]
        
        return self.from_iselection(
            count = remap_info.other_length,
            iselection = [ i for i in raw if i is not None ]
            )
        
        
    def overlap_with(self, other):
        
        return self.from_iselection(
            count = self.count,
            iselection = self.iselection.intersection( other.iselection )
            )
    
    
    def difference_from(self, other):
    
        return self.from_iselection(
            count = self.count,
            iselection = self.iselection.difference( other.iselection )
            )
        
        
    def segments(self):
        
        segments = rsam.split(
            sequence = sorted( self.iselection ),
            consecutivity = lambda left, right: left == right - 1
            )
        
        return [ ( min( s ), max( s ) ) for s in segments ]
    
    
    def complement(self):
        
        return self.from_iselection(
            count = self.count,
            iselection = set( range( self.count ) ).difference( self.iselection ) 
            )
        
    
    def __len__(self):
        
        return len( self.iselection )
    
    
    def __str__(self):
    
        segments = self.segments()
        lines = []
         
        for ( i_seg, ( lower, upper ) ) in enumerate( segments, start = 1 ):
            lines.append( "Segment %s: %s - %s" % ( i_seg, lower, upper ) )
            
        return "\n".join( lines )

            
    @classmethod
    def from_segments(cls, smc, segments):
        
        css = cls( count = len( smc.sequence ) )
        
        for segment in segments:
            css.add(
                segment = smc.sequence_segment(
                    start = segment.pdb_start,
                    stop = segment.pdb_stop
                    )
                )
            
        return css
    
    
    @classmethod
    def from_iselection(cls, count, iselection):
        
        css = cls( count = count )
        
        for posi in iselection:
            css.insert( position = posi )
            
        return css
    
    
    @classmethod
    def from_selection(cls, selection):
        
        css = cls( count = len( selection ) )
            
        for ( posi, is_selected ) in enumerate( selection ):
            if is_selected:
                css.insert( position = posi )
                
        return css
    
    
    @staticmethod
    def common(*selections):
        
        if not selections:
            raise ValueError, "No selections for merging"
        
        count = selections[0].count
        
        if any( count != s.count for s in selections[1:] ):
            raise ValueError, "Incompatible selections"
        
        return count
    
    
    @classmethod
    def union(cls, *selections):
        
        return cls.from_iselection(
            count = cls.common( *selections ),
            iselection = selections[0].iselection.union(
                *[ s.iselection for s in selections[1:] ]
                )
            )
    
    
    @classmethod
    def overlap(cls, *selections):
        
        return cls.from_iselection(
            count = cls.common( *selections ),
            iselection = selections[0].iselection.intersection(
                *[ s.iselection for s in selections[1:] ]
                )
            )
    
    
class MergingStatistics(object):
    """
    Simple statistics on how many times a position was selected
    """
    
    def __init__(self, selections):
        
        self.stats_for = {}
        
        for s in selections:
            for index in s.iselection:
                self.stats_for[ index ] = self.stats_for.setdefault( index, 0 ) + 1
                
                
    def __getitem__(self, key):
        
        if key in self.stats_for:
            return self.stats_for[ key ]
        
        below = max( [ k for k in self.stats_for if k < key ] + [ 0 ] )
        above = min( [ k for k in self.stats_for if key < k ] + [ 0 ] )
        
        if above == 0 and below == 0:
            return 0
        
        elif below == 0:
            return self.stats_for[ above ]
            
        elif above == 0:
            return self.stats_for[ below ]
        
        else:
            return 1.0 / ( above - below ) * (
                ( above - key ) * self.stats_for[ below ]
                + ( key - below ) * self.stats_for[ above ]
                )
        
    
def polish(domain, short = 5, long = 15):
        
    segments = SegmentedSequence(
        segments = list(
            rsam.split(
                sequence = domain.selection,
                consecutivity = lambda left, right: left == right
                )
            )
        )
    
    while not segments.at_right_edge():
        if long <= len( segments.current() ):
            segments.forward()
            continue
        
        left = segments.left()
        right = segments.right()
        current = segments.current()
            
        # medium overhang on terminii
        if ( len( left ) == 0 and long <= len( right ) ) or ( len( right ) == 0 and long <= len( left ) ):
            segments.transform( transformation = lambda s: not s ) # flip
            segments.merge_left()
            segments.merge_right()
            continue
        
        # short different section between two long ones
        if len( current ) < short and ( long <= len( left ) + len( right ) ):
            segments.transform( transformation = lambda s: not s ) # flip
            segments.merge_left()
            segments.merge_right()
            continue
        
        segments.forward()
        
    return ChainSequenceSelection.from_selection( selection = segments.merged() )
    
    
def map_domains_on_target(annotation, alignment, chain, min_domain_length):
    
    smc = SequenceModelCorrespondence(
        sequence = alignment.sequence_strings()[1],
        chain = chain
        )
    sri = SequenceRemapInfo(
        seq_from = alignment.alignments[1],
        seq_to = alignment.alignments[0],
        gap = alignment.gap
        )
    domains = []
    
    for dom in [ d for d in annotation.root.query.domains if d.chain_code.strip() == chain.id ]:
        css = ChainSequenceSelection.from_segments( smc = smc, segments = dom.segments )
        selection = css.remap( remap_info = sri )
        
        if len( selection ) < min_domain_length:
            continue
        
        domains.append( selection )
        
    return domains
        
        
def average_domain_definitions(assignments, cluster_radius, polish_function):
            
    assert assignments
    singles = reduce( operator.add, assignments, [] )
    
    if len( singles ) == 1:
        levels = [ singles ]
        
    else:
        from libtbx import cluster
        cl = cluster.HierarchicalClustering(
            data = singles,
            distance_function = domain_clustering_distance
            )
        levels = cl.getlevel( cluster_radius )
        
    polisheds = [
        polish_function( domain = ChainSequenceSelection.union( *l ) )
        for l in levels
        ]
         
    domain_for = dict(
        reduce(
            operator.add,
            [ [ ( d, m ) for d in l ] for ( l, m ) in zip( levels, polisheds ) ]
            )
        )
    
    alternatives = set()
    
    for domains in assignments:
        alternatives.add( frozenset( [ domain_for[ d ] for d in domains ] ) )
        
    stats_for = dict(
        [ ( m, MergingStatistics( selections = l ) )
            for ( m, l ) in zip( polisheds, levels ) ]
        )
    results = []
    
    for domains in alternatives:
        overlap = ChainSequenceSelection.overlap( *domains )
        stats = [ ( d, stats_for[ d ] ) for d in domains ]
        
        for pos in overlap.iselection:
            best = max( stats, key = lambda p: p[1][ pos ] )[0]
            
            for domain in [ d for d in domains if d != best ]:
                domain.remove( position = pos )
                
        missing = ChainSequenceSelection.union( *domains ).complement()
        results.append( ( domains, missing ) )
                
    return results
                
                
def domain_clustering_distance(left, right):
    
    return 1.0 / ( len( left.overlap_with( other = right ) ) + 1 ) * min(
        len( left.difference_from( other = right ) ),
        len( right.difference_from( other = left ) ),
        )

   
master_params = libtbx.phil.parse(
    """
    input
        .help = "Input files"
    {
        sequence = None
            .help = "Target sequence"
            .type = path
            
        homology = None
            .help = "Homology search performed with target sequence"
            .type = path
            .multiple = True
            
        max_hits = 3
            .help = "Number of hits to extract from homology search"
            .type = int
            .optional = False
    }
    
    configuration
        .help = "Calculation parameters"
    {
        min_domain_length = 10
            .help = "Minimum number of residues in a domain to include in calculations"
            .type = int
            .optional = False
            
        cluster_radius = 0.2
            .help = "Distance between domains to be identified as equivalent"
            .type = float
            .optional = False
            
        polish_short = 5
            .help = "Short segment length for polish algorithm"
            .type = int
            .optional = False
            
        polish_long = 15
            .help = "Long segment length for polish algorithm"
            .type = int
            .optional = False
    }
    """
    )

PROGRAM = "domain_analysis"
VERSION = "0.0.1"

        
def process(sequence, hits, params, logger):
    
    from libtbx.utils import Sorry
    
    logger.info(
        msg = output.underlined( text = "Downloading domain annotations from CATH" )
        )
    assignments = []
    
    from phaser.pipeline import homology
    from phaser.pipeline.proxy import cath
    
    try:
        cath.check()
        
    except RuntimeError, e:
        raise Sorry, "Error: %s" % e
    
    for hit in hits:
        ann = homology.parse_cath_xml(
            data = cath.fetch( identifier = hit.identifier ).data
            )
        
        if not ann or ann.empty():
            logger.info(
                msg = "No annotation available for hit %s_%s" % (
                    hit.identifier,
                    hit.chain,
                    )
                )
            continue
        
        logger.info(
            msg = "Downloaded annotation for %s_%s" % ( hit.identifier, hit.chain )
            )
        assignments.append( ( hit, ann ) )
        
    logger.info( msg = "There are %s annotations available" % len( assignments ) )
    
    if not assignments:
        raise Sorry, "No domain annotations available"
    
    logger.info(
        msg = output.underlined( text = "Mapping domain annotations to target sequence" )
        )
    
    from phaser import proxy
    from phaser.pipeline.proxy import ParameterizedProxy, dbfetch
    
    try:
        dbfetch.check()
        
    except RuntimeError, e:
        raise Sorry, "Error: %s" % e
    
    pdb = proxy.PDBFetch(
        proxy = ParameterizedProxy(
            proxy = dbfetch,
            params = dbfetch.QueryParams.PDB( format = "pdb" )
            )
        )
    pdb_redirections = proxy.get_pdb_redirections()
    
    from phaser.pipeline import expert
    domains = []
    
    for ( hit, ann ) in assignments:
        try:
            root = expert.download_homology_hit_model(
                hit = hit,
                pdb = pdb,
                redirections = pdb_redirections,
                stream = logger.verbose
                )
            
        except RuntimeError, e:
            logger.info(
                msg = "Could not fetch %s_%s: %s" % ( hit.identifier, hit.chain, e )
                )
            continue
        
        alignment = hit.alignment.copy()
        rsam.alignment_stich( sequence = sequence, alignment = alignment, index = 0 )
        
        accepteds = map_domains_on_target(
            annotation = ann,
            alignment = alignment,
            chain = root.models()[0].chains()[0],
            min_domain_length = params.min_domain_length
            )
        
        if len( accepteds ) == 0:
            logger.info(
                msg = "%s_%s: no domains map on target sequence" % (
                    hit.identifier,
                    hit.chain,
                    )
                )
            continue
        
        logger.info(
            msg = "%s_%s: %s domains mapped" % (
                hit.identifier,
                hit.chain,
                len( accepteds ),
                )
            )
        domains.append( accepteds )
        
    if not domains:
        raise Sorry, "No domain annotations map on target structure"
        
    logger.info( msg = output.underlined( text = "Domain analysis" ) )
    
    import functools
    decompositions = average_domain_definitions(
        assignments = domains,
        cluster_radius = params.cluster_radius,
        polish_function = functools.partial(
            polish,
            short = params.polish_short,
            long = params.polish_long,
            )
        )
    
    logger.info(
        msg = "Alternative domain decompositions found: %s" % len( decompositions )
        )
    logger.info( output.blank() )
    
    for ( index, ( domains, missing ) ) in enumerate( decompositions, start = 1 ):
        logger.info(
            msg = output.heading(
                text = "%s. possibility: %s domains" % ( index, len( domains ) )
                )
            )
        logger.indent()
        
        for ( i_dom, dom ) in enumerate( domains, start = 1 ):
            logger.info( msg = "Domain %s:" % i_dom )
            logger.indent()
            logger.info( msg = str( dom ) )
            logger.dedent()
        
        logger.info( msg = "Not annotated:" )
        logger.indent()
        logger.info( msg = str( missing ) )
        logger.dedent()
        logger.dedent()
        
    return decompositions


def run(args, out, verbosity):
    
    logger = output.SingleStream( stream = out, level = verbosity )
    
    logger.info(
        msg = output.banner( text = "%s version %s" % ( PROGRAM, VERSION ) )
        )
    
    from phaser import tbx_utils
    
    accumulator = tbx_utils.PhilAccumulator( master_phil = master_params )
    
    from iotbx import bioinformatics
    
    accumulator.register_file_handler(
        handler = tbx_utils.ExtensionFileHandler(
            extensions = [ ".xml", ".hhr" ],
            template = "input.homology=%s"
            )
        )
    accumulator.register_file_handler(
        handler = tbx_utils.ExtensionFileHandler(
            extensions = bioinformatics.known_sequence_formats(),
            template = "input.sequence=%s"
            )
        )
    
    for argument in args:
        argument.process( accumulator = accumulator )
        
    from libtbx.utils import Sorry
    
    try:
        merged_phil = accumulator.merge()
        
    except RuntimeError, e:
        raise Sorry, e
    
    logger.info( msg = output.underlined( text = "All configuration options:" ) )
    logger.info( msg = merged_phil.as_str() )
    
    try:
        params = merged_phil.extract()
        
    except RuntimeError, e:
        raise Sorry, e
    
    if params.input.sequence is None:
        raise Sorry, "No target sequence defined"
    
    if not params.input.homology:
        raise Sorry, "No homology searches defined"
    
    hits = []
    
    for fname in params.input.homology:
        fobject = tbx_utils.HomologySearchObject.from_file( file_name = fname )
        found = list( fobject.object.hits() )[ : params.input.max_hits ]
        logger.info( msg = "Extracted %s hits from %s" % ( len( found ), fname ) )
        hits.extend( found )
        
    sobject = tbx_utils.SequenceObject.from_file( file_name = params.input.sequence )
    
    if len( sobject.object ) == 0:
        raise Sorry, "No sequences found in %s" % sobject.name
    
    elif 1 < len( sobject.object ):
        logger.warning(
            msg = "Multiple sequences found in %s, using first" % sobject.name
            )
        
    process(
        sequence =  sobject.object[0],
        hits = hits,
        params = params.configuration,
        logger = logger
        )