#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Copyright (C) 2015 National Institutes of Health
#
#    This library is free software; you can redistribute it and/or
#    modify it under the terms of the GNU Lesser General Public
#    License as published by the Free Software Foundation; either
#    version 2.1 of the License, or (at your option) any later version.
#
#    This library is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public
#    License along with this library; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Written by:  Christopher Coletta (github.com/colettace)
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

import numpy as np
from .utils import output_railroad_switch
from .FeatureSpacePredictionExperiment import FeatureSpaceClassificationExperiment


def ordered_uniq( labels ):
    seen = set()
    seen_add = seen.add
    return [ x for x in labels if not (x in seen or seen_add(x) ) ]

#=====================================================================
def RankOrderedPredictedValuesGraph( labels, predicted_values, figsize=(15, 10),
        chart_title=None, fontsize=16, titlefontsize=18, alpha=0.5, s=80, fig=None, ax=None ):
    """Visualizes the distribution of predicted values on a per-class basis.

    Results are sorted in ascending order twice, first by value, then label, resulting in
    a visualization of the distribution of values within a class.

    Requires the package matplotlib to be installed.

    labels - iterable
        Ground truth numeric labels for the sample set.
    predicted_values - iterable
        Predicted values (numbers) for a given set of samples.
    fig - matplotlib.figure.Figure, or None (default):
        If provided, generate axes on the passed-in figure obj, otherwise
        instantiate a new one.

    ax - matplotlib.axes._subplots.AxesSubplot or None (default):
        Plot data on these axes, if provided, otherwise create new axes on figure."""

    import matplotlib.pyplot as plt

    if fig is None and ax is None:
        fig = plt.figure( figsize=figsize )
        ax = fig.add_subplot(111)
    elif fig is None:
        fig = ax.get_figure()
    elif ax is None:
        ax = fig.gca()

    ax.set_title( chart_title, size=titlefontsize )
    ax.set_xlabel( 'Sample Count', size=fontsize )
    ax.set_ylabel( 'Sample Predicted Values', size=fontsize )

    n = len( predicted_values )
    ax.set_xlim( ( - 0.1 * n, 1.1 * n ) )

    predicted_values = np.array( predicted_values )
    class_values = ordered_uniq( labels )
    n_classes = len( class_values )

    _min = min( class_values )
    ampl = max( class_values ) - _min
    class_colors = plt.cm.jet( [ float(val -_min)/ampl for val in class_values ] )

    abscissa_index = 1
    for class_value, class_color in zip( class_values, class_colors ):
        # convert list of bools to the indices that eval to True:
        # why does numpy where return a 1-tuple??
        indices = np.where( [ l == class_value for l in labels ] )[0]
        #print indices
        class_pred_values = sorted( predicted_values[ indices ] )
        x_vals = [ i + abscissa_index for i in xrange( len( class_pred_values ) ) ]
        ax.scatter( x_vals, class_pred_values, c=class_color, marker='o',
                s=s, edgecolor='none', alpha=alpha, label=class_value )
        abscissa_index += len( class_pred_values )

    ax.legend( loc='lower right', fontsize=fontsize )
    return fig

