#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from __future__ import print_function
import sys
import warnings
from functools import reduce
from threading import RLock

if sys.version >= '3':
    basestring = unicode = str
else:
    from itertools import imap as map

from pyspark import since
from pyspark.rdd import RDD, ignore_unicode_prefix
from pyspark.sql.catalog import Catalog
from pyspark.sql.conf import RuntimeConfig
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.readwriter import DataFrameReader
from pyspark.sql.streaming import DataStreamReader
from pyspark.sql.types import Row, DataType, StringType, StructType, _verify_type, \
    _infer_schema, _has_nulltype, _merge_type, _create_converter, _parse_datatype_string
from pyspark.sql.utils import install_exception_handler

__all__ = ["SparkSession"]

def _monkey_patch_RDD(sparkSession):
    def toDF(self, schema=None, sampleRatio=None):
        """
        Converts current :class:`RDD` into a :class:`DataFrame`

        This is a shorthand for ``spark.createDataFrame(rdd, schema, sampleRatio)``

        :param schema: a :class:`pyspark.sql.types.StructType` or list of names of columns
        :param samplingRatio: the sample ratio of rows used for inferring
        :return: a DataFrame

        >>> rdd.toDF().collect()
        [Row(name=u'Alice', age=1)]
        """
        return sparkSession.createDataFrame(self, schema, sampleRatio)

    RDD.toDF = toDF

