import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd

import nose.tools as nt
import numpy.testing as npt
import pandas.util.testing as pdt
from numpy.testing.decorators import skipif
from nose import SkipTest

try:
    import statsmodels.regression.linear_model as smlm
    _no_statsmodels = False
except ImportError:
    _no_statsmodels = True

from . import PlotTestCase
from .. import linearmodels as lm
from .. import algorithms as algo
from .. import utils
from ..palettes import color_palette

rs = np.random.RandomState(0)

class TestLinearPlotter(PlotTestCase):

    rs = np.random.RandomState(77)
    df = pd.DataFrame(dict(x=rs.normal(size=60),
                           d=rs.randint(-2, 3, 60),
                           y=rs.gamma(4, size=60),
                           s=np.tile(list("abcdefghij"), 6)))
    df["z"] = df.y + rs.randn(60)
    df["y_na"] = df.y.copy()
    df.loc[[10, 20, 30], 'y_na'] = np.nan

    def test_establish_variables_from_frame(self):

        p = lm._LinearPlotter()
        p.establish_variables(self.df, x="x", y="y")
        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y, self.df.y)
        pdt.assert_frame_equal(p.data, self.df)

    def test_establish_variables_from_series(self):

        p = lm._LinearPlotter()
        p.establish_variables(None, x=self.df.x, y=self.df.y)
        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y, self.df.y)
        nt.assert_is(p.data, None)

    def test_establish_variables_from_array(self):

        p = lm._LinearPlotter()
        p.establish_variables(None,
                              x=self.df.x.values,
                              y=self.df.y.values)
        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y)
        nt.assert_is(p.data, None)

    def test_establish_variables_from_mix(self):

        p = lm._LinearPlotter()
        p.establish_variables(self.df, x="x", y=self.df.y)
        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y, self.df.y)
        pdt.assert_frame_equal(p.data, self.df)

    def test_establish_variables_from_bad(self):

        p = lm._LinearPlotter()
        with nt.assert_raises(ValueError):
            p.establish_variables(None, x="x", y=self.df.y)

    def test_dropna(self):

        p = lm._LinearPlotter()
        p.establish_variables(self.df, x="x", y_na="y_na")
        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y_na, self.df.y_na)

        p.dropna("x", "y_na")
        mask = self.df.y_na.notnull()
        pdt.assert_series_equal(p.x, self.df.x[mask])
        pdt.assert_series_equal(p.y_na, self.df.y_na[mask])

class TestRegressionPlotter(PlotTestCase):

    rs = np.random.RandomState(49)

    grid = np.linspace(-3, 3, 30)
    n_boot = 100
    bins_numeric = 3
    bins_given = [-1, 0, 1]

    df = pd.DataFrame(dict(x=rs.normal(size=60),
                           d=rs.randint(-2, 3, 60),
                           y=rs.gamma(4, size=60),
                           s=np.tile(list(range(6)), 10)))
    df["z"] = df.y + rs.randn(60)
    df["y_na"] = df.y.copy()

    bw_err = rs.randn(6)[df.s.values] * 2
    df.y += bw_err

    p = 1 / (1 + np.exp(-(df.x * 2 + rs.randn(60))))
    df["c"] = [rs.binomial(1, p_i) for p_i in p]
    df.loc[[10, 20, 30], 'y_na'] = np.nan

    def test_variables_from_frame(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, units="s")

        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y, self.df.y)
        pdt.assert_series_equal(p.units, self.df.s)
        pdt.assert_frame_equal(p.data, self.df)

    def test_variables_from_series(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y, units=self.df.s)

        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y)
        npt.assert_array_equal(p.units, self.df.s)
        nt.assert_is(p.data, None)

    def test_variables_from_mix(self):

        p = lm._RegressionPlotter("x", self.df.y + 1, data=self.df)

        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y + 1)
        pdt.assert_frame_equal(p.data, self.df)

    def test_dropna(self):

        p = lm._RegressionPlotter("x", "y_na", data=self.df)
        nt.assert_equal(len(p.x), pd.notnull(self.df.y_na).sum())

        p = lm._RegressionPlotter("x", "y_na", data=self.df, dropna=False)
        nt.assert_equal(len(p.x), len(self.df.y_na))

    def test_ci(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, ci=95)
        nt.assert_equal(p.ci, 95)
        nt.assert_equal(p.x_ci, 95)

        p = lm._RegressionPlotter("x", "y", data=self.df, ci=95, x_ci=68)
        nt.assert_equal(p.ci, 95)
        nt.assert_equal(p.x_ci, 68)

    @skipif(_no_statsmodels)
    def test_fast_regression(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, n_boot=self.n_boot)

        # Fit with the "fast" function, which just does linear algebra
        yhat_fast, _ = p.fit_fast(self.grid)

        # Fit using the statsmodels function with an OLS model
        yhat_smod, _ = p.fit_statsmodels(self.grid, smlm.OLS)

        # Compare the vector of y_hat values
        npt.assert_array_almost_equal(yhat_fast, yhat_smod)

    @skipif(_no_statsmodels)
    def test_regress_poly(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, n_boot=self.n_boot)

        # Fit an first-order polynomial
        yhat_poly, _ = p.fit_poly(self.grid, 1)

        # Fit using the statsmodels function with an OLS model
        yhat_smod, _ = p.fit_statsmodels(self.grid, smlm.OLS)

        # Compare the vector of y_hat values
        npt.assert_array_almost_equal(yhat_poly, yhat_smod)

    def test_regress_logx(self):

        x = np.arange(1, 10)
        y = np.arange(1, 10)
        grid = np.linspace(1, 10, 100)
        p = lm._RegressionPlotter(x, y, n_boot=self.n_boot)

        yhat_lin, _ = p.fit_fast(grid)
        yhat_log, _ = p.fit_logx(grid)

        nt.assert_greater(yhat_lin[0], yhat_log[0])
        nt.assert_greater(yhat_log[20], yhat_lin[20])
        nt.assert_greater(yhat_lin[90], yhat_log[90])

    @skipif(_no_statsmodels)
    def test_regress_n_boot(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, n_boot=self.n_boot)

        # Fast (linear algebra) version
        _, boots_fast = p.fit_fast(self.grid)
        npt.assert_equal(boots_fast.shape, (self.n_boot, self.grid.size))

        # Slower (np.polyfit) version
        _, boots_poly = p.fit_poly(self.grid, 1)
        npt.assert_equal(boots_poly.shape, (self.n_boot, self.grid.size))

        # Slowest (statsmodels) version
        _, boots_smod = p.fit_statsmodels(self.grid, smlm.OLS)
        npt.assert_equal(boots_smod.shape, (self.n_boot, self.grid.size))

    @skipif(_no_statsmodels)
    def test_regress_without_bootstrap(self):

        p = lm._RegressionPlotter("x", "y", data=self.df,
                                  n_boot=self.n_boot, ci=None)

        # Fast (linear algebra) version
        _, boots_fast = p.fit_fast(self.grid)
        nt.assert_is(boots_fast, None)

        # Slower (np.polyfit) version
        _, boots_poly = p.fit_poly(self.grid, 1)
        nt.assert_is(boots_poly, None)

        # Slowest (statsmodels) version
        _, boots_smod = p.fit_statsmodels(self.grid, smlm.OLS)
        nt.assert_is(boots_smod, None)

    def test_numeric_bins(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y)
        x_binned, bins = p.bin_predictor(self.bins_numeric)
        npt.assert_equal(len(bins), self.bins_numeric)
        npt.assert_array_equal(np.unique(x_binned), bins)

    def test_provided_bins(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y)
        x_binned, bins = p.bin_predictor(self.bins_given)
        npt.assert_array_equal(np.unique(x_binned), self.bins_given)

    def test_bin_results(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y)
        x_binned, bins = p.bin_predictor(self.bins_given)
        nt.assert_greater(self.df.x[x_binned == 0].min(),
                          self.df.x[x_binned == -1].max())
        nt.assert_greater(self.df.x[x_binned == 1].min(),
                          self.df.x[x_binned == 0].max())

    def test_scatter_data(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y)
        x, y = p.scatter_data
        npt.assert_array_equal(x, self.df.x)
        npt.assert_array_equal(y, self.df.y)

        p = lm._RegressionPlotter(self.df.d, self.df.y)
        x, y = p.scatter_data
        npt.assert_array_equal(x, self.df.d)
        npt.assert_array_equal(y, self.df.y)

        p = lm._RegressionPlotter(self.df.d, self.df.y, x_jitter=.1)
        x, y = p.scatter_data
        nt.assert_true((x != self.df.d).any())
        npt.assert_array_less(np.abs(self.df.d - x), np.repeat(.1, len(x)))
        npt.assert_array_equal(y, self.df.y)

        p = lm._RegressionPlotter(self.df.d, self.df.y, y_jitter=.05)
        x, y = p.scatter_data
        npt.assert_array_equal(x, self.df.d)
        npt.assert_array_less(np.abs(self.df.y - y), np.repeat(.1, len(y)))

    def test_estimate_data(self):

        p = lm._RegressionPlotter(self.df.d, self.df.y, x_estimator=np.mean)

        x, y, ci = p.estimate_data

        npt.assert_array_equal(x, np.sort(np.unique(self.df.d)))
        npt.assert_array_almost_equal(y, self.df.groupby("d").y.mean())
        npt.assert_array_less(np.array(ci)[:, 0], y)
        npt.assert_array_less(y, np.array(ci)[:, 1])

    def test_estimate_cis(self):

        # set known good seed to avoid the test stochastically failing
        np.random.seed(123)

        p = lm._RegressionPlotter(self.df.d, self.df.y,
                                  x_estimator=np.mean, ci=95)
        _, _, ci_big = p.estimate_data

        p = lm._RegressionPlotter(self.df.d, self.df.y,
                                  x_estimator=np.mean, ci=50)
        _, _, ci_wee = p.estimate_data
        npt.assert_array_less(np.diff(ci_wee), np.diff(ci_big))

        p = lm._RegressionPlotter(self.df.d, self.df.y,
                                  x_estimator=np.mean, ci=None)
        _, _, ci_nil = p.estimate_data
        npt.assert_array_equal(ci_nil, [None] * len(ci_nil))

    def test_estimate_units(self):

        # Seed the RNG locally
        np.random.seed(345)

        p = lm._RegressionPlotter("x", "y", data=self.df,
                                  units="s", x_bins=3)
        _, _, ci_big = p.estimate_data
        ci_big = np.diff(ci_big, axis=1)

        p = lm._RegressionPlotter("x", "y", data=self.df, x_bins=3)
        _, _, ci_wee = p.estimate_data
        ci_wee = np.diff(ci_wee, axis=1)

        npt.assert_array_less(ci_wee, ci_big)

    def test_partial(self):

        x = self.rs.randn(100)
        y = x + self.rs.randn(100)
        z = x + self.rs.randn(100)

        p = lm._RegressionPlotter(y, z)
        _, r_orig = np.corrcoef(p.x, p.y)[0]

        p = lm._RegressionPlotter(y, z, y_partial=x)
        _, r_semipartial = np.corrcoef(p.x, p.y)[0]
        nt.assert_less(r_semipartial, r_orig)

        p = lm._RegressionPlotter(y, z, x_partial=x, y_partial=x)
        _, r_partial = np.corrcoef(p.x, p.y)[0]
        nt.assert_less(r_partial, r_orig)

    @skipif(_no_statsmodels)
    def test_logistic_regression(self):

        p = lm._RegressionPlotter("x", "c", data=self.df,
                                  logistic=True, n_boot=self.n_boot)
        _, yhat, _ = p.fit_regression(x_range=(-3, 3))
        npt.assert_array_less(yhat, 1)
        npt.assert_array_less(0, yhat)

    @skipif(_no_statsmodels)
    def test_robust_regression(self):

        p_ols = lm._RegressionPlotter("x", "y", data=self.df,
                                      n_boot=self.n_boot)
        _, ols_yhat, _ = p_ols.fit_regression(x_range=(-3, 3))

        p_robust = lm._RegressionPlotter("x", "y", data=self.df,
                                         robust=True, n_boot=self.n_boot)
        _, robust_yhat, _ = p_robust.fit_regression(x_range=(-3, 3))

        nt.assert_equal(len(ols_yhat), len(robust_yhat))

    @skipif(_no_statsmodels)
    def test_lowess_regression(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, lowess=True)
        grid, yhat, err_bands = p.fit_regression(x_range=(-3, 3))

        nt.assert_equal(len(grid), len(yhat))
        nt.assert_is(err_bands, None)

    def test_regression_options(self):

        with nt.assert_raises(ValueError):
            lm._RegressionPlotter("x", "y", data=self.df,
                                  lowess=True, order=2)

        with nt.assert_raises(ValueError):
            lm._RegressionPlotter("x", "y", data=self.df,
                                  lowess=True, logistic=True)

    def test_regression_limits(self):

        f, ax = plt.subplots()
        ax.scatter(self.df.x, self.df.y)
        p = lm._RegressionPlotter("x", "y", data=self.df)
        grid, _, _ = p.fit_regression(ax)
        xlim = ax.get_xlim()
        nt.assert_equal(grid.min(), xlim[0])
        nt.assert_equal(grid.max(), xlim[1])

        p = lm._RegressionPlotter("x", "y", data=self.df, truncate=True)
        grid, _, _ = p.fit_regression()
        nt.assert_equal(grid.min(), self.df.x.min())
        nt.assert_equal(grid.max(), self.df.x.max())

