#!/usr/bin/env ccp4-python

import os, sys, re, glob, shutil, copy, tempfile, gc
import math, re

import mdp
import pandas, numpy

import libtbx.phil
import libtbx.easy_mp
import iotbx.pdb
import mmtbx.tls.tools

from libtbx.utils import Sorry, Failure
from scitbx.array_family import flex

from bamboo.common.path import easy_directory
from bamboo.common.command import CommandManager
from bamboo.common.logs import Log
from bamboo.stats.cluster import generate_group_idxs

from giant.dataset import CrystallographicModel
from giant.structure.b_factors import occupancy_weighted_average_b_factor
from giant.structure.select import protein, backbone, sidechains

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

from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import proj3d

numpy.set_printoptions(threshold=numpy.nan)

EIGHT_PI_SQ = 8*math.pi*math.pi

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

PROGRAM = 'giant.datasets.b_factor_analysis'

DESCRIPTION = """
    Analyse the variation/conservartion of B-factors (under different models) of a set of related structures.
"""

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

blank_arg_prepend = {'.pdb':'pdb=', '.cif':'cif='}

master_phil = libtbx.phil.parse("""
input {
    pdb = None
        .help = "input pdb file"
        .multiple = True
        .type = str
    cif = None
        .type = str
        .multiple = True
    tls_selections = tls_selections.params
        .help = "define the tls groups (used for all structures)"
        .type = path
        .multiple = False
    labelling = filename *foldername
        .type = choice
        .multiple = False
}
output {
    out_dir = multi-dataset-tls-characterisation
        .help = "output directory"
        .type = str
    log_file = multi-dataset-tls-characterisation.log
        .type = str
}
settings {
    cpus = 48
        .type = int
        .multiple = False
}
""")

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

def extract_tls_from_pdb(pdb_file):
    ih = iotbx.pdb.hierarchy.input(pdb_file)
    tls_params = ih.input.extract_tls_params(ih.hierarchy)
    return tls_params

def extract_u_components(input_hierarchy=None, hierarchy=None, input=None):
    assert input_hierarchy or (hierarchy and input)
    if input_hierarchy:
        hierarchy = input_hierarchy.hierarchy
        input = input_hierarchy.input
    else:
        assert hierarchy and input

    cache = hierarchy.atom_selection_cache()
    tls_objs = input.extract_tls_params(hierarchy)
    assert tls_objs.tls_present is True
    tls_prms = [mmtbx.tls.tools.tlso(t=p.t, l=p.l, s=p.s, origin=p.origin) for p in tls_objs.tls_params]
    tls_sels = [cache.selection(p.selection_string) for p in tls_objs.tls_params]

    u_tot = hierarchy.atoms().extract_uij()
    u_tls = mmtbx.tls.tools.u_cart_from_tls(sites_cart=hierarchy.atoms().extract_xyz(),
                                            selections=tls_sels, tlsos=tls_prms)
    u_loc = u_tot - u_tls

    return ADPs(hierarchy=hierarchy, u_tot=u_tot, u_tls=u_tls, u_loc=u_loc)

def identify_tls_boundaries(tls_selections):

    rgx = re.compile("chain (.*) and \(resid (.*) through (.*)\)")

    dividers = {}

    for sel in tls_selections:

        m = rgx.findall(sel)

        if not m:
            print 'No REGEX match found: {}'.format(sel)
            continue

        assert len(m) == 1

        chain, v1, v2 = m[0]
        chain = chain.strip("'")
        v1 = int(v1)
        v2 = int(v2)
        dividers.setdefault(chain, []).extend([v1,v2])

    for k in dividers.keys():
        dividers[k] = sorted(set(dividers[k]))
        boundaries = copy.copy(dividers[k])
        print boundaries
        for v1,v2 in zip(dividers[k][:-1],dividers[k][1:]):
            if abs(v2-v1)==1:
                try: boundaries.remove(v1)
                except: pass
                try: boundaries.remove(v2)
                except: pass
                boundaries.append((v1+v2)/2.0)
        dividers[k] = sorted(boundaries)

    return dividers

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


