"""
Module for formatting output data in Latex.
"""
from abc import ABC, abstractmethod
from typing import Iterator, List, Optional, Sequence, Tuple, Type, Union

import numpy as np

from pandas.core.dtypes.generic import ABCMultiIndex

from pandas.io.formats.format import DataFrameFormatter

def _split_into_full_short_caption(
    caption: Optional[Union[str, Tuple[str, str]]]
) -> Tuple[str, str]:
    """Extract full and short captions from caption string/tuple.

    Parameters
    ----------
    caption : str or tuple, optional
        Either table caption string or tuple (full_caption, short_caption).
        If string is provided, then it is treated as table full caption,
        while short_caption is considered an empty string.

    Returns
    -------
    full_caption, short_caption : tuple
        Tuple of full_caption, short_caption strings.
    """
    if caption:
        if isinstance(caption, str):
            full_caption = caption
            short_caption = ""
        else:
            try:
                full_caption, short_caption = caption
            except ValueError as err:
                msg = "caption must be either a string or a tuple of two strings"
                raise ValueError(msg) from err
    else:
        full_caption = ""
        short_caption = ""
    return full_caption, short_caption

class RowStringConverter(ABC):
    r"""Converter for dataframe rows into LaTeX strings.

    Parameters
    ----------
    formatter : `DataFrameFormatter`
        Instance of `DataFrameFormatter`.
    multicolumn: bool, optional
        Whether to use \multicolumn macro.
    multicolumn_format: str, optional
        Multicolumn format.
    multirow: bool, optional
        Whether to use \multirow macro.

    """

    def __init__(
        self,
        formatter: DataFrameFormatter,
        multicolumn: bool = False,
        multicolumn_format: Optional[str] = None,
        multirow: bool = False,
    ):
        self.fmt = formatter
        self.frame = self.fmt.frame
        self.multicolumn = multicolumn
        self.multicolumn_format = multicolumn_format
        self.multirow = multirow
        self.clinebuf: List[List[int]] = []
        self.strcols = self._get_strcols()
        self.strrows = list(zip(*self.strcols))

    def get_strrow(self, row_num: int) -> str:
        """Get string representation of the row."""
        row = self.strrows[row_num]

        is_multicol = (
            row_num < self.column_levels and self.fmt.header and self.multicolumn
        )

        is_multirow = (
            row_num >= self.header_levels
            and self.fmt.index
            and self.multirow
            and self.index_levels > 1
        )

        is_cline_maybe_required = is_multirow and row_num < len(self.strrows) - 1

        crow = self._preprocess_row(row)

        if is_multicol:
            crow = self._format_multicolumn(crow)
        if is_multirow:
            crow = self._format_multirow(crow, row_num)

        lst = []
        lst.append(" & ".join(crow))
        lst.append(" \\\\")
        if is_cline_maybe_required:
            cline = self._compose_cline(row_num, len(self.strcols))
            lst.append(cline)
        return "".join(lst)

    @property
    def _header_row_num(self) -> int:
        """Number of rows in header."""
        return self.header_levels if self.fmt.header else 0

    @property
    def index_levels(self) -> int:
        """Integer number of levels in index."""
        return self.frame.index.nlevels

    @property
    def column_levels(self) -> int:
        return self.frame.columns.nlevels

    @property
    def header_levels(self) -> int:
        nlevels = self.column_levels
        if self.fmt.has_index_names and self.fmt.show_index_names:
            nlevels += 1
        return nlevels

    def _get_strcols(self) -> List[List[str]]:
        """String representation of the columns."""
        if self.fmt.frame.empty:
            strcols = [[self._empty_info_line]]
        else:
            strcols = self.fmt.get_strcols()

        # reestablish the MultiIndex that has been joined by get_strcols()
        if self.fmt.index and isinstance(self.frame.index, ABCMultiIndex):
            out = self.frame.index.format(
                adjoin=False,
                sparsify=self.fmt.sparsify,
                names=self.fmt.has_index_names,
                na_rep=self.fmt.na_rep,
            )

            # index.format will sparsify repeated entries with empty strings
            # so pad these with some empty space
            def pad_empties(x):
                for pad in reversed(x):
                    if pad:
                        break
                return [x[0]] + [i if i else " " * len(pad) for i in x[1:]]

            gen = (pad_empties(i) for i in out)

            # Add empty spaces for each column level
            clevels = self.frame.columns.nlevels
            out = [[" " * len(i[-1])] * clevels + i for i in gen]

            # Add the column names to the last index column
            cnames = self.frame.columns.names
            if any(cnames):
                new_names = [i if i else "{}" for i in cnames]
                out[self.frame.index.nlevels - 1][:clevels] = new_names

            # Get rid of old multiindex column and add new ones
            strcols = out + strcols[1:]
        return strcols

    @property
    def _empty_info_line(self):
        return (
            f"Empty {type(self.frame).__name__}\n"
            f"Columns: {self.frame.columns}\n"
            f"Index: {self.frame.index}"
        )

    def _preprocess_row(self, row: Sequence[str]) -> List[str]:
        """Preprocess elements of the row."""
        if self.fmt.escape:
            crow = _escape_symbols(row)
        else:
            crow = [x if x else "{}" for x in row]
        if self.fmt.bold_rows and self.fmt.index:
            crow = _convert_to_bold(crow, self.index_levels)
        return crow

    def _format_multicolumn(self, row: List[str]) -> List[str]:
        r"""
        Combine columns belonging to a group to a single multicolumn entry
        according to self.multicolumn_format

        e.g.:
        a &  &  & b & c &
        will become
        \multicolumn{3}{l}{a} & b & \multicolumn{2}{l}{c}
        """
        row2 = row[: self.index_levels]
        ncol = 1
        coltext = ""

        def append_col():
            # write multicolumn if needed
            if ncol > 1:
                row2.append(
                    f"\\multicolumn{{{ncol:d}}}{{{self.multicolumn_format}}}"
                    f"{{{coltext.strip()}}}"
                )
            # don't modify where not needed
            else:
                row2.append(coltext)

        for c in row[self.index_levels :]:
            # if next col has text, write the previous
            if c.strip():
                if coltext:
                    append_col()
                coltext = c
                ncol = 1
            # if not, add it to the previous multicolumn
            else:
                ncol += 1
        # write last column name
        if coltext:
            append_col()
        return row2

    def _format_multirow(self, row: List[str], i: int) -> List[str]:
        r"""
        Check following rows, whether row should be a multirow

        e.g.:     becomes:
        a & 0 &   \multirow{2}{*}{a} & 0 &
          & 1 &     & 1 &
        b & 0 &   \cline{1-2}
                  b & 0 &
        """
        for j in range(self.index_levels):
            if row[j].strip():
                nrow = 1
                for r in self.strrows[i + 1 :]:
                    if not r[j].strip():
                        nrow += 1
                    else:
                        break
                if nrow > 1:
                    # overwrite non-multirow entry
                    row[j] = f"\\multirow{{{nrow:d}}}{{*}}{{{row[j].strip()}}}"
                    # save when to end the current block with \cline
                    self.clinebuf.append([i + nrow - 1, j + 1])
        return row

    def _compose_cline(self, i: int, icol: int) -> str:
        """
        Create clines after multirow-blocks are finished.
        """
        lst = []
        for cl in self.clinebuf:
            if cl[0] == i:
                lst.append(f"\n\\cline{{{cl[1]:d}-{icol:d}}}")
                # remove entries that have been written to buffer
                self.clinebuf = [x for x in self.clinebuf if x[0] != i]
        return "".join(lst)

