#!/usr/bin/env ccp4-python

import os, sys, glob, time, re
import numpy, pandas

import iotbx.pdb

from giant.structure.summary import atomGroupSummary
from giant.xray.symmetry import match_sites_by_symmetry

from pandda.constants import PanddaDatasetFilenames

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

pandda_csv = './pandda/analyses/pandda_analyse_events.csv'
assert os.path.exists(pandda_csv)

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

pandda_table = pandas.read_csv(pandda_csv, sep=',', dtype={'dtag':str})
pandda_table = pandda_table.set_index(['dtag','event_idx'])

event_counts = dict(pandda_table.index.get_level_values('dtag').value_counts())

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

pandda_table['Actual Occupancy'] = -1.0
pandda_table['Refined B-Factor'] = -1.0
new_csv = './pandda_export.csv'

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

top_dir = os.getcwd()
print 'WORKING FROM: {}'.format(top_dir)
            
lig_occs= []

for dir in glob.glob(os.path.join(top_dir,'pandda-export','*')):
    if not os.path.isdir(dir): continue
    
    ###################################################

    print 'LOOKING AT: {}'.format(dir)
    # Name of the dataset should be the directory name
    dtag = os.path.basename(dir)

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

    # Extract only the rows for this dataset
    dataset_info = pandda_table[pandda_table.index.get_level_values('dtag') == dtag]
    if dataset_info.empty: continue
    # Number of events in this dataset
    num_events = event_counts[dtag]
    assert num_events == len(dataset_info)

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

    # Refined structures
    refined_pdb = os.path.join(dir, 'refine.pdb')
    refined_mtz = os.path.join(dir, 'refine.mtz')
    if not (os.path.exists(refined_pdb) and os.path.exists(refined_mtz)):
        print 'REFINED FILES DO NOT EXIST: SKIPPING'
        continue

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

    # Load the structure
    prot_i = iotbx.pdb.hierarchy.input(refined_pdb)
    prot_h = prot_i.hierarchy

    # Extract the ligands from the hierarchy
    ligs = [ag for ag in prot_h.atom_groups() if ag.resname=='LIG']
    print 'DATASET {}: LIGS - {}, EVENTS - {}'.format(dtag, len(ligs), num_events)
    # No modelled ligands... 
    if not ligs: continue

    # Get residues summaries
    lig_sums = [atomGroupSummary(ag) for ag in ligs]

    print '=========================>'
    for ag_sum in lig_sums:
        print 'LIGAND OCCUPANCY: {}'.format(ag_sum.occies_stats.mean)
    print '=========================>'
    for e_index in dataset_info.index:    
        event_info = dataset_info.loc[e_index]
        print 'ESTIMATED OCCUPANCIES: {}'.format(event_info['1-BDC'])
    print '=========================>'

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

    # Extract symmetry properties so we can calculate symmetry pairings
    uc = prot_i.input.crystal_symmetry().unit_cell()
    sg = prot_i.input.crystal_symmetry().space_group()

    # Extract sites for the events and the ligand atoms
    event_sites = dataset_info[['x','y','z']].values.tolist()
    ligand_sites = [ag.atoms().extract_xyz() for ag in ligs]

    # Check to see how we're to match them up
    sym_equiv = match_sites_by_symmetry(ref_sites   = event_sites,
                                        query_sites = ligand_sites,
                                        unit_cell   = uc,
                                        space_group = sg            )

    ###################################################
    
    # Iterate through and match up the sites
    for i_event, e_index in enumerate(dataset_info.index):

        lig_equivs = sym_equiv[i_event]
        if sum(lig_equivs) == 0:
            print 'No Matching ligands for Event {}'.format(e_index)
            continue
        if sum(lig_equivs) > 1:
            print 'More than one matching ligand for Event {}'.format(e_index)
            continue

        for i_lig, lig_ag in enumerate(ligs): 
            # Skip if not equivalent
            if lig_equivs[i_lig] == 0:
                print 'Ligand {} not matched to this site'.format(i_lig)
                continue
            # Skip if ligand associated with more than one event
            if sum(sym_equiv[:,i_lig]) > 1:
                print 'More than one matching event for Ligand {}'.format(i_lig)
                continue

            lig_sum = lig_sums[i_lig]
            pandda_table.set_value(index=e_index, col='Actual Occupancy', value=lig_sum.occies_stats.mean)
            pandda_table.set_value(index=e_index, col='Refined B-Factor', value=round(lig_sum.b_facs_stats.mean,2))

pandda_table.to_csv(new_csv)