#============================================================================
def KernelSmoothedDensityGraph( labels, predicted_values, bw_method=None, figsize=(15, 10),
        chart_title=None, fontsize=16, titlefontsize=18, linewidth=4, fig=None, ax=None ):
    """This graph visualizes the distribution of predicted values generated by classification.
    A kernel-smoothed probability density function is plotted for each image class on
    the same chart allowing comparison of distribution of predicted values amoung image class.

    Requires the packages matplotlib and scipy. Uses scipy.stats.gaussian_kde to
    generate kernel-smoothed probability density functions.

    labels - iterable
        Ground truth numeric labels for the sample set.
    predicted_values - iterable
        Predicted values (numbers) for a given set of samples.
    bw_method -
        Arg passed directly to bw_method arg in scipy.stats.gaussian_kde
    fig - matplotlib.figure.Figure, or None (default):
        If provided, generate axes on the passed-in figure obj, otherwise
        instantiate a new one.
    ax - matplotlib.axes._subplots.AxesSubplot or None (default):
        Plot data on these axes, if provided, otherwise create new axes on figure."""

    import matplotlib.pyplot as plt

    if fig is None and ax is None:
        fig = plt.figure( figsize=figsize )
        ax = fig.add_subplot(111)
    elif fig is None:
        fig = ax.get_figure()
    elif ax is None:
        ax = fig.gca()

    ax.set_title( chart_title, size=titlefontsize )
    ax.set_xlabel( 'Predicted Value', size=fontsize )
    ax.set_ylabel( 'Probability density', size=fontsize )

    predicted_values = np.array( predicted_values )
    labels = np.array( labels )
    class_values = ordered_uniq( labels )
    n_classes = len( class_values )

    _min = min( class_values )
    ampl = max( class_values ) - _min
    class_colors = plt.cm.jet( [ float(val -_min)/ampl for val in class_values ] )

    from scipy.stats import gaussian_kde

    # For now, ignore "FloatingPointError: underflow encountered in exp'"
    np.seterr( under='ignore' )

    for class_value, class_color in zip( class_values, class_colors ):
        # convert list of bools to the indices that eval to True:
        # why does numpy where return a 1-tuple??
        indices = np.where( [ l == class_value for l in labels ] )[0]
        class_pred_vals = predicted_values[ indices ]
        ground_truth_vals = labels[ indices ]

        lobound = class_pred_vals.min()
        hibound = class_pred_vals.max()
        kernel_smoother = gaussian_kde( class_pred_vals, bw_method=bw_method )
        intervals = np.mgrid[ lobound:hibound:100j ]
        density_estimates = kernel_smoother.evaluate( intervals )
        ax.plot( intervals, density_estimates, c=class_color,
            linewidth=linewidth, label=class_value )

    np.seterr( all='raise' )

    ax.legend( fontsize=fontsize )
    return fig

#============================================================================
class _BaseGraph( object ):
    """An abstract base class that is supposed to hold onto objects on which to call
    matplotlib.pyplot API methods."""

    def __init__( self ):

        # general stuff:
        self.chart_title = None
        self.file_name = None
        self.split_result = None

        # pyplot-specific stuff
        self.figure = None
        self.main_axes = None

    def savefig( self, filepath ):
    
        if self.figure == None:
            raise ValueError( 'No figure to save!' )
        self.figure.savefig( filepath )
        print 'Wrote chart "{0}" to file "{1}"'.format( self.chart_title, filepath )
            