class TLSFitterRefiner(object):

    def __init__(self, pdb_file, mtz_file, out_dir, cif_files=[], tag=None, tls_selections=None):

        self.pdb_file = pdb_file
        self.mtz_file = mtz_file
        self.cif_files = cif_files
        self.out_dir = out_dir
        self.tag = tag
        self.tls_selections = []
        self.tls_matrices = None

        self.tls_initial_pdb = os.path.join(self.out_dir, 'initial.pdb')
        self.tls_refined_pdb = os.path.join(self.out_dir, 'refined.pdb')

        self.log = Log(verbose=True)

        if not tls_selections:
            tls_selections = self.determine_tls_groups(pdb_file=pdb_file)

        # Sanitise the tls selections
        for tls in tls_selections:
            if tls.startswith('"') and tls.endswith('"'):
                tls=tls[1:-1]
            assert '\"' not in tls, 'TLS selection cannot include \": {}'.format(tls)
            self.tls_selections.append(tls)

        if not os.path.exists(self.out_dir): os.mkdir(self.out_dir)

    def determine_tls_groups(self, pdb_file):

        self.log.subheading('Determining TLS groups for: {}'.format(pdb_file))

        cmd = CommandManager('phenix.find_tls_groups')
        cmd.add_command_line_arguments(pdb_file)

        cmd.print_settings()
        ret_code = cmd.run()

        if ret_code != 0:
            self.log(cmd.output)
            self.log(cmd.error)
            raise Exception('Failed to determine TLS groups: {}'.format(' '.join(cmd.program)))

        regex = re.compile("refinement\.refine\.adp \{([\s\S]*?)\}")
        tls_command = regex.findall(cmd.output)[0]
        tls_selections = [s.strip() for s in tls_command.split('tls =') if s.strip()]

        self.log.subheading('Identified TLS Selections:')
        for s in tls_selections:
            self.log(s)

        return tls_selections

    def initial_tls_parameters(self):
        """Characterise TLS with phenix.tls"""

        self.log.subheading('Fitting TLS Matrices to selections')
        self.log('writing to output file: {}'.format(self.tls_initial_pdb))

        cmd = CommandManager('phenix.tls')
        cmd.add_command_line_arguments(self.pdb_file)
        cmd.add_command_line_arguments(self.cif_files)
        cmd.add_command_line_arguments('extract_tls=True')
        cmd.add_command_line_arguments([r'selection="{}"'.format(s) for s in self.tls_selections if s is not None])
        cmd.add_command_line_arguments('output_file_name={}'.format(self.tls_initial_pdb))

        cmd.print_settings()
        ret_code = cmd.run()
        cmd.write_output(self.tls_initial_pdb.replace('.pdb', '.log'))

        if ret_code != 0:
            self.log(cmd.output)
            self.log(cmd.error)
            raise Exception('Failed to determine TLS parameters: {}'.format(' '.join(cmd.program)))

        return self.tls_initial_pdb, self.extract_tls_from_pdb(self.tls_initial_pdb)

    def refined_tls_parameters(self):
        """Refine the model with phenix.refine, including the TLS model"""

        out_dir = tempfile.mkdtemp(prefix='tls_refine_')
        out_pre = os.path.join(out_dir, 'tls-refine')

        self.log.subheading('Refining TLS model with phenix.refine')
        self.log('writing to output directory: {}'.format(out_dir))

        cmd = CommandManager('phenix.refine')
        cmd.add_command_line_arguments(self.pdb_file, self.mtz_file)
        cmd.add_command_line_arguments(self.cif_files)
        cmd.add_command_line_arguments('refine.strategy=individual_adp+tls')
        cmd.add_command_line_arguments('main.number_of_macro_cycles=3')
        cmd.add_command_line_arguments([r'refinement.refine.adp.tls="{}"'.format(t) for t in self.tls_selections])
        cmd.add_command_line_arguments('output.prefix={}'.format(out_pre))

        cmd.print_settings()
        ret_code = cmd.run()
        cmd.write_output(out_pre+'.log')

        if ret_code != 0:
            self.log(cmd.output)
            self.log(cmd.error)
            raise Exception('Failed to determine refine model: {}'.format(' '.join(cmd.program)))

        out_pdb = glob.glob(out_pre+'*.pdb')[0]

        self.log('copying output pdb: {} -> {}'.format(out_pdb, self.tls_refined_pdb))
        shutil.copy(out_pdb, self.tls_refined_pdb)

        #os.rmdir(out_dir)

        return self.tls_refined_pdb, self.extract_tls_from_pdb(self.tls_refined_pdb)

    @staticmethod
    def extract_tls_from_pdb(pdb_file):
        return extract_tls_from_pdb(pdb_file)

    def show_tls_params(self, tls_params=None, pdb_file=None):
        if pdb_file: tls_params=self.extract_tls_from_pdb(pdb_file=pdb_file)
        T = tls_params.tls_params[0].t
        L = tls_params.tls_params[0].l
        S = tls_params.tls_params[0].s

        o = ""
        for tls in tls_params.tls_params:
            o += '\n'
            o += 'selection: {}\n'.format(tls.selection_string)
            o += 'origin: {}\n'.format(tls.origin)
            o += 'T: '+str(tls.t)+'\n'
            o += 'L: '+str(tls.l)+'\n'
            o += 'S: '+str(tls.s)+'\n'
        o += '\n'
        self.log(o)