class RowStringIterator(RowStringConverter):
    """Iterator over rows of the header or the body of the table."""

    @abstractmethod
    def __iter__(self) -> Iterator[str]:
        """Iterate over LaTeX string representations of rows."""

class RowHeaderIterator(RowStringIterator):
    """Iterator for the table header rows."""

    def __iter__(self) -> Iterator[str]:
        for row_num in range(len(self.strrows)):
            if row_num < self._header_row_num:
                yield self.get_strrow(row_num)

class RowBodyIterator(RowStringIterator):
    """Iterator for the table body rows."""

    def __iter__(self) -> Iterator[str]:
        for row_num in range(len(self.strrows)):
            if row_num >= self._header_row_num:
                yield self.get_strrow(row_num)

class TableBuilderAbstract(ABC):
    """
    Abstract table builder producing string representation of LaTeX table.

    Parameters
    ----------
    formatter : `DataFrameFormatter`
        Instance of `DataFrameFormatter`.
    column_format: str, optional
        Column format, for example, 'rcl' for three columns.
    multicolumn: bool, optional
        Use multicolumn to enhance MultiIndex columns.
    multicolumn_format: str, optional
        The alignment for multicolumns, similar to column_format.
    multirow: bool, optional
        Use multirow to enhance MultiIndex rows.
    caption: str, optional
        Table caption.
    short_caption: str, optional
        Table short caption.
    label: str, optional
        LaTeX label.
    position: str, optional
        Float placement specifier, for example, 'htb'.
    """

    def __init__(
        self,
        formatter: DataFrameFormatter,
        column_format: Optional[str] = None,
        multicolumn: bool = False,
        multicolumn_format: Optional[str] = None,
        multirow: bool = False,
        caption: Optional[str] = None,
        short_caption: Optional[str] = None,
        label: Optional[str] = None,
        position: Optional[str] = None,
    ):
        self.fmt = formatter
        self.column_format = column_format
        self.multicolumn = multicolumn
        self.multicolumn_format = multicolumn_format
        self.multirow = multirow
        self.caption = caption
        self.short_caption = short_caption
        self.label = label
        self.position = position

    def get_result(self) -> str:
        """String representation of LaTeX table."""
        elements = [
            self.env_begin,
            self.top_separator,
            self.header,
            self.middle_separator,
            self.env_body,
            self.bottom_separator,
            self.env_end,
        ]
        result = "\n".join([item for item in elements if item])
        trailing_newline = "\n"
        result += trailing_newline
        return result

    @property
    @abstractmethod
    def env_begin(self) -> str:
        """Beginning of the environment."""

    @property
    @abstractmethod
    def top_separator(self) -> str:
        """Top level separator."""

    @property
    @abstractmethod
    def header(self) -> str:
        """Header lines."""

    @property
    @abstractmethod
    def middle_separator(self) -> str:
        """Middle level separator."""

    @property
    @abstractmethod
    def env_body(self) -> str:
        """Environment body."""

    @property
    @abstractmethod
    def bottom_separator(self) -> str:
        """Bottom level separator."""

    @property
    @abstractmethod
    def env_end(self) -> str:
        """End of the environment."""

class GenericTableBuilder(TableBuilderAbstract):
    """Table builder producing string representation of LaTeX table."""

    @property
    def header(self) -> str:
        iterator = self._create_row_iterator(over="header")
        return "\n".join(list(iterator))

    @property
    def top_separator(self) -> str:
        return "\\toprule"

    @property
    def middle_separator(self) -> str:
        return "\\midrule" if self._is_separator_required() else ""

    @property
    def env_body(self) -> str:
        iterator = self._create_row_iterator(over="body")
        return "\n".join(list(iterator))

    def _is_separator_required(self) -> bool:
        return bool(self.header and self.env_body)

    @property
    def _position_macro(self) -> str:
        r"""Position macro, extracted from self.position, like [h]."""
        return f"[{self.position}]" if self.position else ""

    @property
    def _caption_macro(self) -> str:
        r"""Caption macro, extracted from self.caption.

        With short caption:
            \caption[short_caption]{caption_string}.

        Without short caption:
            \caption{caption_string}.
        """
        if self.caption:
            return "".join(
                [
                    r"\caption",
                    f"[{self.short_caption}]" if self.short_caption else "",
                    f"{{{self.caption}}}",
                ]
            )
        return ""

    @property
    def _label_macro(self) -> str:
        r"""Label macro, extracted from self.label, like \label{ref}."""
        return f"\\label{{{self.label}}}" if self.label else ""

    def _create_row_iterator(self, over: str) -> RowStringIterator:
        """Create iterator over header or body of the table.

        Parameters
        ----------
        over : {'body', 'header'}
            Over what to iterate.

        Returns
        -------
        RowStringIterator
            Iterator over body or header.
        """
        iterator_kind = self._select_iterator(over)
        return iterator_kind(
            formatter=self.fmt,
            multicolumn=self.multicolumn,
            multicolumn_format=self.multicolumn_format,
            multirow=self.multirow,
        )

    def _select_iterator(self, over: str) -> Type[RowStringIterator]:
        """Select proper iterator over table rows."""
        if over == "header":
            return RowHeaderIterator
        elif over == "body":
            return RowBodyIterator
        else:
            msg = f"'over' must be either 'header' or 'body', but {over} was provided"
            raise ValueError(msg)

