#!/usr/bin/env ccp4-python

import os, sys, glob, shutil
import wget

import pandas
import iotbx.pdb
import libtbx.easy_mp

from bamboo.rscb_utils import get_pdb_index, download_and_organise_pdb_as_necessary

def process_line_wrap(line):

    if not line: return None

    try:
        ret = process_line(line)
        return ret
    except KeyboardInterrupt:
        raise SystemExit()
    except Exception as e:
        print e.__class__, e.message.strip('\n')+'\n'
        raise 

def process_line(line):
    """Download and extract the het molecules from the pdb file"""
    
    pdb_id, pdb_type, pdb_method = line.split('\t')
        
    print 'Processing {}'.format(pdb_id)

    if pdb_type != 'prot':
        return
    elif pdb_method != 'diffraction':
        return

    # Download and organise files as necessary
    pdb_file = download_and_organise_pdb_as_necessary(pdb_id=pdb_id)
   
    try: 
        i = iotbx.pdb.hierarchy.input(file_name=pdb_file)
        h = i.hierarchy
    except:
        print 'Bad PDB File'
        return

    cache = h.atom_selection_cache()
    sel = cache.selection('(not pepnames) and (not water) and (not element H)')
    # Skip if nothing selected
    if sum(sel) == 0: return
    new_h = h.select(sel)
    
    # Extract resolution of file
    res = i.input.get_r_rfree_sigma().high

    sel_c = []

    for lig_rg in new_h.residue_groups():

        # Total occupancy for the residue
        total_occ = 0
        # Mark residue initially to be recorded
        record_residue = True

        # Iterate through conformers of the residue
        for lig_conf in lig_rg.conformers():

            # Skip if single atom
            if lig_conf.atoms().size() < 5: 
                record_residue = False
                break
            # Skip if common small molecule
            if iotbx.pdb.common_residue_names_get_class(lig_conf.only_residue().resname) == 'common_small_molecule': 
                record_residue = False
                break
            # Skip if different occupancies in same residue
            if lig_conf.atoms().extract_occ().standard_deviation_of_the_sample() != 0.0: 
                record_residue = False
                break

            total_occ += lig_conf.atoms().extract_occ()[0]

        # Skip if not to be recorded
        if not record_residue: continue

        # Append ligand object
        sel_c.append((lig_rg.conformers()[0], {'occ':total_occ, 'res':res}))

    if not sel_c: 
        return
   
    ret_tab = pandas.DataFrame(columns=['pdbid','chain','resid','resname','altid','occ','res'])
    for c, info_dict in sel_c:
        ret_tab = ret_tab.append([{ 'pdbid':    pdb_id,
                                    'chain':    c.parent().id,
                                    'resid':    c.only_residue().resid(),
                                    'resname':  c.only_residue().resname,
                                    'altid':    c.altloc,
                                    'occ':      info_dict['occ'],
                                    'res':      info_dict['res']           }], ignore_index=True)
 
    print ret_tab
    return ret_tab

def run():

    pdb_index = get_pdb_index()
    print 'Using PDB file index from {}'.format(pdb_index)

    pdb_entries = open(pdb_index, 'r').read().split('\n')
    num_entries = len(pdb_entries)
    print 'Found {} codes'.format(num_entries)
    
    print '==========================================================>'
    print '===> Extracting HET Molecules'
    print '==========================================================>'

    # Get the HET molecules from the PDB files
    output = libtbx.easy_mp.pool_map(func=process_line_wrap, args=pdb_entries, processes=7)
      
    print '==========================================================>'        
    print 'Extracted molecules'
    print '==========================================================>'        
    print '\n'*2
    print '==========================================================>'        
    print '===> Compiling output'
    print '==========================================================>'        

    out_tab = pandas.DataFrame(columns=['pdbid','chain','resid','resname','altid','occ','res'])
    for ret in output:
        # Nothing returned
        if ret is None: continue
        # Errors
        if isinstance(ret, Exception):
            print ret.__type__, ret.message.strip('\n')+'\n'
            continue       
        # Append to output table 
        out_tab = out_tab.append(ret, ignore_index=True)
        print ret
    # Write output
    out_tab.to_csv('./pdb_het_molecules.csv')

    print '==========================================================>'
    print '===> FINISHED'
    print '==========================================================>'

if __name__=="__main__":
    run()