class RestrainedTLSModelRefiner(object):

    def __init__(self, inp_pdb, inp_mtz, out_pdb='this.pdb', tls_selection_strings=None):

        self.pdb = inp_pdb
        self.mtz = inp_mtz
        self.out = out_pdb
        self.tls = tls_selection_strings

    def refine_model(self):

        out_dir = tempfile.mkdtemp(prefix='tls_refine_')
        out_pre = os.path.join(out_dir, 'tls-refine')

        self.log.subheading('Refining TLS model with phenix.refine')
        self.log('writing to output directory: {}'.format(out_dir))

        cmd = CommandManager('phenix.refine')
        cmd.add_command_line_arguments(self.pdb, self.mtz)
        cmd.add_command_line_arguments('refine.strategy=tls')
        cmd.add_command_line_arguments('main.number_of_macro_cycles=3')
        cmd.add_command_line_arguments([r'refinement.refine.adp.tls="{}"'.format(t) for t in self.tls_selection_strings])
        cmd.add_command_line_arguments('output.prefix={}'.format(out_pre))

        cmd.print_settings()
        ret_code = cmd.run()
        cmd.write_output(out_pre+'.log')

        if ret_code != 0:
            self.log(cmd.output)
            self.log(cmd.error)
            raise Exception('Failed to determine refine model: {}'.format(' '.join(cmd.program)))

        out_pdb = glob.glob(out_pre+'*.pdb')[0]

        self.log('copying output pdb: {} -> {}'.format(out_pdb, self.out_pdb))
        shutil.copy(out_pdb, self.out_pdb)


class ADPs(object):
    def __init__(self, hierarchy, u_tot, u_tls, u_loc):
        self.hierarchy = hierarchy
        self.u_tot = flex.sym_mat3_double(u_tot)
        self.u_tls = flex.sym_mat3_double(u_tls)
        self.u_loc = flex.sym_mat3_double(u_loc)
        self.iso_b_tot = flex.double(EIGHT_PI_SQ * numpy.mean(numpy.array(u_tot)[:,:3], axis=1))
        self.iso_b_tls = flex.double(EIGHT_PI_SQ * numpy.mean(numpy.array(u_tls)[:,:3], axis=1))
        self.iso_b_loc = flex.double(EIGHT_PI_SQ * numpy.mean(numpy.array(u_loc)[:,:3], axis=1))
    def __getitem__(self, item):
        return self.__dict__[item]