class LongTableBuilder(GenericTableBuilder):
    """Concrete table builder for longtable.

    >>> from pandas import DataFrame
    >>> from pandas.io.formats import format as fmt
    >>> df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
    >>> formatter = fmt.DataFrameFormatter(df)
    >>> builder = LongTableBuilder(formatter, caption='a long table',
    ...                            label='tab:long', column_format='lrl')
    >>> table = builder.get_result()
    >>> print(table)
    \\begin{longtable}{lrl}
    \\caption{a long table}
    \\label{tab:long}\\\\
    \\toprule
    {} &  a &   b \\\\
    \\midrule
    \\endfirsthead
    \\caption[]{a long table} \\\\
    \\toprule
    {} &  a &   b \\\\
    \\midrule
    \\endhead
    \\midrule
    \\multicolumn{3}{r}{{Continued on next page}} \\\\
    \\midrule
    \\endfoot
    <BLANKLINE>
    \\bottomrule
    \\endlastfoot
    0 &  1 &  b1 \\\\
    1 &  2 &  b2 \\\\
    \\end{longtable}
    <BLANKLINE>
    """

    @property
    def env_begin(self) -> str:
        first_row = (
            f"\\begin{{longtable}}{self._position_macro}{{{self.column_format}}}"
        )
        elements = [first_row, f"{self._caption_and_label()}"]
        return "\n".join([item for item in elements if item])

    def _caption_and_label(self) -> str:
        if self.caption or self.label:
            double_backslash = "\\\\"
            elements = [f"{self._caption_macro}", f"{self._label_macro}"]
            caption_and_label = "\n".join([item for item in elements if item])
            caption_and_label += double_backslash
            return caption_and_label
        else:
            return ""

    @property
    def middle_separator(self) -> str:
        iterator = self._create_row_iterator(over="header")

        # the content between \endfirsthead and \endhead commands
        # mitigates repeated List of Tables entries in the final LaTeX
        # document when dealing with longtable environments; GH #34360
        elements = [
            "\\midrule",
            "\\endfirsthead",
            f"\\caption[]{{{self.caption}}} \\\\" if self.caption else "",
            self.top_separator,
            self.header,
            "\\midrule",
            "\\endhead",
            "\\midrule",
            f"\\multicolumn{{{len(iterator.strcols)}}}{{r}}"
            "{{Continued on next page}} \\\\",
            "\\midrule",
            "\\endfoot\n",
            "\\bottomrule",
            "\\endlastfoot",
        ]
        if self._is_separator_required():
            return "\n".join(elements)
        return ""

    @property
    def bottom_separator(self) -> str:
        return ""

    @property
    def env_end(self) -> str:
        return "\\end{longtable}"

class RegularTableBuilder(GenericTableBuilder):
    """Concrete table builder for regular table.

    >>> from pandas import DataFrame
    >>> from pandas.io.formats import format as fmt
    >>> df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
    >>> formatter = fmt.DataFrameFormatter(df)
    >>> builder = RegularTableBuilder(formatter, caption='caption', label='lab',
    ...                               column_format='lrc')
    >>> table = builder.get_result()
    >>> print(table)
    \\begin{table}
    \\centering
    \\caption{caption}
    \\label{lab}
    \\begin{tabular}{lrc}
    \\toprule
    {} &  a &   b \\\\
    \\midrule
    0 &  1 &  b1 \\\\
    1 &  2 &  b2 \\\\
    \\bottomrule
    \\end{tabular}
    \\end{table}
    <BLANKLINE>
    """

    @property
    def env_begin(self) -> str:
        elements = [
            f"\\begin{{table}}{self._position_macro}",
            "\\centering",
            f"{self._caption_macro}",
            f"{self._label_macro}",
            f"\\begin{{tabular}}{{{self.column_format}}}",
        ]
        return "\n".join([item for item in elements if item])

    @property
    def bottom_separator(self) -> str:
        return "\\bottomrule"

    @property
    def env_end(self) -> str:
        return "\n".join(["\\end{tabular}", "\\end{table}"])

class TabularBuilder(GenericTableBuilder):
    """Concrete table builder for tabular environment.

    >>> from pandas import DataFrame
    >>> from pandas.io.formats import format as fmt
    >>> df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
    >>> formatter = fmt.DataFrameFormatter(df)
    >>> builder = TabularBuilder(formatter, column_format='lrc')
    >>> table = builder.get_result()
    >>> print(table)
    \\begin{tabular}{lrc}
    \\toprule
    {} &  a &   b \\\\
    \\midrule
    0 &  1 &  b1 \\\\
    1 &  2 &  b2 \\\\
    \\bottomrule
    \\end{tabular}
    <BLANKLINE>
    """

    @property
    def env_begin(self) -> str:
        return f"\\begin{{tabular}}{{{self.column_format}}}"

    @property
    def bottom_separator(self) -> str:
        return "\\bottomrule"

    @property
    def env_end(self) -> str:
        return "\\end{tabular}"

class LatexFormatter:
    r"""
    Used to render a DataFrame to a LaTeX tabular/longtable environment output.

    Parameters
    ----------
    formatter : `DataFrameFormatter`
    longtable : bool, default False
        Use longtable environment.
    column_format : str, default None
        The columns format as specified in `LaTeX table format
        <https://en.wikibooks.org/wiki/LaTeX/Tables>`__ e.g 'rcl' for 3 columns
    multicolumn : bool, default False
        Use \multicolumn to enhance MultiIndex columns.
    multicolumn_format : str, default 'l'
        The alignment for multicolumns, similar to `column_format`
    multirow : bool, default False
        Use \multirow to enhance MultiIndex rows.
    caption : str or tuple, optional
        Tuple (full_caption, short_caption),
        which results in \caption[short_caption]{full_caption};
        if a single string is passed, no short caption will be set.
    label : str, optional
        The LaTeX label to be placed inside ``\label{}`` in the output.
    position : str, optional
        The LaTeX positional argument for tables, to be placed after
        ``\begin{}`` in the output.

    See Also
    --------
    HTMLFormatter
    """

    def __init__(
        self,
        formatter: DataFrameFormatter,
        longtable: bool = False,
        column_format: Optional[str] = None,
        multicolumn: bool = False,
        multicolumn_format: Optional[str] = None,
        multirow: bool = False,
        caption: Optional[Union[str, Tuple[str, str]]] = None,
        label: Optional[str] = None,
        position: Optional[str] = None,
    ):
        self.fmt = formatter
        self.frame = self.fmt.frame
        self.longtable = longtable
        self.column_format = column_format
        self.multicolumn = multicolumn
        self.multicolumn_format = multicolumn_format
        self.multirow = multirow
        self.caption, self.short_caption = _split_into_full_short_caption(caption)
        self.label = label
        self.position = position

    def to_string(self) -> str:
        """
        Render a DataFrame to a LaTeX tabular, longtable, or table/tabular
        environment output.
        """
        return self.builder.get_result()

    @property
    def builder(self) -> TableBuilderAbstract:
        """Concrete table builder.

        Returns
        -------
        TableBuilder
        """
        builder = self._select_builder()
        return builder(
            formatter=self.fmt,
            column_format=self.column_format,
            multicolumn=self.multicolumn,
            multicolumn_format=self.multicolumn_format,
            multirow=self.multirow,
            caption=self.caption,
            short_caption=self.short_caption,
            label=self.label,
            position=self.position,
        )

    def _select_builder(self) -> Type[TableBuilderAbstract]:
        """Select proper table builder."""
        if self.longtable:
            return LongTableBuilder
        if any([self.caption, self.label, self.position]):
            return RegularTableBuilder
        return TabularBuilder

    @property
    def column_format(self) -> Optional[str]:
        """Column format."""
        return self._column_format

    @column_format.setter
    def column_format(self, input_column_format: Optional[str]) -> None:
        """Setter for column format."""
        if input_column_format is None:
            self._column_format = (
                self._get_index_format() + self._get_column_format_based_on_dtypes()
            )
        elif not isinstance(input_column_format, str):
            raise ValueError(
                f"column_format must be str or unicode, "
                f"not {type(input_column_format)}"
            )
        else:
            self._column_format = input_column_format

    def _get_column_format_based_on_dtypes(self) -> str:
        """Get column format based on data type.

        Right alignment for numbers and left - for strings.
        """

        def get_col_type(dtype):
            if issubclass(dtype.type, np.number):
                return "r"
            return "l"

        dtypes = self.frame.dtypes._values
        return "".join(map(get_col_type, dtypes))

    def _get_index_format(self) -> str:
        """Get index column format."""
        return "l" * self.frame.index.nlevels if self.fmt.index else ""

