#!/usr/bin/env ccp4-python

import os, sys, copy, re, shutil

import libtbx.phil
import libtbx.easy_mp

from bamboo.common.command import CommandManager

############################################################################

PROGRAM = 'multi.dimple'

DESCRIPTION = """
    Run dimple on multiple datasets.
"""

############################################################################

blank_arg_prepend = {'.pdb':'reference_pdb=', '.mtz':'mtz='}

master_phil = libtbx.phil.parse("""
input {
    reference_pdb = None
        .type = path
        .multiple = True
    mtz = None
        .type = path
        .multiple = True
}
options {
    use_reference_rfree_flags = False
        .type = bool
}
output {
    overwrite = True
        .type = bool
    dir_suffix = '_dimple'
        .type = str
}
settings {
    cpus = 1
        .type = int
}
""")

############################################################################

def run_dimple(c):
    ret_code = c.run()
    if ret_code != 0:
        print '============================>'
        print 'Dimple returned with an error'
        print '============================>'
        c.print_settings()
        print '============================>'
        print c.output
        print '============================>'
        print c.error
        print '============================>'
    return ret_code

############################################################################

def run(params):

    # Iterate through reference pdb files
    for ref_pdb_file in params.input.reference_pdb:

        # Command-objects for multi-processing
        commands = []

        ref_pdb_file = os.path.abspath(ref_pdb_file)

        # Look for reference mtz if requested
        if params.options.use_reference_rfree_flags:
            ref_mtz_file = ref_pdb_file.replace('.pdb','.mtz')
            assert os.path.exists(ref_mtz_file), 'Reference MTZ does not exist: {}'.format(ref_mtz_file)

        # Create output directory name for this PDB
        output_dirname = os.path.splitext(os.path.basename(ref_pdb_file))[0]+params.output.dir_suffix
        print 'Placing all dimple runs for {} into directories called {}'.format(ref_pdb_file, output_dirname)

        # Iterate through mtz files and run dimple on each
        for mtz_file in params.input.mtz:

            # Create output files and directory
            mtz_file = os.path.abspath(mtz_file)
            output_dir = os.path.join(os.path.dirname(mtz_file), output_dirname)

            # Wipe output directory
            if os.path.exists(output_dir):
                if params.output.overwrite:
                    shutil.rmtree(output_dir)
                else:
                    raise Exception('Output Directory already exists: {}'.format(output_dir))

            # Need to precreate the directory
            os.mkdir(output_dir)

            # Create dimple command object for multi-processing
            cm = CommandManager('dimple')
            if params.options.use_reference_rfree_flags:
                cm.add_command_line_arguments([ '--free-r-flags', ref_mtz_file ])
            cm.add_command_line_arguments([ mtz_file,
                                            ref_pdb_file,
                                            output_dir,
                                        ])
            cm.print_settings()
            commands.append(cm)

        # Run Dimples for this reference pdb file
        print '============================>'
        print 'Running dimple for {} datasets'.format(len(commands))
        returned = libtbx.easy_mp.pool_map(fixed_func=run_dimple, args=commands, processes=params.settings.cpus)

        print '============================>'
        print 'Return code Total: {}'.format(sum(returned))

############################################################################

if __name__=='__main__':
    from giant.jiffies import run_default
    run_default(
        run                 = run,
        master_phil         = master_phil,
        args                = sys.argv[1:],
        blank_arg_prepend   = blank_arg_prepend,
        program             = PROGRAM,
        description         = DESCRIPTION)
