import libtbx.phil
import libtbx.phil.command_line
from libtbx.utils import Sorry

import iotbx.pdb
from iotbx import bioinformatics

import os.path
import logging

# Null logger
class NullHandler(logging.Handler):
    """
    No output
    """
    
    def emit(self, record):
        pass
    
    
# Command-line processing
def process_command_line(args, phil_master):
    
    argument_interpreter = libtbx.phil.command_line.argument_interpreter(
        master_phil = phil_master
        )
    
    phil_objects = []
    files = []
    unknowns = []
    
    for argument in args:
        if os.path.isfile( argument ):
            try:
                file_params = libtbx.phil.parse( file( argument ).read() )
                phil_objects.append( file_params )
            
            except Exception:
                files.append( argument )
                
        elif "=" in argument:
            try:
                command_line_params = argument_interpreter.process(
                    arg = argument
                    )
                phil_objects.append( command_line_params )
            
            except Exception:
                unknowns.append( argument )
                
        else:
            unknowns.append( argument )
            
    return ( phil_objects, files, unknowns )


def group_files_on_extension(files, **kwargs):
    
    sorted = {}
    unknowns = []
    
    for file_name in files:
        for ( type, extensions ) in kwargs.items():
            if any( [ file_name.endswith( ending ) for ending in extensions ] ):
                sorted.setdefault( type, [] ).append( file_name )
                break
        
        else:
            unknowns.append( file_name )
            
    return ( sorted, unknowns )
        

def phil_assign(parameter, value):
    
    return libtbx.phil.parse( "%s=%s" % ( parameter, value ) )


def autoroot(file_name, label):
    
    return "%s_%s" % ( os.path.splitext( file_name )[0], label )


def autoname(file_name, label, extension):
    
    return "%s.%s" % ( autoroot( file_name, label ), extension )


def choice_string(possible, default):
    
    return " ".join(
        [ "%s%s" % ( "*" if n == default else "", n ) for n in possible ]
        )


def get_banner(program, version):
    
    banner = [
        "#" * 79,
        ( "#" + ( "%s %s" % ( program, version ) ).center( 77 ) + "#" ),
        "#" * 79
        ]
    
    return "\n".join( banner )


def format_table(captions, formats, lines):
    
    count = len( captions )
    widths = [ len( c ) for c in captions ]
    wformats = [ "%%%ds" % w for w in widths ]
    table = [ "  ".join( captions ) ]
    
    for l in lines:
        assert len( l ) == count
        table.append(
            "  ".join(
                [ ( wf % ( f % v ) )[:w] for ( f, v, w, wf )
                    in zip( formats, l, widths, wformats ) ]
                )
            )
        
    return "\n".join( table )


def format_table_entries(captions, formats, lines, max_widths = None):
    
    columns = []
    widths = []
    
    if max_widths is None:
        max_widths = [ None ] * len( captions )
    
    for ( c, f, values, mw ) in zip( captions, formats, zip( *lines ), max_widths ):
        converteds = [ c ] + [ f % v for v in values ]
        assert values
        width = max( [ len( cv ) for cv in converteds ] )
        
        if mw is not None:
            width = min( width, mw )
            
        wformat = "%%%ds" % width
        columns.append( [ ( wformat % cv )[ : width ] for cv in converteds ] )
        
    return zip( *columns )


def assemble_table(lines):
    
    table = []
    
    for l in lines:
        table.append( "  ".join( l ) )
        
    return "\n".join( table )


def get_logger(name, out):
    
    logger = logging.getLogger( name = name )
    
    if out:
        handler = logging.StreamHandler( out )
        
    else:
        handler = NullHandler()
        
    handler.setFormatter( logging.Formatter( "%(message)s" ) )
    logger.addHandler( handler )
    logger.setLevel( logging.INFO )
    
    return ( logger, handler )


def get_handler(out):
    
    if out:
        handler = logging.StreamHandler( out )
        
    else:
        handler = NullHandler()
        
    handler.setFormatter( logging.Formatter( "%(message)s" ) )
    
    return handler