class TestRegressionPlots(PlotTestCase):

    rs = np.random.RandomState(56)
    df = pd.DataFrame(dict(x=rs.randn(90),
                           y=rs.randn(90) + 5,
                           z=rs.randint(0, 1, 90),
                           g=np.repeat(list("abc"), 30),
                           h=np.tile(list("xy"), 45),
                           u=np.tile(np.arange(6), 15)))
    bw_err = rs.randn(6)[df.u.values]
    df.y += bw_err

    def test_regplot_basic(self):

        f, ax = plt.subplots()
        lm.regplot("x", "y", self.df)
        nt.assert_equal(len(ax.lines), 1)
        nt.assert_equal(len(ax.collections), 2)

        x, y = ax.collections[0].get_offsets().T
        npt.assert_array_equal(x, self.df.x)
        npt.assert_array_equal(y, self.df.y)

    def test_regplot_selective(self):

        f, ax = plt.subplots()
        ax = lm.regplot("x", "y", self.df, scatter=False, ax=ax)
        nt.assert_equal(len(ax.lines), 1)
        nt.assert_equal(len(ax.collections), 1)
        ax.clear()

        f, ax = plt.subplots()
        ax = lm.regplot("x", "y", self.df, fit_reg=False)
        nt.assert_equal(len(ax.lines), 0)
        nt.assert_equal(len(ax.collections), 1)
        ax.clear()

        f, ax = plt.subplots()
        ax = lm.regplot("x", "y", self.df, ci=None)
        nt.assert_equal(len(ax.lines), 1)
        nt.assert_equal(len(ax.collections), 1)
        ax.clear()

    def test_regplot_scatter_kws_alpha(self):

        f, ax = plt.subplots()
        color = np.array([[0.3, 0.8, 0.5, 0.5]])
        ax = lm.regplot("x", "y", self.df, scatter_kws={'color': color})
        nt.assert_is(ax.collections[0]._alpha, None)
        nt.assert_equal(ax.collections[0]._facecolors[0, 3], 0.5)

        f, ax = plt.subplots()
        color = np.array([[0.3, 0.8, 0.5]])
        ax = lm.regplot("x", "y", self.df, scatter_kws={'color': color})
        nt.assert_equal(ax.collections[0]._alpha, 0.8)

        f, ax = plt.subplots()
        color = np.array([[0.3, 0.8, 0.5]])
        ax = lm.regplot("x", "y", self.df, scatter_kws={'color': color,
                                                        'alpha': 0.4})
        nt.assert_equal(ax.collections[0]._alpha, 0.4)

        f, ax = plt.subplots()
        color = 'r'
        ax = lm.regplot("x", "y", self.df, scatter_kws={'color': color})
        nt.assert_equal(ax.collections[0]._alpha, 0.8)

    def test_regplot_binned(self):

        ax = lm.regplot("x", "y", self.df, x_bins=5)
        nt.assert_equal(len(ax.lines), 6)
        nt.assert_equal(len(ax.collections), 2)

    def test_lmplot_basic(self):

        g = lm.lmplot("x", "y", self.df)
        ax = g.axes[0, 0]
        nt.assert_equal(len(ax.lines), 1)
        nt.assert_equal(len(ax.collections), 2)

        x, y = ax.collections[0].get_offsets().T
        npt.assert_array_equal(x, self.df.x)
        npt.assert_array_equal(y, self.df.y)

    def test_lmplot_hue(self):

        g = lm.lmplot("x", "y", data=self.df, hue="h")
        ax = g.axes[0, 0]

        nt.assert_equal(len(ax.lines), 2)
        nt.assert_equal(len(ax.collections), 4)

    def test_lmplot_markers(self):

        g1 = lm.lmplot("x", "y", data=self.df, hue="h", markers="s")
        nt.assert_equal(g1.hue_kws, {"marker": ["s", "s"]})

        g2 = lm.lmplot("x", "y", data=self.df, hue="h", markers=["o", "s"])
        nt.assert_equal(g2.hue_kws, {"marker": ["o", "s"]})

        with nt.assert_raises(ValueError):
            lm.lmplot("x", "y", data=self.df, hue="h", markers=["o", "s", "d"])

    def test_lmplot_marker_linewidths(self):

        if mpl.__version__ == "1.4.2":
            raise SkipTest

        g = lm.lmplot("x", "y", data=self.df, hue="h",
                      fit_reg=False, markers=["o", "+"])
        c = g.axes[0, 0].collections
        nt.assert_equal(c[0].get_linewidths()[0], 0)
        rclw = mpl.rcParams["lines.linewidth"]
        nt.assert_equal(c[1].get_linewidths()[0], rclw)

    def test_lmplot_facets(self):

        g = lm.lmplot("x", "y", data=self.df, row="g", col="h")
        nt.assert_equal(g.axes.shape, (3, 2))

        g = lm.lmplot("x", "y", data=self.df, col="u", col_wrap=4)
        nt.assert_equal(g.axes.shape, (6,))

        g = lm.lmplot("x", "y", data=self.df, hue="h", col="u")
        nt.assert_equal(g.axes.shape, (1, 6))

    def test_lmplot_hue_col_nolegend(self):

        g = lm.lmplot("x", "y", data=self.df, col="h", hue="h")
        nt.assert_is(g._legend, None)

    def test_lmplot_scatter_kws(self):

        g = lm.lmplot("x", "y", hue="h", data=self.df, ci=None)
        red_scatter, blue_scatter = g.axes[0, 0].collections

        red, blue = color_palette(n_colors=2)
        npt.assert_array_equal(red, red_scatter.get_facecolors()[0, :3])
        npt.assert_array_equal(blue, blue_scatter.get_facecolors()[0, :3])

    def test_residplot(self):

        x, y = self.df.x, self.df.y
        ax = lm.residplot(x, y)

        resid = y - np.polyval(np.polyfit(x, y, 1), x)
        x_plot, y_plot = ax.collections[0].get_offsets().T

        npt.assert_array_equal(x, x_plot)
        npt.assert_array_almost_equal(resid, y_plot)

    @skipif(_no_statsmodels)
    def test_residplot_lowess(self):

        ax = lm.residplot("x", "y", self.df, lowess=True)
        nt.assert_equal(len(ax.lines), 2)

        x, y = ax.lines[1].get_xydata().T
        npt.assert_array_equal(x, np.sort(self.df.x))

    def test_three_point_colors(self):

        x, y = np.random.randn(2, 3)
        ax = lm.regplot(x, y, color=(1, 0, 0))
        color = ax.collections[0].get_facecolors()
        npt.assert_almost_equal(color[0, :3],
                                (1, 0, 0))


import numpy as np

# from matplotlib import pyplot as plt

from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *

class RVLoadHistogram(QGraphicsView):
    '''A histogram for the maximum load across the reachable area'''
    def __init__(self, ik):
        width = 330
        height = 120
        self.scene = QGraphicsScene(0,-15,width,height-15)
        super(RVLoadHistogram, self).__init__(self.scene)
        self.setBackgroundBrush(QBrush(Qt.white))
        self.setRenderHints(QPainter.Antialiasing)
        self.setFrameStyle(0)
        self.setAlignment(Qt.AlignCenter)
        self.setFixedSize(width, height)
        self.setSceneRect(0, 0, width, height)
        self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.scale(1, -1)

        self.subscribers = {
        'mouseEnter' : [],
        'mouseLeave' : [],
        'mouseMove' : []
        }

        self.lines = []
        self.hist = []
        self.edges = []
        self.config = ik.config
        self.update(ik)

        self.setMouseTracking(True)

    def update(self, ik=None):
        if ik is not None:
            self.ik = ik
            self.min_load = self.config['min_load'].value

        for line in self.lines:
            self.scene.removeItem(line)
        self.lines = []

        width = self.width()
        height = self.height()

        loads = np.ma.masked_invalid(self.ik.loads*self.ik.partial_ok)
        loads = np.ma.masked_where(loads == 0, loads).compressed()
        self.hist, self.edges = np.histogram(loads, bins='auto')

        buckets = len(self.hist)
        self.screen_step = width/np.max(self.edges)
        max_count = np.max(self.hist)

        # Display histogram
        for i in range(buckets):
            x = self.edges[i] * self.screen_step
            w = max(1, (self.edges[i+1] - self.edges[i]) * self.screen_step)
            l = (self.edges[i] + self.edges[i + 1]) / 2
            count = self.hist[i]
            if l < self.min_load:
                color = QColor(100,100,100)
            else:
                color = QColor(200, 180, 100)
            # print(count)
            line = self.scene.addLine(x, 5, x, 5 + (height-5) * count/max_count, QPen(color, w))
            self.lines.append(line)
        # Setpoint shows the configuration's minimum load
        setpoint = self.config['min_load'].value * self.screen_step
        line = self.scene.addLine(setpoint, 0, setpoint, height, QPen(QColor(150, 150, 255), 2))
        self.lines.append(line)

    def setMinimumLoad(self, val):
        self.min_load = val
        self.update()

    def subscribe(self, event, function):
        self.subscribers[event].append(function)

    def enterEvent(self, event):
        for func in self.subscribers['mouseEnter']:
            func(event)

    def leaveEvent(self, event):
        self.setMinimumLoad(self.config['min_load'].value)
        for func in self.subscribers['mouseLeave']:
            func(event)

    def mouseMoveEvent(self, event):
        if event.buttons() == Qt.LeftButton:
            self.click(event.pos())
        else:
            pt = self.mapToScene(event.pos())
            self.setMinimumLoad(pt.x()/self.screen_step)
        for func in self.subscribers['mouseMove']:
            func(event)

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.click(event.pos())

    def click(self, pos):
        pt = self.mapToScene(pos)
        self.config['min_load'].value = pt.x()/self.screen_step
        self.config.notifyChange()