#============================================================================
class PredictedValuesGraph( _BaseGraph ):
    """This is a concrete class that can produce two types of graphs that are produced
    from SingleSamplePrediction data stored in a FeatureSpacePrediction."""

    #=================================================================
    def __init__( self, result, name=None, use_averaged_results=True, figsize=(15, 10), fontsize=16, titlefontsize=18 ):
        """Constructor sorts ground truth values contained in FeatureSpacePrediction
        and loads them into self.grouped_coords
        
        use_averaged_results - bool - If this object has averaged results (due to tiling or 
            "per sample" aggregation across splits, use those results instead of 
            individual results."""

        #FIXME: implement user-definable bin edges

        self.split_result = result
        if name is None:
            name = result.name

        self.chart_title = name
        self.figsize = figsize
        self.fontsize = fontsize
        self.titlefontsize = titlefontsize

        gt_vals, pred_vals = result.RankOrderSort( use_averaged_results=use_averaged_results )
        whole_list = zip( gt_vals, pred_vals )

        self.grouped_coords = {}

        if result.test_set.discrete:
            self.class_names = result.test_set.class_names
            self.class_values = result.test_set.interpolation_coefficients
            self.num_classes = result.test_set.num_classes
            for class_val, class_name in zip( self.class_values, self.class_names ):
                self.grouped_coords[ class_name ] = \
                        [ xy for xy in whole_list if xy[0] == class_val ]
        else:
            class_name = result.test_set.name
            self.class_names = [ class_name ]
            self.class_values = [ 1 ]
            self.num_classes = 1
            self.grouped_coords[ class_name ] = whole_list

    #=================================================================
    def show( self ):
        self.figure.show()

    #=====================================================================
    @classmethod
    @output_railroad_switch
    def NewFromHTMLReport( cls, filepath, **kwargs):
        """Helper function to facilitate the fast generation of graphs from C++-generated
        HTML Report files."""

        exp = FeatureSpaceClassificationExperiment.NewFromHTMLReport( filepath )
        exp.GenerateStats()
        exp.PerSampleStatistics( quiet=True )
        newgraphobj = cls( exp, **kwargs )
        return newgraphobj

    #=====================================================================
    def RankOrderedPredictedValuesGraph( self, chart_title=None, fig=None, ax=None ):
        """Calls function wndcharm.visualization.RankOrderedPredictedValuesGraph
        using predicted values scraped from either of 2 objects, a
        wndcharm.FeatureSpacePrediction, or a wndcharm.FeatureSpacePredictionExperiment"""

        xy = [ xy for group_coords in self.grouped_coords.values() for xy in group_coords ]
        xy.sort( key=lambda xy: xy[0] )

        kwargs = {}
        kwargs[ 'labels' ], kwargs[ 'predicted_values' ] = zip( *xy )
        kwargs[ 'chart_title' ] = chart_title
        kwargs[ 'fig' ] = fig
        kwargs[ 'ax' ] = ax
        kwargs[ 'figsize' ] = self.figsize
        kwargs[ 'fontsize' ] = self.fontsize
        kwargs[ 'titlefontsize' ] = self.titlefontsize

        self.figure = RankOrderedPredictedValuesGraph( **kwargs )
        return self
        
    #=====================================================================
    def KernelSmoothedDensityGraph( self, chart_title=None, bw_method=None,
            fig=None, ax=None  ):
        """Calls function wndcharm.visualization.RankOrderedPredictedValuesGraph
        using predicted values scraped from either of 2 objects, a
        wndcharm.FeatureSpacePrediction, or a wndcharm.FeatureSpacePredictionExperiment"""

        xy = [ xy for group_coords in self.grouped_coords.values() for xy in group_coords ]
        xy.sort( key=lambda xy: xy[0] )

        kwargs = {}
        kwargs[ 'labels' ], kwargs[ 'predicted_values' ] = zip( *xy )
        kwargs[ 'chart_title' ] = chart_title
        kwargs[ 'fig' ] = fig
        kwargs[ 'ax' ] = ax
        kwargs[ 'figsize' ] = self.figsize
        kwargs[ 'fontsize' ] = self.fontsize
        kwargs[ 'titlefontsize' ] = self.titlefontsize

        self.figure = KernelSmoothedDensityGraph( **kwargs )
        return self

#============================================================================
class FeatureTimingVersusAccuracyGraph( _BaseGraph ):
    """A cost/benefit analysis of the number of features used and the time it takes to calculate
    that number of features for a single image"""

    #FIXME: Add ability to do the first 50 or 100 features, make the graph, then
    #       ability to resume from where it left off to do the next 50.

    def __init__( self, training_set, feature_weights, test_image_path,
        chart_title=None, max_num_features=300 ):

        self.timing_axes = None
        import time
        timings = []

        from .FeatureSpacePredictionExperiment import FeatureSpaceClassificationExperiment
        from .SingleSamplePrediction import SingleSampleClassification
        from .FeatureSpacePrediction import FeatureSpaceClassification
        from .FeatureVector import FeatureVector

        experiment = FeatureSpaceClassificationExperiment( training_set, training_set, feature_weights )
        for number_of_features_to_use in range( 1, max_num_features + 1 ):

            reduced_ts = None
            reduced_fw = None
            three_timings = []
            # Take the best of 3
            for timing in range( 3 ):
                # Time the creation and classification of a single signature
                t1 = time.time()
                reduced_fw = feature_weights.Threshold( number_of_features_to_use )
                sig = FeatureVector( source_filepath=test_image_path, feature_names=reduced_fw.feature_names ).GenerateFeatures()
                reduced_ts = training_set.FeatureReduce( reduced_fw )
                sig.Normalize( reduced_ts )
        
                result = SingleSampleClassification.NewWND5( reduced_ts, reduced_fw, sig )
                result.Print()
                # FIXME: save intermediates just in case of interruption or parallization
                # result.PickleMe()
                t2 = time.time()
                three_timings.append( t2 - t1 )

            timings.append( min( three_timings ) )

            # now, do a fit-on-fit test to measure classification accuracy
            split_result = FeatureSpaceClassification.NewWND5( reduced_ts, reduced_ts, reduced_fw )
            split_result.Print()
            experiment.individual_results.append( split_result )

        import matplotlib.pyplot as plt

        x_vals = list( range( 1, max_num_features + 1 ) )

        self.figure = plt.figure()
        self.main_axes = self.figure.add_subplot(111)
        if chart_title == None:
            self.chart_title = "Feature timing v. classification accuracy"    
        else:
            self.chart_title = chart_title
        self.main_axes.set_title( self.chart_title )
        self.main_axes.set_xlabel( 'Number of features' )
        self.main_axes.set_ylabel( 'Classification accuracy (%)', color='b' )
        classification_accuracies = \
          [ split_result.classification_accuracy * 100 for split_result in experiment.individual_results ]

        self.main_axes.plot( x_vals, classification_accuracies, color='b', linewidth=2 )
        for tl in self.main_axes.get_yticklabels():
            tl.set_color('b')    

        self.timing_axes = self.main_axes.twinx()
        self.timing_axes.set_ylabel( 'Time to calculate features (s)', color='r' )
        self.timing_axes.plot( x_vals, timings, color='r' )
        for tl in self.timing_axes.get_yticklabels():
            tl.set_color('r')    