def create_output_folder(folder, logger):
    
    if os.path.isdir( folder ):
        logger.info( "Using existing output folder: %s" % folder )
        
    elif os.path.exists( folder ):
        logger.error(
            "Cannot create output folder: existing non-directory"
            )
        raise Sorry, "Existing non-directory: %s" % folder
    
    else:
        logger.info( "Creating output directory: %s" % folder )
        os.mkdir( folder )


def read_in_files(file_names, object_type):
    
    objs = []

    for f in file_names: 
        objs.append( object_type.from_file( file_name = f ) )
        
    return objs


def extension(file_name):
    
    ( name, extension ) = os.path.splitext( file_name )
    return extension
    

PLURAL_FOR_NOUN = {
    "a": "",
    }

def inflect(word, count):
    
    if count == 1:
        return word
    
    if word in PLURAL_FOR_NOUN:
        return PLURAL_FOR_NOUN[ word ]
    
    return "%ss" % word


CONJUGATED_FOR_INFINITIVE = {
    "be": ( "is", "are" ),
    }
    
    
def conjugate(verb, count):
    
    if count == 1:
        if verb == "is":
            return "is"
        
        else:
            return "%ss" % verb
        
    else:
        if verb == "is":
            return "are"
        
        else:
            return verb
        
PHASER_MODEL_IDENTITY_REMARK = "REMARK PHASER ENSEMBLE MODEL %s ID %s" 

def get_phaser_model_identity_remark(mid, identity):
    
    index = 0 if not mid else int( mid )
    return PHASER_MODEL_IDENTITY_REMARK % ( "%d" % index, "%.1f" % identity )


def parse_phaser_model_identity_remark(text):
    
    import re
    match = re.match(
        PHASER_MODEL_IDENTITY_REMARK % ( "\s*(\d+)", "(\d*\.\d+)" ),
        text
        )
    
    if not match:
        raise ValueError, "Incorrect remark format"
    
    mid = match.group( 1 )
    return ( "" if mid == "0" else "%4s" % mid, float( match.group( 2 ) ) )


def dummy_cryst1_record():
    
    from iotbx.pdb import format_cryst1_and_scale_records
    from cctbx import uctbx, sgtbx, crystal
    return format_cryst1_and_scale_records(
        crystal_symmetry = crystal.symmetry(
            unit_cell = uctbx.unit_cell( parameters = [ 1, 1, 1, 90, 90, 90 ] ),
            space_group_info = sgtbx.space_group_info( number = 1 )
            )
        )
    
    
class ChainInfo(object):
    """
    Stores information about chains
    """
    
    def __init__(self, pdb, m_index, c_index, identity):
        
        self.pdb = pdb
        self.m_index = m_index
        self.c_index = c_index
        self.identity = identity
        
        
    @property
    def model(self):
        
        return self.pdb.object.models()[ self.m_index ]
    
    
    @property
    def chain(self):
        
        return self.model.chains()[ self.c_index ]
    
    
    @property
    def identifier(self):
        
        return ( self.pdb.name, self.model.id, self.chain.id, self.c_index )
    
    
    def __str__(self):
        
        return "chain (%s, '%s', '%s')" % (
            self.pdb.name,
            self.model.id,
            self.chain.id
            )
        

class PDBObject(object):
    """
    Opened and parsed PDB file
    """
    
    def __init__(self, root, name, remarks = ()):
        
        self.name = name
        self.object = root
        self.remarks = remarks
        
        
    def match_identity_records(self):
        
        cid_in = {}
        
        for r in self.remarks:
            try:
                ( mid, identity ) = parse_phaser_model_identity_remark( text = r )
                
            except ValueError:
                continue
             
            cid_in.setdefault( mid, [] ).append( identity )
            
        identity_for = {}
            
        for model in self.object.models():
            for ( chain, identity ) in zip( model.chains(), cid_in.get( model.id, [] ) ):
                identity_for[ chain ] = identity
            
        return identity_for
    
    
    def get_chain_infos(self):
        
        infos = []
        identity_for = self.match_identity_records()
        
        for ( m_index, model ) in enumerate( self.object.models() ):
            for ( c_index, chain ) in enumerate( model.chains() ):
                infos.append(
                    ChainInfo(
                        pdb = self,
                        m_index = m_index,
                        c_index = c_index,
                        identity = identity_for.get( chain )
                        )
                    )
                
        return infos
        
        
    @classmethod
    def from_file(cls, file_name):
        
        try:
            input = iotbx.pdb.input(file_name=file_name,
              raise_sorry_if_format_error=True)
            
        except IOError, e:
            raise Sorry, e
        
        return cls(
            root = input.construct_hierarchy(),
            name = file_name,
            remarks = input.remark_section()
            )
        
        
    @classmethod
    def from_gzipped_file(cls, file_name):
        
        import gzip
        
        try:
            data = gzip.GzipFile( file_name ).read()
            
        except IOError, e:
            raise Sorry, e
        
        return cls.from_string(
            pdb_str = data,
            name = file_name
            ) 
    
    
    @classmethod
    def from_string(cls, pdb_str, name):
         
        input = iotbx.pdb.input( lines = pdb_str, source_info = name )
        return cls(
            root = input.construct_hierarchy(),
            name = name,
            remarks = input.remark_section()
            )
        
        
