# LIBTBX_SET_DISPATCHER_NAME phaser.MRage.solutions
from phaser.pipeline import brunett
from phaser import output

FORMAT_FOR = {
    "text": brunett.print_solutions_as_text,
    "xml": brunett.print_solutions_as_xml,
    "script": brunett.print_solutions_as_script,
    }

OUTPUT_WITH = {
    "single": brunett.structure_output_single,
    "ensemble": brunett.structure_output_ensemble,
    "combined": brunett.structure_output_combined,
    "best-single-model": brunett.structure_output_best_single_model,
    }


def show(logger, prefix, selection, args):
    
    FORMAT_FOR[ args.format ]( selection = selection, logger = logger )
    
    
def prepare_data_and_select_structure(logger, selection, index):
    
    brunett.print_solutions_as_text( selection = selection, logger = logger )
    
    logger.info( msg = "Selected: %d" % index )
    logger.info( msg = output.blank() )
    
    if index < 1 or len( selection.solutions ) < index:
        from libtbx.utils import Sorry
        raise Sorry, "Non-existent solution"
    
    structure = selection.solutions[ index - 1 ]
    
    brunett.print_solution_info(
        case = selection.case,
        index = index,
        structure = structure,
        logger = logger
        )
    
    logger.info( msg = output.blank() )
    selection.case.problem.load( logger = logger )
    return structure
    
    
def extract(logger, prefix, selection, args):
    
    structure = prepare_data_and_select_structure(
        logger = logger,
        selection = selection,
        index = args.index,
        )
    
    logger.info( msg = "Regenerating solution..." )
    ( pdb_file, mtz_file ) = brunett.run_llg_job(
        data = selection.case.calculation_data(),
        structure = structure,
        root = "%s_%d" % ( prefix, args.index ),
        stream = logger.info,
        method = OUTPUT_WITH[ args.output ],
        )
    logger.info( msg = "PDB file: %s" % pdb_file )
    logger.info( msg = "Map file: %s" % mtz_file )
    
    logger.info( msg = "Done!" )
    
    if args.coot:
        from phaser import tbx_utils
        tbx_utils.display_in_coot( pdb_file = pdb_file, mtz_file = mtz_file )
        
        
def evaluate(logger, prefix, selection, args):
    
    structure = prepare_data_and_select_structure(
        logger = logger,
        selection = selection,
        index = args.index,
        )
    brunett.calculate_success_probability(
        structure = structure,
        case = selection.case,
        stream = logger.info,
        )
    

import argparse

parser = argparse.ArgumentParser(
    description = "Manipulate MR solution list" 
    )
parser.add_argument(
    "--version",
    action = "version",
    version = "%s" % brunett.VERSION
    )
subparsers = parser.add_subparsers( title = "Valid subcommands" )

# Show functions
parser_show = subparsers.add_parser( "show", help = "show solution list" )
parser_show.set_defaults( func = show )
parser_show.add_argument(
    "file_name",
    metavar = "results-file",
    type = argparse.FileType( "r" ),
    help = "molecular replacement results"
    )
parser_show.add_argument(
    "-f", "--format",
    choices = FORMAT_FOR.keys(),
    help = "Listing format",
    default = "text"
    )

# Extract functions
parser_extract = subparsers.add_parser(
    "extract",
    help = "generate model and map according to selected solution"
    )
parser_extract.set_defaults( func = extract )
parser_extract.add_argument(
    "file_name",
    metavar = "results-file",
    type = argparse.FileType( "r" ),
    help = "molecular replacement results"
    )
parser_extract.add_argument(
    "--index",
    action = "store",
    type = int,
    default = 1,
    help = "index of solution to extract"
    )
parser_extract.add_argument(
    "--coot",
    action = "store_true",
    default = False,
    help = "Load up solution in COOT"
    )
parser_extract.add_argument(
    "-o", "--output",
    choices = OUTPUT_WITH.keys(),
    help = "Output format",
    default = "single"
    )

# Evaluation functions
parser_evaluate = subparsers.add_parser(
    "evaluate",
    help = "evaluate selected solution to determine probability of success"
    )
parser_evaluate.set_defaults( func = evaluate )
parser_evaluate.add_argument(
    "file_name",
    metavar = "results-file",
    type = argparse.FileType( "r" ),
    help = "molecular replacement results"
    )
parser_evaluate.add_argument(
    "--index",
    action = "store",
    type = int,
    default = 1,
    help = "index of solution to extract"
    )

if __name__ == "__main__":
    args = parser.parse_args()
    import sys
    logger = output.SingleStream( stream = sys.stdout )
    import pickle
    ( prefix, selection ) = pickle.load( args.file_name )

    args.func(
        logger = logger,
        prefix = prefix,
        selection = selection,
        args = args
        )