#!/usr/bin/env ccp4-python

import os, sys, shutil, copy

from scipy.interpolate import spline

from libtbx import phil, easy_pickle
from cctbx import sgtbx, maptbx
from scitbx.array_family import flex
from scitbx.math.distributions import normal_distribution

import iotbx.ccp4_map
import iotbx.map_tools

import numpy 

from giant.grid.utils import get_grid_points_within_index_cutoff_of_grid_sites
from giant.xray.maps.bdc import calculate_varying_bdc_correlations, estimate_event_background_correction
from pandda.analyse_graphs import write_occupancy_graph

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

master_phil_def = """
input {
    outdir = ./occupancy-subtracted-maps
        .type = str
    outcsv = False
        .type = bool
    mean_map = None
        .type = path
    query_map = None
        .type = path
    pickle = None
        .type = path
    outprefix = 'event'
        .type = str
    min_occ = 0.00
        .type = float
    max_occ = 1.00
        .type = float
    occ_increment = 0.01
        .type = float
    verbose = False
        .type = bool
}
"""
# Read in the master phil
master_phil = phil.parse(master_phil_def)
# Show defaults and exit
if '--show-defaults' in sys.argv:
    print('\n# DEFAULTS\n')
    master_phil.show()
    sys.exit()
# Copy the args so that we can remove items from the list without affecting sys.argv etc
args = copy.copy(sys.argv[1:])
# Build an interpreter from the master phil
cmd_interpr = master_phil.command_line_argument_interpreter(home_scope="input")
# Look for cmd line arguments
cmd_line_args = [a for a in args if (('=' in a) and not os.path.exists(a))]
# Remove them from the original lists
[args.remove(a) for a in cmd_line_args]
# Parse these arguments
cmd_sources = [cmd_interpr.process(arg=a) for a in cmd_line_args]
# Combine the phils
working_phil = master_phil.fetch(sources=cmd_sources)

# Pull out the python object of the arguments
params = working_phil.extract().input

# Check that a pickle file has been given
assert params.pickle, 'No PICKLE file provided'
assert params.mean_map, 'No MEAN_MAP file provided'
assert params.query_map, 'No QUERY_MAP file provided'

# Create output directory
if not os.path.exists(params.outdir):   os.mkdir(params.outdir)
if params.outcsv:
    params.outcsv = os.path.join(params.outdir, 'all_vals.csv')
    if params.verbose: print 'WRITING TO', params.outcsv
    if not os.path.exists(params.outcsv):
        with open(params.outcsv, 'w') as fh:
            fh.write('xtal, lig_occ, stat_type, stat_val\n')
############################################################################

# Read in the input maps
mean_map = iotbx.ccp4_map.map_reader(params.mean_map)
query_map = iotbx.ccp4_map.map_reader(params.query_map)
# Check that the maps are the same
assert mean_map.data.all() == query_map.data.all()
assert mean_map.data.size() == query_map.data.size()
assert mean_map.unit_cell().is_similar_to(query_map.unit_cell())
# Print map sizes
if params.verbose: print 'ALL', mean_map.data.all() 
if params.verbose: print 'ORIGIN', mean_map.data.origin()
if params.verbose: print 'FOCUS', mean_map.data.focus()

#####################################################################
# Map accessors
grid_accessor = mean_map.data.accessor()

# Extract the map grid points around the ligand
dataset = easy_pickle.load(params.pickle)

# Extract the map uncertainty
try:    map_uncertainty = dataset.child.meta.map_uncertainty
except: map_uncertainty = 'NONE'

# Extract FEATURE REGION
if dataset.events:
    feature_sites = dataset.events[0].cluster.points
    print feature_sites
    # Expand these points by 10 grid points
    tot_sites, max_sites, min_sites = get_grid_points_within_index_cutoff_of_grid_sites(feature_sites, max_grid_dist=3)
    expanded_sites = max_sites
else:
    sys.exit('NO HIT CLUSTERS FOR THIS DATASET (PICKLE)')

# Print
print 'DATASET DETAILS:'
print 'Z  SITES:', len(feature_sites)
print 'Z+ SITES:', len(expanded_sites)
print 'MAP UNCERTAINTY:', map_uncertainty

# Calculate the grid points for the feature region
feature_region = [gp for gp in expanded_sites if mean_map.data[grid_accessor(gp)]]
# Calculate the grid points for the reference region
reference_region = [gp for gp in flex.nested_loop(mean_map.data.all()) if mean_map.data[grid_accessor(gp)]]
    
