#!/usr/bin/env ccp4-python

import os, sys, glob, time
import numpy

from scitbx.array_family import flex

from bamboo.common.path import easy_directory, rel_symlink
from bamboo.common.print import status_bar

from giant.stats.cluster import cluster_data
from giant.stats.utils import calculate_roc_vals

from pandda.analyse import PanddaMultiDatasetAnalyser
from pandda import welcome

from libtbx import easy_pickle

def write(fh, mes):
    fh.write(mes+'\n')
    print mes

def test_pandda_results(summary_files, pandda, ignores=None, output_dir='./validation_output'):
    """If a list of hits is available, test to see whether the pandda identified them"""

    if ignores is None: ignores = []

    output_dir = easy_directory(output_dir)
    output_log = os.path.join(output_dir, 'validation.log')
    print 'Writing Files to {!s}'.format(output_dir)
    print 'Writing Log to {!s}'.format(output_log)

    # Cutoff for how close the "discovered" points have to be to the centroid of the ligand to be correct
    max_dist_cart = 5
    max_dist_grid = max_dist_cart/pandda.reference_grid().grid_spacing()

    try:

        out_fh = open(output_log, 'a')

        # ============================================================================>
        #####
        # Extract list of correct sites (sites that should have been identified)
        #####
        # ============================================================================>

        write(out_fh,'===================================>>>')
        write(out_fh,'LOADING CORRECT LIGANDS FROM FILE')

        # Dictionary of  "correct" sites per xtal - populate with empty list for each crystal
        correct_sites = {}; [correct_sites.setdefault(xtal,[]) for xtal in [d.tag for d in pandda.datasets.all()]]

        for summary in summary_files:
            sites = [line.strip().split(', ') for line in open(summary, 'r').readlines()]
            for xtal, model, label, info, x, y, z in sites:
                if xtal not in correct_sites.keys():
                    continue
                elif info in ignores:
                    write(out_fh, ', '.join(map(str,['REJECTING', xtal, model, label, info, x, y, z])))
                    continue
                else:
                    correct_sites[xtal].append([(float(x), float(y), float(z)), label, info])

        num_search_ligands = sum(map(len,correct_sites.values()))

        write(out_fh,'===================================>>>')
        write(out_fh,'LOOKING FOR {!s} LIGANDS'.format(num_search_ligands))
        
        if num_search_ligands == 0:
            write(out_fh,'NO LIGANDS TO LOOK FOR: EXITING.')
            sys.exit()

        # ============================================================================>

        # Pull out the cluster information from the pandda
        combined_cluster = pandda.process_z_value_clusters()

        # Iterate through the correct ligands and check that there is a cluster near them
        write(out_fh,'===================================>>>')
        write(out_fh,'TESTING STEP 1: Checking Hits have been identified')
        write(out_fh,'===================================>>>')

        # Record missed datasets or missed sites
        missed_datasets = []
        missed_sites = []
        # Minimum distances to the ligands
        minimum_dists = []

        # Go through the datasets with ligands in them
        for d_tag in sorted(correct_sites.keys()):
            write(out_fh,'Checking for the ligands in Dataset {!s}'.format(d_tag))

            # Get the list of sites that we should be detecting
            dataset_ligands = correct_sites[d_tag]

            # Get the dataset handler
            d_handler = pandda.datasets.get(tag=d_tag)
            # Get the clusters of identified points for this dataset
            clust = d_handler.hit_clusters

            # Go through the ligands in the dataset
            for site_cart_native, label, info in dataset_ligands:

                # Map to reference frame            
                if d_handler.local_alignment_transforms():
                    site_cart_ref = d_handler.transform_to_reference([site_cart_native], 'local', d_handler.find_nearest_calpha([site_cart_native]))
                else:
                    site_cart_ref = d_handler.transform_to_reference([site_cart_native], 'global')
                assert len(site_cart_ref) == 1
                site_grid_ref = site_cart_ref.as_double()/pandda.reference_grid().grid_spacing()

                write(out_fh,'\tLooking for {!s} at {!s}'.format(label, list(site_grid_ref)))

                # Check to see if there are any identified points in this dataset
                if clust==None:
                    write(out_fh,"\t\tTHIS DATASET WASN'T IDENTIFIED AS INTERESTING - MISSED IT")
                    missed_datasets.append((d_tag, label))
                    continue

                # Calculate the distances between ligands and the identified clusters
                dists = []

                # Iterate through the clusters in this dataset
                for clust_points in clust.get_points():
                    # Iterate through the points in this cluster
                    for grid_point in clust_points:
                        dists.append((flex.double(grid_point) - flex.double(site_grid_ref)).norm())
                        #print '\t\tHow about {!s}?'.format(grid_point)
                    
                min_dist = min(dists)
                if min_dist < max_dist_grid:
                    write(out_fh,'\t\t\tFOUND IT!')
                else:    
                    write(out_fh,'\t\t\tMISSED IT...')
                    missed_sites.append((d_tag, label))
                write(out_fh, '\t\t\t\tMIN DIST: {!s}'.format(min_dist))

                minimum_dists.append(min_dist*pandda.reference_grid().grid_spacing())

        write(out_fh,'----------------------------------->>>')
        write(out_fh,'MISSED DATASETS: \n\t\t{!s}'.format('\n\t\t'.join(map(str,missed_datasets))))
        write(out_fh,'MISSED LIGANDS:  \n\t\t{!s}'.format('\n\t\t'.join(map(str,missed_sites))))
        write(out_fh,'----------------------------------->>>')
        write(out_fh,'MINIMUM DISTANCES TO LIGAND (A): \n{!s}'.format(map(int,minimum_dists)))

        write(out_fh,'===================================>>>')
        write(out_fh,'Sorting Clusters')

        # Sort the clusters by size
        size_sorted_indices = combined_cluster.sort(sorting_data='sizes', sorting_function=None, decreasing=True)
        # Sort the clusters by mean
        mean_sorted_indices = combined_cluster.sort(sorting_data='values', sorting_function=numpy.mean, decreasing=True)
        # Sort the clusters by max
        max_sorted_indices = combined_cluster.sort(sorting_data='values', sorting_function=max, decreasing=True)

        write(out_fh,'===================================>>>')
        write(out_fh,'TESTING STEP 2: Calculating ROC Curves')

        COMBINED_ROC_RESULTS = {}

        for test_num, (c_indices, rank_vals) in enumerate([(size_sorted_indices, combined_cluster.get_sizes(size_sorted_indices)),
                                                           (mean_sorted_indices, combined_cluster.get_means(mean_sorted_indices)),
                                                           (max_sorted_indices,  combined_cluster.get_maxima(max_sorted_indices))]):

            test_type = ['c_size','z_mean','z_peak'][test_num]