def _escape_symbols(row: Sequence[str]) -> List[str]:
    """Carry out string replacements for special symbols.

    Parameters
    ----------
    row : list
        List of string, that may contain special symbols.

    Returns
    -------
    list
        list of strings with the special symbols replaced.
    """
    return [
        (
            x.replace("\\", "\\textbackslash ")
            .replace("_", "\\_")
            .replace("%", "\\%")
            .replace("$", "\\$")
            .replace("#", "\\#")
            .replace("{", "\\{")
            .replace("}", "\\}")
            .replace("~", "\\textasciitilde ")
            .replace("^", "\\textasciicircum ")
            .replace("&", "\\&")
            if (x and x != "{}")
            else "{}"
        )
        for x in row
    ]

def _convert_to_bold(crow: Sequence[str], ilevels: int) -> List[str]:
    """Convert elements in ``crow`` to bold."""
    return [
        f"\\textbf{{{x}}}" if j < ilevels and x.strip() not in ["", "{}"] else x
        for j, x in enumerate(crow)
    ]

if __name__ == "__main__":
    import doctest

    doctest.testmod()


"""
===================================================
Recursive feature elimination with cross-validation
===================================================

A recursive feature elimination example with automatic tuning of the
number of features selected with cross-validation.
"""
print(__doc__)

import matplotlib.pyplot as plt
from sklearn.svm import SVC
from sklearn.cross_validation import StratifiedKFold
from sklearn.feature_selection import RFECV
from sklearn.datasets import make_classification

# Build a classification task using 3 informative features
X, y = make_classification(n_samples=1000, n_features=25, n_informative=3,
                           n_redundant=2, n_repeated=0, n_classes=8,
                           n_clusters_per_class=1, random_state=0)

# Create the RFE object and compute a cross-validated score.
svc = SVC(kernel="linear")
# The "accuracy" scoring is proportional to the number of correct
# classifications
rfecv = RFECV(estimator=svc, step=1, cv=StratifiedKFold(y, 2),
              scoring='accuracy')
rfecv.fit(X, y)

print("Optimal number of features : %d" % rfecv.n_features_)

# Plot number of features VS. cross-validation scores
plt.figure()
plt.xlabel("Number of features selected")
plt.ylabel("Cross validation score (nb of correct classifications)")
plt.plot(range(1, len(rfecv.grid_scores_) + 1), rfecv.grid_scores_)
plt.show()


#!/usr/bin/env python
"""
refguide_check.py [OPTIONS] [-- ARGS]

Check for a PyWavelets submodule whether the objects in its __all__ dict
correspond to the objects included in the reference guide.

Example of usage::

    $ python refguide_check.py optimize

Note that this is a helper script to be able to check if things are missing;
the output of this script does need to be checked manually.  In some cases
objects are left out of the refguide for a good reason (it's an alias of
another function, or deprecated, or ...)

Another use of this helper script is to check validity of code samples
in docstrings. This is different from doctesting [we do not aim to have
scipy docstrings doctestable!], this is just to make sure that code in
docstrings is valid python::

    $ python refguide_check.py --check_docs optimize

"""
from __future__ import print_function

import sys
import os
import re
import copy
import inspect
import warnings
import doctest
import tempfile
import io
import docutils.core
from docutils.parsers.rst import directives
import shutil
import glob
from doctest import NORMALIZE_WHITESPACE, ELLIPSIS, IGNORE_EXCEPTION_DETAIL
from argparse import ArgumentParser
import numpy as np

# sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'doc',
#                 'sphinxext'))
from numpydoc.docscrape_sphinx import get_doc_object
# Remove sphinx directives that don't run without Sphinx environment
directives._directives.pop('versionadded', None)
directives._directives.pop('versionchanged', None)
directives._directives.pop('moduleauthor', None)
directives._directives.pop('sectionauthor', None)
directives._directives.pop('codeauthor', None)
directives._directives.pop('toctree', None)

BASE_MODULE = "pywt"

PUBLIC_SUBMODULES = []

# Docs for these modules are included in the parent module
OTHER_MODULE_DOCS = {}

# these names are known to fail doctesting and we like to keep it that way
# e.g. sometimes pseudocode is acceptable etc
DOCTEST_SKIPLIST = set([])

# these names are not required to be present in ALL despite being in
# autosummary:: listing
REFGUIDE_ALL_SKIPLIST = []

HAVE_MATPLOTLIB = False

def short_path(path, cwd=None):
    """
    Return relative or absolute path name, whichever is shortest.
    """
    if not isinstance(path, str):
        return path
    if cwd is None:
        cwd = os.getcwd()
    abspath = os.path.abspath(path)
    relpath = os.path.relpath(path, cwd)
    if len(abspath) <= len(relpath):
        return abspath
    return relpath

def find_names(module, names_dict):
    # Refguide entries:
    #
    # - 3 spaces followed by function name, and maybe some spaces, some
    #   dashes, and an explanation; only function names listed in
    #   refguide are formatted like this (mostly, there may be some false
    #   positives)
    #
    # - special directives, such as data and function
    #
    # - (scipy.constants only): quoted list
    #
    patterns = [
        r"^\s\s\s([a-z_0-9A-Z]+)(\s+-+.*)?$",
        r"^\.\. (?:data|function)::\s*([a-z_0-9A-Z]+)\s*$"
    ]

    if module.__name__ == 'scipy.constants':
        patterns += ["^``([a-z_0-9A-Z]+)``"]

    patterns = [re.compile(pattern) for pattern in patterns]
    module_name = module.__name__

    for line in module.__doc__.splitlines():
        res = re.search(r"^\s*\.\. (?:currentmodule|module):: ([a-z0-9A-Z_.]+)\s*$", line)
        if res:
            module_name = res.group(1)
            continue

        for pattern in patterns:
            res = re.match(pattern, line)
            if res is not None:
                name = res.group(1)
                entry = '.'.join([module_name, name])
                names_dict.setdefault(module_name, set()).add(name)
                break