#============================================================================
class HyperparameterOptimizationGraph( _BaseGraph ):
    """Hyper-parameter optimization.
    Graphs the figure of merit as a function of hyperparameter specified by
    method name.."""

    def __init__( self, feature_space ):
        self.feature_space = feature_space
        self.discrete = self.feature_space.discrete

        # Save the intermediate results in case you want to
        # write over them and redraw figure
        self.raw_results = None
        self.lda_results = None

    def GridSearch( self, param='features', param_space=None, param_scale='log',
            show_raw=True, show_lda=True,
            chart_title=None, figsize=(12, 8), y_lim=None, quiet=True,
            use_error_bars=True, text_angle=None, fig=None, ax=None, **kwargs ):
        """Creates graph for Classifier/Regressor figure of merit as a function of
        the parameter user requests, either number of top-ranked features or
        number of samples per class.

        Calls wndcharm.FeatureSpacePredictionExperiment.NumFeaturesGridSearch, which is
        itself a wrapper for NewShuffleSplit, to which kwargs gets passed through.

        Args:
            param - str={'features','samples}
                Indicate whether user wants to grid search over number of top-ranked
                features, or samples.
            param_space - iterable or int or None
                Passed through to NumFeaturesGridSearch()
            param_scale - str={'log', 'linear' }, default='log'
                Passed through to matplotlib.ax.set_xscale()
            show_raw - bool, default=True
                Include in figure results obtained WITHOUT LDA dimensionality reduction
            show_lda - bool, default=True
                Include in figure results obtained WITH LDA dimensionality reduction done
                just before classification in the pipeline.
            y_lim - (float,float), or None
                Sets ymin and ymax of figure. If None, set dynamically using figure data.
            text_angle - int, or None
                Create text labels containing x & y values for all points if not None
            fig - matplotlib.figure.Figure, or None (default):
                If provided, generate axes on the passed-in figure obj, otherwise
                instantiate a new one.
            ax - matplotlib.axes._subplots.AxesSubplot or None (default):
                Plot data on these axes, if provided, otherwise create new axes on figure.

            **kwargs
                Passed through to NewShuffleSplit()
        """

        # dummyproof:
        if 'feature_space' in kwargs:
            if self.feature_space is not kwargs['feature_space']:
                raise ValueError( 'This graph object already has a FeatureSpace associated with it: {}'.format( self.feature_space ) )

        import matplotlib.pyplot as plt

        if fig is None and ax is None:
            fig = plt.figure( figsize=figsize, facecolor='white' )
            ax = fig.add_subplot(111)
        elif fig is None:
            fig = ax.get_figure()
        elif ax is None:
            ax = fig.gca()

        if self.discrete:
            from .FeatureSpacePredictionExperiment import FeatureSpaceClassificationExperiment as Experiment
        else:
            from .FeatureSpacePredictionExperiment import FeatureSpaceRegressionExperiment as Experiment

        if chart_title == None:
            if self.discrete:
                d = '/class'
            else:
                d = ''
            chart_info_str = 'Feature Space=' + self.feature_space.name + '\n'
            if param != 'samples':
                if 'train_size' in kwargs:
                    chart_info_str+= "Train samps" + d + "={}, ".format( kwargs['train_size'] )
            else:
                if 'features_size' in kwargs:
                    f = kwargs['features_size']
                else:
                    f = 0.15
                chart_info_str += "Features size={}, ".format( f )
            if 'test_size' in kwargs:
                chart_info_str += "Test samps" + d + "={}, ".format( kwargs['test_size'] )
            if 'n_iter' in kwargs:
                chart_info_str += "# iter={}".format( kwargs['n_iter'] )
            chart_info_str.rstrip( ', ' )
            # make more room for 3-line title

        if param == 'features':
            GridSearch = Experiment.NumFeaturesGridSearch
            self.axes = ax
            ax.set_xlabel( '# top-ranked features', size=16 )
            if chart_title == None:
                self.chart_title = "Figure of Merit vs. # Top-ranked Features\n" + \
                        chart_info_str
            else:
                self.chart_title = chart_title
        elif param == 'samples':
            GridSearch = Experiment.NumSamplesGridSearch
            self.axes = ax
            ax.set_xlabel( '# samples per class', size=16 )
            if chart_title == None:
                self.chart_title = "Figure of Merit vs. # Samples per class\n" + \
                        chart_info_str
            else:
                self.chart_title = chart_title
        else:
            raise ValueError( 'arg param should be either "features" or "values", not {}'.format( param ))
        ax.set_title( self.chart_title, size=18 )

        data_series = []
        data_series_labels = []

        # Use or non-use of LDA is governed by args to this method, not by kwargs:
        if 'lda' in kwargs:
            del kwargs['lda']
        if 'pre_lda_feature_filter' in kwargs:
            del kwargs['pre_lda_feature_filter']

        if show_raw:
            if self.raw_results is None:
                self.raw_results = GridSearch(
                    feature_space=self.feature_space, param_space=param_space, quiet=quiet,
                    lda=False, **kwargs )
            else:
                # Implement ability to choose different params and add additional runs
                pass
            data_series.append( self.raw_results )
            data_series_labels.append( 'Without LDA' )
        if show_lda:
            if self.lda_results is None:
                self.lda_results = GridSearch(
                    feature_space=self.feature_space, param_space=param_space, quiet=quiet,
                    lda=True, pre_lda_feature_filter=True, **kwargs )
            else:
                pass
            data_series.append( self.lda_results )
            data_series_labels.append( 'With LDA' )

        ax.set_xscale( param_scale )
        ax.set_ylabel( 'Figure of Merit', size=16 )

        # use all_Ys and all_Xs to set xlim and ylim
        all_Ys = []
        all_Xs = []

        C = ['b','y','m' ]
        S = ['o','v','s' ]
        for label, (X, Y, E), c, s in zip( data_series_labels, data_series, C, S ):
            if not use_error_bars or set( E ) == set( [None] ):
                all_Ys.extend( Y )
            else:
                # Include upper-bound
                all_Ys.extend( [ y + e for y, e in zip( Y, E ) ] )
            all_Xs.extend( X )
            ax.plot( X, Y, color=c, linestyle='--', label=label )
            ax.errorbar( X, Y, yerr=E, color=c, marker=s )
            if text_angle is not None:
                for x, y in zip( X, Y ):
                    text = '({},{:0.03})'.format( x, y )
                    ax.annotate( text, xy=(x,y), rotation=float(text_angle), color=c )
                text_angle *= -1

        if y_lim == None:
            y_min = 0
            y_max = 1.05 * max( all_Ys )
            x_min = 0
            x_max = 1.15 * max( all_Xs )
        else:
            y_min, y_max = y_lim

        ax.set_ylim( [ y_min, y_max ] )
        ax.set_xlim( [ x_min, x_max ] )
        ax.legend( fontsize=16, loc='lower right' )

        self.figure = fig
        self.main_axes = ax
        plt.tight_layout()
        return self