class AlignmentObject(object):
    """
    Opened and parsed alignment file
    """
    
    def __init__(self, alignment, name):
        
        if not alignment.alignments:
            raise Sorry, "Empty alignment: %s" % name
        
        self.name = name
        self.object = alignment
        
        
    @classmethod
    def from_file(cls, file_name):
            
        try:
            ali_str = file( file_name ).read()
            
        except IOError, e:
            raise Sorry, "Error while reading %s: %s" % ( file_name, e )
        
        return cls.from_string(
            ali_str = ali_str,
            extension = extension( file_name = file_name ),
            name = file_name
            )
    
    
    @classmethod
    def from_string(cls, ali_str, extension, name):
        
        parser = bioinformatics.alignment_parser_for_extension( extension = extension )
        
        if not parser:
            raise Sorry, "Unknown format: %s" % extension
        
        ( ali, errors ) = parser( ali_str )
        
        if errors:
            raise Sorry, "Incorrect format:\n%s" % ali_str
        
        assert ali
        
        return cls( alignment = ali, name = name )
    
    
class SequenceObject(object):
    """
    Opened and parsed alignment file
    """
    
    def __init__(self, sequences, name):
        
        self.name = name
        self.object = sequences
        
        
    @classmethod
    def from_file(cls, file_name):
            
        try:
            seq_str = file( file_name ).read()
            
        except IOError, e:
            raise Sorry, "Error while reading %s: %s" % ( file_name, e )
        
        return cls.from_string( seq_str = seq_str, name = file_name )
    
    
    @classmethod
    def from_string(cls, seq_str, name):
        
        ( seqs, errors ) = bioinformatics.parse_sequence( data = seq_str )
        
        if errors or not seqs:
            raise Sorry, "Incorrect format:\n%s" % seq_str
        
        for seq in seqs:
            if not seq.name:
                seq.name = os.path.splitext( os.path.basename( name ) )[0]
        
        return cls( sequences = seqs, name = name )
    
    
class HomologySearchObject(object):
    """
    Opened and parsed homology search data
    """
    
    def __init__(self, search, name):
        
        self.name = name
        self.object = search
        
        
    @classmethod
    def from_file(cls, file_name):
        
        data = open( file_name ).read()
    
        if file_name.endswith( ".hhr" ):
            try:
                result = bioinformatics.hhsearch_parser( output = data )
                
            except ValueError, e:
                raise Sorry, e
        
        elif file_name.endswith( ".xml" ):
            from phaser.pipeline import homology
            
            parsers = [
                homology.get_ncbi_blast_parser(),
                homology.get_ebi_blast_parser()
                ]
            
            for p in parsers: 
                try:
                    result = p( data = data )
                    
                except RuntimeError, e:
                    continue
                
                break
            
            else:
                raise Sorry, "No XML parser was able to process this file"
            
        else:
            raise Sorry, "Unknown homology search file"
        
        return cls( search = result, name = file_name )
    
    
class FileArgument(object):
    """
    A file
    """
    
    def __init__(self, file_name):
        
        self.file_name = file_name
        
        
    def process(self, accumulator):
        
        accumulator.process_file( argument = self )
        
        
    def extension(self):
        
        return os.path.splitext( self.file_name )[1]
    
    
    def __str__(self):
        
        return "%s <file>" % self.file_name
        
        
