#!/usr/bin/env ccp4-python

import os, sys, copy, re, glob
import random

import libtbx.phil
import iotbx.pdb, iotbx.mtz

import numpy, pandas

from scitbx.array_family import flex

from bamboo.common.path import easy_directory, rel_symlink
from bamboo import div_bar

bar = div_bar(30)

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

master_phil = libtbx.phil.parse("""
input {
    pdb = None
        .type = str
        .multiple = True
}
output {
    resolution_mean = 1.5
        .type = float
    resolution_std  = 0.1
        .type = float
    out_dir = ./test-data
        .type = path
}
verbose = False
""")

def run(params):

    assert params.input.pdb, 'No PDB Files provided: params.input.pdb'
    for p in params.input.pdb: assert os.path.exists(p), 'PDB File does not exist: {}'.format(p)

    out_dir = easy_directory(params.output.out_dir)
   
    for idx, pdb in enumerate(params.input.pdb):
 
        s = iotbx.pdb.hierarchy.input(pdb)
        s_input     = s.input
        s_hierarchy = s.hierarchy

        num = idx+1
        model_dir = easy_directory(os.path.join(out_dir, 'model_{:06}'.format(num)))
        
        print(bar)
        print('Processing Model {}, {}'.format(num, pdb))

        # Write the output hierarchy
        rel_symlink(pdb, os.path.join(model_dir, 'model.pdb'))

        print('>> Generating structure factors from model'.format(num))
        # Convert the structure to structure factors
        dataset_resolution = random.gauss(mu=params.output.resolution_mean, sigma=params.output.resolution_std)
        xray_s_factors = s_hierarchy.extract_xray_structure(crystal_symmetry=s_input.crystal_symmetry()).structure_factors(d_min=dataset_resolution)
        random_dataset = xray_s_factors.f_calc().as_mtz_dataset(column_root_label='FMODEL')
        random_dataset.mtz_object().write(os.path.join(model_dir, 'model.mtz'))
        

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

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)