class SparkSession(object):
    """The entry point to programming Spark with the Dataset and DataFrame API.

    A SparkSession can be used create :class:`DataFrame`, register :class:`DataFrame` as
    tables, execute SQL over tables, cache tables, and read parquet files.
    To create a SparkSession, use the following builder pattern:

    >>> spark = SparkSession.builder \\
    ...     .master("local") \\
    ...     .appName("Word Count") \\
    ...     .config("spark.some.config.option", "some-value") \\
    ...     .getOrCreate()
    """

    class Builder(object):
        """Builder for :class:`SparkSession`.
        """

        _lock = RLock()
        _options = {}

        @since(2.0)
        def config(self, key=None, value=None, conf=None):
            """Sets a config option. Options set using this method are automatically propagated to
            both :class:`SparkConf` and :class:`SparkSession`'s own configuration.

            For an existing SparkConf, use `conf` parameter.

            >>> from pyspark.conf import SparkConf
            >>> SparkSession.builder.config(conf=SparkConf())
            <pyspark.sql.session...

            For a (key, value) pair, you can omit parameter names.

            >>> SparkSession.builder.config("spark.some.config.option", "some-value")
            <pyspark.sql.session...

            :param key: a key name string for configuration property
            :param value: a value for configuration property
            :param conf: an instance of :class:`SparkConf`
            """
            with self._lock:
                if conf is None:
                    self._options[key] = str(value)
                else:
                    for (k, v) in conf.getAll():
                        self._options[k] = v
                return self

        @since(2.0)
        def master(self, master):
            """Sets the Spark master URL to connect to, such as "local" to run locally, "local[4]"
            to run locally with 4 cores, or "spark://master:7077" to run on a Spark standalone
            cluster.

            :param master: a url for spark master
            """
            return self.config("spark.master", master)

        @since(2.0)
        def appName(self, name):
            """Sets a name for the application, which will be shown in the Spark web UI.

            If no application name is set, a randomly generated name will be used.

            :param name: an application name
            """
            return self.config("spark.app.name", name)

        @since(2.0)
        def enableHiveSupport(self):
            """Enables Hive support, including connectivity to a persistent Hive metastore, support
            for Hive serdes, and Hive user-defined functions.
            """
            return self.config("spark.sql.catalogImplementation", "hive")

        @since(2.0)
        def getOrCreate(self):
            """Gets an existing :class:`SparkSession` or, if there is no existing one, creates a
            new one based on the options set in this builder.

            This method first checks whether there is a valid global default SparkSession, and if
            yes, return that one. If no valid global default SparkSession exists, the method
            creates a new SparkSession and assigns the newly created SparkSession as the global
            default.

            >>> s1 = SparkSession.builder.config("k1", "v1").getOrCreate()
            >>> s1.conf.get("k1") == s1.sparkContext.getConf().get("k1") == "v1"
            True

            In case an existing SparkSession is returned, the config options specified
            in this builder will be applied to the existing SparkSession.

            >>> s2 = SparkSession.builder.config("k2", "v2").getOrCreate()
            >>> s1.conf.get("k1") == s2.conf.get("k1")
            True
            >>> s1.conf.get("k2") == s2.conf.get("k2")
            True
            """
            with self._lock:
                from pyspark.context import SparkContext
                from pyspark.conf import SparkConf
                session = SparkSession._instantiatedContext
                if session is None:
                    sparkConf = SparkConf()
                    for key, value in self._options.items():
                        sparkConf.set(key, value)
                    sc = SparkContext.getOrCreate(sparkConf)
                    # This SparkContext may be an existing one.
                    for key, value in self._options.items():
                        # we need to propagate the confs
                        # before we create the SparkSession. Otherwise, confs like
                        # warehouse path and metastore url will not be set correctly (
                        # these confs cannot be changed once the SparkSession is created).
                        sc._conf.set(key, value)
                    session = SparkSession(sc)
                for key, value in self._options.items():
                    session._jsparkSession.sessionState().conf().setConfString(key, value)
                for key, value in self._options.items():
                    session.sparkContext._conf.set(key, value)
                return session

    builder = Builder()

    _instantiatedContext = None

    @ignore_unicode_prefix
    def __init__(self, sparkContext, jsparkSession=None):
        """Creates a new SparkSession.

        >>> from datetime import datetime
        >>> spark = SparkSession(sc)
        >>> allTypes = sc.parallelize([Row(i=1, s="string", d=1.0, l=1,
        ...     b=True, list=[1, 2, 3], dict={"s": 0}, row=Row(a=1),
        ...     time=datetime(2014, 8, 1, 14, 1, 5))])
        >>> df = allTypes.toDF()
        >>> df.createOrReplaceTempView("allTypes")
        >>> spark.sql('select i+1, d+1, not b, list[1], dict["s"], time, row.a '
        ...            'from allTypes where b and i > 0').collect()
        [Row((i + CAST(1 AS BIGINT))=2, (d + CAST(1 AS DOUBLE))=2.0, (NOT b)=False, list[1]=2, \
            dict[s]=0, time=datetime.datetime(2014, 8, 1, 14, 1, 5), a=1)]
        >>> df.rdd.map(lambda x: (x.i, x.s, x.d, x.l, x.b, x.time, x.row.a, x.list)).collect()
        [(1, u'string', 1.0, 1, True, datetime.datetime(2014, 8, 1, 14, 1, 5), 1, [1, 2, 3])]
        """
        from pyspark.sql.context import SQLContext
        self._sc = sparkContext
        self._jsc = self._sc._jsc
        self._jvm = self._sc._jvm
        if jsparkSession is None:
            jsparkSession = self._jvm.SparkSession(self._jsc.sc())
        self._jsparkSession = jsparkSession
        self._jwrapped = self._jsparkSession.sqlContext()
        self._wrapped = SQLContext(self._sc, self, self._jwrapped)
        _monkey_patch_RDD(self)
        install_exception_handler()
        if SparkSession._instantiatedContext is None:
            SparkSession._instantiatedContext = self

    @since(2.0)
    def newSession(self):
        """
        Returns a new SparkSession as new session, that has separate SQLConf,
        registered temporary views and UDFs, but shared SparkContext and
        table cache.
        """
        return self.__class__(self._sc, self._jsparkSession.newSession())

    @property
    @since(2.0)
    def sparkContext(self):
        """Returns the underlying :class:`SparkContext`."""
        return self._sc

    @property
    @since(2.0)
    def version(self):
        """The version of Spark on which this application is running."""
        return self._jsparkSession.version()

    @property
    @since(2.0)
    def conf(self):
        """Runtime configuration interface for Spark.

        This is the interface through which the user can get and set all Spark and Hadoop
        configurations that are relevant to Spark SQL. When getting the value of a config,
        this defaults to the value set in the underlying :class:`SparkContext`, if any.
        """
        if not hasattr(self, "_conf"):
            self._conf = RuntimeConfig(self._jsparkSession.conf())
        return self._conf

    @property
    @since(2.0)
    def catalog(self):
        """Interface through which the user may create, drop, alter or query underlying
        databases, tables, functions etc.
        """
        if not hasattr(self, "_catalog"):
            self._catalog = Catalog(self)
        return self._catalog

    @property
    @since(2.0)
    def udf(self):
        """Returns a :class:`UDFRegistration` for UDF registration.

        :return: :class:`UDFRegistration`
        """
        from pyspark.sql.context import UDFRegistration
        return UDFRegistration(self._wrapped)

    @since(2.0)
    def range(self, start, end=None, step=1, numPartitions=None):
        """
        Create a :class:`DataFrame` with single :class:`pyspark.sql.types.LongType` column named
        ``id``, containing elements in a range from ``start`` to ``end`` (exclusive) with
        step value ``step``.

        :param start: the start value
        :param end: the end value (exclusive)
        :param step: the incremental step (default: 1)
        :param numPartitions: the number of partitions of the DataFrame
        :return: :class:`DataFrame`

        >>> spark.range(1, 7, 2).collect()
        [Row(id=1), Row(id=3), Row(id=5)]

        If only one argument is specified, it will be used as the end value.

        >>> spark.range(3).collect()
        [Row(id=0), Row(id=1), Row(id=2)]
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism

        if end is None:
            jdf = self._jsparkSession.range(0, int(start), int(step), int(numPartitions))
        else:
            jdf = self._jsparkSession.range(int(start), int(end), int(step), int(numPartitions))

        return DataFrame(jdf, self._wrapped)

    def _inferSchemaFromList(self, data):
        """
        Infer schema from list of Row or tuple.

        :param data: list of Row or tuple
        :return: :class:`pyspark.sql.types.StructType`
        """
        if not data:
            raise ValueError("can not infer schema from empty dataset")
        first = data[0]
        if type(first) is dict:
            warnings.warn("inferring schema from dict is deprecated,"
                          "please use pyspark.sql.Row instead")
        schema = reduce(_merge_type, map(_infer_schema, data))
        if _has_nulltype(schema):
            raise ValueError("Some of types cannot be determined after inferring")
        return schema

    def _inferSchema(self, rdd, samplingRatio=None):
        """
        Infer schema from an RDD of Row or tuple.

        :param rdd: an RDD of Row or tuple
        :param samplingRatio: sampling ratio, or no sampling (default)
        :return: :class:`pyspark.sql.types.StructType`
        """
        first = rdd.first()
        if not first:
            raise ValueError("The first row in RDD is empty, "
                             "can not infer schema")
        if type(first) is dict:
            warnings.warn("Using RDD of dict to inferSchema is deprecated. "
                          "Use pyspark.sql.Row instead")

        if samplingRatio is None:
            schema = _infer_schema(first)
            if _has_nulltype(schema):
                for row in rdd.take(100)[1:]:
                    schema = _merge_type(schema, _infer_schema(row))
                    if not _has_nulltype(schema):
                        break
                else:
                    raise ValueError("Some of types cannot be determined by the "
                                     "first 100 rows, please try again with sampling")
        else:
            if samplingRatio < 0.99:
                rdd = rdd.sample(False, float(samplingRatio))
            schema = rdd.map(_infer_schema).reduce(_merge_type)
        return schema

    def _createFromRDD(self, rdd, schema, samplingRatio):
        """
        Create an RDD for DataFrame from an existing RDD, returns the RDD and schema.
        """
        if schema is None or isinstance(schema, (list, tuple)):
            struct = self._inferSchema(rdd, samplingRatio)
            converter = _create_converter(struct)
            rdd = rdd.map(converter)
            if isinstance(schema, (list, tuple)):
                for i, name in enumerate(schema):
                    struct.fields[i].name = name
                    struct.names[i] = name
            schema = struct

        elif not isinstance(schema, StructType):
            raise TypeError("schema should be StructType or list or None, but got: %s" % schema)

        # convert python objects to sql data
        rdd = rdd.map(schema.toInternal)
        return rdd, schema

    def _createFromLocal(self, data, schema):
        """
        Create an RDD for DataFrame from a list or pandas.DataFrame, returns
        the RDD and schema.
        """
        # make sure data could consumed multiple times
        if not isinstance(data, list):
            data = list(data)

        if schema is None or isinstance(schema, (list, tuple)):
            struct = self._inferSchemaFromList(data)
            converter = _create_converter(struct)
            data = map(converter, data)
            if isinstance(schema, (list, tuple)):
                for i, name in enumerate(schema):
                    struct.fields[i].name = name
                    struct.names[i] = name
            schema = struct

        elif not isinstance(schema, StructType):
            raise TypeError("schema should be StructType or list or None, but got: %s" % schema)

        # convert python objects to sql data
        data = [schema.toInternal(row) for row in data]
        return self._sc.parallelize(data), schema

    @since(2.0)
    @ignore_unicode_prefix
    def createDataFrame(self, data, schema=None, samplingRatio=None, verifySchema=True):
        """
        Creates a :class:`DataFrame` from an :class:`RDD`, a list or a :class:`pandas.DataFrame`.

        When ``schema`` is a list of column names, the type of each column
        will be inferred from ``data``.

        When ``schema`` is ``None``, it will try to infer the schema (column names and types)
        from ``data``, which should be an RDD of :class:`Row`,
        or :class:`namedtuple`, or :class:`dict`.

        When ``schema`` is :class:`pyspark.sql.types.DataType` or a datatype string, it must match
        the real data, or an exception will be thrown at runtime. If the given schema is not
        :class:`pyspark.sql.types.StructType`, it will be wrapped into a
        :class:`pyspark.sql.types.StructType` as its only field, and the field name will be "value",
        each record will also be wrapped into a tuple, which can be converted to row later.

        If schema inference is needed, ``samplingRatio`` is used to determined the ratio of
        rows used for schema inference. The first row will be used if ``samplingRatio`` is ``None``.

        :param data: an RDD of any kind of SQL data representation(e.g. row, tuple, int, boolean,
            etc.), or :class:`list`, or :class:`pandas.DataFrame`.
        :param schema: a :class:`pyspark.sql.types.DataType` or a datatype string or a list of
            column names, default is ``None``.  The data type string format equals to
            :class:`pyspark.sql.types.DataType.simpleString`, except that top level struct type can
            omit the ``struct<>`` and atomic types use ``typeName()`` as their format, e.g. use
            ``byte`` instead of ``tinyint`` for :class:`pyspark.sql.types.ByteType`. We can also use
            ``int`` as a short name for ``IntegerType``.
        :param samplingRatio: the sample ratio of rows used for inferring
        :param verifySchema: verify data types of every row against schema.
        :return: :class:`DataFrame`

        .. versionchanged:: 2.1
           Added verifySchema.

        >>> l = [('Alice', 1)]
        >>> spark.createDataFrame(l).collect()
        [Row(_1=u'Alice', _2=1)]
        >>> spark.createDataFrame(l, ['name', 'age']).collect()
        [Row(name=u'Alice', age=1)]

        >>> d = [{'name': 'Alice', 'age': 1}]
        >>> spark.createDataFrame(d).collect()
        [Row(age=1, name=u'Alice')]

        >>> rdd = sc.parallelize(l)
        >>> spark.createDataFrame(rdd).collect()
        [Row(_1=u'Alice', _2=1)]
        >>> df = spark.createDataFrame(rdd, ['name', 'age'])
        >>> df.collect()
        [Row(name=u'Alice', age=1)]

        >>> from pyspark.sql import Row
        >>> Person = Row('name', 'age')
        >>> person = rdd.map(lambda r: Person(*r))
        >>> df2 = spark.createDataFrame(person)
        >>> df2.collect()
        [Row(name=u'Alice', age=1)]

        >>> from pyspark.sql.types import *
        >>> schema = StructType([
        ...    StructField("name", StringType(), True),
        ...    StructField("age", IntegerType(), True)])
        >>> df3 = spark.createDataFrame(rdd, schema)
        >>> df3.collect()
        [Row(name=u'Alice', age=1)]

        >>> spark.createDataFrame(df.toPandas()).collect()  # doctest: +SKIP
        [Row(name=u'Alice', age=1)]
        >>> spark.createDataFrame(pandas.DataFrame([[1, 2]])).collect()  # doctest: +SKIP
        [Row(0=1, 1=2)]

        >>> spark.createDataFrame(rdd, "a: string, b: int").collect()
        [Row(a=u'Alice', b=1)]
        >>> rdd = rdd.map(lambda row: row[1])
        >>> spark.createDataFrame(rdd, "int").collect()
        [Row(value=1)]
        >>> spark.createDataFrame(rdd, "boolean").collect() # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
            ...
        Py4JJavaError: ...
        """
        if isinstance(data, DataFrame):
            raise TypeError("data is already a DataFrame")

        if isinstance(schema, basestring):
            schema = _parse_datatype_string(schema)

        try:
            import pandas
            has_pandas = True
        except Exception:
            has_pandas = False
        if has_pandas and isinstance(data, pandas.DataFrame):
            if schema is None:
                schema = [str(x) for x in data.columns]
            data = [r.tolist() for r in data.to_records(index=False)]

        verify_func = _verify_type if verifySchema else lambda _, t: True
        if isinstance(schema, StructType):
            def prepare(obj):
                verify_func(obj, schema)
                return obj
        elif isinstance(schema, DataType):
            dataType = schema
            schema = StructType().add("value", schema)

            def prepare(obj):
                verify_func(obj, dataType)
                return obj,
        else:
            if isinstance(schema, list):
                schema = [x.encode('utf-8') if not isinstance(x, str) else x for x in schema]
            prepare = lambda obj: obj

        if isinstance(data, RDD):
            rdd, schema = self._createFromRDD(data.map(prepare), schema, samplingRatio)
        else:
            rdd, schema = self._createFromLocal(map(prepare, data), schema)
        jrdd = self._jvm.SerDeUtil.toJavaArray(rdd._to_java_object_rdd())
        jdf = self._jsparkSession.applySchemaToPythonRDD(jrdd.rdd(), schema.json())
        df = DataFrame(jdf, self._wrapped)
        df._schema = schema
        return df

    @ignore_unicode_prefix
    @since(2.0)
    def sql(self, sqlQuery):
        """Returns a :class:`DataFrame` representing the result of the given query.

        :return: :class:`DataFrame`

        >>> df.createOrReplaceTempView("table1")
        >>> df2 = spark.sql("SELECT field1 AS f1, field2 as f2 from table1")
        >>> df2.collect()
        [Row(f1=1, f2=u'row1'), Row(f1=2, f2=u'row2'), Row(f1=3, f2=u'row3')]
        """
        return DataFrame(self._jsparkSession.sql(sqlQuery), self._wrapped)

    @since(2.0)
    def table(self, tableName):
        """Returns the specified table as a :class:`DataFrame`.

        :return: :class:`DataFrame`

        >>> df.createOrReplaceTempView("table1")
        >>> df2 = spark.table("table1")
        >>> sorted(df.collect()) == sorted(df2.collect())
        True
        """
        return DataFrame(self._jsparkSession.table(tableName), self._wrapped)

    @property
    @since(2.0)
    def read(self):
        """
        Returns a :class:`DataFrameReader` that can be used to read data
        in as a :class:`DataFrame`.

        :return: :class:`DataFrameReader`
        """
        return DataFrameReader(self._wrapped)

    @property
    @since(2.0)
    def readStream(self):
        """
        Returns a :class:`DataStreamReader` that can be used to read data streams
        as a streaming :class:`DataFrame`.

        .. note:: Experimental.

        :return: :class:`DataStreamReader`
        """
        return DataStreamReader(self._wrapped)

    @property
    @since(2.0)
    def streams(self):
        """Returns a :class:`StreamingQueryManager` that allows managing all the
        :class:`StreamingQuery` StreamingQueries active on `this` context.

        .. note:: Experimental.

        :return: :class:`StreamingQueryManager`
        """
        from pyspark.sql.streaming import StreamingQueryManager
        return StreamingQueryManager(self._jsparkSession.streams())

    @since(2.0)
    def stop(self):
        """Stop the underlying :class:`SparkContext`.
        """
        self._sc.stop()
        SparkSession._instantiatedContext = None

    @since(2.0)
    def __enter__(self):
        """
        Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax.
        """
        return self

    @since(2.0)
    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax.

        Specifically stop the SparkSession on exit of the with block.
        """
        self.stop()

def _test():
    import os
    import doctest
    from pyspark.context import SparkContext
    from pyspark.sql import Row
    import pyspark.sql.session

    os.chdir(os.environ["SPARK_HOME"])

    globs = pyspark.sql.session.__dict__.copy()
    sc = SparkContext('local[4]', 'PythonTest')
    globs['sc'] = sc
    globs['spark'] = SparkSession(sc)
    globs['rdd'] = rdd = sc.parallelize(
        [Row(field1=1, field2="row1"),
         Row(field1=2, field2="row2"),
         Row(field1=3, field2="row3")])
    globs['df'] = rdd.toDF()
    (failure_count, test_count) = doctest.testmod(
        pyspark.sql.session, globs=globs,
        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE)
    globs['sc'].stop()
    if failure_count:
        exit(-1)

if __name__ == "__main__":
    _test()


import numpy as np
import matplotlib
import matplotlib.colors as colors
import matplotlib.patches as patches
import matplotlib.mathtext as mathtext
import matplotlib.pyplot as plt
import matplotlib.artist as artist
import matplotlib.image as image

