from __future__ import division

from phaser import tbx_utils
from phaser import sculptor
from phaser import ncs
from phaser import proxy

import iotbx.pdb
import iotbx.pdb.fetch
from iotbx import bioinformatics
import libtbx.phil
from libtbx.utils import Sorry
from scitbx.array_family import flex

import os
import cStringIO

PROGRAM = "Model preparation"

PHIL_MASTER = libtbx.phil.parse(
    """input
        .help = "Input files"
    {
        homology_search
            .help = "Homology search"
        {
            file_name = None
                .help = "File name"
                .optional = False
                .type = path
                
            max_hits = None
                .help = "Max number of hits"
                .optional = True
                .type = int
        }
        
        remove_alternate_conformations = False
            .help = "Remove alternate conformations"
            .type = bool
            .optional = False
            
        max_model_count = 1
            .help = "Max models to keep when processing an entry with multiple structure"
            .type = int
            .optional = True
    }
    
    output
        .help = "Output options"
    {
        folder = .
            .help = "Output folder; None is current folder"
            .type = path
            .optional = False
    }
    
    sculptor
        .help = "Sculptor keywords"
    {
        %(sculptor)s
        %(other)s
    }
    """ % {
        "sculptor": sculptor.PHIL_SCULPTOR,
        "other": sculptor.PHIL_OTHER,
        }
    )

def process_command_line(args, logger):
    
    ( phil_objects, files, unknowns ) = tbx_utils.process_command_line(
        args = args,
        phil_master = PHIL_MASTER
        )
    
    if unknowns:
        usage( logger = logger )
        raise Sorry, "Unknown keyword(s): %s" % ", ".join( unknowns )
    
    # Identify files
    ( files_for, unknowns ) = tbx_utils.group_files_on_extension(
        files = files,
        homology = [ ".hhr", ".xml" ]
        )
    
    if unknowns:
        usage( logger = logger )
        raise Sorry, "Unknown file(s): %s" % ", ".join( unknowns )
    
    if "homology" in files_for:
        hss = files_for[ "homology" ]
        count = len( hss )
        logger.info(
            "Guessing: %s %s %s" % (
                ", ".join( hss ),
                tbx_utils.conjugate( verb = "is", count = count ),
                tbx_utils.inflect( word = "homology search file", count = count )
                )
            )
        phil_objects.extend(
            [ libtbx.phil.parse( "%s=%s" % ( "input.homology_search.file_name", f ) )
                for f in hss ]
            )
        
    # Merge phil sources
    merged_phil = PHIL_MASTER.fetch( sources = phil_objects )
    
    logger.info( "" )
    logger.info( "All configuration options:\n" )
    logger.info( merged_phil.as_str() )
    
    params = merged_phil.extract()
    
    return params 


def usage(logger):
    
    logger.info( "Prepares MR models for MR_Rosetta" )
    logger.info( "Required input:" )
    logger.info( "hhpred_file - results of homology search in hhpred format" )