"""
Various tools for extracting signal components from a fit of the amplitude
distribution
"""

from . import pdf
from .Classdef import Statfit
import numpy as np
import time
import random
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from lmfit import minimize, Parameters, report_fit

def param0(sample, method='basic'):
    """Estimate initial parameters for HK fitting

    Arguments
    ---------
    sample : sequence
        amplitudes

    Keywords
    --------
    method : string
        method to compute the initial parameters
    """
    if method == 'basic':
        a = np.nanmean(sample)
        s = np.nanstd(sample)
        mu = 1.
    return {'a':a, 's':s, 'mu':mu}

def lmfit(sample, fit_model='hk', bins='auto', p0 = None,
          xtol=1e-4, ftol=1e-4):
    """Lmfit

    Arguments
    ---------
    sample : sequence
        amplitudes between 0 and 1.

    Keywords
    --------
    fit_model : string
        name of the function (in pdf module) to use for the fit
    bins : string
        method to compute the bin width (inherited from numpy.histogram)
    p0 : dict
        Initial parameters. If None, estimated automatically.
    xtol : float
        ??
    ftol : float
        ??

    Return
    ------
    A Statfit Class
    """
    start = time.time()
    winsize = len(sample)
    bad = False

    #--------------------------------------------------------------------------
    # Clean sample
    #--------------------------------------------------------------------------
    sample = np.array(sample)
    sample = sample[np.isfinite(sample)]
    if len(sample) == 0:
        bad = True
        sample = np.zeros(10)+1

    #--------------------------------------------------------------------------
    # Make the histogram
    #--------------------------------------------------------------------------
#    n, edges, patches = hist(sample, bins=bins, normed=True)
    n, edges = np.histogram(sample, bins=bins, density=True)
#    plt.clf()

    x = ((np.roll(edges, -1) + edges)/2.)[0:-1]

    #--------------------------------------------------------------------------
    # Initial Parameters for the fit
    #--------------------------------------------------------------------------
    if p0 is None:
        p0 = param0(sample)

    prm0 = Parameters()
    #     (Name,    Value,                 Vary,   Min,    Max,    Expr)
    prm0.add('a',     p0['a'],              True,   0,  1,      None)
    prm0.add('s',     p0['s'],              True,   0,  1,      None)
    prm0.add('mu',    p0['mu'],             True,   .5, 10,    None)
    prm0.add('pt',    np.average(sample)**2,False,  0,  1,      'a**2+2*s**2*mu')

    #if fit_model == 'hk':
    #    # From [Dutt and Greenleaf. 1994, eq.14]
    #    prm0.add('a4',    np.average(sample)**4,False,  0,  1,
    #             '8*(1+1/mu)*s**4 + 8*s**2*s**2 + a**4')

    #--------------------------------------------------------------------------
    # Fit
    #--------------------------------------------------------------------------
    pdf2use = getattr(pdf, fit_model)

    # use 'lbfgs' fit if error with 'leastsq' fit
    try:
        p = minimize(pdf2use, prm0, args=(x, n), method='leastsq',
            xtol=xtol, ftol=ftol)
    except KeyboardInterrupt:
        raise
    except:
        print('!! Error with LEASTSQ fit, use L-BFGS-B instead')
        p = minimize(pdf2use, prm0, args=(x, n), method='lbfgs')

    #--------------------------------------------------------------------------
    # Output
    #--------------------------------------------------------------------------
    elapsed = time.time() - start

    values = {}

    # Create values dict For lmfit >0.9.0 compatibility since it is no longer
    # in the minimize output
    for i in p.params.keys():
        values[i] = p.params[i].value

    # Results
    result = Statfit(sample, pdf2use, values, p.params,
             p.chisqr, p.redchi, elapsed, p.nfev, p.message, p.success,
             p.residual, x, n, edges, bins=bins)

    # Identify bad results
    if bad is True:
        result.success = False
        result.values['a'] = 0
        result.values['s'] = 0
        result.values['mu'] = 0
        result.values['pt'] = 0
        result.chisqr = 0
        result.redchi = 0
        result.message = 'No valid data in the sample'
        result.residual = 0

    return result


"""
This file provides a way to obtain thermodynamic quantities from an
interpolation of available QMC solutions
"""

import numpy as np

import matplotlib.pyplot as plt
import matplotlib

from matplotlib import rc
rc('font', **{'family':'serif'})
rc('text', usetex=True)

import glob
import os

import ldaconf
basedir = ldaconf.basedir

from scipy.spatial import Delaunay
from scipy.interpolate import CloughTocher2DInterpolator, LinearNDInterpolator
from scipy.interpolate.interpnd import _ndim_coords_from_arrays

import logging
# create logger
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())
#logger.disabled = True

def get_qty_mu( dat, mu, MUCOL, COL, **kwargs ):
    # Control the interpolation between availble
    # density points here
    #~qtyinterp = 'nearest'
    qtyinterp = 'linear'
    msg = kwargs.get('msg', None)

    DENSCOL = 1
    ENTRCOL = 2
    SPICOL  = 3
    CMPRCOL = 4

    if COL == SPICOL:
        default_minus = 1.0
        default_plus  = 0.0
    elif COL == ENTRCOL:
        default_minus = 0.0
        default_plus  = 0.0
    elif COL == DENSCOL:
        default_minus = 0.0
        default_plus  = 2.0
    elif COL == CMPRCOL:
        default_minus = 0.0
        default_plus  = 0.0
    else:
        raise ValueError("Column not defined: COL = {:d}".format(COL) )

    CAREFUL = kwargs.get('careful', True)
    if CAREFUL and (mu < -10. or mu > 60.):
        CAREFUL = False

    if qtyinterp == 'nearest':
        index = np.argmin( np.abs(dat[:, MUCOL] - mu ))
        qtyresult = dat[index,COL]

    else:
        # find the two closest chemical potentials that
        # stride the point
        mudat = dat[:,MUCOL]

        verbose = False

        if np.all(mu < mudat):
            qtyresult = default_minus

            if COL == DENSCOL or COL == ENTRCOL:
               if verbose:
                   print "QTY=", COL,
                   print "===>>> mu={:0.2f} ".format(mu), msg
               if  dat[:,DENSCOL].min() < 0.1 :
                   qtyresult = default_minus
               elif CAREFUL:
                   return 'out-of-bounds'
               #print "====>>> BE CAREFUL :  Using default density" + \
               #      " n=%.2f"%default_minus  + \
               #      " at mu={:0.2f}  ".format(mu),
               #if msg is not None:
               #    print msg
               #raise ValueError('density error')

        elif np.all( mu > mudat):
            qtyresult = default_plus

            if COL == DENSCOL or COL == ENTRCOL:
               if verbose:
                   print "QTY=", COL,
                   print "====>>> mu={:0.2f} ".format(mu), msg
               if  dat[:,DENSCOL].max() > 1.9 :
                   qtyresult = default_plus
               elif CAREFUL:
                   return 'out-of-bounds'
               #print "====>>> BE CAREFUL :  Using default density" + \
               #      " n=%.2f"%default_plus  + \
               #      " at mu={:0.2f}  ".format(mu),
               #if msg is not None:
               #    print msg
               #raise ValueError('density error')

        else:
            # since the mu's are ordered we can do:
            index0 = np.where( mudat <=mu )[0][-1]
            index1 = np.where( mudat > mu )[0][0]

            qty0 = dat[ index0, COL ]
            qty1 = dat[ index1, COL ]

            mu0 = dat[ index0, MUCOL ]
            mu1 = dat[ index1, MUCOL ]

            qtyresult =  qty0 +  (mu-mu0) * (qty1-qty0) / (mu1-mu0)

    return qtyresult

                #print
                #print " mu = ", mu
                #print "index0 = ", index0
                #print "index1 = ", index1
                #print "Doing linear interpolation for the qty"
                #print " mu0 = ", mu0
                #print " mu1 = ", mu1
                #print "qty0 = ", qty0
                #print "qty1 = ", qty1
                #print "qtyresult = ", qtyresult