class ItemProperties:
    def __init__(self, fontsize=14, labelcolor='black', bgcolor='yellow',
                 alpha=1.0):
        self.fontsize = fontsize
        self.labelcolor = labelcolor
        self.bgcolor = bgcolor
        self.alpha = alpha

        self.labelcolor_rgb = colors.colorConverter.to_rgb(labelcolor)
        self.bgcolor_rgb = colors.colorConverter.to_rgb(bgcolor)

class MenuItem(artist.Artist):
    parser = mathtext.MathTextParser("Bitmap")
    padx = 5
    pady = 5
    def __init__(self, fig, labelstr, props=None, hoverprops=None,
                 on_select=None):
        artist.Artist.__init__(self)

        self.set_figure(fig)
        self.labelstr = labelstr

        if props is None:
            props = ItemProperties()

        if hoverprops is None:
            hoverprops = ItemProperties()

        self.props = props
        self.hoverprops = hoverprops

        self.on_select = on_select

        x, self.depth = self.parser.to_mask(
            labelstr, fontsize=props.fontsize, dpi=fig.dpi)

        if props.fontsize!=hoverprops.fontsize:
            raise NotImplementedError(
                        'support for different font sizes not implemented')

        self.labelwidth = x.shape[1]
        self.labelheight = x.shape[0]

        self.labelArray = np.zeros((x.shape[0], x.shape[1], 4))
        self.labelArray[:, :, -1] = x/255.

        self.label = image.FigureImage(fig, origin='upper')
        self.label.set_array(self.labelArray)

        # we'll update these later
        self.rect = patches.Rectangle((0,0), 1,1)

        self.set_hover_props(False)

        fig.canvas.mpl_connect('button_release_event', self.check_select)

    def check_select(self, event):
        over, junk = self.rect.contains(event)
        if not over:
            return

        if self.on_select is not None:
            self.on_select(self)

    def set_extent(self, x, y, w, h):
        print x, y, w, h
        self.rect.set_x(x)
        self.rect.set_y(y)
        self.rect.set_width(w)
        self.rect.set_height(h)

        self.label.ox = x+self.padx
        self.label.oy = y-self.depth+self.pady/2.

        self.rect._update_patch_transform()
        self.hover = False

    def draw(self, renderer):
        self.rect.draw(renderer)
        self.label.draw(renderer)

    def set_hover_props(self, b):
        if b:
            props = self.hoverprops
        else:
            props = self.props

        r, g, b = props.labelcolor_rgb
        self.labelArray[:, :, 0] = r
        self.labelArray[:, :, 1] = g
        self.labelArray[:, :, 2] = b
        self.label.set_array(self.labelArray)
        self.rect.set(facecolor=props.bgcolor, alpha=props.alpha)

    def set_hover(self, event):
        'check the hover status of event and return true if status is changed'
        b,junk = self.rect.contains(event)

        changed = (b != self.hover)

        if changed:
            self.set_hover_props(b)

        self.hover = b
        return changed

class Menu:

    def __init__(self, fig, menuitems):
        self.figure = fig
        fig.suppressComposite = True

        self.menuitems = menuitems
        self.numitems = len(menuitems)

        maxw = max([item.labelwidth for item in menuitems])
        maxh = max([item.labelheight for item in menuitems])

        totalh = self.numitems*maxh + (self.numitems+1)*2*MenuItem.pady

        x0 = 100
        y0 = 400

        width = maxw + 2*MenuItem.padx
        height = maxh+MenuItem.pady

        for item in menuitems:
            left = x0
            bottom = y0-maxh-MenuItem.pady

            item.set_extent(left, bottom, width, height)

            fig.artists.append(item)
            y0 -= maxh + MenuItem.pady

        fig.canvas.mpl_connect('motion_notify_event', self.on_move)

    def on_move(self, event):
        draw = False
        for item in self.menuitems:
            draw = item.set_hover(event)
            if draw:
                self.figure.canvas.draw()
                break

fig = plt.figure()
fig.subplots_adjust(left=0.3)
props = ItemProperties(labelcolor='black', bgcolor='yellow',
                       fontsize=15, alpha=0.2)
hoverprops = ItemProperties(labelcolor='white', bgcolor='blue',
                            fontsize=15, alpha=0.2)

menuitems = []
for label in ('open', 'close', 'save', 'save as', 'quit'):
    def on_select(item):
        print 'you selected', item.labelstr
    item = MenuItem(fig, label, props=props, hoverprops=hoverprops,
                    on_select=on_select)
    menuitems.append(item)

menu = Menu(fig, menuitems)
plt.show()


# -*- coding: utf-8 -*-
from __future__ import division, print_function

from functools import partial

import warnings
import numpy as np
from pandas import Series, isnull
from pandas.types.common import is_integer_dtype
import pandas.core.nanops as nanops
import pandas.util.testing as tm

use_bn = nanops._USE_BOTTLENECK