#           if test_type != 'z_peak':
#               write(out_fh,'Skipping test: {!s}'.format(test_type))
#               continue

            # Output ROC information
            ROC_OUTPUT = []

            write(out_fh,'===================================>>>')
            write(out_fh,'Ranking by: {!s}'.format(test_type))
            write(out_fh,'===================================>>>')

            # List to see when a ligand is identified - avoids multiple identifications
            identified_ligands = []
            identified_idx = 1

            # Calculate ROC Curves etc - Iterate through clusters and check if near ligand
            for c_rank, c_index in enumerate(c_indices):
                # Pull out information to identify the cluster
                d_tag, c_num = combined_cluster.get_keys([c_index])[0]
                #c_centroid = combined_cluster.get_centroids([c_index])[0]
                write(out_fh,'Checking Dataset {:4}, Cluster {:4}, Rank: {:4}, Val: {:6.2f}'.format(d_tag, c_num, c_rank, rank_vals[c_rank]))

                # Reset
                is_a_ligand = 0
                skip = False

                # Check if there is a ligand in the dataset
                if d_tag not in correct_sites.keys():
                    pass
                else:

                    # Get the ligands present in the dataset
                    dataset_ligands = correct_sites[d_tag]

                    for l_num, (site_cart_native, label, info) in enumerate(dataset_ligands):
                        # Mark whether this ligand has been found
                        ligand_found = 0

                        # Map to reference frame            
                        if d_handler.local_alignment_transforms():
                            site_cart_ref = d_handler.transform_to_reference([site_cart_native], 'local', d_handler.find_nearest_calpha([site_cart_native]))
                        else:
                            site_cart_ref = d_handler.transform_to_reference([site_cart_native], 'global')
                        assert len(site_cart_ref) == 1
                        site_grid_ref = site_cart_ref.as_double()/pandda.reference_grid().grid_spacing()
                
                        # Iterate through and see if this overlaps with the ligand
                        for c_point in combined_cluster.get_points([c_index])[0]:
                            # Calculate distance to ligand
                            dist = (flex.double(site_grid_ref) - flex.double(c_point)).norm()
                            # Check if correct
                            if dist < max_dist_grid:
                                ligand_found = 1
                                break

                        # If it's a ligand, check to see if it's been previously identified
                        if ligand_found:
                            if (d_tag, l_num) in identified_ligands:
                                write(out_fh,'\t. - Ligand already found: Dataset {!s}'.format(d_tag))
                                print 'THIS SHOULDN\'T HAPPEN!!!!'
                                skip = True
                            else:
                                is_a_ligand = 1
                                write(out_fh,"\t{!s} - Ligand found: Dataset {!s}".format(identified_idx, d_tag))
                                identified_idx += 1
                                identified_ligands.append((d_tag, l_num))

                if not skip:
                    ROC_OUTPUT.append({'d_tag': d_tag, 'c_num': c_num, 'is_a_ligand': is_a_ligand, 'rank': c_rank, 'val': rank_vals[c_rank]})

            # Calculate the ROC Curve
            correct_class = [t['is_a_ligand'] for t in ROC_OUTPUT]
            rank_vals     = [t['val'] for t in ROC_OUTPUT]
            ROC_RESULTS = calculate_roc_vals(correct_class=correct_class, rank_vals=rank_vals)
            COMBINED_ROC_RESULTS[test_type] = ROC_RESULTS

            write(out_fh,'Ligands Identified: {!s}'.format([l[0] for l in identified_ligands]))
        
        write(out_fh,'----------------------------------->>>')
        write(out_fh,'SUMMARIES')
        write(out_fh,'----------------------------------->>>')
        write(out_fh,'MISSED DATASETS: {!s}\n\t\t{!s}'.format(len(missed_datasets), '\n\t\t'.join(map(str,missed_datasets))))
        write(out_fh,'MISSED LIGANDS: {!s}\n\t\t{!s}'.format(len(missed_sites), '\n\t\t'.join(map(str,missed_sites))))
        write(out_fh,'----------------------------------->>>')
        write(out_fh,'MINIMUM DISTANCES TO LIGAND (A): \n{!s}'.format(map(int,minimum_dists)))
        write(out_fh,'----------------------------------->>>')
        write(out_fh,'IDENTIFIED LIGANDS: {!s}\n\t\t{!s}'.format(len(identified_ligands), '\n\t\t'.join(map(str,[(l[0], correct_sites[l[0]][l[1]][1]) for l in identified_ligands]))))

        # Load plotting
        try:
            import matplotlib
            # Setup so that we can write without a display connected
            matplotlib.interactive(0)
            default_backend, validate_function = matplotlib.defaultParams['backend']
            from matplotlib import pyplot
            output_graphs = True
        except:
            output_graphs = False

        # OUTPUT ROC CURVE
        if output_graphs:
        
            write(out_fh,'----------------------------------->>>')
            write(out_fh,'PLOTTING ROC GRAPHS')

            for test_type in COMBINED_ROC_RESULTS.keys():

                roc = COMBINED_ROC_RESULTS[test_type]

                fig = pyplot.figure()
                pyplot.title('ROC CURVE FOR {!s}'.format(test_type))
                pyplot.plot(roc['VALS'], roc['SENS'], marker='o', linestyle='-', color='k', label='Sensitivity')
                pyplot.plot(roc['VALS'], roc['SPEC'], marker='o', linestyle='-', color='r', label='Specificity')
                pyplot.plot(roc['VALS'], roc['PREC'], marker='o', linestyle='-', color='g', label='Precision')
                pyplot.plot(roc['VALS'], roc['ACC'],  marker='o', linestyle='-', color='b', label='Accuracy')
                pyplot.xlabel('SCORE CUTOFF ({!s})'.format(test_type))
                pyplot.ylabel('ROC RESPONSE')