def find_closest_qmc( U=8, T=0.67, mu=4.0, **kwargs):
    """
    This function finds the closest values of U and T in the QMC data
    that straddle the values U and T given as arguments.

    """

    nUs = 4
    nTs = 3
    ALLPTS = kwargs.get('ALLPTS', False)

    # select which quantity will be returned, options are
    # spi and entropy
    QTY = kwargs.get('QTY', 'spi' )

    if QTY == 'spi':
        datadir = basedir + 'COMB_Final_Spi/'
    elif QTY == 'entropy':
        datadir = basedir + 'COMB_Final_Entr/'
    elif QTY == 'density':
        datadir = basedir + 'COMB_Final_Spi/'
    elif QTY == 'kappa':
        datadir = basedir + 'COMB_Final_Spi/'
    else:
        raise ValueError('Quantity not defined:' + str(QTY) )

    fname = datadir + 'U*'
    us = [ float(u.split('/U')[-1]) for u in glob.glob(fname) ]
    du = [ np.abs(U-u) for u in us ]
    index = np.argsort(du)

    if ALLPTS:
        Ulist0 = range(len(index))
    else:
        Ulist0 = range( nUs )

    us = [ us[index[i]] for i in Ulist0]
    #print us
    #print du
    #print index
    #print "Closest Us = ", us

    datfiles = []
    for u in us:

        # For the Spi and Stheta data
        if QTY == 'spi' or QTY == 'density' or QTY == 'kappa':
            fname = datadir + 'U{U:02d}/T*dat'.format(U=int(u))
            fs = sorted(glob.glob(fname))
            Ts = [ float(f.split('T')[1].split('.dat')[0]) for f in fs ]
        elif QTY=='entropy':
            fname = datadir + 'U{U:02d}/S*dat'.format(U=int(u))
            fs = sorted(glob.glob(fname))
            Ts = [ float(f.split('S')[1].split('.dat')[0]) for f in fs ]

        Ts_g = [] ; Ts_l = [];
        for t in Ts:
            if t > T:
                Ts_g.append(t)
            else:
                Ts_l.append(t)

        order_g = np.argsort( [ np.abs( T -t ) for t in Ts_g ] )
        order_l = np.argsort( [ np.abs( T -t ) for t in Ts_l ] )

        try:
            Tpts = [ Ts_g[ order_g[0]] , Ts_l[ order_l[0]]   ]
        except:
            #print
            #print "problem adding U=",u, "T=",Ts
            #print "available T data does not stride the point"
            #print "T  =", T
            #print "Ts =", Ts
            #print "will add nearest Ts nevertheless"
            Tpts = [  ]
            #raise ValueError("QMC data not available.")

        dT = [ np.abs( T - t) for t in Ts ]
        index = np.argsort(dT)

        if ALLPTS:
            Tlist0 = range(len(Ts))
        else:
            Tlist0 = range( min(nTs , len(Ts)))

        for i in Tlist0:
            Tnew = Ts[index[i]]
            if Tnew not in Tpts:
                Tpts.append(Tnew)
        for Tpt in Tpts:
            index = Ts.index( Tpt )
            try:
                datfiles.append( [ fs[ index ], u, Ts[index] ] )
            except:
                print "problem adding U=",u, "T=",Ts
                raise

        # Need to make sure that selected T values stride both
        # sides of the point

        #print
        #print u
        #print Ts
        #print dT
        #print index
        #print fs

#        for i in range(min(3, len(Ts))):
#            try:
#                datfiles.append( [ fs[index[i]], u, Ts[index[i]] ] )
#            except:
#                print "problem adding U=",u, "T=",Ts
#                raise
#
            #datfiles.append( [ fs[index[1]], u, Ts[index[1]] ] )

    #print datfiles
    MUCOL   = 0
    DENSCOL = 1
    ENTRCOL = 2
    SPICOL  = 3
    CMPRCOL = 4

    if QTY == 'spi':
        COL = SPICOL
    elif QTY == 'entropy':
        COL = ENTRCOL
    elif QTY == 'density':
        COL = DENSCOL
    elif QTY == 'kappa':
        COL = CMPRCOL

    msg0 = 'U={:0.2f}, T={:0.2f}'.format(U,T)

    logger.debug("number of nearby points  = " +  str(len(datfiles)))
    basedat = []
    basedaterr = []
    datserr = []
    for mm, f in enumerate(datfiles):
        # f[0] is the datafile name
        # f[1] is U
        # f[2] is T

        radius = kwargs.get('radius', np.nan )
        msg = 'U={:0.2f}, T={:0.2f}'.format(U,T) + \
               ' mu={:0.2f}, r={:0.2f}, Upt={:0.3f}, Tpt={:0.3f}'.\
               format(mu, radius, f[1], f[2])

        try:
            dat = np.loadtxt(f[0])
            spival = get_qty_mu( dat, mu, MUCOL, COL, msg=msg )

            # Toggle the false here to plot all of the out of bounds
            if spival == 'out-of-bounds':
                #spival_symmetry =
                logger.info('qty is out of bounds')
                basedaterr.append( [f[1], f[2], np.nan] )
                datserr.append( dat )

                if False:
                    fig = plt.figure( figsize=(3.5,3.5))
                    gs = matplotlib.gridspec.GridSpec( 1,1 ,\
                            left=0.15, right=0.96, bottom=0.12, top=0.88)
                    ax = fig.add_subplot( gs[0] )
                    ax.grid(alpha=0.5)
                    ax.plot( dat[:,MUCOL], dat[:,COL], '.-')
                    ax.axvline( mu )
                    ax.text( 0.5, 1.05, msg, ha='center', va='bottom', \
                        transform=ax.transAxes, fontsize=6.)
                    if matplotlib.get_backend() == 'agg':
                        fig.savefig('err_mu_%02d.png'%mm, dpi=200)
                        plt.close(fig)
                    else:
                        plt.show()
                        plt.close(fig)
                continue
            else:
                basedat.append( [f[1], f[2], spival] )
        except Exception as e :
            print "Failed to get data from file = ", f

            # toggle plotting, not implemented yet:
            if True:
                fig = plt.figure( figsize=(3.5,3.5))
                gs = matplotlib.gridspec.GridSpec( 1,1 ,\
                        left=0.15, right=0.96, bottom=0.12, top=0.88)
                ax = fig.add_subplot( gs[0] )
                ax.grid(alpha=0.5)
                ax.plot( dat[:,MUCOL], dat[:,COL], '.-')
                ax.axvline( mu )
                ax.text( 0.5, 1.05, msg, ha='center', va='bottom', \
                    transform=ax.transAxes)
                if matplotlib.get_backend() == 'agg':
                    fig.savefig('err_mu_%02d.png'%mm, dpi=200)
                else:
                    plt.show()
            raise e
    logger.debug("number of nearby valid points  = " +  str(len(basedat)))

    error = False
    points = None

    # MAKE THE TRIANGULATION
    basedat =   np.array(basedat)

    Us = np.unique(basedat[:,0] )
    Ts = np.unique(basedat[:,1] )

    validTriang = not ( len(Us) ==1  or len(Ts) ==  1 )
    #print "#Us={:d}, #Ts={:d}".format( len(Us), len(Ts) )
    #print msg

    if validTriang:
        points = _ndim_coords_from_arrays(( basedat[:,0] , basedat[:,1]))
        #print "Closest dat = ", basedat
        #finterp = CloughTocher2DInterpolator(points, basedat[:,2])
        finterp = LinearNDInterpolator( points, basedat[:,2] )
    else:

        logerr = 'not enough finterp points, QTY=%s'%QTY + '\n' + msg + '\n' \
                  + "number of basedat pts = " + str(len(basedat))
        print basedat
        print "len Us = ", len(Us)
        print "len Ts = ", len(Ts)
        print "len 'out-of-bounds' = ", len( basedaterr )
        if len( basedaterr ) > 0:
            for bb, bdaterr in enumerate(basedaterr):
                msgbb = 'U={:0.2f}, T={:0.2f}'.format(U,T) +\
                   ' mu={:0.2f}, r={:0.2f}, Upt={:0.3f}, Tpt={:0.3f}'.\
               format(mu, radius, basedaterr[bb][0], basedaterr[bb][1] )
                daterr = datserr[bb]
                fig = plt.figure( figsize=(3.5,3.5))
                gs = matplotlib.gridspec.GridSpec( 1,1 ,\
                        left=0.15, right=0.96, bottom=0.12, top=0.88)
                ax = fig.add_subplot( gs[0] )
                ax.grid(alpha=0.5)
                ax.plot( daterr[:,MUCOL], daterr[:,COL], '.-')
                ax.axvline( mu )
                ax.text( 0.5, 1.05, msgbb, ha='center', va='bottom', \
                    transform=ax.transAxes, fontsize=6.)
                if matplotlib.get_backend() == 'agg':
                    fig.savefig('err_mu_%02d.png'%bb, dpi=200)
                    plt.close(fig)
                else:
                    plt.show()
                    plt.close(fig)
        logger.exception(logerr)
        raise ValueError('finterp')

    if points == None:
        logger.warning( "points object is None"  )

    if error == False:
        try:
            result = finterp( U,T )
            if np.isnan(result):
                if U >= 30.0 and U <=32.5:
                    result = finterp( 29.99, T )
                    logger.warning(" qmc: U={:0.1f} replaced to U=29.99 ".\
                                     format(U) )
            if np.isnan(result):
                raise Exception("\n!!!! qmc: Invalid result, QTY:%s!!!!\n"%QTY \
                        + msg0)
        except Exception as e:
            if kwargs.get('error_nan', False):
                return np.nan
            else:
                error = True
                logger.exception("Invalid QTY result!")

    if error == False:
        if result >= 8. and QTY == 'spi' :
            print " Obtained Spi > 8. : U={:0.2f}, T={:0.2f}, mu={:0.2f}".\
                   format( U, T, mu ),
            print "  ==> Spi={:0.2f}".format(float(result))
            error = True
        elif result >=4. and QTY == 'entropy':
            print " Obtained Ent > 4. : U={:0.2f}, T={:0.2f}, mu={:0.2f}".\
                   format( U, T, mu ),
            print "  ==> Result={:0.2f}".format(float(result))
            error = True

    logger.debug("error status = " + str(error))
    if error or kwargs.get('showinterp',False):
        logger.debug("Inside error if statement...")

        if kwargs.get('error_nan', False):
            pass
            #return np.nan

        #print "Interp points:"
        #print basedat
        if len(basedat) == 0 and len(basedaterr) > 0 :

            basedaterr =   np.array(basedaterr)
            Userr = np.unique(basedaterr[:,0] )
            Tserr = np.unique(basedaterr[:,1] )
            validTriangerr = not ( len(Userr) ==1  or len(Tserr) ==  1 )

            points = _ndim_coords_from_arrays(( basedaterr[:,0] , basedaterr[:,1]))
            tri = Delaunay(points)

        else:
            tri = Delaunay(points)
        fig = plt.figure( figsize=(3.5,3.5))
        gs = matplotlib.gridspec.GridSpec( 1,1 ,\
                left=0.15, right=0.96, bottom=0.12, top=0.88)
        ax = fig.add_subplot( gs[0] )
        ax.grid(alpha=0.5)
        ax.triplot(points[:,0], points[:,1], tri.simplices.copy())
        ax.plot(points[:,0], points[:,1], 'o')
        ax.plot( U, T, 'o', ms=6., color='red')
        xlim = ax.get_xlim()
        dx = (xlim[1]-xlim[0])/10.
        ax.set_xlim( xlim[0]-dx, xlim[1]+dx )
        ylim = ax.get_ylim()
        dy = (ylim[1]-ylim[0])/10.
        ax.set_ylim( ylim[0]-dy, ylim[1]+dy )
        ax.set_xlabel('$U/t$')
        ax.set_ylabel('$T/t$',rotation=0,labelpad=8)

        tt = kwargs.get('title_text','')
        ax.set_title( tt + '$U/t={:.2f}$'.format(U) + \
                      ',\ \ ' + '$T/t={:.2f}$'.format(T), \
                ha='center', va='bottom', fontsize=10)

        save_err = kwargs.get('save_err',None)
        if save_err is not None:
            print "Saving png."
            fig.savefig( save_err, dpi=300)
        if matplotlib.get_backend() == 'agg':
            fig.savefig('err.png', dpi=200)
            print "Saved error to err.png"
        else:
            plt.show()
        if not kwargs.get('single', False):
            raise ValueError("Could not interpolate using QMC data.")
        if ALLPTS:
            if 'savepath' in kwargs.keys():
                fig.savefig( kwargs.get('savepath',None) , dpi=300)
        if error:
            raise

    return result