class TestnanopsDataFrame(tm.TestCase):

    def setUp(self):
        np.random.seed(11235)
        nanops._USE_BOTTLENECK = False

        self.arr_shape = (11, 7, 5)

        self.arr_float = np.random.randn(*self.arr_shape)
        self.arr_float1 = np.random.randn(*self.arr_shape)
        self.arr_complex = self.arr_float + self.arr_float1 * 1j
        self.arr_int = np.random.randint(-10, 10, self.arr_shape)
        self.arr_bool = np.random.randint(0, 2, self.arr_shape) == 0
        self.arr_str = np.abs(self.arr_float).astype('S')
        self.arr_utf = np.abs(self.arr_float).astype('U')
        self.arr_date = np.random.randint(0, 20000,
                                          self.arr_shape).astype('M8[ns]')
        self.arr_tdelta = np.random.randint(0, 20000,
                                            self.arr_shape).astype('m8[ns]')

        self.arr_nan = np.tile(np.nan, self.arr_shape)
        self.arr_float_nan = np.vstack([self.arr_float, self.arr_nan])
        self.arr_float1_nan = np.vstack([self.arr_float1, self.arr_nan])
        self.arr_nan_float1 = np.vstack([self.arr_nan, self.arr_float1])
        self.arr_nan_nan = np.vstack([self.arr_nan, self.arr_nan])

        self.arr_inf = self.arr_float * np.inf
        self.arr_float_inf = np.vstack([self.arr_float, self.arr_inf])
        self.arr_float1_inf = np.vstack([self.arr_float1, self.arr_inf])
        self.arr_inf_float1 = np.vstack([self.arr_inf, self.arr_float1])
        self.arr_inf_inf = np.vstack([self.arr_inf, self.arr_inf])

        self.arr_nan_inf = np.vstack([self.arr_nan, self.arr_inf])
        self.arr_float_nan_inf = np.vstack([self.arr_float, self.arr_nan,
                                            self.arr_inf])
        self.arr_nan_float1_inf = np.vstack([self.arr_float, self.arr_inf,
                                             self.arr_nan])
        self.arr_nan_nan_inf = np.vstack([self.arr_nan, self.arr_nan,
                                          self.arr_inf])
        self.arr_obj = np.vstack([self.arr_float.astype(
            'O'), self.arr_int.astype('O'), self.arr_bool.astype(
                'O'), self.arr_complex.astype('O'), self.arr_str.astype(
                    'O'), self.arr_utf.astype('O'), self.arr_date.astype('O'),
            self.arr_tdelta.astype('O')])

        with np.errstate(invalid='ignore'):
            self.arr_nan_nanj = self.arr_nan + self.arr_nan * 1j
            self.arr_complex_nan = np.vstack([self.arr_complex,
                                              self.arr_nan_nanj])

            self.arr_nan_infj = self.arr_inf * 1j
            self.arr_complex_nan_infj = np.vstack([self.arr_complex,
                                                   self.arr_nan_infj])

        self.arr_float_2d = self.arr_float[:, :, 0]
        self.arr_float1_2d = self.arr_float1[:, :, 0]
        self.arr_complex_2d = self.arr_complex[:, :, 0]
        self.arr_int_2d = self.arr_int[:, :, 0]
        self.arr_bool_2d = self.arr_bool[:, :, 0]
        self.arr_str_2d = self.arr_str[:, :, 0]
        self.arr_utf_2d = self.arr_utf[:, :, 0]
        self.arr_date_2d = self.arr_date[:, :, 0]
        self.arr_tdelta_2d = self.arr_tdelta[:, :, 0]

        self.arr_nan_2d = self.arr_nan[:, :, 0]
        self.arr_float_nan_2d = self.arr_float_nan[:, :, 0]
        self.arr_float1_nan_2d = self.arr_float1_nan[:, :, 0]
        self.arr_nan_float1_2d = self.arr_nan_float1[:, :, 0]
        self.arr_nan_nan_2d = self.arr_nan_nan[:, :, 0]
        self.arr_nan_nanj_2d = self.arr_nan_nanj[:, :, 0]
        self.arr_complex_nan_2d = self.arr_complex_nan[:, :, 0]

        self.arr_inf_2d = self.arr_inf[:, :, 0]
        self.arr_float_inf_2d = self.arr_float_inf[:, :, 0]
        self.arr_nan_inf_2d = self.arr_nan_inf[:, :, 0]
        self.arr_float_nan_inf_2d = self.arr_float_nan_inf[:, :, 0]
        self.arr_nan_nan_inf_2d = self.arr_nan_nan_inf[:, :, 0]

        self.arr_float_1d = self.arr_float[:, 0, 0]
        self.arr_float1_1d = self.arr_float1[:, 0, 0]
        self.arr_complex_1d = self.arr_complex[:, 0, 0]
        self.arr_int_1d = self.arr_int[:, 0, 0]
        self.arr_bool_1d = self.arr_bool[:, 0, 0]
        self.arr_str_1d = self.arr_str[:, 0, 0]
        self.arr_utf_1d = self.arr_utf[:, 0, 0]
        self.arr_date_1d = self.arr_date[:, 0, 0]
        self.arr_tdelta_1d = self.arr_tdelta[:, 0, 0]

        self.arr_nan_1d = self.arr_nan[:, 0, 0]
        self.arr_float_nan_1d = self.arr_float_nan[:, 0, 0]
        self.arr_float1_nan_1d = self.arr_float1_nan[:, 0, 0]
        self.arr_nan_float1_1d = self.arr_nan_float1[:, 0, 0]
        self.arr_nan_nan_1d = self.arr_nan_nan[:, 0, 0]
        self.arr_nan_nanj_1d = self.arr_nan_nanj[:, 0, 0]
        self.arr_complex_nan_1d = self.arr_complex_nan[:, 0, 0]

        self.arr_inf_1d = self.arr_inf.ravel()
        self.arr_float_inf_1d = self.arr_float_inf[:, 0, 0]
        self.arr_nan_inf_1d = self.arr_nan_inf[:, 0, 0]
        self.arr_float_nan_inf_1d = self.arr_float_nan_inf[:, 0, 0]
        self.arr_nan_nan_inf_1d = self.arr_nan_nan_inf[:, 0, 0]

    def tearDown(self):
        nanops._USE_BOTTLENECK = use_bn

    def check_results(self, targ, res, axis, check_dtype=True):
        res = getattr(res, 'asm8', res)
        res = getattr(res, 'values', res)

        # timedeltas are a beast here
        def _coerce_tds(targ, res):
            if hasattr(targ, 'dtype') and targ.dtype == 'm8[ns]':
                if len(targ) == 1:
                    targ = targ[0].item()
                    res = res.item()
                else:
                    targ = targ.view('i8')
            return targ, res

        try:
            if axis != 0 and hasattr(
                    targ, 'shape') and targ.ndim and targ.shape != res.shape:
                res = np.split(res, [targ.shape[0]], axis=0)[0]
        except:
            targ, res = _coerce_tds(targ, res)

        try:
            tm.assert_almost_equal(targ, res, check_dtype=check_dtype)
        except:

            # handle timedelta dtypes
            if hasattr(targ, 'dtype') and targ.dtype == 'm8[ns]':
                targ, res = _coerce_tds(targ, res)
                tm.assert_almost_equal(targ, res, check_dtype=check_dtype)
                return

            # There are sometimes rounding errors with
            # complex and object dtypes.
            # If it isn't one of those, re-raise the error.
            if not hasattr(res, 'dtype') or res.dtype.kind not in ['c', 'O']:
                raise
            # convert object dtypes to something that can be split into
            # real and imaginary parts
            if res.dtype.kind == 'O':
                if targ.dtype.kind != 'O':
                    res = res.astype(targ.dtype)
                else:
                    try:
                        res = res.astype('c16')
                    except:
                        res = res.astype('f8')
                    try:
                        targ = targ.astype('c16')
                    except:
                        targ = targ.astype('f8')
            # there should never be a case where numpy returns an object
            # but nanops doesn't, so make that an exception
            elif targ.dtype.kind == 'O':
                raise
            tm.assert_almost_equal(targ.real, res.real,
                                   check_dtype=check_dtype)
            tm.assert_almost_equal(targ.imag, res.imag,
                                   check_dtype=check_dtype)

    def check_fun_data(self, testfunc, targfunc, testarval, targarval,
                       targarnanval, check_dtype=True, **kwargs):
        for axis in list(range(targarval.ndim)) + [None]:
            for skipna in [False, True]:
                targartempval = targarval if skipna else targarnanval
                try:
                    targ = targfunc(targartempval, axis=axis, **kwargs)
                    res = testfunc(testarval, axis=axis, skipna=skipna,
                                   **kwargs)
                    self.check_results(targ, res, axis,
                                       check_dtype=check_dtype)
                    if skipna:
                        res = testfunc(testarval, axis=axis, **kwargs)
                        self.check_results(targ, res, axis,
                                           check_dtype=check_dtype)
                    if axis is None:
                        res = testfunc(testarval, skipna=skipna, **kwargs)
                        self.check_results(targ, res, axis,
                                           check_dtype=check_dtype)
                    if skipna and axis is None:
                        res = testfunc(testarval, **kwargs)
                        self.check_results(targ, res, axis,
                                           check_dtype=check_dtype)
                except BaseException as exc:
                    exc.args += ('axis: %s of %s' % (axis, testarval.ndim - 1),
                                 'skipna: %s' % skipna, 'kwargs: %s' % kwargs)
                    raise

        if testarval.ndim <= 1:
            return

        try:
            testarval2 = np.take(testarval, 0, axis=-1)
            targarval2 = np.take(targarval, 0, axis=-1)
            targarnanval2 = np.take(targarnanval, 0, axis=-1)
        except ValueError:
            return
        self.check_fun_data(testfunc, targfunc, testarval2, targarval2,
                            targarnanval2, check_dtype=check_dtype, **kwargs)

    def check_fun(self, testfunc, targfunc, testar, targar=None,
                  targarnan=None, **kwargs):
        if targar is None:
            targar = testar
        if targarnan is None:
            targarnan = testar
        testarval = getattr(self, testar)
        targarval = getattr(self, targar)
        targarnanval = getattr(self, targarnan)
        try:
            self.check_fun_data(testfunc, targfunc, testarval, targarval,
                                targarnanval, **kwargs)
        except BaseException as exc:
            exc.args += ('testar: %s' % testar, 'targar: %s' % targar,
                         'targarnan: %s' % targarnan)
            raise

    def check_funs(self, testfunc, targfunc, allow_complex=True,
                   allow_all_nan=True, allow_str=True, allow_date=True,
                   allow_tdelta=True, allow_obj=True, **kwargs):
        self.check_fun(testfunc, targfunc, 'arr_float', **kwargs)
        self.check_fun(testfunc, targfunc, 'arr_float_nan', 'arr_float',
                       **kwargs)
        self.check_fun(testfunc, targfunc, 'arr_int', **kwargs)
        self.check_fun(testfunc, targfunc, 'arr_bool', **kwargs)
        objs = [self.arr_float.astype('O'), self.arr_int.astype('O'),
                self.arr_bool.astype('O')]

        if allow_all_nan:
            self.check_fun(testfunc, targfunc, 'arr_nan', **kwargs)

        if allow_complex:
            self.check_fun(testfunc, targfunc, 'arr_complex', **kwargs)
            self.check_fun(testfunc, targfunc, 'arr_complex_nan',
                           'arr_complex', **kwargs)
            if allow_all_nan:
                self.check_fun(testfunc, targfunc, 'arr_nan_nanj', **kwargs)
            objs += [self.arr_complex.astype('O')]

        if allow_str:
            self.check_fun(testfunc, targfunc, 'arr_str', **kwargs)
            self.check_fun(testfunc, targfunc, 'arr_utf', **kwargs)
            objs += [self.arr_str.astype('O'), self.arr_utf.astype('O')]

        if allow_date:
            try:
                targfunc(self.arr_date)
            except TypeError:
                pass
            else:
                self.check_fun(testfunc, targfunc, 'arr_date', **kwargs)
                objs += [self.arr_date.astype('O')]

        if allow_tdelta:
            try:
                targfunc(self.arr_tdelta)
            except TypeError:
                pass
            else:
                self.check_fun(testfunc, targfunc, 'arr_tdelta', **kwargs)
                objs += [self.arr_tdelta.astype('O')]

        if allow_obj:
            self.arr_obj = np.vstack(objs)
            # some nanops handle object dtypes better than their numpy
            # counterparts, so the numpy functions need to be given something
            # else
            if allow_obj == 'convert':
                targfunc = partial(self._badobj_wrap, func=targfunc,
                                   allow_complex=allow_complex)
            self.check_fun(testfunc, targfunc, 'arr_obj', **kwargs)

    def check_funs_ddof(self,
                        testfunc,
                        targfunc,
                        allow_complex=True,
                        allow_all_nan=True,
                        allow_str=True,
                        allow_date=False,
                        allow_tdelta=False,
                        allow_obj=True, ):
        for ddof in range(3):
            try:
                self.check_funs(testfunc, targfunc, allow_complex,
                                allow_all_nan, allow_str, allow_date,
                                allow_tdelta, allow_obj, ddof=ddof)
            except BaseException as exc:
                exc.args += ('ddof %s' % ddof, )
                raise

    def _badobj_wrap(self, value, func, allow_complex=True, **kwargs):
        if value.dtype.kind == 'O':
            if allow_complex:
                value = value.astype('c16')
            else:
                value = value.astype('f8')
        return func(value, **kwargs)

    def test_nanany(self):
        self.check_funs(nanops.nanany, np.any, allow_all_nan=False,
                        allow_str=False, allow_date=False, allow_tdelta=False)

    def test_nanall(self):
        self.check_funs(nanops.nanall, np.all, allow_all_nan=False,
                        allow_str=False, allow_date=False, allow_tdelta=False)

    def test_nansum(self):
        self.check_funs(nanops.nansum, np.sum, allow_str=False,
                        allow_date=False, allow_tdelta=True, check_dtype=False)

    def test_nanmean(self):
        self.check_funs(nanops.nanmean, np.mean, allow_complex=False,
                        allow_obj=False, allow_str=False, allow_date=False,
                        allow_tdelta=True)

    def test_nanmean_overflow(self):
        # GH 10155
        # In the previous implementation mean can overflow for int dtypes, it
        # is now consistent with numpy

        # numpy < 1.9.0 is not computing this correctly
        from distutils.version import LooseVersion
        if LooseVersion(np.__version__) >= '1.9.0':
            for a in [2 ** 55, -2 ** 55, 20150515061816532]:
                s = Series(a, index=range(500), dtype=np.int64)
                result = s.mean()
                np_result = s.values.mean()
                self.assertEqual(result, a)
                self.assertEqual(result, np_result)
                self.assertTrue(result.dtype == np.float64)

    def test_returned_dtype(self):

        dtypes = [np.int16, np.int32, np.int64, np.float32, np.float64]
        if hasattr(np, 'float128'):
            dtypes.append(np.float128)

        for dtype in dtypes:
            s = Series(range(10), dtype=dtype)
            group_a = ['mean', 'std', 'var', 'skew', 'kurt']
            group_b = ['min', 'max']
            for method in group_a + group_b:
                result = getattr(s, method)()
                if is_integer_dtype(dtype) and method in group_a:
                    self.assertTrue(
                        result.dtype == np.float64,
                        "return dtype expected from %s is np.float64, "
                        "got %s instead" % (method, result.dtype))
                else:
                    self.assertTrue(
                        result.dtype == dtype,
                        "return dtype expected from %s is %s, "
                        "got %s instead" % (method, dtype, result.dtype))

    def test_nanmedian(self):
        with warnings.catch_warnings(record=True):
            self.check_funs(nanops.nanmedian, np.median, allow_complex=False,
                            allow_str=False, allow_date=False,
                            allow_tdelta=True, allow_obj='convert')

    def test_nanvar(self):
        self.check_funs_ddof(nanops.nanvar, np.var, allow_complex=False,
                             allow_str=False, allow_date=False,
                             allow_tdelta=True, allow_obj='convert')

    def test_nanstd(self):
        self.check_funs_ddof(nanops.nanstd, np.std, allow_complex=False,
                             allow_str=False, allow_date=False,
                             allow_tdelta=True, allow_obj='convert')

    def test_nansem(self):
        tm.skip_if_no_package('scipy.stats')
        tm._skip_if_scipy_0_17()
        from scipy.stats import sem
        self.check_funs_ddof(nanops.nansem, sem, allow_complex=False,
                             allow_str=False, allow_date=False,
                             allow_tdelta=True, allow_obj='convert')

    def _minmax_wrap(self, value, axis=None, func=None):
        res = func(value, axis)
        if res.dtype.kind == 'm':
            res = np.atleast_1d(res)
        return res

    def test_nanmin(self):
        func = partial(self._minmax_wrap, func=np.min)
        self.check_funs(nanops.nanmin, func, allow_str=False, allow_obj=False)

    def test_nanmax(self):
        func = partial(self._minmax_wrap, func=np.max)
        self.check_funs(nanops.nanmax, func, allow_str=False, allow_obj=False)

    def _argminmax_wrap(self, value, axis=None, func=None):
        res = func(value, axis)
        nans = np.min(value, axis)
        nullnan = isnull(nans)
        if res.ndim:
            res[nullnan] = -1
        elif (hasattr(nullnan, 'all') and nullnan.all() or
              not hasattr(nullnan, 'all') and nullnan):
            res = -1
        return res

    def test_nanargmax(self):
        func = partial(self._argminmax_wrap, func=np.argmax)
        self.check_funs(nanops.nanargmax, func, allow_str=False,
                        allow_obj=False, allow_date=True, allow_tdelta=True)

    def test_nanargmin(self):
        func = partial(self._argminmax_wrap, func=np.argmin)
        if tm.sys.version_info[0:2] == (2, 6):
            self.check_funs(nanops.nanargmin, func, allow_date=True,
                            allow_tdelta=True, allow_str=False,
                            allow_obj=False)
        else:
            self.check_funs(nanops.nanargmin, func, allow_str=False,
                            allow_obj=False)

    def _skew_kurt_wrap(self, values, axis=None, func=None):
        if not isinstance(values.dtype.type, np.floating):
            values = values.astype('f8')
        result = func(values, axis=axis, bias=False)
        # fix for handling cases where all elements in an axis are the same
        if isinstance(result, np.ndarray):
            result[np.max(values, axis=axis) == np.min(values, axis=axis)] = 0
            return result
        elif np.max(values) == np.min(values):
            return 0.
        return result

    def test_nanskew(self):
        tm.skip_if_no_package('scipy.stats')
        tm._skip_if_scipy_0_17()
        from scipy.stats import skew
        func = partial(self._skew_kurt_wrap, func=skew)
        self.check_funs(nanops.nanskew, func, allow_complex=False,
                        allow_str=False, allow_date=False, allow_tdelta=False)

    def test_nankurt(self):
        tm.skip_if_no_package('scipy.stats')
        tm._skip_if_scipy_0_17()
        from scipy.stats import kurtosis
        func1 = partial(kurtosis, fisher=True)
        func = partial(self._skew_kurt_wrap, func=func1)
        self.check_funs(nanops.nankurt, func, allow_complex=False,
                        allow_str=False, allow_date=False, allow_tdelta=False)

    def test_nanprod(self):
        self.check_funs(nanops.nanprod, np.prod, allow_str=False,
                        allow_date=False, allow_tdelta=False)

    def check_nancorr_nancov_2d(self, checkfun, targ0, targ1, **kwargs):
        res00 = checkfun(self.arr_float_2d, self.arr_float1_2d, **kwargs)
        res01 = checkfun(self.arr_float_2d, self.arr_float1_2d,
                         min_periods=len(self.arr_float_2d) - 1, **kwargs)
        tm.assert_almost_equal(targ0, res00)
        tm.assert_almost_equal(targ0, res01)

        res10 = checkfun(self.arr_float_nan_2d, self.arr_float1_nan_2d,
                         **kwargs)
        res11 = checkfun(self.arr_float_nan_2d, self.arr_float1_nan_2d,
                         min_periods=len(self.arr_float_2d) - 1, **kwargs)
        tm.assert_almost_equal(targ1, res10)
        tm.assert_almost_equal(targ1, res11)

        targ2 = np.nan
        res20 = checkfun(self.arr_nan_2d, self.arr_float1_2d, **kwargs)
        res21 = checkfun(self.arr_float_2d, self.arr_nan_2d, **kwargs)
        res22 = checkfun(self.arr_nan_2d, self.arr_nan_2d, **kwargs)
        res23 = checkfun(self.arr_float_nan_2d, self.arr_nan_float1_2d,
                         **kwargs)
        res24 = checkfun(self.arr_float_nan_2d, self.arr_nan_float1_2d,
                         min_periods=len(self.arr_float_2d) - 1, **kwargs)
        res25 = checkfun(self.arr_float_2d, self.arr_float1_2d,
                         min_periods=len(self.arr_float_2d) + 1, **kwargs)
        tm.assert_almost_equal(targ2, res20)
        tm.assert_almost_equal(targ2, res21)
        tm.assert_almost_equal(targ2, res22)
        tm.assert_almost_equal(targ2, res23)
        tm.assert_almost_equal(targ2, res24)
        tm.assert_almost_equal(targ2, res25)

    def check_nancorr_nancov_1d(self, checkfun, targ0, targ1, **kwargs):
        res00 = checkfun(self.arr_float_1d, self.arr_float1_1d, **kwargs)
        res01 = checkfun(self.arr_float_1d, self.arr_float1_1d,
                         min_periods=len(self.arr_float_1d) - 1, **kwargs)
        tm.assert_almost_equal(targ0, res00)
        tm.assert_almost_equal(targ0, res01)

        res10 = checkfun(self.arr_float_nan_1d, self.arr_float1_nan_1d,
                         **kwargs)
        res11 = checkfun(self.arr_float_nan_1d, self.arr_float1_nan_1d,
                         min_periods=len(self.arr_float_1d) - 1, **kwargs)
        tm.assert_almost_equal(targ1, res10)
        tm.assert_almost_equal(targ1, res11)

        targ2 = np.nan
        res20 = checkfun(self.arr_nan_1d, self.arr_float1_1d, **kwargs)
        res21 = checkfun(self.arr_float_1d, self.arr_nan_1d, **kwargs)
        res22 = checkfun(self.arr_nan_1d, self.arr_nan_1d, **kwargs)
        res23 = checkfun(self.arr_float_nan_1d, self.arr_nan_float1_1d,
                         **kwargs)
        res24 = checkfun(self.arr_float_nan_1d, self.arr_nan_float1_1d,
                         min_periods=len(self.arr_float_1d) - 1, **kwargs)
        res25 = checkfun(self.arr_float_1d, self.arr_float1_1d,
                         min_periods=len(self.arr_float_1d) + 1, **kwargs)
        tm.assert_almost_equal(targ2, res20)
        tm.assert_almost_equal(targ2, res21)
        tm.assert_almost_equal(targ2, res22)
        tm.assert_almost_equal(targ2, res23)
        tm.assert_almost_equal(targ2, res24)
        tm.assert_almost_equal(targ2, res25)

    def test_nancorr(self):
        targ0 = np.corrcoef(self.arr_float_2d, self.arr_float1_2d)[0, 1]
        targ1 = np.corrcoef(self.arr_float_2d.flat,
                            self.arr_float1_2d.flat)[0, 1]
        self.check_nancorr_nancov_2d(nanops.nancorr, targ0, targ1)
        targ0 = np.corrcoef(self.arr_float_1d, self.arr_float1_1d)[0, 1]
        targ1 = np.corrcoef(self.arr_float_1d.flat,
                            self.arr_float1_1d.flat)[0, 1]
        self.check_nancorr_nancov_1d(nanops.nancorr, targ0, targ1,
                                     method='pearson')

    def test_nancorr_pearson(self):
        targ0 = np.corrcoef(self.arr_float_2d, self.arr_float1_2d)[0, 1]
        targ1 = np.corrcoef(self.arr_float_2d.flat,
                            self.arr_float1_2d.flat)[0, 1]
        self.check_nancorr_nancov_2d(nanops.nancorr, targ0, targ1,
                                     method='pearson')
        targ0 = np.corrcoef(self.arr_float_1d, self.arr_float1_1d)[0, 1]
        targ1 = np.corrcoef(self.arr_float_1d.flat,
                            self.arr_float1_1d.flat)[0, 1]
        self.check_nancorr_nancov_1d(nanops.nancorr, targ0, targ1,
                                     method='pearson')

    def test_nancorr_kendall(self):
        tm.skip_if_no_package('scipy.stats')
        from scipy.stats import kendalltau
        targ0 = kendalltau(self.arr_float_2d, self.arr_float1_2d)[0]
        targ1 = kendalltau(self.arr_float_2d.flat, self.arr_float1_2d.flat)[0]
        self.check_nancorr_nancov_2d(nanops.nancorr, targ0, targ1,
                                     method='kendall')
        targ0 = kendalltau(self.arr_float_1d, self.arr_float1_1d)[0]
        targ1 = kendalltau(self.arr_float_1d.flat, self.arr_float1_1d.flat)[0]
        self.check_nancorr_nancov_1d(nanops.nancorr, targ0, targ1,
                                     method='kendall')

    def test_nancorr_spearman(self):
        tm.skip_if_no_package('scipy.stats')
        from scipy.stats import spearmanr
        targ0 = spearmanr(self.arr_float_2d, self.arr_float1_2d)[0]
        targ1 = spearmanr(self.arr_float_2d.flat, self.arr_float1_2d.flat)[0]
        self.check_nancorr_nancov_2d(nanops.nancorr, targ0, targ1,
                                     method='spearman')
        targ0 = spearmanr(self.arr_float_1d, self.arr_float1_1d)[0]
        targ1 = spearmanr(self.arr_float_1d.flat, self.arr_float1_1d.flat)[0]
        self.check_nancorr_nancov_1d(nanops.nancorr, targ0, targ1,
                                     method='spearman')

    def test_nancov(self):
        targ0 = np.cov(self.arr_float_2d, self.arr_float1_2d)[0, 1]
        targ1 = np.cov(self.arr_float_2d.flat, self.arr_float1_2d.flat)[0, 1]
        self.check_nancorr_nancov_2d(nanops.nancov, targ0, targ1)
        targ0 = np.cov(self.arr_float_1d, self.arr_float1_1d)[0, 1]
        targ1 = np.cov(self.arr_float_1d.flat, self.arr_float1_1d.flat)[0, 1]
        self.check_nancorr_nancov_1d(nanops.nancov, targ0, targ1)

    def check_nancomp(self, checkfun, targ0):
        arr_float = self.arr_float
        arr_float1 = self.arr_float1
        arr_nan = self.arr_nan
        arr_nan_nan = self.arr_nan_nan
        arr_float_nan = self.arr_float_nan
        arr_float1_nan = self.arr_float1_nan
        arr_nan_float1 = self.arr_nan_float1

        while targ0.ndim:
            try:
                res0 = checkfun(arr_float, arr_float1)
                tm.assert_almost_equal(targ0, res0)

                if targ0.ndim > 1:
                    targ1 = np.vstack([targ0, arr_nan])
                else:
                    targ1 = np.hstack([targ0, arr_nan])
                res1 = checkfun(arr_float_nan, arr_float1_nan)
                tm.assert_numpy_array_equal(targ1, res1, check_dtype=False)

                targ2 = arr_nan_nan
                res2 = checkfun(arr_float_nan, arr_nan_float1)
                tm.assert_numpy_array_equal(targ2, res2, check_dtype=False)
            except Exception as exc:
                exc.args += ('ndim: %s' % arr_float.ndim, )
                raise

            try:
                arr_float = np.take(arr_float, 0, axis=-1)
                arr_float1 = np.take(arr_float1, 0, axis=-1)
                arr_nan = np.take(arr_nan, 0, axis=-1)
                arr_nan_nan = np.take(arr_nan_nan, 0, axis=-1)
                arr_float_nan = np.take(arr_float_nan, 0, axis=-1)
                arr_float1_nan = np.take(arr_float1_nan, 0, axis=-1)
                arr_nan_float1 = np.take(arr_nan_float1, 0, axis=-1)
                targ0 = np.take(targ0, 0, axis=-1)
            except ValueError:
                break

    def test_nangt(self):
        targ0 = self.arr_float > self.arr_float1
        self.check_nancomp(nanops.nangt, targ0)

    def test_nange(self):
        targ0 = self.arr_float >= self.arr_float1
        self.check_nancomp(nanops.nange, targ0)

    def test_nanlt(self):
        targ0 = self.arr_float < self.arr_float1
        self.check_nancomp(nanops.nanlt, targ0)

    def test_nanle(self):
        targ0 = self.arr_float <= self.arr_float1
        self.check_nancomp(nanops.nanle, targ0)

    def test_naneq(self):
        targ0 = self.arr_float == self.arr_float1
        self.check_nancomp(nanops.naneq, targ0)

    def test_nanne(self):
        targ0 = self.arr_float != self.arr_float1
        self.check_nancomp(nanops.nanne, targ0)

    def check_bool(self, func, value, correct, *args, **kwargs):
        while getattr(value, 'ndim', True):
            try:
                res0 = func(value, *args, **kwargs)
                if correct:
                    self.assertTrue(res0)
                else:
                    self.assertFalse(res0)
            except BaseException as exc:
                exc.args += ('dim: %s' % getattr(value, 'ndim', value), )
                raise
            if not hasattr(value, 'ndim'):
                break
            try:
                value = np.take(value, 0, axis=-1)
            except ValueError:
                break

    def test__has_infs(self):
        pairs = [('arr_complex', False), ('arr_int', False),
                 ('arr_bool', False), ('arr_str', False), ('arr_utf', False),
                 ('arr_complex', False), ('arr_complex_nan', False),
                 ('arr_nan_nanj', False), ('arr_nan_infj', True),
                 ('arr_complex_nan_infj', True)]
        pairs_float = [('arr_float', False), ('arr_nan', False),
                       ('arr_float_nan', False), ('arr_nan_nan', False),
                       ('arr_float_inf', True), ('arr_inf', True),
                       ('arr_nan_inf', True), ('arr_float_nan_inf', True),
                       ('arr_nan_nan_inf', True)]

        for arr, correct in pairs:
            val = getattr(self, arr)
            try:
                self.check_bool(nanops._has_infs, val, correct)
            except BaseException as exc:
                exc.args += (arr, )
                raise

        for arr, correct in pairs_float:
            val = getattr(self, arr)
            try:
                self.check_bool(nanops._has_infs, val, correct)
                self.check_bool(nanops._has_infs, val.astype('f4'), correct)
                self.check_bool(nanops._has_infs, val.astype('f2'), correct)
            except BaseException as exc:
                exc.args += (arr, )
                raise

    def test__isfinite(self):
        pairs = [('arr_complex', False), ('arr_int', False),
                 ('arr_bool', False), ('arr_str', False), ('arr_utf', False),
                 ('arr_complex', False), ('arr_complex_nan', True),
                 ('arr_nan_nanj', True), ('arr_nan_infj', True),
                 ('arr_complex_nan_infj', True)]
        pairs_float = [('arr_float', False), ('arr_nan', True),
                       ('arr_float_nan', True), ('arr_nan_nan', True),
                       ('arr_float_inf', True), ('arr_inf', True),
                       ('arr_nan_inf', True), ('arr_float_nan_inf', True),
                       ('arr_nan_nan_inf', True)]

        func1 = lambda x: np.any(nanops._isfinite(x).ravel())

        # TODO: unused?
        # func2 = lambda x: np.any(nanops._isfinite(x).values.ravel())

        for arr, correct in pairs:
            val = getattr(self, arr)
            try:
                self.check_bool(func1, val, correct)
            except BaseException as exc:
                exc.args += (arr, )
                raise

        for arr, correct in pairs_float:
            val = getattr(self, arr)
            try:
                self.check_bool(func1, val, correct)
                self.check_bool(func1, val.astype('f4'), correct)
                self.check_bool(func1, val.astype('f2'), correct)
            except BaseException as exc:
                exc.args += (arr, )
                raise

    def test__bn_ok_dtype(self):
        self.assertTrue(nanops._bn_ok_dtype(self.arr_float.dtype, 'test'))
        self.assertTrue(nanops._bn_ok_dtype(self.arr_complex.dtype, 'test'))
        self.assertTrue(nanops._bn_ok_dtype(self.arr_int.dtype, 'test'))
        self.assertTrue(nanops._bn_ok_dtype(self.arr_bool.dtype, 'test'))
        self.assertTrue(nanops._bn_ok_dtype(self.arr_str.dtype, 'test'))
        self.assertTrue(nanops._bn_ok_dtype(self.arr_utf.dtype, 'test'))
        self.assertFalse(nanops._bn_ok_dtype(self.arr_date.dtype, 'test'))
        self.assertFalse(nanops._bn_ok_dtype(self.arr_tdelta.dtype, 'test'))
        self.assertFalse(nanops._bn_ok_dtype(self.arr_obj.dtype, 'test'))

