#!/usr/bin/env ccp4-python

import os, sys, copy, glob

import numpy, pandas, json

import libtbx.phil, libtbx.easy_mp
import iotbx.pdb

from libtbx.utils import Sorry, Failure

from bamboo.common.logs import Log
from bamboo.common.path import easy_directory, rel_symlink
from bamboo.common.command import CommandManager

from giant.structure.select import protein, backbone, sidechains
from giant.structure.formatting import PhenixSelection

import matplotlib
matplotlib.interactive(False)
from matplotlib import pyplot
pyplot.switch_backend('agg')
pyplot.interactive(0)

numpy.set_printoptions(threshold=numpy.nan)

from IPython import embed

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

PROGRAM = 'pandemic.adp_sweep'

DESCRIPTION = """
    Fit a series of B-factor models to a series of B-factor refinements
"""

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

blank_arg_prepend = {None:'refinement_dir=', '.pdb':'pdb='}

master_phil = libtbx.phil.parse("""
input {
    pdb = None
        .help = "input pdb files for refinement"
        .multiple = True
        .type = str
    cif = None
        .type = str
    refinement_dir = None
        .help = "input directory, containing different B-factor refinements"
        .optional = False
        .type = str
    labelling = filename *foldername
        .type = choice
}
output {
    out_dir = b-factor-fitting-sweep
        .help = "output directory"
        .type = str
    out_script = run_adp_sweeps.sh
        .help = "output script to run jobs"
        .type = str
}
options {
    tls_selections = tls_selections.params
        .type = path
        .multiple = False
    refinements = *isotropic tls *anisotropic
        .type = choice(multi=True)
    fitting_groups = *tls *chain *residue backbone_sidechain
        .type = choice(multi=True)
}
parameterisation {
    tbd = True
}
settings {
    max_cpus_per_job = 1
        .type = int
        .multiple = False
    concurrent_jobs = 1
        .type = int
        .multiple = False
}
""", process_includes=True)

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

def wrapper_run(arg):
    return arg.run()

def validate_parameters(params):
    pass

def run_refinements(params, log=None):

    if log is None: log = Log(verbose=True)
    log.heading('Creating B-factor refinement job')

    cmd = CommandManager('giant.multi.adp_refine')
    cmd.add_command_line_arguments(params.input.pdb)
    cmd.add_command_line_arguments(params.input.cif)
    cmd.add_command_line_arguments(r'tls_selections={}'.format(params.options.tls_selections))
    cmd.add_command_line_arguments(r'labelling={}'.format(params.input.labelling))
    cmd.add_command_line_arguments(r'b_factor_models={}'.format('+'.join(params.options.refinements)))
    cmd.add_command_line_arguments(r'out_dir={}'.format(params.input.refinement_dir))
    cmd.add_command_line_arguments(r'log_file={}.log'.format(params.input.dir))
    cmd.add_command_line_arguments(r'cpus={}'.format(params.settings.max_cpus_per_job*params.settings.concurrent_jobs))
    cmd.print_settings()
    ret = cmd.run()
    if ret != 0:
        raise Exception('Failed in refinement')