#feature_region=[gp for i,gp in enumerate(flex.nested_loop(mean_map.data.all())) if expanded_bool[i]],
#reference_region=[gp for i,gp in enumerate(flex.nested_loop(mean_map.data.all())) if sel_bool[i]],

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

lig_occs, feature_corrs, global_corrs = calculate_varying_bdc_correlations(
    ref_map=mean_map.data.as_double(),
    query_map=query_map.data.as_double(),
    feature_region=feature_region,
    reference_region=reference_region,
    min_remain=params.min_occ,
    max_remain=params.max_occ,
    bdc_increment=params.occ_increment)

# Calculate occupancy by correlation discrepancy
corr_discrep = [global_corrs[i] - feature_corrs[i] for i in range(len(feature_corrs))]
max_corr_discrep = max(corr_discrep)
lig_occ_est = lig_occs[corr_discrep.index(max_corr_discrep)]
print 'CORRELATION-DISCREPANCY OCCUPANCY ESTIMATE:', lig_occ_est

# Calculate correlation gradients
feature_gradients = [feature_corrs[i+1] - feature_corrs[i-1] for i in range(1, len(feature_corrs)-1)]
global_gradients =  [global_corrs[i+1] - global_corrs[i-1] for i in range(1, len(global_corrs)-1)]
corr_gradient_discrep = [feature_gradients[i] - global_gradients[i] for i in range(0, len(feature_gradients))]
gradient_occs = lig_occs[1:-1]
# Calculate occupancy by correlation discrepancy
max_corr_grad_discrep = max(corr_gradient_discrep)
lig_occ_est_grad = gradient_occs[corr_gradient_discrep.index(max_corr_grad_discrep)]
print '   GRADIENT-DISCREPANCY OCCUPANCY ESTIMATE:', lig_occ_est_grad

occ_est = estimate_event_background_correction(
    ref_map=mean_map.data.as_double(),
    query_map=query_map.data.as_double(),
    feature_region=feature_region,
    reference_region=reference_region,
    min_remain=params.min_occ,
    max_remain=params.max_occ,
    bdc_increment=params.occ_increment,
    method='value')
print 'FUNCTION:', occ_est
occ_est = estimate_event_background_correction(
    ref_map=mean_map.data.as_double(),
    query_map=query_map.data.as_double(),
    feature_region=feature_region,
    reference_region=reference_region,
    min_remain=params.min_occ,
    max_remain=params.max_occ,
    bdc_increment=params.occ_increment,
    method='gradient')
print 'FUNCTION:', occ_est

#####################################################################
#for i_occ, lig_occ in enumerate(all_lig_occs):
#    
#    outfile = params.outprefix + '-lig-occ-{!s}.ccp4'.format(lig_occ)
#    outpath = os.path.join(params.outdir, outfile) 
#    if params.verbose: print outpath
#
#    # "Un-pad" the map
#    gridding_first = mean_map.data.origin()
#    gridding_last = tuple(flex.int(query_map.data.all())-2)
#    diff_data = maptbx.copy(diff_data, first=gridding_first, last=gridding_last)
# 
#    # Reducing the size of the grid by 1 seems to be neccessary                                
#    iotbx.map_tools.write_ccp4_map( sites_cart = None, 
#                                    unit_cell = query_map.unit_cell(), 
#                                    map_data = diff_data, 
#                                    n_real = diff_data.all(), 
#                                    file_name = outpath, 
#                                    buffer=0 )
#
#    iotbx.ccp4_map.write_ccp4_map(  file_name = outpath,
#                                    unit_cell = query_map.unit_cell(),
#                                    space_group = sgtbx.space_group_info(number=query_map.space_group_number).group(),
#                                    gridding_first = gridding_first,
#                                    gridding_last = gridding_last,
#                                    map_data = diff_data,
#                                    labels = None )

#####################################################################
if params.outcsv:
    print 'WRITING CSV:', params.outcsv
    outstr = '' 
    for i_occ, lig_occ in enumerate(lig_occs):
        outstr += ', '.join(map(str, [params.outprefix, lig_occ, 'CORRELATION-DISCREPANCY', corr_discrep[i_occ]])) + '\n'
    for i_occ, lig_occ in enumerate(gradient_occs):
        outstr += ', '.join(map(str, [params.outprefix, lig_occ, 'CORRELATION-GRADIENT-DISCREPANCY', corr_gradient_discrep[i_occ]])) + '\n'
                                    # label, lig_occ, stat_type, stat_val
    with open(params.outcsv, 'a') as fh:    
        fh.write(outstr)
#####################################################################