class TestEnsureNumeric(tm.TestCase):

    def test_numeric_values(self):
        # Test integer
        self.assertEqual(nanops._ensure_numeric(1), 1, 'Failed for int')
        # Test float
        self.assertEqual(nanops._ensure_numeric(1.1), 1.1, 'Failed for float')
        # Test complex
        self.assertEqual(nanops._ensure_numeric(1 + 2j), 1 + 2j,
                         'Failed for complex')

    def test_ndarray(self):
        # Test numeric ndarray
        values = np.array([1, 2, 3])
        self.assertTrue(np.allclose(nanops._ensure_numeric(values), values),
                        'Failed for numeric ndarray')

        # Test object ndarray
        o_values = values.astype(object)
        self.assertTrue(np.allclose(nanops._ensure_numeric(o_values), values),
                        'Failed for object ndarray')

        # Test convertible string ndarray
        s_values = np.array(['1', '2', '3'], dtype=object)
        self.assertTrue(np.allclose(nanops._ensure_numeric(s_values), values),
                        'Failed for convertible string ndarray')

        # Test non-convertible string ndarray
        s_values = np.array(['foo', 'bar', 'baz'], dtype=object)
        self.assertRaises(ValueError, lambda: nanops._ensure_numeric(s_values))

    def test_convertable_values(self):
        self.assertTrue(np.allclose(nanops._ensure_numeric('1'), 1.0),
                        'Failed for convertible integer string')
        self.assertTrue(np.allclose(nanops._ensure_numeric('1.1'), 1.1),
                        'Failed for convertible float string')
        self.assertTrue(np.allclose(nanops._ensure_numeric('1+1j'), 1 + 1j),
                        'Failed for convertible complex string')

    def test_non_convertable_values(self):
        self.assertRaises(TypeError, lambda: nanops._ensure_numeric('foo'))
        self.assertRaises(TypeError, lambda: nanops._ensure_numeric({}))
        self.assertRaises(TypeError, lambda: nanops._ensure_numeric([]))