def get_all_dict(module):
    """Return a copy of the __all__ dict with irrelevant items removed."""
    if hasattr(module, "__all__"):
        all_dict = copy.deepcopy(module.__all__)
    else:
        all_dict = copy.deepcopy(dir(module))
        all_dict = [name for name in all_dict
                    if not name.startswith("_")]
    for name in ['absolute_import', 'division', 'print_function']:
        try:
            all_dict.remove(name)
        except ValueError:
            pass

    # Modules are almost always private; real submodules need a separate
    # run of refguide_check.
    all_dict = [name for name in all_dict
                if not inspect.ismodule(getattr(module, name, None))]

    deprecated = []
    not_deprecated = []
    for name in all_dict:
        f = getattr(module, name, None)
        if callable(f) and is_deprecated(f):
            deprecated.append(name)
        else:
            not_deprecated.append(name)

    others = set(dir(module)).difference(set(deprecated)).difference(set(not_deprecated))

    return not_deprecated, deprecated, others

def compare(all_dict, others, names, module_name):
    """Return sets of objects only in __all__, refguide, or completely missing."""
    only_all = set()
    for name in all_dict:
        if name not in names:
            only_all.add(name)

    only_ref = set()
    missing = set()
    for name in names:
        if name not in all_dict:
            for pat in REFGUIDE_ALL_SKIPLIST:
                if re.match(pat, module_name + '.' + name):
                    if name not in others:
                        missing.add(name)
                    break
            else:
                only_ref.add(name)

    return only_all, only_ref, missing

def is_deprecated(f):
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("error")
        try:
            f(**{"not a kwarg": None})
        except DeprecationWarning:
            return True
        except:
            pass
        return False

def check_items(all_dict, names, deprecated, others, module_name, dots=True):
    num_all = len(all_dict)
    num_ref = len(names)

    output = ""

    output += "Non-deprecated objects in __all__: %i\n" % num_all
    output += "Objects in refguide: %i\n\n" % num_ref

    only_all, only_ref, missing = compare(all_dict, others, names, module_name)
    dep_in_ref = set(only_ref).intersection(deprecated)
    only_ref = set(only_ref).difference(deprecated)

    if len(dep_in_ref) > 0:
        output += "Deprecated objects in refguide::\n\n"
        for name in sorted(deprecated):
            output += "    " + name + "\n"

    if len(only_all) == len(only_ref) == len(missing) == 0:
        if dots:
            output_dot('.')
        return [(None, True, output)]
    else:
        if len(only_all) > 0:
            output += "ERROR: objects in %s.__all__ but not in refguide::\n\n" % module_name
            for name in sorted(only_all):
                output += "    " + name + "\n"

        if len(only_ref) > 0:
            output += "ERROR: objects in refguide but not in %s.__all__::\n\n" % module_name
            for name in sorted(only_ref):
                output += "    " + name + "\n"

        if len(missing) > 0:
            output += "ERROR: missing objects::\n\n"
            for name in sorted(missing):
                output += "    " + name + "\n"

        if dots:
            output_dot('F')
        return [(None, False, output)]

def validate_rst_syntax(text, name, dots=True):
    if text is None:
        if dots:
            output_dot('E')
        return False, "ERROR: %s: no documentation" % (name,)

    ok_unknown_items = set([
        'mod', 'currentmodule', 'autosummary', 'data',
        'obj', 'versionadded', 'versionchanged', 'module', 'class',
        'ref', 'func', 'toctree', 'moduleauthor',
        'sectionauthor', 'codeauthor', 'eq',
    ])

    # Run through docutils
    error_stream = io.StringIO()

    def resolve(name, is_label=False):
        return ("http://foo", name)

    token = '<RST-VALIDATE-SYNTAX-CHECK>'

    docutils.core.publish_doctree(
        text, token,
        settings_overrides = dict(halt_level=5,
                                  traceback=True,
                                  default_reference_context='title-reference',
                                  default_role='emphasis',
                                  link_base='',
                                  resolve_name=resolve,
                                  stylesheet_path='',
                                  raw_enabled=0,
                                  file_insertion_enabled=0,
                                  warning_stream=error_stream))

    # Print errors, disregarding unimportant ones
    error_msg = error_stream.getvalue()
    errors = error_msg.split(token)
    success = True
    output = ""

    for error in errors:
        lines = error.splitlines()
        if not lines:
            continue

        m = re.match(r'.*Unknown (?:interpreted text role|directive type) "(.*)".*$', lines[0])
        if m:
            if m.group(1) in ok_unknown_items:
                continue

        m = re.match(r'.*Error in "math" directive:.*unknown option: "label"', " ".join(lines), re.S)
        if m:
            continue

        output += name + lines[0] + "::\n    " + "\n    ".join(lines[1:]).rstrip() + "\n"
        success = False

    if not success:
        output += "    " + "-"*72 + "\n"
        for lineno, line in enumerate(text.splitlines()):
            output += "    %-4d    %s\n" % (lineno+1, line)
        output += "    " + "-"*72 + "\n\n"

    if dots:
        output_dot('.' if success else 'F')
    return success, output

def output_dot(msg='.', stream=sys.stderr):
    stream.write(msg)
    stream.flush()

def check_rest(module, names, dots=True):
    """
    Check reStructuredText formatting of docstrings

    Returns: [(name, success_flag, output), ...]
    """

    try:
        skip_types = (dict, str, unicode, float, int)
    except NameError:
        # python 3
        skip_types = (dict, str, float, int)

    results = []

    if module.__name__[6:] not in OTHER_MODULE_DOCS:
        results += [(module.__name__,) +
                    validate_rst_syntax(inspect.getdoc(module),
                                        module.__name__, dots=dots)]

    for name in names:
        full_name = module.__name__ + '.' + name
        obj = getattr(module, name, None)

        if obj is None:
            results.append((full_name, False, "%s has no docstring" % (full_name,)))
            continue
        elif isinstance(obj, skip_types):
            continue

        if inspect.ismodule(obj):
            text = inspect.getdoc(obj)
        else:
            try:
                text = str(get_doc_object(obj))
            except:
                import traceback
                results.append((full_name, False,
                                "Error in docstring format!\n" +
                                traceback.format_exc()))
                continue

        m = re.search("([\x00-\x09\x0b-\x1f])", text)
        if m:
            msg = ("Docstring contains a non-printable character %r! "
                   "Maybe forgot r\"\"\"?" % (m.group(1),))
            results.append((full_name, False, msg))
            continue

        try:
            src_file = short_path(inspect.getsourcefile(obj))
        except TypeError:
            src_file = None

        if src_file:
            file_full_name = src_file + ':' + full_name
        else:
            file_full_name = full_name

        results.append((full_name,) +
                       validate_rst_syntax(text, file_full_name, dots=dots))

    return results