import re, inspect, textwrap, pydoc
import sphinx
from docscrape import NumpyDocString, FunctionDoc, ClassDoc

class SphinxDocString(NumpyDocString):
    def __init__(self, docstring, config={}):
        self.use_plots = config.get('use_plots', False)
        NumpyDocString.__init__(self, docstring, config=config)

    # string conversion routines
    def _str_header(self, name, symbol='`'):
        return ['.. rubric:: ' + name, '']

    def _str_field_list(self, name):
        return [':' + name + ':']

    def _str_indent(self, doc, indent=4):
        out = []
        for line in doc:
            out += [' '*indent + line]
        return out

    def _str_signature(self):
        return ['']
        if self['Signature']:
            return ['``%s``' % self['Signature']] + ['']
        else:
            return ['']

    def _str_summary(self):
        return self['Summary'] + ['']

    def _str_extended_summary(self):
        return self['Extended Summary'] + ['']

    def _str_param_list(self, name):
        out = []
        if self[name]:
            out += self._str_field_list(name)
            out += ['']
            for param,param_type,desc in self[name]:
                out += self._str_indent(['**%s** : %s' % (param.strip(),
                                                          param_type)])
                out += ['']
                out += self._str_indent(desc,8)
                out += ['']
        return out

    @property
    def _obj(self):
        if hasattr(self, '_cls'):
            return self._cls
        elif hasattr(self, '_f'):
            return self._f
        return None

    def _str_member_list(self, name):
        """
        Generate a member listing, autosummary:: table where possible,
        and a table where not.

        """
        out = []
        if self[name]:
            out += ['.. rubric:: %s' % name, '']
            prefix = getattr(self, '_name', '')

            if prefix:
                prefix = '~%s.' % prefix

            autosum = []
            others = []
            for param, param_type, desc in self[name]:
                param = param.strip()
                if not self._obj or hasattr(self._obj, param):
                    autosum += ["   %s%s" % (prefix, param)]
                else:
                    others.append((param, param_type, desc))

            if autosum:
                out += ['.. autosummary::', '   :toctree:', '']
                out += autosum

            if others:
                maxlen_0 = max([len(x[0]) for x in others])
                maxlen_1 = max([len(x[1]) for x in others])
                hdr = "="*maxlen_0 + "  " + "="*maxlen_1 + "  " + "="*10
                fmt = '%%%ds  %%%ds  ' % (maxlen_0, maxlen_1)
                n_indent = maxlen_0 + maxlen_1 + 4
                out += [hdr]
                for param, param_type, desc in others:
                    out += [fmt % (param.strip(), param_type)]
                    out += self._str_indent(desc, n_indent)
                out += [hdr]
            out += ['']
        return out

    def _str_section(self, name):
        out = []
        if self[name]:
            out += self._str_header(name)
            out += ['']
            content = textwrap.dedent("\n".join(self[name])).split("\n")
            out += content
            out += ['']
        return out

    def _str_see_also(self, func_role):
        out = []
        if self['See Also']:
            see_also = super(SphinxDocString, self)._str_see_also(func_role)
            out = ['.. seealso::', '']
            out += self._str_indent(see_also[2:])
        return out

    def _str_warnings(self):
        out = []
        if self['Warnings']:
            out = ['.. warning::', '']
            out += self._str_indent(self['Warnings'])
        return out

    def _str_index(self):
        idx = self['index']
        out = []
        if len(idx) == 0:
            return out

        out += ['.. index:: %s' % idx.get('default','')]
        for section, references in idx.iteritems():
            if section == 'default':
                continue
            elif section == 'refguide':
                out += ['   single: %s' % (', '.join(references))]
            else:
                out += ['   %s: %s' % (section, ','.join(references))]
        return out

    def _str_references(self):
        out = []
        if self['References']:
            out += self._str_header('References')
            if isinstance(self['References'], str):
                self['References'] = [self['References']]
            out.extend(self['References'])
            out += ['']
            # Latex collects all references to a separate bibliography,
            # so we need to insert links to it
            if sphinx.__version__ >= "0.6":
                out += ['.. only:: latex','']
            else:
                out += ['.. latexonly::','']
            items = []
            for line in self['References']:
                m = re.match(r'.. \[([a-z0-9._-]+)\]', line, re.I)
                if m:
                    items.append(m.group(1))
            out += ['   ' + ", ".join(["[%s]_" % item for item in items]), '']
        return out

    def _str_examples(self):
        examples_str = "\n".join(self['Examples'])

        if (self.use_plots and 'import matplotlib' in examples_str
                and 'plot::' not in examples_str):
            out = []
            out += self._str_header('Examples')
            out += ['.. plot::', '']
            out += self._str_indent(self['Examples'])
            out += ['']
            return out
        else:
            return self._str_section('Examples')

    def __str__(self, indent=0, func_role="obj"):
        out = []
        out += self._str_signature()
        out += self._str_index() + ['']
        out += self._str_summary()
        out += self._str_extended_summary()
        for param_list in ('Parameters', 'Returns', 'Raises'):
            out += self._str_param_list(param_list)
        out += self._str_warnings()
        out += self._str_see_also(func_role)
        out += self._str_section('Notes')
        out += self._str_references()
        out += self._str_examples()
        for param_list in ('Attributes', 'Methods'):
            out += self._str_member_list(param_list)
        out = self._str_indent(out,indent)
        return '\n'.join(out)

class SphinxFunctionDoc(SphinxDocString, FunctionDoc):
    def __init__(self, obj, doc=None, config={}):
        self.use_plots = config.get('use_plots', False)
        FunctionDoc.__init__(self, obj, doc=doc, config=config)

class SphinxClassDoc(SphinxDocString, ClassDoc):
    def __init__(self, obj, doc=None, func_doc=None, config={}):
        self.use_plots = config.get('use_plots', False)
        ClassDoc.__init__(self, obj, doc=doc, func_doc=None, config=config)

class SphinxObjDoc(SphinxDocString):
    def __init__(self, obj, doc=None, config={}):
        self._f = obj
        SphinxDocString.__init__(self, doc, config=config)

def get_doc_object(obj, what=None, doc=None, config={}):
    if what is None:
        if inspect.isclass(obj):
            what = 'class'
        elif inspect.ismodule(obj):
            what = 'module'
        elif callable(obj):
            what = 'function'
        else:
            what = 'object'
    if what == 'class':
        return SphinxClassDoc(obj, func_doc=SphinxFunctionDoc, doc=doc,
                              config=config)
    elif what in ('function', 'method'):
        return SphinxFunctionDoc(obj, doc=doc, config=config)
    else:
        if doc is None:
            doc = pydoc.getdoc(obj)
        return SphinxObjDoc(obj, doc, config=config)


# -*- coding: utf-8 -*-
"""Sphinx directive to support embedded IPython code.

This directive allows pasting of entire interactive IPython sessions, prompts
and all, and their code will actually get re-executed at doc build time, with
all prompts renumbered sequentially. It also allows you to input code as a pure
python input by giving the argument python to the directive. The output looks
like an interactive ipython section.

To enable this directive, simply list it in your Sphinx ``conf.py`` file
(making sure the directory where you placed it is visible to sphinx, as is
needed for all Sphinx directives).

By default this directive assumes that your prompts are unchanged IPython ones,
but this can be customized. The configurable options that can be placed in
conf.py are

ipython_savefig_dir:
    The directory in which to save the figures. This is relative to the
    Sphinx source directory. The default is `html_static_path`.
ipython_rgxin:
    The compiled regular expression to denote the start of IPython input
    lines. The default is re.compile('In \[(\d+)\]:\s?(.*)\s*'). You
    shouldn't need to change this.
ipython_rgxout:
    The compiled regular expression to denote the start of IPython output
    lines. The default is re.compile('Out\[(\d+)\]:\s?(.*)\s*'). You
    shouldn't need to change this.
ipython_promptin:
    The string to represent the IPython input prompt in the generated ReST.
    The default is 'In [%d]:'. This expects that the line numbers are used
    in the prompt.
ipython_promptout:

    The string to represent the IPython prompt in the generated ReST. The
    default is 'Out [%d]:'. This expects that the line numbers are used
    in the prompt.

ToDo
----

- Turn the ad-hoc test() function into a real test suite.
- Break up ipython-specific functionality from matplotlib stuff into better
  separated code.

Authors
-------

- John D Hunter: orignal author.
- Fernando Perez: refactoring, documentation, cleanups, port to 0.11.
- VáclavŠmilauer <eudoxos-AT-arcig.cz>: Prompt generalizations.
- Skipper Seabold, refactoring, cleanups, pure python addition
"""

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

# Stdlib
import cStringIO
import os
import re
import sys
import tempfile
import ast

# To keep compatibility with various python versions
try:
    from hashlib import md5
except ImportError:
    from md5 import md5

# Third-party
try:
    import matplotlib
    matplotlib.use('Agg')
except ImportError:
    print "Couldn't find matplotlib"

import sphinx
from docutils.parsers.rst import directives
from docutils import nodes
from sphinx.util.compat import Directive

# Our own
from IPython import Config, InteractiveShell
from IPython.core.profiledir import ProfileDir
from IPython.utils import io

#-----------------------------------------------------------------------------
# Globals
#-----------------------------------------------------------------------------
# for tokenizing blocks
COMMENT, INPUT, OUTPUT =  range(3)

#-----------------------------------------------------------------------------
# Functions and class declarations
#-----------------------------------------------------------------------------
def block_parser(part, rgxin, rgxout, fmtin, fmtout):
    """
    part is a string of ipython text, comprised of at most one
    input, one ouput, comments, and blank lines.  The block parser
    parses the text into a list of::

      blocks = [ (TOKEN0, data0), (TOKEN1, data1), ...]

    where TOKEN is one of [COMMENT | INPUT | OUTPUT ] and
    data is, depending on the type of token::

      COMMENT : the comment string

      INPUT: the (DECORATOR, INPUT_LINE, REST) where
         DECORATOR: the input decorator (or None)
         INPUT_LINE: the input as string (possibly multi-line)
         REST : any stdout generated by the input line (not OUTPUT)

      OUTPUT: the output string, possibly multi-line
    """

    block = []
    lines = part.split('\n')
    N = len(lines)
    i = 0
    decorator = None
    while 1:

        if i==N:
            # nothing left to parse -- the last line
            break

        line = lines[i]
        i += 1
        line_stripped = line.strip()
        if line_stripped.startswith('#'):
            block.append((COMMENT, line))
            continue

        if line_stripped.startswith('@'):
            # we're assuming at most one decorator -- may need to
            # rethink
            decorator = line_stripped
            continue

        # does this look like an input line?
        matchin = rgxin.match(line)
        if matchin:
            lineno, inputline = int(matchin.group(1)), matchin.group(2)

            # the ....: continuation string
            continuation = '   %s:'%''.join(['.']*(len(str(lineno))+2))
            Nc = len(continuation)
            # input lines can continue on for more than one line, if
            # we have a '\' line continuation char or a function call
            # echo line 'print'.  The input line can only be
            # terminated by the end of the block or an output line, so
            # we parse out the rest of the input line if it is
            # multiline as well as any echo text

            rest = []
            while i<N:

                # look ahead; if the next line is blank, or a comment, or
                # an output line, we're done

                nextline = lines[i]
                matchout = rgxout.match(nextline)
                #print "nextline=%s, continuation=%s, starts=%s"%(nextline, continuation, nextline.startswith(continuation))
                if matchout or nextline.startswith('#'):
                    break
                elif nextline.startswith(continuation):
                    inputline += '\n' + nextline[Nc:]
                else:
                    rest.append(nextline)
                i+= 1

            block.append((INPUT, (decorator, inputline, '\n'.join(rest))))
            continue

        # if it looks like an output line grab all the text to the end
        # of the block
        matchout = rgxout.match(line)
        if matchout:
            lineno, output = int(matchout.group(1)), matchout.group(2)
            if i<N-1:
                output = '\n'.join([output] + lines[i:])

            block.append((OUTPUT, output))
            break

    return block