class MultiDatasetBFactorAnalysis(object):

    _BFAC_ATTRS = ['tot',       'tls',       'loc' ]
    _BFAC_NAMES = ['total',     'tls',       'residual'  ]
    _FILT_FUNCS = [ protein,     backbone,    sidechains ]
    _FILT_NAMES = ['protein',   'backbone',  'sidechains']

    def __init__(self, models):

        self.models = models
        self.adps = [extract_u_components(input=m.input, hierarchy=m.hierarchy) for m in self.models]

        self.log = Log(verbose=True)

    def write_b_factor_plots(self, boundaries=None):

        data_dict = {}

        # Which atoms to look at (protein/backbone/sidechain)
        for i_filt in range(3):

            filt_func = self._FILT_FUNCS[i_filt]
            filt_name = self._FILT_NAMES[i_filt]

            self.log.subheading('Calculating B-factor plots for {}'.format(filt_name))
            self.log('Iterating through {} structures...'.format(len(self.models)))

            for i_strc, (mdl, adps) in enumerate(zip(self.models, self.adps)):

                sys.stdout.write('\n'*(not i_strc%50)+'>'); sys.stdout.flush()

                # Which bfactors to add to structure (Total/TLS/Residual)
                for i_bfac in range(3):

                    bfac_attr = self._BFAC_ATTRS[i_bfac]
                    bfac_name = self._BFAC_NAMES[i_bfac]

                    h = mdl.hierarchy.deep_copy()
                    h.atoms().set_b(adps['iso_b_'+bfac_attr])
                    h.atoms().set_uij(adps['u_'+bfac_attr])

                    file_name = os.path.splitext(mdl.filename)[0]+'-'+filt_name+'-'+bfac_name+'.pdb'

                    h.write_pdb_file(file_name=file_name,
                                     crystal_symmetry=mdl.input.crystal_symmetry())

                    for chn in protein(h).chains():

                        chn_h = iotbx.pdb.hierarchy.new_hierarchy_from_chain(chn)
                        av_bs = [( rg.resseq_as_int(), occupancy_weighted_average_b_factor(rg.atoms()) ) for rg in filt_func(chn_h).residue_groups()]

                        data_dict.setdefault(filt_name,{}).setdefault(chn.id,{}).setdefault(bfac_name,[]).append(av_bs)

            sys.stdout.write('\n')

            self.log.subheading('Generating plots')

            for chn_id in sorted(data_dict[filt_name].keys()):

                self.log('Writing images for function {} chain {}'.format(filt_name, chn_id))

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

                fig, axes = pyplot.subplots(3, sharex=True)

                for ax, bfac_name in zip(axes, self._BFAC_NAMES):
                    ax.set_title(bfac_name)
                    for d in data_dict[filt_name][chn_id][bfac_name]:
                        d = zip(*d)
                        line, = ax.plot(d[0], d[1], 'b-', alpha=1.0/(1.0+5.0*math.log(len(self.models))), linewidth=1)
                    if boundaries:
                        [ax.axvline(x=x, ymin=0.0, ymax=1.0, linewidth=2, color='r') for x in boundaries[chn_id]]
                axes[-1].set_xlabel('Residue Number', color='k', size=16)
                axes[ 1].set_ylabel('Isotropic B-factor', color='k', size=16)
                pyplot.subplots_adjust()
                pyplot.savefig('{}-{}.png'.format(chn_id, filt_name), dpi=300)
                pyplot.close(fig)

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

                fig, axes = pyplot.subplots(len(self._BFAC_NAMES)+1, sharex=True)

                ax_1_3 = axes[0:len(self._BFAC_NAMES)]
                ax_4 = axes[-1]
                ax_4.set_title('RMS B-factor')

                for ax, bfac_name in zip(ax_1_3, self._BFAC_NAMES):
                    ax.set_title(bfac_name)

                    # Collate all B-factors from the datasets
                    all_data = []; [all_data.extend(d) for d in data_dict[filt_name][chn_id][bfac_name]]
                    res_nums, b_factors = zip(*all_data)

                    # Group B-factors by residue number and make boxplot
                    x_vals, y_vals = zip(*[[i, [b_factors[j] for j in idxs]] for i,idxs in generate_group_idxs(res_nums)])
                    x_labs = [str(x) if x%100==0 else '' for i,x in enumerate(x_vals)]
                    ax.boxplot(x=y_vals, positions=x_vals, labels=x_vals, sym='')

                    # Calculation the RMSD B-factor for each residue
                    y_stds = map(numpy.std, y_vals)
                    line, = ax_4.plot(x_vals, y_stds, linewidth=1, label=bfac_name)

                    if boundaries:
                        [ax.axvline(x=x, ymin=0.0, ymax=1.0, linewidth=2, color='r') for x in boundaries[chn_id]]

                if boundaries:
                    [ax_4.axvline(x=x, ymin=0.0, ymax=1.0, linewidth=2, color='r') for x in boundaries[chn_id]]

                axes[-1].set_xlabel('Residue Number', color='k', size=16)

                x_ticks_maj = range(int(round(min(x_vals),-2)), int(round(max(x_vals),-2)+1), 100)
                x_ticks_min = range(int(round(min(x_vals),-1)), int(round(max(x_vals),-1)+1), 10)
                axes[-1].set_xticks(x_ticks_maj)
                axes[-1].set_xticks(x_ticks_min, minor=True)
                axes[-1].set_xticklabels(map(str,x_ticks_maj))
                axes[ 1].set_ylabel('Isotropic B-factor', color='k', size=16)
                pyplot.legend(loc='center right',
                              bbox_to_anchor=(0.99, 0.25),
                              bbox_transform=pyplot.gcf().transFigure,
                              fontsize=8,
                              ncol=1)
                pyplot.subplots_adjust(hspace=0.4)
                pyplot.savefig('{}-{}-boxplot.png'.format(chn_id, filt_name), dpi=300)
                pyplot.close(fig)

    def write_median_b_factors(self):

        b_factors = [numpy.array(a.iso_b_loc) for a in self.adps]
        b_sizes = set(map(len, b_factors))
        b_factors = numpy.array(b_factors)

        if len(b_sizes) > 1:
            raise Sorry('Structures have different numbers of atoms -- cannot proceed. Sizes: {}'.format(b_sizes))

        med_b_factors = flex.double(numpy.median(b_factors, axis=0).tolist())

        g_zero = (med_b_factors>0)
        l_zero = (med_b_factors<0)
        med_b_factors.set_selected(l_zero, 0)

        for model in self.models:

            med_filename = model.filename.replace('.pdb','.median.pdb')

            mdl_hierarchy = model.hierarchy.deep_copy()
            mdl_b_facs = mdl_hierarchy.atoms().extract_b()
            mdl_b_facs.set_selected(g_zero, 0)

            mdl_hierarchy.atoms().set_b(med_b_factors+mdl_b_facs)
            print 'Writing {}'.format(med_filename)
            mdl_hierarchy.write_pdb_file(med_filename,
                                         anisou=False,
                                         crystal_symmetry=model.input.crystal_symmetry())

        from IPython import embed; embed()