def find_chains_identical_to(reference, others, logger):
    
    logger.info(
        "Finding chains identical in sequence to chain %s" % reference.id
        )

    identicals = set( [ reference.id ] )
    
    if not others.atoms():
      return identicals

    ref = ncs.ReferenceChain(
        chain = reference,
        macromolecule_type = ncs.MacromoleculeType.protein
        )
    
    for chain in others.models()[0].chains():
        logger.info( "Comparing with chain %s" % chain.id )
        ( overlap, offset ) = ref.determine_numbering_shift( chain = chain )
        min_accepted = max( 10, len( chain.residue_groups() ) // 2 )
        logger.info( "Minimum acceptable HSSP length is %s" % min_accepted )
        logger.info( "Length of longest HSSP found: %s" % len( overlap ) )
        
        if len( overlap ) < min_accepted:
            logger.info( "The longest HSSP is too short" )
            logger.info( "Discarding this chain" )
            continue
        
        logger.info( "Performing alignment using numbering shift %s" % offset )
        chain_ali = ncs.ChainAlignment(
            reference = reference,
            aligned = chain,
            offset = offset
            )
        
        min_sequence_identity = 0.8
        seqid = chain_ali.get_sequence_identity()
        logger.info( "Sequence identity between the two chains: %s" % seqid )
        
        if seqid < min_sequence_identity:
            logger.info( "Sequence identity too low" )
            logger.info( "Discarding this chain" )
            continue
        
        logger.info( "This chain passes all tests" )
        identicals.add( chain.id )

    return identicals


def run(args, out = None):
    
    ( logger, handler ) = tbx_utils.get_logger( name = PROGRAM, out = out )
    
    try:
        logger.info(
            tbx_utils.get_banner( program = PROGRAM, version = "0.0.2" )
            )
        params = process_command_line( args = args, logger = logger )
        
        if not params.input.homology_search.file_name:
            usage( logger = logger )
            raise Sorry, "No homology search specified"
        
        logger.info(
            "Reading homology search file: %s" % params.input.homology_search.file_name
            )
        
        hsobj = tbx_utils.HomologySearchObject.from_file(
          file_name = params.input.homology_search.file_name
          )
        hss = hsobj.object
        
        if params.input.homology_search.max_hits:
            logger.info(
                "Considering only the first %s hits" % params.input.homology_search.max_hits
                )
            hss.restrict( max_count = params.input.homology_search.max_hits )
        
        tbx_utils.create_output_folder(
            folder = params.output.folder,
            logger = logger
            )
        
        logger.info( "Setting up sculptor..." )
        sculptors = [
            (
                "protocol",
                "",
                sculptor.get_sculptor_object( 
                    deletion = params.sculptor.macromolecule.deletion,
                    polishing = params.sculptor.macromolecule.polishing,
                    renumber = params.sculptor.macromolecule.renumber,
                    pruning = params.sculptor.macromolecule.pruning,
                    completion = params.sculptor.macromolecule.completion,
                    rename = params.sculptor.macromolecule.rename,
                    bfactor = params.sculptor.macromolecule.bfactor
                    ),
                )
            ]
        discarder = sculptor.get_discarder_object( params = params.sculptor.hetero )
        factory = sculptor.chain_sample_factory( processors = sculptors, logger = logger )
        
        if params.input.max_model_count is not None:
            if params.input.max_model_count <= 0:
                raise Sorry, "max_model_count: positive number expected"
        
        logger.info( "Reading list of obsoleted entries..." )
        redirections = proxy.get_pdb_redirections()
        logger.info( "Read %s entries" % redirections.count() )
        
        logger.info( "Processing hits..." )
        processed = []
        
        for ( index, hit ) in enumerate( hss.hits() ):
            logger.info( "Hit #%s: %s chain %s" % ( index, hit.identifier, hit.chain ) )
            logger.info( "Annotation:" )
            logger.info( hit.annotation )
            
            if redirections.retracted( identifier = hit.identifier ):
                logger.info( "This entry has been retracted." )
                logger.info( "Continue with next..." )
                continue
            
            
            if redirections.obsoleted( identifier = hit.identifier ):
                logger.info( "This entry has been obsoleted." )
                repl = redirections.replacement_for( identifier = hit.identifier )
                assert repl
                logger.info( "Replacement entry: %s" % repl )
                hit.identifier = repl
            
            
            logger.info( "Retrieving file..." )
            logstream = cStringIO.StringIO()
            try:
                stream = iotbx.pdb.fetch.fetch(
                    id = hit.identifier,
                    log = logstream,
                    )
                root = iotbx.pdb.input(
                    lines = stream.read(),
                    source_info = "rcsb",
                    ) \
                    .construct_hierarchy()
                
            except RuntimeError, e:
                logger.warning( "WARNING: download problem (%s)" % e )
                fetchlog = logstream.getvalue()
                
                if fetchlog:
                    logger.warning( "Fetch log:" )
                    logger.warning( fetchlog )
                  
                logger.warning( "Continue with next..." )
                continue
            
            if params.input.remove_alternate_conformations:
                logger.info( "Removing alternate conformations..." )
                sculptor.remove_alternate_conformations( root = root )
                
            if params.input.max_model_count:
                logger.info(
                    "Discarding models above requested max %s if any" % params.input.max_model_count
                    )
                sculptor.restrict_model_count(
                    root = root,
                    max_count = params.input.max_model_count
                    )
            
            if not root.atoms():
                logger.warning( "No atoms in PDB file: %s" % hit.identifier )
                logger.warning( "Continue with next..." )
                continue
            
            # Discard non-protein
            logger.info( "Selecting reference chain %s" % hit.chain )
            asc = root.atom_selection_cache()
            ref_r = root.select(
                asc.selection( "pepnames and chain '%s'" % hit.chain ),
                True,
                )
            
            if not ref_r.atoms():
                logger.warning( "No atoms found for chain %s" % hit.chain )
                logger.warning( "Continue with next..." )
                continue
            
            if 1 < len( ref_r.models()[0].chains() ):
                logger.warning( "Multiple chains with chain_id %s" % hit.chain )
                logger.warning( "Using only first" )
                
            reference = ref_r.models()[0].chains()[0]
            
            if not ncs.MacromoleculeType.protein.recognize_chain(
                chain = reference,
                confidence = 0.8,
                ):
                logger.warning( "Reference chain not recognized as protein" )
                logger.warning( "Continue with next..." )
                continue
            
            identical_chains = find_chains_identical_to(
                reference = reference,
                others = root.select(
                    asc.selection( "pepnames and not chain '%s'" % hit.chain ),
                    True,
                    ),
                logger = logger,
                )
            logger.info(
                "Identical chains: %s" % ", ".join( identical_chains )
                )
            written = []
            
            for c in sorted( identical_chains ):
                logger.info( "Processing chain %s" % c )
                croot = root.select( asc.selection( "chain '%s'" % c ), True )
                
                ali = hit.alignment
                logger.info( "Alignment:" )
                logger.info( str( ali ) )
                
                logger.info( "Running sculptor..." )
                ofiles = sculptor.process_single_pdb(
                    pdb = tbx_utils.PDBObject(
                        root = croot,
                        name = "%s_%s" % ( hit.identifier, c )
                        ),
                    factory = factory,
                    alignments = [
                        tbx_utils.AlignmentObject(
                            alignment = ali,
                            name = "Hit %s" % index
                            )
                        ],
                    sequence_for = {},
                    sculptors = sculptors,
                    discarder = discarder,
                    min_length = params.sculptor.min_hssp_length,
                    min_fraction = params.sculptor.min_matching_fraction,
                    prefix = os.path.join(
                        params.output.folder,
                        "%s_%s" % ( hit.identifier, c )
                        ),
                    suffix = "sculpt",
                    output_method = sculptor.write_pdb_output,
                    logger = logger
                    )
                assert len( ofiles ) == 1
                written.extend( ofiles )
            
            processed.append( written )
            logger.info( "Finished processing hit %s" % hit.identifier )
            logger.info( "Written models: %s" % ", ".join( written ) )
            
        logger.info(
            "Processing complete for %s" % params.input.homology_search.file_name
            )
        
    finally:
        logger.removeHandler( handler )
        handler.flush()
        handler.close()

    return processed