### Doctest helpers ####

# the namespace to run examples in
DEFAULT_NAMESPACE = {'np': np}

# the namespace to do checks in
CHECK_NAMESPACE = {
      'np': np,
      'assert_allclose': np.testing.assert_allclose,
      'assert_equal': np.testing.assert_equal,
      # recognize numpy repr's
      'array': np.array,
      'matrix': np.matrix,
      'int64': np.int64,
      'uint64': np.uint64,
      'int8': np.int8,
      'int32': np.int32,
      'float64': np.float64,
      'dtype': np.dtype,
      'nan': np.nan,
      'NaN': np.nan,
      'inf': np.inf,
      'Inf': np.inf, }

class DTRunner(doctest.DocTestRunner):
    DIVIDER = "\n"

    def __init__(self, item_name, checker=None, verbose=None, optionflags=0):
        self._item_name = item_name
        doctest.DocTestRunner.__init__(self, checker=checker, verbose=verbose,
                                       optionflags=optionflags)

    def _report_item_name(self, out, new_line=False):
        if self._item_name is not None:
            if new_line:
                out("\n")
            self._item_name = None

    def report_start(self, out, test, example):
        self._checker._source = example.source
        return doctest.DocTestRunner.report_start(self, out, test, example)

    def report_success(self, out, test, example, got):
        if self._verbose:
            self._report_item_name(out, new_line=True)
        return doctest.DocTestRunner.report_success(
            self, out, test, example, got)

    def report_unexpected_exception(self, out, test, example, exc_info):
        self._report_item_name(out)
        return doctest.DocTestRunner.report_unexpected_exception(
            self, out, test, example, exc_info)

    def report_failure(self, out, test, example, got):
        self._report_item_name(out)
        return doctest.DocTestRunner.report_failure(self, out, test,
                                                    example, got)

class Checker(doctest.OutputChecker):
    obj_pattern = re.compile('at 0x[0-9a-fA-F]+>')
    vanilla = doctest.OutputChecker()
    rndm_markers = {'# random', '# Random', '#random', '#Random', "# may vary"}
    stopwords = {'plt.', '.hist', '.show', '.ylim', '.subplot(',
                 'set_title', 'imshow', 'plt.show', 'ax.axis', 'plt.plot(',
                 '.bar(', '.title', '.ylabel', '.xlabel', 'set_ylim',
                 'set_xlim', '# reformatted'}

    def __init__(self, parse_namedtuples=True, ns=None, atol=1e-8, rtol=1e-2):
        self.parse_namedtuples = parse_namedtuples
        self.atol, self.rtol = atol, rtol
        if ns is None:
            self.ns = dict(CHECK_NAMESPACE)
        else:
            self.ns = ns

    def check_output(self, want, got, optionflags):
        # cut it short if they are equal
        if want == got:
            return True

        # skip stopwords in source
        if any(word in self._source for word in self.stopwords):
            return True

        # skip random stuff
        if any(word in want for word in self.rndm_markers):
            return True

        # skip function/object addresses
        if self.obj_pattern.search(got):
            return True

        # ignore comments (e.g. signal.freqresp)
        if want.lstrip().startswith("#"):
            return True

        # try the standard doctest
        try:
            if self.vanilla.check_output(want, got, optionflags):
                return True
        except Exception:
            pass

        # OK then, convert strings to objects
        try:
            a_want = eval(want, dict(self.ns))
            a_got = eval(got, dict(self.ns))
        except:
            if not self.parse_namedtuples:
                return False
            # suppose that "want"  is a tuple, and "got" is smth like
            # MoodResult(statistic=10, pvalue=0.1).
            # Then convert the latter to the tuple (10, 0.1),
            # and then compare the tuples.
            try:
                num = len(a_want)
                regex = ('[\w\d_]+\(' +
                         ', '.join(['[\w\d_]+=(.+)']*num) +
                         '\)')
                grp = re.findall(regex, got.replace('\n', ' '))
                if len(grp) > 1:  # no more than one for now
                    return False
                # fold it back to a tuple
                got_again = '(' + ', '.join(grp[0]) + ')'
                return self.check_output(want, got_again, optionflags)
            except Exception:
                return False

        # ... and defer to numpy
        try:
            return self._do_check(a_want, a_got)
        except Exception:
            # heterog tuple, eg (1, np.array([1., 2.]))
           try:
                return all(self._do_check(w, g) for w, g in zip(a_want, a_got))
           except (TypeError, ValueError):
                return False

    def _do_check(self, want, got):
        # This should be done exactly as written to correctly handle all of
        # numpy-comparable objects, strings, and heterogenous tuples
        try:
            if want == got:
                return True
        except Exception:
            pass
        return np.allclose(want, got, atol=self.atol, rtol=self.rtol)

def _run_doctests(tests, full_name, verbose, doctest_warnings):
    """Run modified doctests for the set of `tests`.

    Returns: list of [(success_flag, output), ...]
    """
    flags = NORMALIZE_WHITESPACE | ELLIPSIS | IGNORE_EXCEPTION_DETAIL
    runner = DTRunner(full_name, checker=Checker(), optionflags=flags,
                      verbose=verbose)

    output = []
    success = True
    def out(msg):
        output.append(msg)

    class MyStderr(object):
        """Redirect stderr to the current stdout"""
        def write(self, msg):
            if doctest_warnings:
                sys.stdout.write(msg)
            else:
                out(msg)

    # Run tests, trying to restore global state afterward
    old_printoptions = np.get_printoptions()
    old_errstate = np.seterr()
    old_stderr = sys.stderr
    cwd = os.getcwd()
    tmpdir = tempfile.mkdtemp()
    sys.stderr = MyStderr()
    try:
        os.chdir(tmpdir)

        # try to ensure random seed is NOT reproducible
        np.random.seed(None)

        for t in tests:
            t.filename = short_path(t.filename, cwd)
            fails, successes = runner.run(t, out=out)
            if fails > 0:
                success = False
    finally:
        sys.stderr = old_stderr
        os.chdir(cwd)
        shutil.rmtree(tmpdir)
        np.set_printoptions(**old_printoptions)
        np.seterr(**old_errstate)

    return success, output