class TestNanvarFixedValues(tm.TestCase):

    # xref GH10242

    def setUp(self):
        # Samples from a normal distribution.
        self.variance = variance = 3.0
        self.samples = self.prng.normal(scale=variance ** 0.5, size=100000)

    def test_nanvar_all_finite(self):
        samples = self.samples
        actual_variance = nanops.nanvar(samples)
        tm.assert_almost_equal(actual_variance, self.variance,
                               check_less_precise=2)

    def test_nanvar_nans(self):
        samples = np.nan * np.ones(2 * self.samples.shape[0])
        samples[::2] = self.samples

        actual_variance = nanops.nanvar(samples, skipna=True)
        tm.assert_almost_equal(actual_variance, self.variance,
                               check_less_precise=2)

        actual_variance = nanops.nanvar(samples, skipna=False)
        tm.assert_almost_equal(actual_variance, np.nan, check_less_precise=2)

    def test_nanstd_nans(self):
        samples = np.nan * np.ones(2 * self.samples.shape[0])
        samples[::2] = self.samples

        actual_std = nanops.nanstd(samples, skipna=True)
        tm.assert_almost_equal(actual_std, self.variance ** 0.5,
                               check_less_precise=2)

        actual_std = nanops.nanvar(samples, skipna=False)
        tm.assert_almost_equal(actual_std, np.nan,
                               check_less_precise=2)

    def test_nanvar_axis(self):
        # Generate some sample data.
        samples_norm = self.samples
        samples_unif = self.prng.uniform(size=samples_norm.shape[0])
        samples = np.vstack([samples_norm, samples_unif])

        actual_variance = nanops.nanvar(samples, axis=1)
        tm.assert_almost_equal(actual_variance, np.array(
            [self.variance, 1.0 / 12]), check_less_precise=2)

    def test_nanvar_ddof(self):
        n = 5
        samples = self.prng.uniform(size=(10000, n + 1))
        samples[:, -1] = np.nan  # Force use of our own algorithm.

        variance_0 = nanops.nanvar(samples, axis=1, skipna=True, ddof=0).mean()
        variance_1 = nanops.nanvar(samples, axis=1, skipna=True, ddof=1).mean()
        variance_2 = nanops.nanvar(samples, axis=1, skipna=True, ddof=2).mean()

        # The unbiased estimate.
        var = 1.0 / 12
        tm.assert_almost_equal(variance_1, var,
                               check_less_precise=2)

        # The underestimated variance.
        tm.assert_almost_equal(variance_0, (n - 1.0) / n * var,
                               check_less_precise=2)

        # The overestimated variance.
        tm.assert_almost_equal(variance_2, (n - 1.0) / (n - 2.0) * var,
                               check_less_precise=2)

    def test_ground_truth(self):
        # Test against values that were precomputed with Numpy.
        samples = np.empty((4, 4))
        samples[:3, :3] = np.array([[0.97303362, 0.21869576, 0.55560287
                                     ], [0.72980153, 0.03109364, 0.99155171],
                                    [0.09317602, 0.60078248, 0.15871292]])
        samples[3] = samples[:, 3] = np.nan

        # Actual variances along axis=0, 1 for ddof=0, 1, 2
        variance = np.array([[[0.13762259, 0.05619224, 0.11568816
                               ], [0.20643388, 0.08428837, 0.17353224],
                              [0.41286776, 0.16857673, 0.34706449]],
                             [[0.09519783, 0.16435395, 0.05082054
                               ], [0.14279674, 0.24653093, 0.07623082],
                              [0.28559348, 0.49306186, 0.15246163]]])

        # Test nanvar.
        for axis in range(2):
            for ddof in range(3):
                var = nanops.nanvar(samples, skipna=True, axis=axis, ddof=ddof)
                tm.assert_almost_equal(var[:3], variance[axis, ddof])
                self.assertTrue(np.isnan(var[3]))

        # Test nanstd.
        for axis in range(2):
            for ddof in range(3):
                std = nanops.nanstd(samples, skipna=True, axis=axis, ddof=ddof)
                tm.assert_almost_equal(std[:3], variance[axis, ddof] ** 0.5)
                self.assertTrue(np.isnan(std[3]))

    def test_nanstd_roundoff(self):
        # Regression test for GH 10242 (test data taken from GH 10489). Ensure
        # that variance is stable.
        data = Series(766897346 * np.ones(10))
        for ddof in range(3):
            result = data.std(ddof=ddof)
            self.assertEqual(result, 0.0)

    @property
    def prng(self):
        return np.random.RandomState(1234)