class MultiDatasetTLSAnalysis(object):

    _t_name = ['T11','T22','T33','T12','T13','T23']
    _l_name = ['L11','L22','L33','L12','L13','L23']
    _s_name = ['S11','S12','S13','S21','S22','S23','S31','S32','S33']
    _all_names = _t_name+_l_name+_s_name

    pca_groups = [('T',   _t_name),
                  ('L',   _l_name),
                  ('S',   _s_name),
                  ('TLS', _all_names)]

    def __init__(self, models, out_dir='tls-analysis', csv_base='tls-params-'):

        self.out_dir = easy_directory(out_dir)

        self.csv_base = os.path.join(self.out_dir, os.path.basename(csv_base))
        self.tables = {}

        self.log = Log(verbose=True)

        for m in models:
            self.add(m)

    def add(self, model):

        tls_params = extract_tls_from_pdb(model.filename)

        for tls_fit in tls_params.tls_params:
            # Extract data table for this selection
            tls_table = self.tables.setdefault(tls_fit.selection_string, pandas.DataFrame(columns=self._t_name+self._l_name+self._s_name))

            tls_table.loc[model.tag]=None
            tls_table.loc[model.tag,self._t_name] = tls_fit.t
            tls_table.loc[model.tag,self._l_name] = tls_fit.l
            tls_table.loc[model.tag,self._s_name] = tls_fit.s

    def show(self):
        for selection in sorted(self.tables.keys()):
            self.log.subheading(selection)
            print self.tables[selection]

    def write(self):
        with open(self.csv_base+'selections.log', 'w') as csv_log:
            for i, selection in enumerate(sorted(self.tables.keys())):
                n = i+1
                csv_log.write('selection {:03d} : {}\n'.format(n, selection))
                self.tables[selection].to_csv(self.csv_base+'selection-{:03d}'.format(n)+'.csv')

    def make_plots(self):

        from IPython import embed; embed()

        raise SystemExit()

    def run_pca(self):

        # Iterate through the different TLS selections
        for tls_sel_str in sorted(self.tables.keys()):

            sel_table = self.tables[tls_sel_str]

            # Analyse different parts of the TLS parameterisation
            for group_name, group_sel in self.pca_groups:

                self.log.subheading('{}   -   ({})'.format(group_name, ','.join(group_sel)))

                sel_data = sel_table[group_sel].values.astype('float64')

                pca = mdp.nodes.PCANode(reduce=True)
                pca.train(sel_data)
                pca.stop_training()

                # Bolt on a couple of other statistics
                pca.d_perc = 100.0*pca.d/pca.total_variance
                pca.component_corr_to_avg = numpy.corrcoef(pca.avg, pca.get_recmatrix())[0][1:]

                # Summary
                self.log('The PCA was trained on {} {} matrices (of {} parameters each)'.format(pca.tlen, group_name, pca.input_dim))
                self.log('The resultant PCA is formed of {} principal components'.format(pca.output_dim))
                self.log('...which explain {}% of the observed variance'.format(100.0*pca.explained_variance))
                self.log('...and which individually explain')
                self.log('\t{}'.format('\n\t'.join(['{:7.3f}\t({:7.3f}%)'.format(v,p) for v,p in zip(pca.d,pca.d_perc)])))
                self.log('...of the variance.')
                self.log.bar()
                self.log('The average TLS parameters are')
                self.log('\t'+'\n\t'.join(['{} : {:8.3f}'.format(*v) for v in zip(group_sel,pca.avg.tolist()[0])]))
                self.log.bar()
                self.log('The correlations between the average TLS model and each of the principle components is:')
                self.log('\t{}'.format('\n\t'.join(['{:8.3f}'.format(s) for s in pca.component_corr_to_avg])))

        #        print 'projection matrix: '
        #        print pca.get_projmatrix()

                proj_vals = pca.execute(sel_data)

                print proj_vals.shape

        #        self.plot_bar(pca.d)
        #        self.plot_2d(proj_vals[:,:2])

                from IPython import embed; embed()

                #help(pca)

    def plot_2d(self, data, block=True):
        fig = pyplot.figure(figsize=(8,8))
        ax = fig.add_subplot(111)
        pyplot.rcParams['legend.fontsize'] = 10
        ax.plot(data[:,0], data[:,1], 'o', markersize=8, color='blue', alpha=0.5)
        ax.set_aspect('equal')
        pyplot.show(block)

    def plot_3d(self, data, block=True):
        fig = pyplot.figure(figsize=(8,8))
        ax = fig.add_subplot(111, projection='3d')
        pyplot.rcParams['legend.fontsize'] = 10
        ax.plot(data[:,0], data[:,1], data[:,2], 'o', markersize=8, color='blue', alpha=0.5)
        ax.set_aspect('equal')
        pyplot.show(block)

    def plot_bar(self, vals, block=True):
        fig = pyplot.figure(figsize=(8,8))
        ax = fig.add_subplot(111)
        pyplot.rcParams['legend.fontsize'] = 10
        ax.plot(range(1,len(vals)+1), vals, 'o', markersize=8, color='blue', alpha=0.5)
        pyplot.show(block)