class EmbeddedSphinxShell(object):
    """An embedded IPython instance to run inside Sphinx"""

    def __init__(self):

        self.cout = cStringIO.StringIO()

        # Create config object for IPython
        config = Config()
        config.Global.display_banner = False
        config.Global.exec_lines = ['import numpy as np',
                                    'from pylab import *'
                                    ]
        config.InteractiveShell.autocall = False
        config.InteractiveShell.autoindent = False
        config.InteractiveShell.colors = 'NoColor'

        # create a profile so instance history isn't saved
        tmp_profile_dir = tempfile.mkdtemp(prefix='profile_')
        profname = 'auto_profile_sphinx_build'
        pdir = os.path.join(tmp_profile_dir,profname)
        profile = ProfileDir.create_profile_dir(pdir)

        # Create and initialize ipython, but don't start its mainloop
        IP = InteractiveShell.instance(config=config, profile_dir=profile)
        # io.stdout redirect must be done *after* instantiating InteractiveShell
        io.stdout = self.cout
        io.stderr = self.cout

        # For debugging, so we can see normal output, use this:
        #from IPython.utils.io import Tee
        #io.stdout = Tee(self.cout, channel='stdout') # dbg
        #io.stderr = Tee(self.cout, channel='stderr') # dbg

        # Store a few parts of IPython we'll need.
        self.IP = IP
        self.user_ns = self.IP.user_ns
        self.user_global_ns = self.IP.user_global_ns

        self.input = ''
        self.output = ''

        self.is_verbatim = False
        self.is_doctest = False
        self.is_suppress = False

        # on the first call to the savefig decorator, we'll import
        # pyplot as plt so we can make a call to the plt.gcf().savefig
        self._pyplot_imported = False

    def clear_cout(self):
        self.cout.seek(0)
        self.cout.truncate(0)

    def process_input_line(self, line, store_history=True):
        """process the input, capturing stdout"""
        #print "input='%s'"%self.input
        stdout = sys.stdout
        splitter = self.IP.input_splitter
        try:
            sys.stdout = self.cout
            splitter.push(line)
            more = splitter.push_accepts_more()
            if not more:
                source_raw = splitter.source_raw_reset()[1]
                self.IP.run_cell(source_raw, store_history=store_history)
        finally:
            sys.stdout = stdout

    def process_image(self, decorator):
        """
        # build out an image directive like
        # .. image:: somefile.png
        #    :width 4in
        #
        # from an input like
        # savefig somefile.png width=4in
        """
        savefig_dir = self.savefig_dir
        source_dir = self.source_dir
        saveargs = decorator.split(' ')
        filename = saveargs[1]
        # insert relative path to image file in source
        outfile = os.path.relpath(os.path.join(savefig_dir,filename),
                    source_dir)

        imagerows = ['.. image:: %s'%outfile]

        for kwarg in saveargs[2:]:
            arg, val = kwarg.split('=')
            arg = arg.strip()
            val = val.strip()
            imagerows.append('   :%s: %s'%(arg, val))

        image_file = os.path.basename(outfile) # only return file name
        image_directive = '\n'.join(imagerows)
        return image_file, image_directive

    # Callbacks for each type of token
    def process_input(self, data, input_prompt, lineno):
        """Process data block for INPUT token."""
        decorator, input, rest = data
        image_file = None
        image_directive = None
        #print 'INPUT:', data  # dbg
        is_verbatim = decorator=='@verbatim' or self.is_verbatim
        is_doctest = decorator=='@doctest' or self.is_doctest
        is_suppress = decorator=='@suppress' or self.is_suppress
        is_savefig = decorator is not None and \
                     decorator.startswith('@savefig')

        input_lines = input.split('\n')
        if len(input_lines) > 1:
            if input_lines[-1] != "":
                input_lines.append('') # make sure there's a blank line
                                       # so splitter buffer gets reset

        continuation = '   %s:'%''.join(['.']*(len(str(lineno))+2))
        Nc = len(continuation)

        if is_savefig:
            image_file, image_directive = self.process_image(decorator)

        ret = []
        is_semicolon = False

        for i, line in enumerate(input_lines):
            if line.endswith(';'):
                is_semicolon = True

            if i==0:
                # process the first input line
                if is_verbatim:
                    self.process_input_line('')
                    self.IP.execution_count += 1 # increment it anyway
                else:
                    # only submit the line in non-verbatim mode
                    self.process_input_line(line, store_history=True)
                formatted_line = '%s %s'%(input_prompt, line)
            else:
                # process a continuation line
                if not is_verbatim:
                    self.process_input_line(line, store_history=True)

                formatted_line = '%s %s'%(continuation, line)

            if not is_suppress:
                ret.append(formatted_line)

        if not is_suppress and len(rest.strip()) and is_verbatim:
            # the "rest" is the standard output of the
            # input, which needs to be added in
            # verbatim mode
            ret.append(rest)

        self.cout.seek(0)
        output = self.cout.read()
        if not is_suppress and not is_semicolon:
            ret.append(output)
        elif is_semicolon: # get spacing right
            ret.append('')

        self.cout.truncate(0)
        return (ret, input_lines, output, is_doctest, image_file,
                    image_directive)
        #print 'OUTPUT', output  # dbg

    def process_output(self, data, output_prompt,
                       input_lines, output, is_doctest, image_file):
        """Process data block for OUTPUT token."""
        if is_doctest:
            submitted = data.strip()
            found = output
            if found is not None:
                found = found.strip()

                # XXX - fperez: in 0.11, 'output' never comes with the prompt
                # in it, just the actual output text.  So I think all this code
                # can be nuked...

                # the above comment does not appear to be accurate... (minrk)

                ind = found.find(output_prompt)
                if ind<0:
                    e='output prompt="%s" does not match out line=%s' % \
                       (output_prompt, found)
                    raise RuntimeError(e)
                found = found[len(output_prompt):].strip()

                if found!=submitted:
                    e = ('doctest failure for input_lines="%s" with '
                         'found_output="%s" and submitted output="%s"' %
                         (input_lines, found, submitted) )
                    raise RuntimeError(e)
                #print 'doctest PASSED for input_lines="%s" with found_output="%s" and submitted output="%s"'%(input_lines, found, submitted)

    def process_comment(self, data):
        """Process data fPblock for COMMENT token."""
        if not self.is_suppress:
            return [data]

    def save_image(self, image_file):
        """
        Saves the image file to disk.
        """
        self.ensure_pyplot()
        command = 'plt.gcf().savefig("%s")'%image_file
        #print 'SAVEFIG', command  # dbg
        self.process_input_line('bookmark ipy_thisdir', store_history=False)
        self.process_input_line('cd -b ipy_savedir', store_history=False)
        self.process_input_line(command, store_history=False)
        self.process_input_line('cd -b ipy_thisdir', store_history=False)
        self.process_input_line('bookmark -d ipy_thisdir', store_history=False)
        self.clear_cout()

    def process_block(self, block):
        """
        process block from the block_parser and return a list of processed lines
        """
        ret = []
        output = None
        input_lines = None
        lineno = self.IP.execution_count

        input_prompt = self.promptin%lineno
        output_prompt = self.promptout%lineno
        image_file = None
        image_directive = None

        for token, data in block:
            if token==COMMENT:
                out_data = self.process_comment(data)
            elif token==INPUT:
                (out_data, input_lines, output, is_doctest, image_file,
                    image_directive) = \
                          self.process_input(data, input_prompt, lineno)
            elif token==OUTPUT:
                out_data = \
                    self.process_output(data, output_prompt,
                                        input_lines, output, is_doctest,
                                        image_file)
            if out_data:
                ret.extend(out_data)

        # save the image files
        if image_file is not None:
            self.save_image(image_file)

        return ret, image_directive

    def ensure_pyplot(self):
        if self._pyplot_imported:
            return
        self.process_input_line('import matplotlib.pyplot as plt',
                                store_history=False)

    def process_pure_python(self, content):
        """
        content is a list of strings. it is unedited directive conent

        This runs it line by line in the InteractiveShell, prepends
        prompts as needed capturing stderr and stdout, then returns
        the content as a list as if it were ipython code
        """
        output = []
        savefig = False # keep up with this to clear figure
        multiline = False # to handle line continuation
        multiline_start = None
        fmtin = self.promptin

        ct = 0

        for lineno, line in enumerate(content):

            line_stripped = line.strip()
            if not len(line):
                output.append(line)
                continue

            # handle decorators
            if line_stripped.startswith('@'):
                output.extend([line])
                if 'savefig' in line:
                    savefig = True # and need to clear figure
                continue

            # handle comments
            if line_stripped.startswith('#'):
                output.extend([line])
                continue

            # deal with lines checking for multiline
            continuation  = u'   %s:'% ''.join(['.']*(len(str(ct))+2))
            if not multiline:
                modified = u"%s %s" % (fmtin % ct, line_stripped)
                output.append(modified)
                ct += 1
                try:
                    ast.parse(line_stripped)
                    output.append(u'')
                except Exception: # on a multiline
                    multiline = True
                    multiline_start = lineno
            else: # still on a multiline
                modified = u'%s %s' % (continuation, line)
                output.append(modified)
                try:
                    mod = ast.parse(
                            '\n'.join(content[multiline_start:lineno+1]))
                    if isinstance(mod.body[0], ast.FunctionDef):
                        # check to see if we have the whole function
                        for element in mod.body[0].body:
                            if isinstance(element, ast.Return):
                                multiline = False
                    else:
                        output.append(u'')
                        multiline = False
                except Exception:
                    pass

            if savefig: # clear figure if plotted
                self.ensure_pyplot()
                self.process_input_line('plt.clf()', store_history=False)
                self.clear_cout()
                savefig = False

        return output