def create_sweep_tls_group_choice(input_pdbs, out_dir, params, log=None):

    if log is None: log = Log(verbose=True)

    out_dir = easy_directory(out_dir)

    # Write all pdb files to an input file to pass to program
    pdb_eff = os.path.join(out_dir, 'input_pdbs.eff')
    with open(pdb_eff, 'w') as fh:
        fh.write('\n'.join([r'pdb="{}"'.format(p) for p in input_pdbs]))

    jobs = []

    base_manager = CommandManager("pandemic.adp")
    base_manager.add_command_line_arguments(r'labelling=foldername')
    base_manager.add_command_line_arguments(pdb_eff)

    # Load first structure for creating TLS groups
    h = protein(iotbx.pdb.hierarchy.input(input_pdbs[0]).hierarchy)

    for mode in params.options.fitting_groups:
        # Compile command object
        cmd = copy.deepcopy(base_manager)
        jobs.append(cmd)
        # Create output directory
        sweep_dir = easy_directory(os.path.join(out_dir, 'group_{}'.format(mode)))
        cmd.add_command_line_arguments(r'out_dir={}'.format(sweep_dir))
        # Create TLS groups for fitting
        tls_groups = []
        if   mode == "chain":
            for chain in h.chains():
                #tls_groups.append("chain '{}'".format(chain.id))
                tls_groups.append(PhenixSelection.format(chain))
        elif mode == "residue":
            for rg in h.residue_groups():
                tls_groups.append(PhenixSelection.format(rg))
        elif mode == "tls":
            for gp in open(params.options.tls_selections, 'r').read().strip().split('\n'):
                tls_groups.append(gp)
        elif mode == "backbone_sidechain":
            for ag in backbone(h).atom_groups():
                if ag.resname in ['ALA','GLY']:
                    tls_groups.append(PhenixSelection.format(ag))
                else:
                    tls_groups.append(PhenixSelection.format(ag)+" and (name C or name CA or name N or name O)")
            for ag in sidechains(h).atom_groups():
                if ag.resname in ['ALA','GLY']:
                    continue
                else:
                    tls_groups.append(PhenixSelection.format(ag)+" and not (name C or name CA or name N or name O)")
        assert tls_groups, 'No TLS groups found: mode={}'.format(mode)

        # Add groups to a params file, and pass file to the command object
        tls_eff = os.path.join(sweep_dir, 'fitting_groups.eff')
        with open(tls_eff, 'w') as fh:
            fh.write('\n'.join([r'tls_group="{}"'.format(g) for g in sorted(tls_groups)]))
        cmd.add_command_line_arguments(tls_eff)
    
        cmd.n_cpus = min(params.settings.max_cpus_per_job, len(tls_groups))
        cmd.add_command_line_arguments(r'cpus={}'.format(cmd.n_cpus))

    assert len(params.options.fitting_groups) == len(jobs)

    for mode, cmd in zip(params.options.fitting_groups, jobs):
        log.subheading(mode)
        cmd.print_settings()

    return jobs

def create_sweep_refinement_type(params, log=None):

    if log is None: log = Log(verbose=True)

    jobs = []

    for ref_type in params.options.refinements:

        log.heading('Creating jobs to analyse {}-refined structures'.format(ref_type))

        pdbs = sorted(glob.glob(os.path.join(params.input.refinement_dir, '*', '*-{}.pdb'.format(ref_type))))
        log('Found {} files.'.format(len(pdbs)))
        out_dir = easy_directory(os.path.join(params.output.out_dir, ref_type))
        log('Outputting to directory: {}'.format(out_dir))

        jobs.extend(create_sweep_tls_group_choice(input_pdbs=pdbs, out_dir=out_dir, params=params))

    return jobs

def write_sweeps(jobs, params, log=None):

    if log is None: log = Log(verbose=True)

    log.heading('Writing commands for {} jobs'.format(len(jobs)))

    with open(params.output.out_script, 'w') as fh:
        fh.write('#!/usr/bin/env bash\n\n')
        for i, j in enumerate(jobs):
            log.subheading('Job {} of {}'.format(i+1, len(jobs)))
            log(str(j))
            job_file = os.path.splitext(params.output.out_script)[0]+'-job-{}.sh'.format(i+1)
            log('Writing command to {}'.format(job_file))
            with open(job_file, 'w') as jf:
                cmd = j.as_command()
                jf.write('#!/usr/bin/env bash\n\n')
                jf.write(cmd+'\n')
#            fh.write('qsub -pe orte {} {}\n'.format(j.n_cpus, os.path.abspath(job_file)))
            fh.write('qsub {}\n'.format(os.path.abspath(job_file)))

def run_sweeps(jobs, params, log=None):

    if log is None: log = Log(verbose=True)

    log.heading('Running {} jobs'.format(len(jobs)))

    if params.settings.concurrent_jobs > 1:
        for i, j in enumerate(jobs):
            log.subheading('Job {} of {}'.format(i+1, len(jobs)))
            log(str(j))
        log.subheading('Running {} jobs ({} at a time)'.format(len(jobs),params.settings.concurrent_jobs))
        ret = libtbx.easy_mp.pool_map(processes=params.settings.concurrent_jobs, func=wrapper_run, args=jobs, chunksize=1)
    else:
        for i, j in enumerate(jobs):
            log.subheading('Job {} of {}'.format(i+1, len(jobs)))
            log(str(j))
            ret = wrapper_run(j)

#    for i, j in enumerate(ret):
#        j.write_output('job-{}.log'.format(i+1))

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

def run(params):

    params.output.out_dir = os.path.abspath(params.output.out_dir)
    params.input.refinement_dir = os.path.abspath(params.input.refinement_dir)

    easy_directory(params.output.out_dir)

    validate_parameters(params)

    if params.input.pdb: run_refinements(params)

    jobs = create_sweep_refinement_type(params)

    if params.output.out_script is not None:
        write_sweeps(jobs, params)
    else:
        run_sweeps(jobs, params)

    #embed()

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

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)