class TestNanskewFixedValues(tm.TestCase):

    # xref GH 11974

    def setUp(self):
        # Test data + skewness value (computed with scipy.stats.skew)
        self.samples = np.sin(np.linspace(0, 1, 200))
        self.actual_skew = -0.1875895205961754

    def test_constant_series(self):
        # xref GH 11974
        for val in [3075.2, 3075.3, 3075.5]:
            data = val * np.ones(300)
            skew = nanops.nanskew(data)
            self.assertEqual(skew, 0.0)

    def test_all_finite(self):
        alpha, beta = 0.3, 0.1
        left_tailed = self.prng.beta(alpha, beta, size=100)
        self.assertLess(nanops.nanskew(left_tailed), 0)

        alpha, beta = 0.1, 0.3
        right_tailed = self.prng.beta(alpha, beta, size=100)
        self.assertGreater(nanops.nanskew(right_tailed), 0)

    def test_ground_truth(self):
        skew = nanops.nanskew(self.samples)
        self.assertAlmostEqual(skew, self.actual_skew)

    def test_axis(self):
        samples = np.vstack([self.samples,
                             np.nan * np.ones(len(self.samples))])
        skew = nanops.nanskew(samples, axis=1)
        tm.assert_almost_equal(skew, np.array([self.actual_skew, np.nan]))

    def test_nans(self):
        samples = np.hstack([self.samples, np.nan])
        skew = nanops.nanskew(samples, skipna=False)
        self.assertTrue(np.isnan(skew))

    def test_nans_skipna(self):
        samples = np.hstack([self.samples, np.nan])
        skew = nanops.nanskew(samples, skipna=True)
        tm.assert_almost_equal(skew, self.actual_skew)

    @property
    def prng(self):
        return np.random.RandomState(1234)

class TestNankurtFixedValues(tm.TestCase):

    # xref GH 11974

    def setUp(self):
        # Test data + kurtosis value (computed with scipy.stats.kurtosis)
        self.samples = np.sin(np.linspace(0, 1, 200))
        self.actual_kurt = -1.2058303433799713

    def test_constant_series(self):
        # xref GH 11974
        for val in [3075.2, 3075.3, 3075.5]:
            data = val * np.ones(300)
            kurt = nanops.nankurt(data)
            self.assertEqual(kurt, 0.0)

    def test_all_finite(self):
        alpha, beta = 0.3, 0.1
        left_tailed = self.prng.beta(alpha, beta, size=100)
        self.assertLess(nanops.nankurt(left_tailed), 0)

        alpha, beta = 0.1, 0.3
        right_tailed = self.prng.beta(alpha, beta, size=100)
        self.assertGreater(nanops.nankurt(right_tailed), 0)

    def test_ground_truth(self):
        kurt = nanops.nankurt(self.samples)
        self.assertAlmostEqual(kurt, self.actual_kurt)

    def test_axis(self):
        samples = np.vstack([self.samples,
                             np.nan * np.ones(len(self.samples))])
        kurt = nanops.nankurt(samples, axis=1)
        tm.assert_almost_equal(kurt, np.array([self.actual_kurt, np.nan]))

    def test_nans(self):
        samples = np.hstack([self.samples, np.nan])
        kurt = nanops.nankurt(samples, skipna=False)
        self.assertTrue(np.isnan(kurt))

    def test_nans_skipna(self):
        samples = np.hstack([self.samples, np.nan])
        kurt = nanops.nankurt(samples, skipna=True)
        tm.assert_almost_equal(kurt, self.actual_kurt)

    @property
    def prng(self):
        return np.random.RandomState(1234)

def test_int64_add_overflow():
    # see gh-14068
    msg = "Overflow in int64 addition"
    m = np.iinfo(np.int64).max
    n = np.iinfo(np.int64).min

    with tm.assertRaisesRegexp(OverflowError, msg):
        nanops._checked_add_with_arr(np.array([m, m]), m)
    with tm.assertRaisesRegexp(OverflowError, msg):
        nanops._checked_add_with_arr(np.array([m, m]), np.array([m, m]))
    with tm.assertRaisesRegexp(OverflowError, msg):
        nanops._checked_add_with_arr(np.array([n, n]), n)
    with tm.assertRaisesRegexp(OverflowError, msg):
        nanops._checked_add_with_arr(np.array([n, n]), np.array([n, n]))
    with tm.assertRaisesRegexp(OverflowError, msg):
        nanops._checked_add_with_arr(np.array([m, n]), np.array([n, n]))
    with tm.assertRaisesRegexp(OverflowError, msg):
        with tm.assert_produces_warning(RuntimeWarning):
            nanops._checked_add_with_arr(np.array([m, m]),
                                         np.array([np.nan, m]))

if __name__ == '__main__':
    import nose
    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure', '-s'
                         ], exit=False)


# pylint: disable=W0201

from statsmodels.compat.python import iteritems, string_types, range
import numpy as np
from statsmodels.tools.decorators import cache_readonly
import pandas as pd

from . import var_model as _model
from . import util
from . import plotting

FULL_SAMPLE = 0
ROLLING = 1
EXPANDING = 2

def _get_window_type(window_type):
    if window_type in (FULL_SAMPLE, ROLLING, EXPANDING):
        return window_type
    elif isinstance(window_type, string_types):
        window_type_up = window_type.upper()

        if window_type_up in ('FULL SAMPLE', 'FULL_SAMPLE'):
            return FULL_SAMPLE
        elif window_type_up == 'ROLLING':
            return ROLLING
        elif window_type_up == 'EXPANDING':
            return EXPANDING

    raise Exception('Unrecognized window type: %s' % window_type)

class DynamicVAR(object):
    """
    Estimates time-varying vector autoregression (VAR(p)) using
    equation-by-equation least squares

    Parameters
    ----------
    data : pandas.DataFrame
    lag_order : int, default 1
    window : int
    window_type : {'expanding', 'rolling'}
    min_periods : int or None
        Minimum number of observations to require in window, defaults to window
        size if None specified
    trend : {'c', 'nc', 'ct', 'ctt'}
        TODO

    Returns
    -------
    **Attributes**:

    coefs : WidePanel
        items : coefficient names
        major_axis : dates
        minor_axis : VAR equation names
    """
    def __init__(self, data, lag_order=1, window=None, window_type='expanding',
                 trend='c', min_periods=None):
        self.lag_order = lag_order

        self.names = list(data.columns)
        self.neqs = len(self.names)

        self._y_orig = data

        # TODO: deal with trend
        self._x_orig = _make_lag_matrix(data, lag_order)
        self._x_orig['intercept'] = 1

        (self.y, self.x, self.x_filtered, self._index,
         self._time_has_obs) = _filter_data(self._y_orig, self._x_orig)

        self.lag_order = lag_order
        self.trendorder = util.get_trendorder(trend)

        self._set_window(window_type, window, min_periods)

    def _set_window(self, window_type, window, min_periods):
        self._window_type = _get_window_type(window_type)

        if self._is_rolling:
            if window is None:
                raise Exception('Must pass window when doing rolling '
                                'regression')

            if min_periods is None:
                min_periods = window
        else:
            window = len(self.x)
            if min_periods is None:
                min_periods = 1

        self._window = int(window)
        self._min_periods = min_periods

    @cache_readonly
    def T(self):
        """
        Number of time periods in results
        """
        return len(self.result_index)

    @property
    def nobs(self):
        # Stub, do I need this?
        data = dict((eq, r.nobs) for eq, r in iteritems(self.equations))
        return pd.DataFrame(data)

    @cache_readonly
    def equations(self):
        eqs = {}
        for col, ts in iteritems(self.y):
            model = pd.ols(y=ts, x=self.x, window=self._window,
                           window_type=self._window_type,
                           min_periods=self._min_periods)

            eqs[col] = model

        return eqs

    @cache_readonly
    def coefs(self):
        """
        Return dynamic regression coefficients as WidePanel
        """
        data = {}
        for eq, result in iteritems(self.equations):
            data[eq] = result.beta

        panel = pd.WidePanel.fromDict(data)

        # Coefficient names become items
        return panel.swapaxes('items', 'minor')

    @property
    def result_index(self):
        return self.coefs.major_axis

    @cache_readonly
    def _coefs_raw(self):
        """
        Reshape coefficients to be more amenable to dynamic calculations

        Returns
        -------
        coefs : (time_periods x lag_order x neqs x neqs)
        """
        coef_panel = self.coefs.copy()
        del coef_panel['intercept']

        coef_values = coef_panel.swapaxes('items', 'major').values
        coef_values = coef_values.reshape((len(coef_values),
                                           self.lag_order,
                                           self.neqs, self.neqs))

        return coef_values

    @cache_readonly
    def _intercepts_raw(self):
        """
        Similar to _coefs_raw, return intercept values in easy-to-use matrix
        form

        Returns
        -------
        intercepts : (T x K)
        """
        return self.coefs['intercept'].values

    @cache_readonly
    def resid(self):
        data = {}
        for eq, result in iteritems(self.equations):
            data[eq] = result.resid

        return pd.DataFrame(data)

    def forecast(self, steps=1):
        """
        Produce dynamic forecast

        Parameters
        ----------
        steps

        Returns
        -------
        forecasts : pandas.DataFrame
        """
        output = np.empty((self.T - steps, self.neqs))

        y_values = self.y.values
        y_index_map = dict((d, idx) for idx, d in enumerate(self.y.index))
        result_index_map = dict((d, idx) for idx, d in enumerate(self.result_index))

        coefs = self._coefs_raw
        intercepts = self._intercepts_raw

        # can only produce this many forecasts
        forc_index = self.result_index[steps:]
        for i, date in enumerate(forc_index):
            # TODO: check that this does the right thing in weird cases...
            idx = y_index_map[date] - steps
            result_idx = result_index_map[date] - steps

            y_slice = y_values[:idx]

            forcs = _model.forecast(y_slice, coefs[result_idx],
                                    intercepts[result_idx], steps)

            output[i] = forcs[-1]

        return pd.DataFrame(output, index=forc_index, columns=self.names)

    def plot_forecast(self, steps=1, figsize=(10, 10)):
        """
        Plot h-step ahead forecasts against actual realizations of time
        series. Note that forecasts are lined up with their respective
        realizations.

        Parameters
        ----------
        steps :
        """
        import matplotlib.pyplot as plt

        fig, axes = plt.subplots(figsize=figsize, nrows=self.neqs,
                                 sharex=True)

        forc = self.forecast(steps=steps)
        dates = forc.index

        y_overlay = self.y.reindex(dates)

        for i, col in enumerate(forc.columns):
            ax = axes[i]

            y_ts = y_overlay[col]
            forc_ts = forc[col]

            y_handle = ax.plot(dates, y_ts.values, 'k.', ms=2)
            forc_handle = ax.plot(dates, forc_ts.values, 'k-')

        fig.legend((y_handle, forc_handle), ('Y', 'Forecast'))
        fig.autofmt_xdate()

        fig.suptitle('Dynamic %d-step forecast' % steps)

        # pretty things up a bit
        plotting.adjust_subplots(bottom=0.15, left=0.10)
        plt.draw_if_interactive()

    @property
    def _is_rolling(self):
        return self._window_type == ROLLING

    @cache_readonly
    def r2(self):
        """Returns the r-squared values."""
        data = dict((eq, r.r2) for eq, r in iteritems(self.equations))
        return pd.DataFrame(data)

