#!/usr/bin/env python
__author__ = "Jai Ram Rideout"
__copyright__ = "Copyright 2007-2016, The Cogent Project"
__credits__ = ["Jai Ram Rideout"]
__license__ = "GPL"
__version__ = "1.9"
__maintainer__ = "Jai Ram Rideout"
__email__ = "jai.rideout@gmail.com"
__status__ = "Production"

"""Tests public and private functions in the distribution_plots module."""
from warnings import filterwarnings

filterwarnings("ignore", "More than 20 figures")

from matplotlib import use
use('Agg', warn=False)

from StringIO import StringIO
import sys
import matplotlib.colors as colors
from matplotlib.pyplot import boxplot
from numpy import array
from cogent.draw.distribution_plots import _validate_input,\
    _get_distribution_markers, _validate_x_values, _create_plot,\
    _calc_data_point_locations, _set_axes_options, generate_box_plots,\
    generate_comparative_plots, _calc_data_point_ticks, _plot_bar_data,\
    _plot_scatter_data, _plot_box_data, _color_box_plot,\
    _create_legend, _set_figure_size
from cogent.util.unit_test import TestCase, main

class DistributionPlotsTests(TestCase):
    """Tests of the distribution_plots module."""
    def setUp(self):
        """Create some data to be used in the tests."""
        # Test null data list.
        self.Null = None

        # Test empty data list.
        self.Empty = []

        # Test nested empty data list.
        self.EmptyNested = [[]]

        # Test nested empty data list (for bar/scatter plots).
        self.EmptyDeeplyNested = [[[]]]

        # Test invalid number of samples in data list (for bar/scatter plots).
        self.InvalidNumSamples = [[[1, 2, 3, 4, 5]],
                                  [[4, 5, 6, 7, 8], [2, 3, 2]],
                                  [[4, 7, 10, 33, 32, 6, 7, 8]]]

        # Test valid data with one sample (for bar/scatter plots).
        self.ValidSingleSampleData = [[[1, 2, 3, 4, 5]],
                                      [[4, 5, 6, 7, 8]],
                                      [[4, 7, 10, 33, 32, 6, 7, 8]]]

        # Test valid data with three samples and four data points
        # (for bar/scatter plots).
        self.ValidTypicalData = [[[1.0, 2, 3.5, 5], [2, 3, 5, 6], [2, 3, 8]],
                                 [[4, 7, 8], [8, 9, 10, 11], [9.0, 4, 1, 1]],
                                 [[4, 33, 32, 6, 8], [5, 4, 8, 13], [1, 1, 2]],
                                 [[2, 2, 2, 2], [3, 9, 8], [2, 1, 6, 7, 4, 5]]]

        # Test typical data to be plotted by the boxplot function.
        self.ValidTypicalBoxData = [[3.4, 10, 11.67, 12.0, 2, 2, 99.99],
                                    [2.3, 4, 5, 88, 9, 10, 11, 1, 0, 3, -8],
                                    [2, 9, 7, 5, 6]]

    def test_validate_input_null(self):
        """_validate_input() should raise a ValueError if null data is passed
        to it."""
        self.assertRaises(ValueError, _validate_input,
                          self.Null, None, None, None)

    def test_validate_input_empty(self):
        """_validate_input() should raise a ValueError if empty data is passed
        to it."""
        self.assertRaises(ValueError, _validate_input,
                          self.Empty, None, None, None)

    def test_validate_input_empty_nested(self):
        """_validate_input() should raise a ValueError if empty nested data is
        passed to it."""
        self.assertRaises(ValueError, _validate_input,
                          self.EmptyNested, None, None, None)

    def test_validate_input_empty_deeply_nested(self):
        """_validate_input() should pass for deeply nested empty data."""
        num_points, num_samples = _validate_input(self.EmptyDeeplyNested,
                                                  None, None, None)
        self.assertEqual(num_points, 1)
        self.assertEqual(num_samples, 1)

    def test_validate_input_invalid_num_samples(self):
        """_validate_input() should raise a ValueError if an inconsistent
        number of samples in included in the data."""
        self.assertRaises(ValueError, _validate_input,
                          self.InvalidNumSamples, None, None, None)

    def test_validate_x_values_invalid_x_values(self):
        """_validate_x_values() should raise a ValueError on an invalid number
        of x_values."""
        self.assertRaises(ValueError, _validate_x_values,
                          [1, 2, 3, 4], ["T0", "T1", "T2"],
                          len(self.ValidSingleSampleData))

    def test_validate_x_values_invalid_x_tick_labels(self):
        """_validate_x_values() should raise a ValueError on an invalid number
        of x_tick_labels."""
        self.assertRaises(ValueError, _validate_x_values,
                          None, ["T0"], len(self.ValidSingleSampleData))

    def test_validate_x_values_nonnumber_x_values(self):
        """_validate_x_values() should raise a ValueError on x_values that
        aren't numbers."""
        self.assertRaises(ValueError, _validate_x_values,
                ["foo", 2, 3], None, len(self.ValidSingleSampleData))

    def test_validate_x_values_valid_x_values(self):
        """_validate_x_values() should not throw an exception."""
        _validate_x_values([1, 2.0, 3], None, 3)

    def test_validate_input_invalid_data_point_names(self):
        """_validate_input() should raise a ValueError on data_point_names that
        are an invalid length."""
        self.assertRaises(ValueError, _validate_input,
                self.ValidSingleSampleData, None, ["T0", "T1"], None)

    def test_validate_input_invalid_sample_names(self):
        """_validate_input() should raise a ValueError on sample_names that are
        an invalid length."""
        self.assertRaises(ValueError, _validate_input,
                self.ValidSingleSampleData, None, None, ["Men", "Women"])

    def test_validate_input_all_valid_input(self):
        """_validate_input() should return valid information about the data
        without throwing an exception."""
        self.assertEqual(_validate_input(self.ValidTypicalData, [1, 3, 4, 8],
                                         ["T0", "T1", "T2", "T3"],
                                         ["Infants", "Children", "Teens"]),
                                         (4, 3))

    def test_get_distribution_markers_null_marker_list(self):
        """_get_distribution_markers() should return a list of predefined
        matplotlib markers."""
        self.assertEqual(_get_distribution_markers('colors', None, 5),
                ['b', 'g', 'r', 'c', 'm'])

    def test_get_distribution_markers_empty_marker_list(self):
        """_get_distribution_markers() should return a list of predefined
        matplotlib markers."""
        self.assertEqual(_get_distribution_markers('colors', None, 4),
                ['b', 'g', 'r', 'c'])

    def test_get_distribution_markers_insufficient_markers(self):
        """_get_distribution_markers() should return a wrapped list of
        predefined markers."""
        # Save stdout and replace it with something that will capture the print
        # statement. Note: this code was taken from here:
        # http://stackoverflow.com/questions/4219717/how-to-assert-output-
        #     with-nosetest-unittest-in-python/4220278#4220278
        saved_stdout = sys.stdout
        try:
            out = StringIO()
            sys.stdout = out
            self.assertEqual(_get_distribution_markers('colors', None, 10),
                ['b', 'g', 'r', 'c', 'm', 'y', 'w', 'b', 'g', 'r'])
            self.assertEqual(_get_distribution_markers('symbols',
                ['^', '>', '<'], 5), ['^', '>', '<', '^', '>'])
            output = out.getvalue().strip()
            self.assertEqual(output, "There are not enough markers to "
                    "uniquely represent each distribution in your dataset. "
                    "You may want to provide a list of markers that is at "
                    "least as large as the number of distributions in your "
                    "dataset.\nThere are not enough markers to "
                    "uniquely represent each distribution in your dataset. "
                    "You may want to provide a list of markers that is at "
                    "least as large as the number of distributions in your "
                    "dataset.")
        finally:
            sys.stdout = saved_stdout

    def test_get_distribution_markers_bad_marker_type(self):
        """_get_distribution_markers() should raise a ValueError."""
        self.assertRaises(ValueError, _get_distribution_markers, 'shapes', [],
                3)

    def test_get_distribution_markers_zero_markers(self):
        """_get_distribution_markers() should return an empty list."""
        self.assertEqual(_get_distribution_markers('symbols', None, 0), [])
        self.assertEqual(_get_distribution_markers('symbols', ['^'], 0), [])

    def test_create_plot(self):
        """_create_plot() should return a tuple containing a Figure and
        Axes."""
        fig, ax = _create_plot()
        self.assertEqual(fig.__class__.__name__, "Figure")
        self.assertEqual(ax.__class__.__name__, "AxesSubplot")

    def test_plot_bar_data(self):
        """_plot_bar_data() should return a list of Rectangle objects."""
        fig, ax = _create_plot()
        result = _plot_bar_data(ax, [1, 2, 3], 'red', 0.5, 3.75, 1.5, 'stdv')
        self.assertEqual(result[0].__class__.__name__, "Rectangle")
        self.assertEqual(len(result), 1)
        self.assertFloatEqual(result[0].get_width(), 0.5)
        self.assertFloatEqual(result[0].get_facecolor(), (1.0, 0.0, 0.0, 1.0))
        self.assertFloatEqual(result[0].get_height(), 2.0)

        fig, ax = _create_plot()
        result = _plot_bar_data(ax, [1, 2, 3], 'red', 0.5, 3.75, 1.5, 'sem')
        self.assertEqual(result[0].__class__.__name__, "Rectangle")
        self.assertEqual(len(result), 1)
        self.assertFloatEqual(result[0].get_width(), 0.5)
        self.assertFloatEqual(result[0].get_facecolor(), (1.0, 0.0, 0.0, 1.0))
        self.assertFloatEqual(result[0].get_height(), 2.0)

    def test_plot_bar_data_bad_error_bar_type(self):
        """_plot_bar_data() should raise an exception on bad error bar type."""
        fig, ax = _create_plot()
        self.assertRaises(ValueError, _plot_bar_data, ax, [1, 2, 3], 'red',
                0.5, 3.75, 1.5, 'var')

    def test_plot_bar_data_empty(self):
        """_plot_bar_data() should not error when given empty list of data,
        but should not plot anything."""
        fig, ax = _create_plot()
        result = _plot_bar_data(ax, [], 'red', 0.5, 3.75, 1.5, 'stdv')
        self.assertTrue(result is None)

        fig, ax = _create_plot()
        result = _plot_bar_data(ax, [], 'red', 0.5, 3.75, 1.5, 'sem')
        self.assertTrue(result is None)

    def test_plot_scatter_data_empty(self):
        """_plot_scatter_data() should not error when given empty list of data,
        but should not plot anything."""
        fig, ax = _create_plot()
        result = _plot_scatter_data(ax, [], '^', 0.77, 1, 1.5, 'stdv')
        self.assertTrue(result is None)

    def test_plot_box_data(self):
        """_plot_box_data() should return a dictionary for Line2D's."""
        fig, ax = _create_plot()
        result = _plot_box_data(ax, [0, 0, 7, 8, -3, 44], 'blue', 0.33, 55,
                1.5, 'stdv')
        self.assertEqual(type(result), dict)
        # make sure the values are not empty
        for key in ['boxes', 'medians', 'whiskers', 'fliers', 'caps']:
            val = result[key]
            self.assertTrue(len(val), 0)

    def test_plot_box_data_empty(self):
        """_plot_box_data() should not error when given empty list of data,
        but should not plot anything."""
        fig, ax = _create_plot()
        result = _plot_box_data(ax, [], 'blue', 0.33, 55, 1.5, 'stdv')
        # the original test was poor as it was testing a property of matplotlib
        # I'm simply exercising this here. These capabilities will be dropped
        # in a future version.
        for key in ['boxes', 'medians', 'whiskers', 'fliers', 'caps']:
            val = result[key]

    def test_calc_data_point_locations_invalid_widths(self):
        """_calc_data_point_locations() should raise a ValueError
        exception when it encounters bad widths."""
        self.assertRaises(ValueError, _calc_data_point_locations, [1, 2, 3],
                3, 2, -2, 0.5)
        self.assertRaises(ValueError, _calc_data_point_locations, [1, 2, 3],
                3, 2, 2, -0.5)

    def test_calc_data_point_locations_default_spacing(self):
        """_calc_data_point_locations() should return an array containing
        the x-axis locations for each data point, evenly spaced from 1..n."""
        locs = _calc_data_point_locations(None, 4, 2, 0.25, 0.5)
        self.assertEqual(locs, array([1.0, 2.0, 3.0, 4.0]))

    def test_calc_data_point_locations_custom_spacing(self):
        """_calc_data_point_locations() should return an array containing
        the x-axis locations for each data point, spaced according to a custom
        spacing scheme."""
        locs = _calc_data_point_locations([3, 4, 10, 12], 4, 2, 0.25, 0.75)
        self.assertEqual(locs, array([3.75, 5.0, 12.5, 15.0]))

    def test_calc_data_point_ticks(self):
        """_calc_data_point_ticks() should return an array containing the
        x-axis locations for each data point tick."""
        ticks = _calc_data_point_ticks(array([1, 5, 9, 11]), 1, 0.5, False)
        self.assertFloatEqual(ticks, array([1.25, 5.25, 9.25, 11.25]))

        ticks = _calc_data_point_ticks(array([0]), 3, 0.5, False)
        self.assertFloatEqual(ticks, array([0.75]))

    def test_set_axes_options(self):
        """_set_axes_options() should set the labels on the axes and not raise
        any exceptions."""
        fig, ax = _create_plot()
        _set_axes_options(ax, "Plot Title", "x-axis label", "y-axis label",
                          x_tick_labels=["T0", "T1"])
        self.assertEqual(ax.get_title(), "Plot Title")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(ax.get_xticklabels()[0].get_text(), "T0")
        self.assertEqual(ax.get_xticklabels()[1].get_text(), "T1")

    def test_set_axes_options_ylim(self):
        """_set_axes_options() should set the y-axis limits."""
        fig, ax = _create_plot()
        _set_axes_options(ax, "Plot Title", "x-axis label", "y-axis label",
                          x_tick_labels=["T0", "T1", "T2"], y_min=0, y_max=1)
        self.assertEqual(ax.get_title(), "Plot Title")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(ax.get_xticklabels()[0].get_text(), "T0")
        self.assertEqual(ax.get_xticklabels()[1].get_text(), "T1")
        self.assertFloatEqual(ax.get_ylim(), [0, 1])

    def test_set_axes_options_bad_ylim(self):
        """_set_axes_options() should raise an exception when given non-numeric
        y limits."""
        fig, ax = _create_plot()
        self.assertRaises(ValueError, _set_axes_options, ax, "Plot Title",
                          "x-axis label", "y-axis label",
                          x_tick_labels=["T0", "T1", "T2"], y_min='car',
                          y_max=30)

    def test_create_legend(self):
        """_create_box_plot_legend() should create a legend on valid input."""
        fig, ax = _create_plot()
        _create_legend(ax, ['b', 'r'], ['dist1', 'dist2'], 'colors')
        self.assertEqual(len(ax.get_legend().get_texts()), 2)

        fig, ax = _create_plot()
        _create_legend(ax, ['^', '<', '>'], ['dist1', 'dist2', 'dist3'],
                'symbols')
        self.assertEqual(len(ax.get_legend().get_texts()), 3)

    def test_create_legend_invalid_input(self):
        """Test raises error on bad input."""
        fig, ax = _create_plot()
        self.assertRaises(ValueError, _create_legend, ax,
                ['^', '<', '>'], ['dist1', 'dist2'], 'symbols')
        self.assertRaises(ValueError, _create_legend, ax, ['^', '<', '>'],
                ['dist1', 'dist2', 'dist3'], 'foo')

    def test_generate_box_plots(self):
        """generate_box_plots() should return a valid Figure object."""
        fig = generate_box_plots(self.ValidTypicalBoxData, [1, 4, 10],
                                 ["Data 1", "Data 2", "Data 3"], "Test",
                                 "x-axis label", "y-axis label")
        ax = fig.get_axes()[0]
        self.assertEqual(ax.get_title(), "Test")
        self.assertEqual(ax.get_xlabel(), "x-axis label")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(len(ax.get_xticklabels()), 3)
        self.assertFloatEqual(ax.get_xticks(), [1, 4, 10])

    def test_generate_comparative_plots_bar(self):
        """generate_comparative_plots() should return a valid barchart Figure
        object."""
        fig = generate_comparative_plots('bar', self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], ['b', 'r', 'g'],
                "x-axis label", "y-axis label", "Test")
        ax = fig.get_axes()[0]
        self.assertEqual(ax.get_title(), "Test")
        self.assertEqual(ax.get_xlabel(), "x-axis label")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(len(ax.get_xticklabels()), 4)
        self.assertFloatEqual(ax.get_xticks(), [2.3, 7.4, 17.6, 19.3])

    def test_generate_comparative_plots_insufficient_colors(self):
        """generate_comparative_plots() should work even when there aren't
        enough colors. We should capture a print statement that warns the
        users."""
        saved_stdout = sys.stdout
        try:
            out = StringIO()
            sys.stdout = out
            generate_comparative_plots('bar', self.ValidTypicalData,
                    [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                    ["Infants", "Children", "Teens"], ['b', 'r'],
                    "x-axis label", "y-axis label", "Test")
            output = out.getvalue().strip()
            self.assertEqual(output, "There are not enough markers to "
                    "uniquely represent each distribution in your dataset. "
                    "You may want to provide a list of markers that is at "
                    "least as large as the number of distributions in your "
                    "dataset.")
        finally:
            sys.stdout = saved_stdout

    def test_generate_comparative_plots_scatter(self):
        """generate_comparative_plots() should return a valid scatterplot
        Figure object."""
        fig = generate_comparative_plots('scatter', self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], ['^', '>', '<'],
                "x-axis label", "y-axis label", "Test")
        ax = fig.get_axes()[0]
        self.assertEqual(ax.get_title(), "Test")
        self.assertEqual(ax.get_xlabel(), "x-axis label")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(len(ax.get_xticklabels()), 4)
        self.assertFloatEqual(ax.get_xticks(), [2.1, 7.2, 17.4, 19.1])

    def test_generate_comparative_plots_insufficient_symbols(self):
        """generate_comparative_plots() should work even when there aren't
        enough symbols. We should capture a print statement that warns the
        users."""
        saved_stdout = sys.stdout
        try:
            out = StringIO()
            sys.stdout = out
            generate_comparative_plots('scatter', self.ValidTypicalData,
                    [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                    ["Infants", "Children", "Teens"], ['^'],
                    "x-axis label", "y-axis label", "Test")
            output = out.getvalue().strip()
            self.assertEqual(output, "There are not enough markers to "
                    "uniquely represent each distribution in your dataset. "
                    "You may want to provide a list of markers that is at "
                    "least as large as the number of distributions in your "
                    "dataset.")
        finally:
            sys.stdout = saved_stdout

    def test_generate_comparative_plots_empty_marker_list(self):
        """generate_comparative_plots() should use the predefined list of
        markers if an empty list is provided by the user."""
        generate_comparative_plots('scatter', self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], [],
                "x-axis label", "y-axis label", "Test")

    def test_generate_comparative_plots_box(self):
        """generate_comparative_plots() should return a valid boxplot Figure
        object."""
        fig = generate_comparative_plots('box', self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], ['b', 'g', 'y'],
                "x-axis label", "y-axis label", "Test")
        ax = fig.get_axes()[0]
        self.assertEqual(ax.get_title(), "Test")
        self.assertEqual(ax.get_xlabel(), "x-axis label")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(len(ax.get_xticklabels()), 4)
        self.assertFloatEqual(ax.get_xticks(), [2.1, 7.2, 17.4, 19.1])

    def test_generate_comparative_plots_error(self):
        """generate_comparative_plots() should raise a ValueError for an
        invalid plot type."""
        self.assertRaises(ValueError, generate_comparative_plots, 'pie',
                self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], ['b', 'g', 'y'],
                "x-axis label", "y-axis label", "Test")

    def test_color_box_plot(self):
        """_color_box_plot() should not throw an exception when passed the
        proper input."""
        fig, ax = _create_plot()
        box_plot = boxplot(self.ValidTypicalBoxData)
        _color_box_plot(ax, box_plot, 'blue')

    def test_set_figure_size(self):
        """Test setting a valid figure size."""
        fig, ax = _create_plot()
        _set_axes_options(ax, 'foo', 'x_foo', 'y_foo',
                          x_tick_labels=['foofoofoo', 'barbarbar'],
                          x_tick_labels_orientation='vertical')
        _set_figure_size(fig, 3, 4)
        self.assertFloatEqual(fig.get_size_inches(), (3, 4))

    def test_set_figure_size_defaults(self):
        """Test setting a figure size using matplotlib defaults."""
        fig, ax = _create_plot()
        _set_axes_options(ax, 'foo', 'x_foo', 'y_foo',
                          x_tick_labels=['foofoofoo', 'barbarbar'],
                          x_tick_labels_orientation='vertical')
        orig_fig_size = fig.get_size_inches()
        _set_figure_size(fig)
        self.assertFloatEqual(fig.get_size_inches(), orig_fig_size)

    def test_set_figure_size_invalid(self):
        """Test setting a figure size using invalid dimensions."""
        fig, ax = _create_plot()
        _set_axes_options(ax, 'foo', 'x_foo', 'y_foo',
                          x_tick_labels=['foofoofoo', 'barbarbar'],
                          x_tick_labels_orientation='vertical')
        orig_fig_size = fig.get_size_inches()
        _set_figure_size(fig, -1, 0)
        self.assertFloatEqual(fig.get_size_inches(), orig_fig_size)

    def test_set_figure_size_long_labels(self):
        """Test setting a figure size that has really long labels."""
        saved_stdout = sys.stdout
        try:
            out = StringIO()
            sys.stdout = out

            fig, ax = _create_plot()
            _set_axes_options(ax, 'foo', 'x_foo', 'y_foo',
                              x_tick_labels=['foofoofooooooooooooooooooooooooo'
                              'ooooooooooooooooooooooooooooooooooooooooooooooo'
                              'ooooooooooooooooooooo', 'barbarbar'],
                              x_tick_labels_orientation='vertical')
            _set_figure_size(fig, 3, 3)
            self.assertFloatEqual(fig.get_size_inches(), (3, 3))
            output = out.getvalue().strip()
#            self.assertEqual(output,
#            "Warning: could not automatically resize plot to make room for "
#            "axes labels and plot title. This can happen if the labels or "
#            "title are extremely long and the plot size is too small. Your "
#            "plot may have its labels and/or title cut-off. To fix this, "
#            "try increasing the plot's size (in inches) and try again.")
        finally:
            sys.stdout = saved_stdout


if __name__ == '__main__':
    main()