def check_doctests(module, verbose, ns=None,
                   dots=True, doctest_warnings=False):
    """Check code in docstrings of the module's public symbols.

    Returns: list of [(item_name, success_flag, output), ...]
    """
    if ns is None:
        ns = dict(DEFAULT_NAMESPACE)

    # Loop over non-deprecated items
    results = []

    for name in get_all_dict(module)[0]:
        full_name = module.__name__ + '.' + name

        if full_name in DOCTEST_SKIPLIST:
            continue

        try:
            obj = getattr(module, name)
        except AttributeError:
            import traceback
            results.append((full_name, False,
                            "Missing item!\n" +
                            traceback.format_exc()))
            continue

        finder = doctest.DocTestFinder()
        try:
            tests = finder.find(obj, name, globs=dict(ns))
        except:
            import traceback
            results.append((full_name, False,
                            "Failed to get doctests!\n" +
                            traceback.format_exc()))
            continue

        success, output = _run_doctests(tests, full_name, verbose,
                                        doctest_warnings)

        if dots:
            output_dot('.' if success else 'F')

        results.append((full_name, success, "".join(output)))

        if HAVE_MATPLOTLIB:
            import matplotlib.pyplot as plt
            plt.close('all')

    return results

def check_doctests_testfile(fname, verbose, ns=None,
                            dots=True, doctest_warnings=False):
    """Check code in a text file.

    Mimic `check_doctests` above, differing mostly in test discovery.
    (which is borrowed from stdlib's doctest.testfile here,
     https://github.com/python-git/python/blob/master/Lib/doctest.py)

    Returns: list of [(item_name, success_flag, output), ...]

    Notes
    -----

    We also try to weed out pseudocode:
    * We maintain a list of exceptions which signal pseudocode,
    * We split the text file into "blocks" of code separated by empty lines
      and/or intervening text.
    * If a block contains a marker, the whole block is then assumed to be
      pseudocode. It is then not being doctested.

    The rationale is that typically, the text looks like this:

    blah
    <BLANKLINE>
    >>> from numpy import some_module   # pseudocode!
    >>> func = some_module.some_function
    >>> func(42)                  # still pseudocode
    146
    <BLANKLINE>
    blah
    <BLANKLINE>
    >>> 2 + 3        # real code, doctest it
    5

    """
    results = []

    if ns is None:
        ns = dict(DEFAULT_NAMESPACE)

    _, short_name = os.path.split(fname)
    if short_name in DOCTEST_SKIPLIST:
        return results

    full_name = fname
    text = open(fname).read()

    PSEUDOCODE = set(['some_function', 'some_module', 'import example',
                      'ctypes.CDLL',     # likely need compiling, skip it
                      'integrate.nquad(func,'  # ctypes integrate tutotial
    ])

    # split the text into "blocks" and try to detect and omit pseudocode blocks.
    parser = doctest.DocTestParser()
    good_parts = []
    for part in text.split('\n\n'):
        tests = parser.get_doctest(part, ns, fname, fname, 0)
        if any(word in ex.source for word in PSEUDOCODE
                                 for ex in tests.examples):
            # omit it
            pass
        else:
            # `part` looks like a good code, let's doctest it
            good_parts += [part]

    # Reassemble the good bits and doctest them:
    good_text = '\n\n'.join(good_parts)
    tests = parser.get_doctest(good_text, ns, fname, fname, 0)
    success, output = _run_doctests([tests], full_name, verbose,
                                    doctest_warnings)

    if dots:
        output_dot('.' if success else 'F')

    results.append((full_name, success, "".join(output)))

    if HAVE_MATPLOTLIB:
        import matplotlib.pyplot as plt
        plt.close('all')

    return results

def init_matplotlib():
    global HAVE_MATPLOTLIB

    try:
        import matplotlib
        matplotlib.use('Agg')
        HAVE_MATPLOTLIB = True
    except ImportError:
        HAVE_MATPLOTLIB = False

def main(argv):
    parser = ArgumentParser(usage=__doc__.lstrip())
    parser.add_argument("module_names", metavar="SUBMODULES", default=[],
                        nargs='*',
                        help="Submodules to check (default: all public)")
    parser.add_argument("--doctests", action="store_true",
                        help="Run also doctests")
    parser.add_argument("-v", "--verbose", action="count", default=0)
    parser.add_argument("--doctest-warnings", action="store_true",
                        help="Enforce warning checking for doctests")
    parser.add_argument("--skip-examples", action="store_true",
                        help="Skip running doctests in the examples.")
    args = parser.parse_args(argv)

    modules = []
    names_dict = {}

    if args.module_names:
        args.skip_examples = True
    else:
        args.module_names = list(PUBLIC_SUBMODULES)

    os.environ['SCIPY_PIL_IMAGE_VIEWER'] = 'true'

    module_names = list(args.module_names)
    for name in list(module_names):
        if name in OTHER_MODULE_DOCS:
            name = OTHER_MODULE_DOCS[name]
            if name not in module_names:
                module_names.append(name)

    for submodule_name in module_names:
        module_name = BASE_MODULE + '.' + submodule_name
        __import__(module_name)
        module = sys.modules[module_name]

        if submodule_name not in OTHER_MODULE_DOCS:
            find_names(module, names_dict)

        if submodule_name in args.module_names:
            modules.append(module)

    dots = True
    success = True
    results = []

    print("Running checks for %d modules:" % (len(modules),))

    if args.doctests or not args.skip_examples:
        init_matplotlib()

    for module in modules:
        if dots:
            if module is not modules[0]:
                sys.stderr.write(' ')
            sys.stderr.write(module.__name__ + ' ')
            sys.stderr.flush()

        all_dict, deprecated, others = get_all_dict(module)
        names = names_dict.get(module.__name__, set())

        mod_results = []
        mod_results += check_items(all_dict, names, deprecated, others, module.__name__)
        mod_results += check_rest(module, set(names).difference(deprecated),
                                  dots=dots)
        if args.doctests:
            mod_results += check_doctests(module, (args.verbose >= 2), dots=dots,
                                          doctest_warnings=args.doctest_warnings)

        for v in mod_results:
            assert isinstance(v, tuple), v

        results.append((module, mod_results))

    if dots:
        sys.stderr.write("\n")
        sys.stderr.flush()

    if not args.skip_examples:
        examples_path = os.path.join(
            os.getcwd(), 'doc', 'source', 'regression', '*.rst')
        print('\nChecking examples files at %s:' % examples_path)
        for filename in sorted(glob.glob(examples_path)):
            if dots:
                sys.stderr.write('\n')
                sys.stderr.write(os.path.split(filename)[1] + ' ')
                sys.stderr.flush()

            examples_results = check_doctests_testfile(
                filename, (args.verbose >= 2), dots=dots,
                doctest_warnings=args.doctest_warnings)

            def scratch(): pass        # stub out a "module", see below
            scratch.__name__ = filename
            results.append((scratch, examples_results))

        if dots:
            sys.stderr.write("\n")
            sys.stderr.flush()

    # Report results
    all_success = True

    for module, mod_results in results:
        success = all(x[1] for x in mod_results)
        all_success = all_success and success

        if success and args.verbose == 0:
            continue

        print("")
        print("=" * len(module.__name__))
        print(module.__name__)
        print("=" * len(module.__name__))
        print("")

        for name, success, output in mod_results:
            if name is None:
                if not success or args.verbose >= 1:
                    print(output.strip())
                    print("")
            elif not success or (args.verbose >= 2 and output.strip()):
                print(name)
                print("-"*len(name))
                print("")
                print(output.strip())
                print("")

    if all_success:
        print("\nOK: refguide and doctests checks passed!")
        sys.exit(0)
    else:
        print("\nERROR: refguide or doctests have errors")
        sys.exit(1)