class IpythonDirective(Directive):

    has_content = True
    required_arguments = 0
    optional_arguments = 4 # python, suppress, verbatim, doctest
    final_argumuent_whitespace = True
    option_spec = { 'python': directives.unchanged,
                    'suppress' : directives.flag,
                    'verbatim' : directives.flag,
                    'doctest' : directives.flag,
                  }

    shell = EmbeddedSphinxShell()

    def get_config_options(self):
        # contains sphinx configuration variables
        config = self.state.document.settings.env.config

        # get config variables to set figure output directory
        confdir = self.state.document.settings.env.app.confdir
        savefig_dir = config.ipython_savefig_dir
        source_dir = os.path.dirname(self.state.document.current_source)
        if savefig_dir is None:
            savefig_dir = config.html_static_path
        if isinstance(savefig_dir, list):
            savefig_dir = savefig_dir[0] # safe to assume only one path?
        savefig_dir = os.path.join(confdir, savefig_dir)

        # get regex and prompt stuff
        rgxin     = config.ipython_rgxin
        rgxout    = config.ipython_rgxout
        promptin  = config.ipython_promptin
        promptout = config.ipython_promptout

        return savefig_dir, source_dir, rgxin, rgxout, promptin, promptout

    def setup(self):
        # reset the execution count if we haven't processed this doc
        #NOTE: this may be borked if there are multiple seen_doc tmp files
        #check time stamp?
        seen_docs = [i for i in os.listdir(tempfile.tempdir)
            if i.startswith('seen_doc')]
        if seen_docs:
            fname = os.path.join(tempfile.tempdir, seen_docs[0])
            docs = open(fname).read().split('\n')
            if not self.state.document.current_source in docs:
                self.shell.IP.history_manager.reset()
                self.shell.IP.execution_count = 1
        else: # haven't processed any docs yet
            docs = []

        # get config values
        (savefig_dir, source_dir, rgxin,
                rgxout, promptin, promptout) = self.get_config_options()

        # and attach to shell so we don't have to pass them around
        self.shell.rgxin = rgxin
        self.shell.rgxout = rgxout
        self.shell.promptin = promptin
        self.shell.promptout = promptout
        self.shell.savefig_dir = savefig_dir
        self.shell.source_dir = source_dir

        # setup bookmark for saving figures directory

        self.shell.process_input_line('bookmark ipy_savedir %s'%savefig_dir,
                                      store_history=False)
        self.shell.clear_cout()

        # write the filename to a tempfile because it's been "seen" now
        if not self.state.document.current_source in docs:
            fd, fname = tempfile.mkstemp(prefix="seen_doc", text=True)
            fout = open(fname, 'a')
            fout.write(self.state.document.current_source+'\n')
            fout.close()

        return rgxin, rgxout, promptin, promptout

    def teardown(self):
        # delete last bookmark
        self.shell.process_input_line('bookmark -d ipy_savedir',
                                      store_history=False)
        self.shell.clear_cout()

    def run(self):
        debug = False

        #TODO, any reason block_parser can't be a method of embeddable shell
        # then we wouldn't have to carry these around
        rgxin, rgxout, promptin, promptout = self.setup()

        options = self.options
        self.shell.is_suppress = 'suppress' in options
        self.shell.is_doctest = 'doctest' in options
        self.shell.is_verbatim = 'verbatim' in options

        # handle pure python code
        if 'python' in self.arguments:
            content = self.content
            self.content = self.shell.process_pure_python(content)

        parts = '\n'.join(self.content).split('\n\n')

        lines = ['.. code-block:: ipython','']
        figures = []

        for part in parts:

            block = block_parser(part, rgxin, rgxout, promptin, promptout)

            if len(block):
                rows, figure = self.shell.process_block(block)
                for row in rows:
                    lines.extend(['   %s'%line for line in row.split('\n')])

                if figure is not None:
                    figures.append(figure)

        #text = '\n'.join(lines)
        #figs = '\n'.join(figures)

        for figure in figures:
            lines.append('')
            lines.extend(figure.split('\n'))
            lines.append('')

        #print lines
        if len(lines)>2:
            if debug:
                print '\n'.join(lines)
            else: #NOTE: this raises some errors, what's it for?
                #print 'INSERTING %d lines'%len(lines)
                self.state_machine.insert_input(
                    lines, self.state_machine.input_lines.source(0))

        text = '\n'.join(lines)
        txtnode = nodes.literal_block(text, text)
        txtnode['language'] = 'ipython'
        #imgnode = nodes.image(figs)

        # cleanup
        self.teardown()

        return []#, imgnode]

# Enable as a proper Sphinx directive
def setup(app):
    setup.app = app

    app.add_directive('ipython', IpythonDirective)
    app.add_config_value('ipython_savefig_dir', None, True)
    app.add_config_value('ipython_rgxin',
                         re.compile('In \[(\d+)\]:\s?(.*)\s*'), True)
    app.add_config_value('ipython_rgxout',
                         re.compile('Out\[(\d+)\]:\s?(.*)\s*'), True)
    app.add_config_value('ipython_promptin', 'In [%d]:', True)
    app.add_config_value('ipython_promptout', 'Out[%d]:', True)

# Simple smoke test, needs to be converted to a proper automatic test.
def test():

    examples = [
        r"""
In [9]: pwd
Out[9]: '/home/jdhunter/py4science/book'

In [10]: cd bookdata/
/home/jdhunter/py4science/book/bookdata

In [2]: from pylab import *

In [2]: ion()

In [3]: im = imread('stinkbug.png')

@savefig mystinkbug.png width=4in
In [4]: imshow(im)
Out[4]: <matplotlib.image.AxesImage object at 0x39ea850>

""",
        r"""

In [1]: x = 'hello world'

# string methods can be
# used to alter the string
@doctest
In [2]: x.upper()
Out[2]: 'HELLO WORLD'

@verbatim
In [3]: x.st<TAB>
x.startswith  x.strip
""",
    r"""

In [130]: url = 'http://ichart.finance.yahoo.com/table.csv?s=CROX\
   .....: &d=9&e=22&f=2009&g=d&a=1&br=8&c=2006&ignore=.csv'

In [131]: print url.split('&')
['http://ichart.finance.yahoo.com/table.csv?s=CROX', 'd=9', 'e=22', 'f=2009', 'g=d', 'a=1', 'b=8', 'c=2006', 'ignore=.csv']

In [60]: import urllib

""",
    r"""\

In [133]: import numpy.random

@suppress
In [134]: numpy.random.seed(2358)

@doctest
In [135]: numpy.random.rand(10,2)
Out[135]:
array([[ 0.64524308,  0.59943846],
       [ 0.47102322,  0.8715456 ],
       [ 0.29370834,  0.74776844],
       [ 0.99539577,  0.1313423 ],
       [ 0.16250302,  0.21103583],
       [ 0.81626524,  0.1312433 ],
       [ 0.67338089,  0.72302393],
       [ 0.7566368 ,  0.07033696],
       [ 0.22591016,  0.77731835],
       [ 0.0072729 ,  0.34273127]])

""",

    r"""
In [106]: print x
jdh

In [109]: for i in range(10):
   .....:     print i
   .....:
   .....:
0
1
2
3
4
5
6
7
8
9
""",

        r"""

In [144]: from pylab import *

In [145]: ion()

# use a semicolon to suppress the output
@savefig test_hist.png width=4in
In [151]: hist(np.random.randn(10000), 100);

@savefig test_plot.png width=4in
In [151]: plot(np.random.randn(10000), 'o');
   """,

        r"""
# use a semicolon to suppress the output
In [151]: plt.clf()

@savefig plot_simple.png width=4in
In [151]: plot([1,2,3])

@savefig hist_simple.png width=4in
In [151]: hist(np.random.randn(10000), 100);

""",
     r"""
# update the current fig
In [151]: ylabel('number')

In [152]: title('normal distribution')

@savefig hist_with_text.png
In [153]: grid(True)

        """,
        ]
    # skip local-file depending first example:
    examples = examples[1:]

    #ipython_directive.DEBUG = True  # dbg
    #options = dict(suppress=True)  # dbg
    options = dict()
    for example in examples:
        content = example.split('\n')
        ipython_directive('debug', arguments=None, options=options,
                          content=content, lineno=0,
                          content_offset=None, block_text=None,
                          state=None, state_machine=None,
                          )

# Run test suite as a script
if __name__=='__main__':
    if not os.path.isdir('_static'):
        os.mkdir('_static')
    test()
    print 'All OK? Check figures in _static/'


#!/usr/bin/env python
# -*- encoding: utf-8 -*-
#
# This file is auto-generated by h2o-3/h2o-bindings/bin/gen_python.py
# Copyright 2016 H2O.ai;  Apache License Version 2.0 (see LICENSE for details)
#
from __future__ import absolute_import, division, print_function, unicode_literals

from h2o.estimators.estimator_base import H2OEstimator
from h2o.exceptions import H2OValueError
from h2o.frame import H2OFrame
from h2o.utils.typechecks import assert_is_type, Enum, numeric