#                pyplot.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)
                pyplot.legend()
                # Apply tight layout to prevent overlaps
                pyplot.tight_layout()
                # Save
                pyplot.savefig(os.path.join(output_dir, 'roc-cutoff-{!s}.png'.format(test_type)))
                pyplot.close(fig)

                fig = pyplot.figure()
                pyplot.title('ROC CURVE FOR {!s}'.format(test_type))
                pyplot.plot(roc['FPR'], roc['TPR'], marker='o', linestyle='-', color='g')
                pyplot.xlabel('FPR')
                pyplot.ylabel('TPR')
                # Apply tight layout to prevent overlaps
                pyplot.tight_layout()
                # Save
                pyplot.savefig(os.path.join(output_dir, 'roc-tpr-fpr-{!s}.png'.format(test_type)))
                pyplot.close(fig)

        # ============================================================================>
        # ======================================>
        # Manual Analyses - These only need processing once
        # ======================================>
        # ============================================================================>
        analyses_dir = output_dir
        # ============================================================================>

        if 0:
            write(out_fh,'===================================>>>')
            write(out_fh,'Calculating Deviations of C-alphas between structures')

            rms = lambda vals: numpy.sqrt(numpy.mean(numpy.abs(vals)**2))
            norm = lambda vals: numpy.sqrt(numpy.sum(numpy.abs(vals)**2))

            # Pull all c-alpha sites for each structure
            all_sites = numpy.array([d.transform_points_to_reference(d.get_calpha_sites()) for d in pandda.datasets.all()])
            # Calculate the mean x,y,z for each c-alpha
            mean_sites = numpy.mean(all_sites, axis=0)
            # Differences from the mean
            diff_sites = all_sites - mean_sites
            # Euclidean norms of the distances moved
            diff_norms = numpy.apply_along_axis(norm, axis=2, arr=diff_sites)

            with open(os.path.join(analyses_dir,'calpha_variation.csv'), 'w') as fh:
                for row in diff_norms:
                    out_list = row.round(3).tolist()
                    out_line = ', '.join(map(str,out_list)) + '\n'
                    fh.write(out_line)

            write(out_fh,'Largest deviation from the mean site: {!s}'.format(diff_norms.max()))
            write(out_fh,'Average deviation from the mean site: {!s}'.format(diff_norms.mean()))

        # ============================================================================>

        if 0:
            write(out_fh,'===================================>>>')
            write(out_fh,'Clustering the Refined Structures')

            distance_matrix = []
            for d1 in pandda.datasets.all():
               distance_matrix.append([d1.transform_points_to_reference(d1.get_calpha_sites()).rms_difference(d2.transform_points_to_reference(d2.get_calpha_sites())) for d2 in pandda.datasets.all()])

            distance_matrix = numpy.array(distance_matrix)

            with open(os.path.join(analyses_dir,'calpha_distance_matrix.csv'), 'w') as fh:
                for row in distance_matrix:
                    out_list = row.round(3).tolist()
                    out_line = ', '.join(map(str,out_list)) + '\n'
                    fh.write(out_line)

        # ============================================================================>

    except:
        write(out_fh,'FAILURE DURING TESTING')
        raise
    finally:
        out_fh.close()

    return 0