class DynamicPanelVAR(DynamicVAR):
    """
    Dynamic (time-varying) panel vector autoregression using panel ordinary
    least squares

    Parameters
    ----------
    """
    def __init__(self, data, lag_order=1, window=None, window_type='expanding',
                 trend='c', min_periods=None):
        self.lag_order = lag_order
        self.neqs = len(data.columns)

        self._y_orig = data

        # TODO: deal with trend
        self._x_orig = _make_lag_matrix(data, lag_order)
        self._x_orig['intercept'] = 1

        (self.y, self.x, self.x_filtered, self._index,
         self._time_has_obs) = _filter_data(self._y_orig, self._x_orig)

        self.lag_order = lag_order
        self.trendorder = util.get_trendorder(trend)

        self._set_window(window_type, window, min_periods)

def _filter_data(lhs, rhs):
    """
    Data filtering routine for dynamic VAR

    lhs : DataFrame
        original data
    rhs : DataFrame
        lagged variables

    Returns
    -------

    """
    def _has_all_columns(df):
        return np.isfinite(df.values).sum(1) == len(df.columns)

    rhs_valid = _has_all_columns(rhs)
    if not rhs_valid.all():
        pre_filtered_rhs = rhs[rhs_valid]
    else:
        pre_filtered_rhs = rhs

    index = lhs.index.union(rhs.index)
    if not index.equals(rhs.index) or not index.equals(lhs.index):
        rhs = rhs.reindex(index)
        lhs = lhs.reindex(index)

        rhs_valid = _has_all_columns(rhs)

    lhs_valid = _has_all_columns(lhs)
    valid = rhs_valid & lhs_valid

    if not valid.all():
        filt_index = rhs.index[valid]
        filtered_rhs = rhs.reindex(filt_index)
        filtered_lhs = lhs.reindex(filt_index)
    else:
        filtered_rhs, filtered_lhs = rhs, lhs

    return filtered_lhs, filtered_rhs, pre_filtered_rhs, index, valid

def _make_lag_matrix(x, lags):
    data = {}
    columns = []
    for i in range(1, 1 + lags):
        lagstr = 'L%d.'% i
        lag = x.shift(i).rename(columns=lambda c: lagstr + c)
        data.update(lag._series)
        columns.extend(lag.columns)

    return pd.DataFrame(data, columns=columns)

class Equation(object):
    """
    Stub, estimate one equation
    """

    def __init__(self, y, x):
        pass

if __name__ == '__main__':
    import pandas.util.testing as ptest

    ptest.N = 500
    data = ptest.makeTimeDataFrame().cumsum(0)

    var = DynamicVAR(data, lag_order=2, window_type='expanding')
    var2 = DynamicVAR(data, lag_order=2, window=10,
                      window_type='rolling')


"""
===============================================================
Model selection with Probabilistic PCA and Factor Analysis (FA)
===============================================================

Probabilistic PCA and Factor Analysis are probabilistic models.
The consequence is that the likelihood of new data can be used
for model selection and covariance estimation.
Here we compare PCA and FA with cross-validation on low rank data corrupted
with homoscedastic noise (noise variance
is the same for each feature) or heteroscedastic noise (noise variance
is the different for each feature). In a second step we compare the model
likelihood to the likelihoods obtained from shrinkage covariance estimators.

One can observe that with homoscedastic noise both FA and PCA succeed
in recovering the size of the low rank subspace. The likelihood with PCA
is higher than FA in this case. However PCA fails and overestimates
the rank when heteroscedastic noise is present. Under appropriate
circumstances the low rank models are more likely than shrinkage models.

The automatic estimation from
Automatic Choice of Dimensionality for PCA. NIPS 2000: 598-604
by Thomas P. Minka is also compared.

"""

# Authors: Alexandre Gramfort
#          Denis A. Engemann
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt
from scipy import linalg

from sklearn.decomposition import PCA, FactorAnalysis
from sklearn.covariance import ShrunkCovariance, LedoitWolf
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import GridSearchCV

print(__doc__)

###############################################################################
# Create the data

n_samples, n_features, rank = 1000, 50, 10
sigma = 1.
rng = np.random.RandomState(42)
U, _, _ = linalg.svd(rng.randn(n_features, n_features))
X = np.dot(rng.randn(n_samples, rank), U[:, :rank].T)

# Adding homoscedastic noise
X_homo = X + sigma * rng.randn(n_samples, n_features)

# Adding heteroscedastic noise
sigmas = sigma * rng.rand(n_features) + sigma / 2.
X_hetero = X + rng.randn(n_samples, n_features) * sigmas

###############################################################################
# Fit the models

n_components = np.arange(0, n_features, 5)  # options for n_components

def compute_scores(X):
    pca = PCA(svd_solver='full')
    fa = FactorAnalysis()

    pca_scores, fa_scores = [], []
    for n in n_components:
        pca.n_components = n
        fa.n_components = n
        pca_scores.append(np.mean(cross_val_score(pca, X)))
        fa_scores.append(np.mean(cross_val_score(fa, X)))

    return pca_scores, fa_scores

def shrunk_cov_score(X):
    shrinkages = np.logspace(-2, 0, 30)
    cv = GridSearchCV(ShrunkCovariance(), {'shrinkage': shrinkages})
    return np.mean(cross_val_score(cv.fit(X).best_estimator_, X))

def lw_score(X):
    return np.mean(cross_val_score(LedoitWolf(), X))

for X, title in [(X_homo, 'Homoscedastic Noise'),
                 (X_hetero, 'Heteroscedastic Noise')]:
    pca_scores, fa_scores = compute_scores(X)
    n_components_pca = n_components[np.argmax(pca_scores)]
    n_components_fa = n_components[np.argmax(fa_scores)]

    pca = PCA(svd_solver='full', n_components='mle')
    pca.fit(X)
    n_components_pca_mle = pca.n_components_

    print("best n_components by PCA CV = %d" % n_components_pca)
    print("best n_components by FactorAnalysis CV = %d" % n_components_fa)
    print("best n_components by PCA MLE = %d" % n_components_pca_mle)

    plt.figure()
    plt.plot(n_components, pca_scores, 'b', label='PCA scores')
    plt.plot(n_components, fa_scores, 'r', label='FA scores')
    plt.axvline(rank, color='g', label='TRUTH: %d' % rank, linestyle='-')
    plt.axvline(n_components_pca, color='b',
                label='PCA CV: %d' % n_components_pca, linestyle='--')
    plt.axvline(n_components_fa, color='r',
                label='FactorAnalysis CV: %d' % n_components_fa,
                linestyle='--')
    plt.axvline(n_components_pca_mle, color='k',
                label='PCA MLE: %d' % n_components_pca_mle, linestyle='--')

    # compare with other covariance estimators
    plt.axhline(shrunk_cov_score(X), color='violet',
                label='Shrunk Covariance MLE', linestyle='-.')
    plt.axhline(lw_score(X), color='orange',
                label='LedoitWolf MLE' % n_components_pca_mle, linestyle='-.')

    plt.xlabel('nb of components')
    plt.ylabel('CV scores')
    plt.legend(loc='lower right')
    plt.title(title)

plt.show()


"""
===============================
Nearest Centroid Classification
===============================

Sample usage of Nearest Centroid classification.
It will plot the decision boundaries for each class.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn import datasets
from sklearn.neighbors import NearestCentroid

n_neighbors = 15

# import some data to play with
iris = datasets.load_iris()
X = iris.data[:, :2]  # we only take the first two features. We could
                      # avoid this ugly slicing by using a two-dim dataset
y = iris.target

h = .02  # step size in the mesh

# Create color maps
cmap_light = ListedColormap(['#FFAAAA', '#AAFFAA', '#AAAAFF'])
cmap_bold = ListedColormap(['#FF0000', '#00FF00', '#0000FF'])

for shrinkage in [None, 0.1]:
    # we create an instance of Neighbours Classifier and fit the data.
    clf = NearestCentroid(shrink_threshold=shrinkage)
    clf.fit(X, y)
    y_pred = clf.predict(X)
    print(shrinkage, np.mean(y == y_pred))
    # Plot the decision boundary. For that, we will assign a color to each
    # point in the mesh [x_min, m_max]x[y_min, y_max].
    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                         np.arange(y_min, y_max, h))
    Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])

    # Put the result into a color plot
    Z = Z.reshape(xx.shape)
    plt.figure()
    plt.pcolormesh(xx, yy, Z, cmap=cmap_light)

    # Plot also the training points
    plt.scatter(X[:, 0], X[:, 1], c=y, cmap=cmap_bold)
    plt.title("3-Class classification (shrink_threshold=%r)"
              % shrinkage)
    plt.axis('tight')

plt.show()