def wrapper_run(tls_fit):
    if os.path.exists(tls_fit.tls_initial_pdb) and os.path.exists(tls_fit.tls_refined_pdb):
        tls_fit.log('tls parameters already fitted and refined: {}'.format(tls_fit.tag))
        return tls_fit
    tls_fit.log.heading('Processing: {}'.format(tls_fit.tag))
    if not os.path.exists(tls_fit.tls_initial_pdb):
        pdb1 = tls_fit.initial_tls_parameters()
    if not os.path.exists(tls_fit.tls_refined_pdb):
        pdb2 = tls_fit.refined_tls_parameters()
    return tls_fit

def run(params):

    log = Log(verbose=True)
    log.heading('Validating input parameters')

    out_dir = easy_directory(params.output.out_dir)

    if os.path.exists(params.input.tls_selections):
        tls_selections = open(params.input.tls_selections, 'r').read().strip().split('\n')
        log('Using existing TLS selections:')
        log('\t'+'\n\t'.join(tls_selections))
    else:
        tls_selections = None

    log.heading('Fitting TLS Parameters')

    all_fits = []

    for p in params.input.pdb:

        if params.input.labelling == 'foldername':
            tag = os.path.basename(os.path.dirname(p))
        elif params.input.labelling == 'filename':
            tag = os.path.basename(os.path.splitext(p)[0])

        fit = TLSFitterRefiner(pdb_file=p,
                               mtz_file=p.replace('.pdb', '.mtz'),
                               cif_files=params.input.cif,
                               out_dir=os.path.join(out_dir,tag+'-tls'),
                               tag=tag,
                               tls_selections=tls_selections)

        if tls_selections is None:
            tls_selections=fit.tls_selections

            assert not os.path.exists(params.input.tls_selections)
            with open(params.input.tls_selections, 'w') as fh: fh.write('\n'.join(tls_selections))

        all_fits.append(fit)

    all_fits = libtbx.easy_mp.pool_map(fixed_func=wrapper_run, args=all_fits, processes=params.settings.cpus, chunksize=1)

    tls_boundaries = identify_tls_boundaries(tls_selections)

    log.heading('Extracting models for refined structures')
    all_models = [CrystallographicModel.from_file(f.tls_refined_pdb).label(tag=f.tag) for f in all_fits]
    log('{} models extracted'.format(len(all_models)))

    # ===================================++>
    # FILTERING
    # ===================================++>
    reject = []
    for m in all_models:
        if not m.hierarchy.is_similar_hierarchy(all_models[0].hierarchy):
            print '{} is not the same hierarchy as the first model'.format(m.tag)
            reject.append(m)
    all_models = [m for m in all_models if m not in reject]
    log('{} models after filtering'.format(len(all_models)))
    if len(all_models) == 1:
        print 'Only one model after filtering -- cannot proceed'

    # ===================================++>
    # BFACTOR ANALYSIS
    # ===================================++>

    log.heading('Performing B-factor analysis')
    multi_bfa = MultiDatasetBFactorAnalysis(models=all_models)
    multi_bfa.write_b_factor_plots(boundaries=tls_boundaries)
    multi_bfa.write_median_b_factors()

    # ===================================++>
    # TLS ANALYSIS
    # ===================================++>

    log.heading('Performing TLS analysis')
    multi_tls = MultiDatasetTLSAnalysis(models=all_models)

#    multi_tls.show()
    multi_tls.write()
    multi_tls.run_pca()



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

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)