# ============================================================================>
#
#   COMMAND LINE RUN
#
# ============================================================================>

if __name__ == '__main__':

    welcome(os.getlogin())

    print 'PROCESSING ARGUMENTS'
    # Find summary files that have been supplied
    summary_files = [a for a in sys.argv if os.path.isfile(a) and a.endswith('.summary')]
    if not summary_files:
        sys.exit('NO ".summary" FILES PROVIDED')

    # Find pickles that have been supplied
    pickles = [a for a in sys.argv if os.path.isfile(a) and a.endswith('.pickle')]
    if not pickles:
        sys.exit('NO ".pickle" FILES PROVIDED')

    print 'LOADING PICKLED OBJECTS'
    pickle_objs = [easy_pickle.load(p) for p in pickles]
    potential_panddas = [o for o in pickle_objs if isinstance(o, PanddaMultiDatasetAnalyser)]
    if not potential_panddas:
        sys.exit('NO PANDDAS PROVIDED')
    elif len(potential_panddas) != 1:
        sys.exit('MORE THAN ONE PANDDA SUPPLIED')

    print 'IGNORING MOLECULES'
    ignores = [a[7:].replace('\'','').replace('\"','') for a in sys.argv if a.startswith('ignore=')]

    pandda = potential_panddas[0]

    print 'RUNNING TESTING'
    test_pandda_results(summary_files=summary_files, pandda=pandda, ignores=ignores)