write_occupancy_graph(x_values=lig_occs, global_values=global_corrs, local_values=feature_corrs, diff_values=corr_discrep, out_file='new_occupancy.png')

#####################################################################
import matplotlib
matplotlib.interactive(0)
from matplotlib import pyplot
#####################################################################
fig, (axis_1_1, axis_2_1) = pyplot.subplots(2, sharex=True)
#####################################################################
# LINEAR CORRELATION
# 1st y-axis on first subplot
# lig_occs, feature_corrs, global_corrs, corr_discrep
line_1_1, = axis_1_1.plot(lig_occs, global_corrs, 'g--', label='GLOBAL-CORR')
#for tl in axis_1.get_yticklabels():    tl.set_color('g')
line_1_2, = axis_1_1.plot(lig_occs, feature_corrs, 'k--', label='FEATURE-CORR')
axis_1_1.set_ylabel('MEAN-MAP CORRELATIONS', color='k')
axis_1_1.set_ylim((-1, 1))
#####################################################################
# 1st y-axis on second subplot
line_2_1, = axis_2_1.plot(lig_occs, corr_discrep, 'b-', label='DISCREP-CORR')
axis_2_1.set_ylabel('CORRELATION DIFFERENCE', color='b')
# Plot line at the maximum
line_2_2, = axis_2_1.plot([lig_occ_est,lig_occ_est],[-1,1], 'k-')
text_2_1 = axis_2_1.text(0.03+lig_occ_est, -0.5, str(lig_occ_est))
# 2nd y-axis on second subplot
axis_2_2 = axis_2_1.twinx()
line_2_3, = axis_2_2.plot(gradient_occs, corr_gradient_discrep, 'r-', label='DISCREP-CORR-GRAD')
# Plot line at the maximum
line_2_4, = axis_2_2.plot([lig_occ_est_grad,lig_occ_est_grad],[-0.1,0.1], 'k-')
text_2_2 = axis_2_2.text(0.03+lig_occ_est_grad, -0.1, str(lig_occ_est_grad))
# Joint legend
axis_1_1.legend(handles=[line_1_1, line_1_2, line_2_1, line_2_3], loc=1)
#####################################################################






## DIFF MAP RMS
#line_2, = axis_2.plot(all_lig_occs, all_occ_rmsss, 'b-', label='RMS DEV')
#line_3, = axis_2.plot(all_lig_occs, moving_window_ave, 'k-*', label='RMS DEV - SMOOTH')
#axis_2.set_xlabel('LIGAND OCCUPANCY')
#axis_2.set_ylabel('RMS DEV', color='b')
#axis_2.set_ylim((-1.2,0.5))
## Second y-axis on second subplot
#axis_2_2 = axis_2.twinx()
#line_4, = axis_2_2.plot(all_lig_occs, moving_window_std, 'g--', label='RMS DEV LOCAL RMS')
#axis_2_2.set_ylabel('RMS DEV VARIATION', color='g')
#axis_2_2.set_ylim((0, 0.2))
## Joint legend
##axis_2.legend(handles=[line_2, line_3, line_4], loc=1)
##for t2 in axis_2.get_yticklabels():    t2.set_color('b')
######################################################################
## Proxy-Z maps
#line_5, = axis_3.plot(all_lig_occs, norm_all_occ_wuncs, 'r-', label='W-UNCS')
#axis_3.set_xlabel('LIGAND OCCUPANCY')
#axis_3.set_ylabel('Weighted Uncertainties', color='r')
#axis_3.set_ylim((0.9, 1.1))
## Second y-axis on third subplot
#axis_3_2 = axis_3.twinx()
#line_6, = axis_3_2.plot(all_lig_occs, norm_all_occ_uuncs, 'r-', label='U-UNCS')
#axis_3_2.set_ylabel('Un-weighted Uncertainties', color='r')
#axis_3_2.set_ylim((-1, 12))
## Joint legend
#axis_3.legend(handles=[line_5, line_6], loc=1)
#####################################################################
axis_1_1.set_title('MAP QUALITY V LIGAND OCCUPANCY')
axis_2_1.set_xlabel('LIGAND OCCUPANCY')
fig.subplots_adjust(hspace=0)
pyplot.setp([a.get_xticklabels() for a in fig.axes[:-1]], visible=False)
pyplot.setp([axis_2_1.get_xticklabels()], visible=True)
#####################################################################
# Apply tight layout to prevent overlaps
#pyplot.tight_layout()
# Save
pyplot.savefig(os.path.join(params.outdir, '{!s}-lig_occupancy_graph.png'.format(params.outprefix)))
pyplot.close(fig)
#####################################################################