if __name__ == '__main__':
    main(argv=sys.argv[1:])


from test_utils import PySparkTestCase
from sparkts.timeseriesrdd import *
from sparkts.timeseriesrdd import _TimeSeriesSerializer
from sparkts.datetimeindex import *
import pandas as pd
import numpy as np
from unittest import TestCase
from io import BytesIO
from pyspark.sql import SQLContext

class TimeSeriesSerializerTestCase(TestCase):
    def test_times_series_serializer(self):
        serializer = _TimeSeriesSerializer()
        stream = BytesIO()
        series = [('abc', np.array([4.0, 4.0, 5.0])), ('123', np.array([1.0, 2.0, 3.0]))]
        serializer.dump_stream(iter(series), stream)
        stream.seek(0)
        reconstituted = list(serializer.load_stream(stream))
        self.assertEquals(reconstituted[0][0], series[0][0])
        self.assertEquals(reconstituted[1][0], series[1][0])
        self.assertTrue((reconstituted[0][1] == series[0][1]).all())
        self.assertTrue((reconstituted[1][1] == series[1][1]).all())

class TimeSeriesRDDTestCase(PySparkTestCase):
    def test_time_series_rdd(self):
        freq = DayFrequency(1, self.sc)
        start = '2015-04-09'
        dt_index = uniform(start, periods=10, freq=freq, sc=self.sc)
        vecs = [np.arange(0, 10), np.arange(10, 20), np.arange(20, 30)]
        rdd = self.sc.parallelize(vecs).map(lambda x: (str(x[0]), x))
        tsrdd = TimeSeriesRDD(dt_index, rdd)
        self.assertEquals(tsrdd.count(), 3)

        contents = tsrdd.collectAsMap()
        self.assertEquals(len(contents), 3)
        self.assertTrue((contents["0"] == np.arange(0, 10)).all())
        self.assertTrue((contents["10"] == np.arange(10, 20)).all())
        self.assertTrue((contents["20"] == np.arange(20, 30)).all())

        subslice = tsrdd['2015-04-10':'2015-04-15']
        self.assertEquals(subslice.index(), uniform('2015-04-10', periods=6, freq=freq, sc=self.sc))
        contents = subslice.collectAsMap()
        self.assertEquals(len(contents), 3)
        self.assertTrue((contents["0"] == np.arange(1, 7)).all())
        self.assertTrue((contents["10"] == np.arange(11, 17)).all())
        self.assertTrue((contents["20"] == np.arange(21, 27)).all())

    def test_to_instants(self):
        vecs = [np.arange(x, x + 4) for x in np.arange(0, 20, 4)]
        labels = ['a', 'b', 'c', 'd', 'e']
        start = '2015-4-9'
        dt_index = uniform(start, periods=4, freq=DayFrequency(1, self.sc), sc=self.sc)
        rdd = self.sc.parallelize(zip(labels, vecs), 3)
        tsrdd = TimeSeriesRDD(dt_index, rdd)
        samples = tsrdd.to_instants().collect()
        target_dates = ['2015-4-9', '2015-4-10', '2015-4-11', '2015-4-12']
        self.assertEquals([x[0] for x in samples], [pd.Timestamp(x) for x in target_dates])
        self.assertTrue((samples[0][1] == np.arange(0, 20, 4)).all())
        self.assertTrue((samples[1][1] == np.arange(1, 20, 4)).all())
        self.assertTrue((samples[2][1] == np.arange(2, 20, 4)).all())
        self.assertTrue((samples[3][1] == np.arange(3, 20, 4)).all())

    def test_to_observations(self):
        sql_ctx = SQLContext(self.sc)
        vecs = [np.arange(x, x + 4) for x in np.arange(0, 20, 4)]
        labels = ['a', 'b', 'c', 'd', 'e']
        start = '2015-4-9'
        dt_index = uniform(start, periods=4, freq=DayFrequency(1, self.sc), sc=self.sc)
        print(dt_index._jdt_index.size())
        rdd = self.sc.parallelize(zip(labels, vecs), 3)
        tsrdd = TimeSeriesRDD(dt_index, rdd)

        obsdf = tsrdd.to_observations_dataframe(sql_ctx)
        tsrdd_from_df = time_series_rdd_from_observations( \
            dt_index, obsdf, 'timestamp', 'key', 'value')

        ts1 = tsrdd.collect()
        ts1.sort(key = lambda x: x[0])
        ts2 = tsrdd_from_df.collect()
        ts2.sort(key = lambda x: x[0])
        self.assertTrue(all([pair[0][0] == pair[1][0] and (pair[0][1] == pair[1][1]).all() \
            for pair in zip(ts1, ts2)]))

        df1 = obsdf.collect()
        df1.sort(key = lambda x: x.value)
        df2 = tsrdd_from_df.to_observations_dataframe(sql_ctx).collect()
        df2.sort(key = lambda x: x.value)
        self.assertEquals(df1, df2)

    def test_filter(self):
        vecs = [np.arange(x, x + 4) for x in np.arange(0, 20, 4)]
        labels = ['a', 'b', 'c', 'd', 'e']
        start = '2015-4-9'
        dt_index = uniform(start, periods=4, freq=DayFrequency(1, self.sc), sc=self.sc)
        rdd = self.sc.parallelize(zip(labels, vecs), 3)
        tsrdd = TimeSeriesRDD(dt_index, rdd)
        filtered = tsrdd.filter(lambda x: x[0] == 'a' or x[0] == 'b')
        self.assertEquals(filtered.count(), 2)
        # assert it has TimeSeriesRDD functionality:
        filtered['2015-04-10':'2015-04-15'].count()

    def test_to_pandas_series_rdd(self):
        vecs = [np.arange(x, x + 4) for x in np.arange(0, 20, 4)]
        labels = ['a', 'b', 'c', 'd', 'e']
        start = '2015-4-9'
        dt_index = uniform(start, periods=4, freq=DayFrequency(1, self.sc), sc=self.sc)
        rdd = self.sc.parallelize(zip(labels, vecs), 3)
        tsrdd = TimeSeriesRDD(dt_index, rdd)

        series_arr = tsrdd.to_pandas_series_rdd().collect()

        pd_index = dt_index.to_pandas_index()
        self.assertEquals(len(vecs), len(series_arr))
        for i in xrange(len(vecs)):
            self.assertEquals(series_arr[i][0], labels[i])
            self.assertTrue(pd.Series(vecs[i], pd_index).equals(series_arr[i][1]))