class PhilArgument(object):
    """
    A phil read from a file
    """
    
    def __init__(self, phil):
        
        self.phil = phil
         
         
    def process(self, accumulator):
        
        accumulator.process_phil( argument = self )
    
    
    def __str__(self):
        
        return "Phil:\n%s" % self.phil.as_str()
        
    
    @classmethod
    def from_string(cls, string):
        
        try:
            phil = libtbx.phil.parse( string )
            
        except Exception:
            raise ValueError, "Incorrect PHIL format"
        
        return cls( phil = phil )
    
    
    @classmethod
    def from_file(cls, file_name):
        
        return cls.from_string( string = open( file_name ).read() )
    
    
    @classmethod
    def command_line_argument(cls, interpreter, argument): 
        
        try:
            phil = interpreter.process( arg = argument )
            
        except Exception:
            raise ValueError, "Incorrect PHIL command line argument"
        
        return cls( phil = phil )
    
    
class UnknownArgument(object):
    """
    An argument that is unknown
    """
    
    def __init__(self, argument):
        
        self.argument = argument
        
        
    def process(self, accumulator):
        
        accumulator.process_unknown( argument = self )
    
    
    def __str__(self):
        
        return self.argument
        
        
class PhilArgumentFactory(object):
    
    def __init__(self, master_phil):
        
        self.interpreter = libtbx.phil.command_line.argument_interpreter(
            master_phil = master_phil
            )
        
        
    def __call__(self, argument):
        
        if os.path.isfile( argument ):
            try:
                phil = PhilArgument.from_file( file_name = argument )
            
            except ValueError:
                phil = FileArgument( file_name = argument )
                
        elif "=" in argument:
            try:
                phil = PhilArgument.command_line_argument(
                    interpreter = self.interpreter,
                    argument = argument
                    )
            
            except ValueError:
                phil = UnknownArgument( argument = argument )
                
        else:
            phil = UnknownArgument( argument = argument )
                
        return phil
    
    
class PhilAccumulator(object):
    """
    Unifies Phil command line handling (using the visitor pattern)
    """
    
    def __init__(self, master_phil):
        
        self.master_phil = master_phil
        self.phils = []
        
        self.file_handlers = []
        self.unknown_handlers = []
        
        
    def register_file_handler(self, handler):
        
        self.file_handlers.append( handler )
        
        
    def process_file(self, argument):
        
        for handler in self.file_handlers:
            try:
                result = handler( argument = argument, phils = self.phils )
                
            except RuntimeError:
                continue
            
            break
            
        else:
            self.process_unknown( argument = argument )
            
            
    def process_phil(self, argument):
        
        self.phils.append( argument.phil )
        
        
    def process_unknown(self, argument):
        
        for handler in self.unknown_handlers:
            try:
                handler( argument = argument, phils = self.phils )
                
            except RuntimeError:
                continue
            
            break
            
        else:
            raise Sorry, "Unknown argument: %s" % argument
    
    
    def merge(self):
        
        return self.master_phil.fetch( sources = self.phils )


class ExtensionFileHandler(object):
    
    def __init__(self, extensions, template):
        
        self.extensions = extensions
        self.template = template
        
        
    def __call__(self, argument, phils):
        
        if argument.extension() in self.extensions:
            return phils.append(
                libtbx.phil.parse( self.template % argument.file_name )
                )
        
        raise RuntimeError, "Cannot process argument: %s" % argument


def display_in_coot(pdb_file, mtz_file):
    
    SCRIPT_TEMPLATE = (
        "read_pdb( '%s' ); graphics_to_ca_representation(0);"
        + "set_show_symmetry_master( 1 ); symmetry_as_calphas( 0, 1 );"
        + "auto_read_make_and_draw_maps( '%s' )"
        )
    
    from subprocess import Popen
    import os
    env = os.environ.copy()
    del env [ "LD_LIBRARY_PATH" ]
    p=Popen(
        ( "coot", "--python", "-c", SCRIPT_TEMPLATE % ( pdb_file, mtz_file ) ),
        env = env,
        )
    p.communicate()