class H2OSingularValueDecompositionEstimator(H2OEstimator):
    """
    Singular Value Decomposition

    """

    algo = "svd"
    supervised_learning = False

    def __init__(self,
                 model_id=None,  # type: Optional[Union[None, str, H2OEstimator]]
                 training_frame=None,  # type: Optional[Union[None, str, H2OFrame]]
                 validation_frame=None,  # type: Optional[Union[None, str, H2OFrame]]
                 ignored_columns=None,  # type: Optional[List[str]]
                 ignore_const_cols=True,  # type: bool
                 score_each_iteration=False,  # type: bool
                 transform="none",  # type: Literal["none", "standardize", "normalize", "demean", "descale"]
                 svd_method="gram_s_v_d",  # type: Literal["gram_s_v_d", "power", "randomized"]
                 nv=1,  # type: int
                 max_iterations=1000,  # type: int
                 seed=-1,  # type: int
                 keep_u=True,  # type: bool
                 u_name=None,  # type: Optional[str]
                 use_all_factor_levels=True,  # type: bool
                 max_runtime_secs=0.0,  # type: float
                 export_checkpoints_dir=None,  # type: Optional[str]
                 ):
        """
        :param model_id: Destination id for this model; auto-generated if not specified.
               Defaults to ``None``.
        :type model_id: Union[None, str, H2OEstimator], optional
        :param training_frame: Id of the training data frame.
               Defaults to ``None``.
        :type training_frame: Union[None, str, H2OFrame], optional
        :param validation_frame: Id of the validation data frame.
               Defaults to ``None``.
        :type validation_frame: Union[None, str, H2OFrame], optional
        :param ignored_columns: Names of columns to ignore for training.
               Defaults to ``None``.
        :type ignored_columns: List[str], optional
        :param ignore_const_cols: Ignore constant columns.
               Defaults to ``True``.
        :type ignore_const_cols: bool
        :param score_each_iteration: Whether to score during each iteration of model training.
               Defaults to ``False``.
        :type score_each_iteration: bool
        :param transform: Transformation of training data
               Defaults to ``"none"``.
        :type transform: Literal["none", "standardize", "normalize", "demean", "descale"]
        :param svd_method: Method for computing SVD (Caution: Randomized is currently experimental and unstable)
               Defaults to ``"gram_s_v_d"``.
        :type svd_method: Literal["gram_s_v_d", "power", "randomized"]
        :param nv: Number of right singular vectors
               Defaults to ``1``.
        :type nv: int
        :param max_iterations: Maximum iterations
               Defaults to ``1000``.
        :type max_iterations: int
        :param seed: RNG seed for k-means++ initialization
               Defaults to ``-1``.
        :type seed: int
        :param keep_u: Save left singular vectors?
               Defaults to ``True``.
        :type keep_u: bool
        :param u_name: Frame key to save left singular vectors
               Defaults to ``None``.
        :type u_name: str, optional
        :param use_all_factor_levels: Whether first factor level is included in each categorical expansion
               Defaults to ``True``.
        :type use_all_factor_levels: bool
        :param max_runtime_secs: Maximum allowed runtime in seconds for model training. Use 0 to disable.
               Defaults to ``0.0``.
        :type max_runtime_secs: float
        :param export_checkpoints_dir: Automatically export generated models to this directory.
               Defaults to ``None``.
        :type export_checkpoints_dir: str, optional
        """
        super(H2OSingularValueDecompositionEstimator, self).__init__()
        self._parms = {}
        self._id = self._parms['model_id'] = model_id
        self.training_frame = training_frame
        self.validation_frame = validation_frame
        self.ignored_columns = ignored_columns
        self.ignore_const_cols = ignore_const_cols
        self.score_each_iteration = score_each_iteration
        self.transform = transform
        self.svd_method = svd_method
        self.nv = nv
        self.max_iterations = max_iterations
        self.seed = seed
        self.keep_u = keep_u
        self.u_name = u_name
        self.use_all_factor_levels = use_all_factor_levels
        self.max_runtime_secs = max_runtime_secs
        self.export_checkpoints_dir = export_checkpoints_dir
        self._parms["_rest_version"] = 99

    @property
    def training_frame(self):
        """
        Id of the training data frame.

        Type: ``Union[None, str, H2OFrame]``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator()
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("training_frame")

    @training_frame.setter
    def training_frame(self, training_frame):
        self._parms["training_frame"] = H2OFrame._validate(training_frame, 'training_frame')

    @property
    def validation_frame(self):
        """
        Id of the validation data frame.

        Type: ``Union[None, str, H2OFrame]``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> train, valid = arrests.split_frame(ratios=[.8])
        >>> fit_h2o = H2OSingularValueDecompositionEstimator()
        >>> fit_h2o.train(x=list(range(4)),
        ...               training_frame=train,
        ...               validation_frame=valid)
        >>> fit_h2o
        """
        return self._parms.get("validation_frame")

    @validation_frame.setter
    def validation_frame(self, validation_frame):
        self._parms["validation_frame"] = H2OFrame._validate(validation_frame, 'validation_frame')

    @property
    def ignored_columns(self):
        """
        Names of columns to ignore for training.

        Type: ``List[str]``.
        """
        return self._parms.get("ignored_columns")

    @ignored_columns.setter
    def ignored_columns(self, ignored_columns):
        assert_is_type(ignored_columns, None, [str])
        self._parms["ignored_columns"] = ignored_columns

    @property
    def ignore_const_cols(self):
        """
        Ignore constant columns.

        Type: ``bool``, defaults to ``True``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(ignore_const_cols=False,
        ...                                                  nv=4)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("ignore_const_cols")

    @ignore_const_cols.setter
    def ignore_const_cols(self, ignore_const_cols):
        assert_is_type(ignore_const_cols, None, bool)
        self._parms["ignore_const_cols"] = ignore_const_cols

    @property
    def score_each_iteration(self):
        """
        Whether to score during each iteration of model training.

        Type: ``bool``, defaults to ``False``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(nv=4,
        ...                                                  score_each_iteration=True)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("score_each_iteration")

    @score_each_iteration.setter
    def score_each_iteration(self, score_each_iteration):
        assert_is_type(score_each_iteration, None, bool)
        self._parms["score_each_iteration"] = score_each_iteration

    @property
    def transform(self):
        """
        Transformation of training data

        Type: ``Literal["none", "standardize", "normalize", "demean", "descale"]``, defaults to ``"none"``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(nv=4,
        ...                                                  transform="standardize",
        ...                                                  max_iterations=2000)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("transform")

    @transform.setter
    def transform(self, transform):
        assert_is_type(transform, None, Enum("none", "standardize", "normalize", "demean", "descale"))
        self._parms["transform"] = transform

    @property
    def svd_method(self):
        """
        Method for computing SVD (Caution: Randomized is currently experimental and unstable)

        Type: ``Literal["gram_s_v_d", "power", "randomized"]``, defaults to ``"gram_s_v_d"``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(svd_method="power")
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("svd_method")

    @svd_method.setter
    def svd_method(self, svd_method):
        assert_is_type(svd_method, None, Enum("gram_s_v_d", "power", "randomized"))
        self._parms["svd_method"] = svd_method

    @property
    def nv(self):
        """
        Number of right singular vectors

        Type: ``int``, defaults to ``1``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(nv=4,
        ...                                                  transform="standardize",
        ...                                                  max_iterations=2000)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("nv")

    @nv.setter
    def nv(self, nv):
        assert_is_type(nv, None, int)
        self._parms["nv"] = nv

    @property
    def max_iterations(self):
        """
        Maximum iterations

        Type: ``int``, defaults to ``1000``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(nv=4,
        ...                                                  transform="standardize",
        ...                                                  max_iterations=2000)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("max_iterations")

    @max_iterations.setter
    def max_iterations(self, max_iterations):
        assert_is_type(max_iterations, None, int)
        self._parms["max_iterations"] = max_iterations

    @property
    def seed(self):
        """
        RNG seed for k-means++ initialization

        Type: ``int``, defaults to ``-1``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(nv=4, seed=-3)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("seed")

    @seed.setter
    def seed(self, seed):
        assert_is_type(seed, None, int)
        self._parms["seed"] = seed

    @property
    def keep_u(self):
        """
        Save left singular vectors?

        Type: ``bool``, defaults to ``True``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(keep_u=False)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("keep_u")

    @keep_u.setter
    def keep_u(self, keep_u):
        assert_is_type(keep_u, None, bool)
        self._parms["keep_u"] = keep_u

    @property
    def u_name(self):
        """
        Frame key to save left singular vectors

        Type: ``str``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(u_name="fit_h2o")
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o.u_name
        >>> fit_h2o
        """
        return self._parms.get("u_name")

    @u_name.setter
    def u_name(self, u_name):
        assert_is_type(u_name, None, str)
        self._parms["u_name"] = u_name

    @property
    def use_all_factor_levels(self):
        """
        Whether first factor level is included in each categorical expansion

        Type: ``bool``, defaults to ``True``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(use_all_factor_levels=False)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("use_all_factor_levels")

    @use_all_factor_levels.setter
    def use_all_factor_levels(self, use_all_factor_levels):
        assert_is_type(use_all_factor_levels, None, bool)
        self._parms["use_all_factor_levels"] = use_all_factor_levels

    @property
    def max_runtime_secs(self):
        """
        Maximum allowed runtime in seconds for model training. Use 0 to disable.

        Type: ``float``, defaults to ``0.0``.

        :examples:

        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(nv=4,
        ...                                                  transform="standardize",
        ...                                                  max_runtime_secs=25)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> fit_h2o
        """
        return self._parms.get("max_runtime_secs")

    @max_runtime_secs.setter
    def max_runtime_secs(self, max_runtime_secs):
        assert_is_type(max_runtime_secs, None, numeric)
        self._parms["max_runtime_secs"] = max_runtime_secs

    @property
    def export_checkpoints_dir(self):
        """
        Automatically export generated models to this directory.

        Type: ``str``.

        :examples:

        >>> import tempfile
        >>> from os import listdir
        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> checkpoints_dir = tempfile.mkdtemp()
        >>> fit_h2o = H2OSingularValueDecompositionEstimator(export_checkpoints_dir=checkpoints_dir,
        ...                                                  seed=-5)
        >>> fit_h2o.train(x=list(range(4)), training_frame=arrests)
        >>> len(listdir(checkpoints_dir))
        """
        return self._parms.get("export_checkpoints_dir")

    @export_checkpoints_dir.setter
    def export_checkpoints_dir(self, export_checkpoints_dir):
        assert_is_type(export_checkpoints_dir, None, str)
        self._parms["export_checkpoints_dir"] = export_checkpoints_dir

    def init_for_pipeline(self):
        """
        Returns H2OSVD object which implements fit and transform method to be used in sklearn.Pipeline properly.
        All parameters defined in self.__params, should be input parameters in H2OSVD.__init__ method.

        :returns: H2OSVD object

        :examples:

        >>> from h2o.transforms.preprocessing import H2OScaler
        >>> from h2o.estimators import H2ORandomForestEstimator
        >>> from h2o.estimators import H2OSingularValueDecompositionEstimator
        >>> from sklearn.pipeline import Pipeline
        >>> arrests = h2o.import_file("https://s3.amazonaws.com/h2o-public-test-data/smalldata/pca_test/USArrests.csv")
        >>> pipe = Pipeline([("standardize", H2OScaler()),
        ...                  ("svd", H2OSingularValueDecompositionEstimator(nv=3).init_for_pipeline()),
        ...                  ("rf", H2ORandomForestEstimator(seed=42,ntrees=50))])
        >>> pipe.fit(arrests[1:], arrests[0])
        """
        import inspect
        from h2o.transforms.decomposition import H2OSVD
        # check which parameters can be passed to H2OSVD init
        var_names = list(dict(inspect.getmembers(H2OSVD.__init__.__code__))['co_varnames'])
        parameters = {k: v for k, v in self._parms.items() if k in var_names}
        return H2OSVD(**parameters)


"""
The :mod:`sklearn.neighbors` module implements the k-nearest neighbors
algorithm.
"""

from .ball_tree import BallTree
from .kd_tree import KDTree
from .dist_metrics import DistanceMetric
from .graph import kneighbors_graph, radius_neighbors_graph
from .unsupervised import NearestNeighbors
from .classification import KNeighborsClassifier, RadiusNeighborsClassifier
from .regression import KNeighborsRegressor, RadiusNeighborsRegressor
from .nearest_centroid import NearestCentroid
from .kde import KernelDensity
from .approximate import LSHForest
from .lof import LocalOutlierFactor

__all__ = ['BallTree',
           'DistanceMetric',
           'KDTree',
           'KNeighborsClassifier',
           'KNeighborsRegressor',
           'NearestCentroid',
           'NearestNeighbors',
           'RadiusNeighborsClassifier',
           'RadiusNeighborsRegressor',
           'kneighbors_graph',
           'radius_neighbors_graph',
           'KernelDensity',
           'LSHForest',
           'LocalOutlierFactor']