from copy import copy
from collections import namedtuple

import datetime
import numpy as np
import pandas as pd

from syscore.genutils import transfer_object_attributes
from syscore.pdutils import make_df_from_list_of_named_tuple
from syscore.objects import header, table, body_text, arg_not_supplied, missing_data

from sysdata.data_blob import dataBlob
from sysproduction.data.orders import dataOrders
from sysproduction.data.instruments import diagInstruments

from sysproduction.reporting.risk_report import  get_current_annualised_stdev_for_instrument

def trades_info(
    data=arg_not_supplied,
    calendar_days_back=1,
    end_date=arg_not_supplied,
    start_date=arg_not_supplied,
):
    """
    Report on system status

    :param: data blob
    :return: list of formatted output items
    """
    if data is arg_not_supplied:
        data = dataBlob()

    if end_date is arg_not_supplied:
        end_date = datetime.datetime.now()

    if start_date is arg_not_supplied:
        start_date = end_date - datetime.timedelta(days=calendar_days_back)

    results_object = get_trades_report_data(
        data, start_date=start_date, end_date=end_date
    )
    formatted_output = format_trades_data(results_object)

    return formatted_output

def get_trades_report_data(data, start_date, end_date):

    broker_orders = get_recent_broker_orders(data, start_date, end_date)
    if len(broker_orders) == 0:
        empty_df = pd.DataFrame()
        results_object = dict(overview=empty_df)
        return results_object

    overview = broker_orders[
        [
            "instrument_code",
            "strategy_name",
            "contract_date",
            "fill_datetime",
            "fill",
            "filled_price",
        ]
    ]

    delays = create_delay_df(broker_orders)
    raw_slippage = create_raw_slippage_df(broker_orders)
    vol_slippage = create_vol_norm_slippage_df(raw_slippage, data)
    cash_slippage = create_cash_slippage_df(raw_slippage, data)

    summary_dict = {}
    item_list = [
        "delay",
        "bid_ask",
        "execution",
        "versus_limit",
        "versus_parent_limit",
        "total_trading",
    ]
    detailed_raw_results = get_stats_for_slippage_groups(
        raw_slippage, item_list)
    summary_dict.update(detailed_raw_results)

    item_list = [
        "delay_vol",
        "bid_ask_vol",
        "execution_vol",
        "versus_limit_vol",
        "versus_parent_limit_vol",
        "total_trading_vol",
    ]
    detailed_vol_results = get_stats_for_slippage_groups(
        vol_slippage, item_list)
    summary_dict.update(detailed_vol_results)

    item_list = [
        "delay_cash",
        "bid_ask_cash",
        "execution_cash",
        "versus_limit_cash",
        "versus_parent_limit_cash",
        "total_trading_cash",
    ]
    detailed_cash_results = get_stats_for_slippage_groups(
        cash_slippage, item_list)
    summary_dict.update(detailed_cash_results)

    results_object = dict(
        overview=overview,
        delays=delays,
        raw_slippage=raw_slippage,
        vol_slippage=vol_slippage,
        cash_slippage=cash_slippage,
        summary_dict=summary_dict,
    )

    return results_object

def format_trades_data(results_object):
    """
    Put the results into a printable format

    :param results_dict: dict, keys are different segments
    :return:
    """

    formatted_output = []

    formatted_output.append(
        header("Trades report produced on %s" % (str(datetime.datetime.now())))
    )

    if len(results_object["overview"]) == 0:
        formatted_output.append(body_text("No trades in relevant period"))

        return formatted_output

    table1_df = results_object["overview"]
    table1 = table("Broker orders", table1_df)
    formatted_output.append(table1)

    table2_df = results_object["delays"]
    table2 = table("Delays", table2_df)
    formatted_output.append(table2)

    table3_df = results_object["raw_slippage"]
    table3 = table("Slippage (ticks per lot)", table3_df)
    formatted_output.append(table3)

    table4_df = results_object["vol_slippage"]
    table4 = table(
        "Slippage (normalised by annual vol, BP of annual SR)",
        table4_df)
    formatted_output.append(table4)

    table5_df = results_object["cash_slippage"]
    table5 = table("Slippage (In base currency)", table5_df)
    formatted_output.append(table5)

    summary_results_dict = results_object["summary_dict"]
    for summary_table_name, summary_table_item in summary_results_dict.items():
        summary_table = table(
            "Summary %s" %
            summary_table_name,
            summary_table_item)
        formatted_output.append(summary_table)

    return formatted_output

tradesData = namedtuple(
    "tradesData",
    [
        "order_id",
        "instrument_code",
        "strategy_name",
        "contract_date",
        "fill",
        "filled_price",
        "mid_price",
        "side_price",
        "offside_price",
        "parent_reference_price",  # from contract order
        "parent_reference_datetime",  # from instrument order
        "submit_datetime",
        "fill_datetime",
        "limit_price",
        "trade",
        "buy_or_sell",
        "parent_limit_price",
        "commission",
    ],
)

data = dataBlob()

def get_recent_broker_orders(data, start_date, end_date):
    data_orders = dataOrders(data)
    order_id_list = data_orders.get_historic_broker_order_ids_in_date_range(
        start_date, end_date
    )
    orders_as_list = [get_tuple_object_from_order_id(
        data, order_id) for order_id in order_id_list]
    pdf = make_df_from_list_of_named_tuple(tradesData, orders_as_list)

    return pdf

def get_tuple_object_from_order_id(data, order_id):
    data_orders = dataOrders(data)
    order = data_orders.get_historic_broker_order_from_order_id_with_execution_data(
        order_id)
    tuple_object = transfer_object_attributes(tradesData, order)

    return tuple_object

def create_delay_df(broker_orders):
    delay_data_as_list = [
        delay_row(
            broker_orders.iloc[irow]) for irow in range(
            len(broker_orders))]
    delay_data_df = pd.concat(delay_data_as_list, axis=1)
    delay_data_df = delay_data_df.transpose()
    delay_data_df.index = broker_orders.index

    return delay_data_df

def delay_row(order_row):
    submit_minus_generated, filled_minus_submit = delay_calculations_for_order_row(
        order_row)
    new_order_row = copy(order_row)
    new_order_row = new_order_row[
        [
            "instrument_code",
            "strategy_name",
            "parent_reference_datetime",
            "submit_datetime",
            "fill_datetime",
        ]
    ]
    new_order_row = new_order_row.append(
        pd.Series(
            [submit_minus_generated, filled_minus_submit],
            index=["submit_minus_generated", "filled_minus_submit"],
        )
    )

    return new_order_row

def delay_calculations_for_order_row(order_row):

    submit_minus_generated = delay_calc(
        order_row.parent_reference_datetime, order_row.submit_datetime
    )

    filled_minus_submit = delay_calc(
        order_row.submit_datetime,
        order_row.fill_datetime)

    return submit_minus_generated, filled_minus_submit

def delay_calc(first_time, second_time):
    if first_time is None or second_time is None:
        return np.nan

    time_diff = second_time - first_time
    time_diff_seconds = time_diff.total_seconds()

    if time_diff_seconds < 0:
        return np.nan

    return time_diff_seconds

def create_raw_slippage_df(broker_orders):
    raw_slippage_data_as_list = [
        raw_slippage_row(
            broker_orders.iloc[irow]) for irow in range(
            len(broker_orders))]
    raw_slippage_df = pd.concat(raw_slippage_data_as_list, axis=1)
    raw_slippage_df = raw_slippage_df.transpose()
    raw_slippage_df.index = broker_orders.index

    return raw_slippage_df

def raw_slippage_row(order_row):
    (
        delay,
        bid_ask,
        execution,
        versus_limit,
        versus_parent_limit,
        total_trading,
    ) = price_calculations_for_order_row(order_row)
    new_order_row = copy(order_row)
    new_order_row = new_order_row[
        [
            "instrument_code",
            "strategy_name",
            "trade",
            "parent_reference_price",
            "parent_limit_price",
            "mid_price",
            "side_price",
            "offside_price",
            "limit_price",
            "filled_price",
        ]
    ]
    new_order_row = new_order_row.append(
        pd.Series(
            [
                delay,
                bid_ask,
                execution,
                versus_limit,
                versus_parent_limit,
                total_trading,
            ],
            index=[
                "delay",
                "bid_ask",
                "execution",
                "versus_limit",
                "versus_parent_limit",
                "total_trading",
            ],
        )
    )

    return new_order_row

def price_calculations_for_order_row(order_row):
    buying_multiplier = order_row.buy_or_sell

    # Following are always floats: parent_reference_price, limit_price,
    # calculated_mid_price, calculated_side_price, fill_price
    delay = price_slippage(
        buying_multiplier,
        order_row.parent_reference_price,
        order_row.mid_price,
    )

    bid_ask = price_slippage(
        buying_multiplier,
        order_row.mid_price,
        order_row.side_price,
    )

    execution = price_slippage(
        buying_multiplier,
        order_row.side_price,
        order_row.filled_price,
    )

    total_trading = bid_ask + execution

    versus_limit = price_slippage(
        buying_multiplier,
        order_row.limit_price,
        order_row.filled_price)

    versus_parent_limit = price_slippage(
        buying_multiplier,
        order_row.parent_limit_price,
        order_row.filled_price,
    )

    return delay, bid_ask, execution, versus_limit, versus_parent_limit, total_trading

def price_slippage(buying_multiplier, first_price, second_price):
    # Slippage is always negative (bad) positive (good)
    # This will return a negative number if second price is adverse versus
    # first price
    if first_price is None or second_price is None:
        return np.nan

    # 1 if buying, -1 if selling
    # if buying, want second price to be lower than first
    # if selling, want second price to be higher than first
    slippage = buying_multiplier * (first_price - second_price)
    return slippage

def create_cash_slippage_df(raw_slippage, data):
    # What does this slippage mean in money terms

    cash_slippage_data_as_list = [
        cash_slippage_row(raw_slippage.iloc[irow], data)
        for irow in range(len(raw_slippage))
    ]
    cash_slippage_df = pd.concat(cash_slippage_data_as_list, axis=1)
    cash_slippage_df = cash_slippage_df.transpose()
    cash_slippage_df.index = raw_slippage.index

    return cash_slippage_df

def cash_slippage_row(slippage_row, data):
    # rewrite
    (
        delay_cash,
        bid_ask_cash,
        execution_cash,
        versus_limit_cash,
        versus_parent_limit_cash,
        total_trading_cash,
        value_of_price_point,
    ) = cash_calculations_for_slippage_row(slippage_row, data)
    new_slippage_row = copy(slippage_row)
    new_slippage_row = new_slippage_row[
        [
            "instrument_code",
            "strategy_name",
            "trade",
        ]
    ]
    new_slippage_row = new_slippage_row.append(
        pd.Series(
            [
                value_of_price_point,
                delay_cash,
                bid_ask_cash,
                execution_cash,
                versus_limit_cash,
                versus_parent_limit_cash,
                total_trading_cash,
            ],
            index=[
                "value_of_price_point",
                "delay_cash",
                "bid_ask_cash",
                "execution_cash",
                "versus_limit_cash",
                "versus_parent_limit_cash",
                "total_trading_cash",
            ],
        )
    )

    return new_slippage_row

def cash_calculations_for_slippage_row(slippage_row, data):
    # What's a tick worth in base currency?
    diag_instruments = diagInstruments(data)
    value_of_price_point = diag_instruments.get_point_size_base_currency(
        slippage_row.instrument_code
    )
    input_items = [
        "delay",
        "bid_ask",
        "execution",
        "versus_limit",
        "versus_parent_limit",
        "total_trading",
    ]
    output = [value_of_price_point * slippage_row[input_name]
              for input_name in input_items]

    return tuple(output + [value_of_price_point])

def create_vol_norm_slippage_df(raw_slippage, data):
    # What does this slippage mean in vol normalised terms
    for irow in range(len(raw_slippage)):
        vol_slippage_row(raw_slippage.iloc[irow], data)

    vol_slippage_data_as_list = [
        vol_slippage_row(raw_slippage.iloc[irow], data)
        for irow in range(len(raw_slippage))
    ]
    vol_slippage_df = pd.concat(vol_slippage_data_as_list, axis=1)
    vol_slippage_df = vol_slippage_df.transpose()
    vol_slippage_df.index = raw_slippage.index

    return vol_slippage_df

def vol_slippage_row(slippage_row, data):
    # rewrite
    (
        vol_delay,
        vol_bid_ask,
        vol_execution,
        vol_versus_limit,
        vol_versus_parent_limit,
        total_trading_vol,
        last_annual_vol,
    ) = vol_calculations_for_slippage_row(slippage_row, data)
    new_slippage_row = copy(slippage_row)
    new_slippage_row = new_slippage_row[
        [
            "instrument_code",
            "strategy_name",
            "trade",
        ]
    ]
    new_slippage_row = new_slippage_row.append(
        pd.Series(
            [
                last_annual_vol,
                vol_delay,
                vol_bid_ask,
                vol_execution,
                vol_versus_limit,
                vol_versus_parent_limit,
                total_trading_vol,
            ],
            index=[
                "last_annual_vol",
                "delay_vol",
                "bid_ask_vol",
                "execution_vol",
                "versus_limit_vol",
                "versus_parent_limit_vol",
                "total_trading_vol",
            ],
        )
    )

    return new_slippage_row

def vol_calculations_for_slippage_row(slippage_row, data):

    last_annual_vol = get_last_annual_vol_for_slippage_row(slippage_row, data)

    input_items = [
        "delay",
        "bid_ask",
        "execution",
        "versus_limit",
        "versus_parent_limit",
        "total_trading",
    ]
    output = [10000 * slippage_row[input_name] /
              last_annual_vol for input_name in input_items]

    return tuple(output + [last_annual_vol])

def get_last_annual_vol_for_slippage_row(slippage_row, data):
    instrument_code = slippage_row.instrument_code
    last_annual_vol = get_current_annualised_stdev_for_instrument(data,
        instrument_code)

    return last_annual_vol

def get_stats_for_slippage_groups(df_to_process, item_list):
    results = {}
    for item_name in item_list:

        sum_data = df_to_process.groupby(
            ["strategy_name", "instrument_code"]).agg({item_name: "sum"})
        count_data = df_to_process.groupby(
            ["strategy_name", "instrument_code"]).agg({item_name: "count"})
        avg_data = sum_data / count_data

        try:
            std = df_to_process.groupby(
                ["strategy_name", "instrument_code"]).agg({item_name: "std"})
        except pd.core.base.DataError:
            # not enough items to calculate standard deviation
            std = np.nan

        lower_range = avg_data + (-2 * std)
        upper_range = avg_data + (2 * std)

        results[item_name + " Sum"] = sum_data
        results[item_name + " Count"] = count_data
        results[item_name + " Mean"] = avg_data
        results[item_name + " Lower range"] = lower_range
        results[item_name + " Upper range"] = upper_range

        total_sum_data = df_to_process.groupby(["strategy_name"]).agg(
            {item_name: "sum"}
        )

        results[item_name + " Total Sum"] = total_sum_data

    return results


from statsmodels.compat.python import iteritems, StringIO
import warnings

from statsmodels.formula.api import ols
from statsmodels.formula.formulatools import make_hypotheses_matrices
from statsmodels.tools import add_constant
from statsmodels.datasets.longley import load, load_pandas

import numpy.testing as npt
from statsmodels.tools.testing import assert_equal
from numpy.testing.utils import WarningManager

longley_formula = 'TOTEMP ~ GNPDEFL + GNP + UNEMP + ARMED + POP + YEAR'

class CheckFormulaOLS(object):

    @classmethod
    def setupClass(cls):
        cls.data = load()

    def test_endog_names(self):
        assert self.model.endog_names == 'TOTEMP'

    def test_exog_names(self):
        assert self.model.exog_names == ['Intercept', 'GNPDEFL', 'GNP',
                                             'UNEMP', 'ARMED', 'POP', 'YEAR']
    def test_design(self):
        npt.assert_equal(self.model.exog,
                         add_constant(self.data.exog, prepend=True))

    def test_endog(self):
        npt.assert_equal(self.model.endog, self.data.endog)

    def test_summary(self):
        # smoke test
        warn_ctx = WarningManager()
        warn_ctx.__enter__()
        try:
            warnings.filterwarnings("ignore",
                                    "kurtosistest only valid for n>=20")
            self.model.fit().summary()
        finally:
            warn_ctx.__exit__()

class TestFormulaPandas(CheckFormulaOLS):
    @classmethod
    def setupClass(cls):
        data = load_pandas().data
        cls.model = ols(longley_formula, data)
        super(TestFormulaPandas, cls).setupClass()

class TestFormulaDict(CheckFormulaOLS):
    @classmethod
    def setupClass(cls):
        data = dict((k, v.tolist()) for k, v in iteritems(load_pandas().data))
        cls.model = ols(longley_formula, data)
        super(TestFormulaDict, cls).setupClass()

class TestFormulaRecArray(CheckFormulaOLS):
    @classmethod
    def setupClass(cls):
        data = load().data
        cls.model = ols(longley_formula, data)
        super(TestFormulaRecArray, cls).setupClass()

def test_tests():
    formula = 'TOTEMP ~ GNPDEFL + GNP + UNEMP + ARMED + POP + YEAR'
    dta = load_pandas().data
    results = ols(formula, dta).fit()
    test_formula = '(GNPDEFL = GNP), (UNEMP = 2), (YEAR/1829 = 1)'
    LC = make_hypotheses_matrices(results, test_formula)
    R = LC.coefs
    Q = LC.constants
    npt.assert_almost_equal(R, [[0, 1, -1, 0, 0, 0, 0],
                               [0, 0 , 0, 1, 0, 0, 0],
                               [0, 0, 0, 0, 0, 0, 1./1829]], 8)
    npt.assert_array_equal(Q, [[0],[2],[1]])

def test_formula_labels():
    # make sure labels pass through patsy as expected
    # data(Duncan) from car in R
    dta = StringIO(""""type" "income" "education" "prestige"\n"accountant" "prof" 62 86 82\n"pilot" "prof" 72 76 83\n"architect" "prof" 75 92 90\n"author" "prof" 55 90 76\n"chemist" "prof" 64 86 90\n"minister" "prof" 21 84 87\n"professor" "prof" 64 93 93\n"dentist" "prof" 80 100 90\n"reporter" "wc" 67 87 52\n"engineer" "prof" 72 86 88\n"undertaker" "prof" 42 74 57\n"lawyer" "prof" 76 98 89\n"physician" "prof" 76 97 97\n"welfare.worker" "prof" 41 84 59\n"teacher" "prof" 48 91 73\n"conductor" "wc" 76 34 38\n"contractor" "prof" 53 45 76\n"factory.owner" "prof" 60 56 81\n"store.manager" "prof" 42 44 45\n"banker" "prof" 78 82 92\n"bookkeeper" "wc" 29 72 39\n"mail.carrier" "wc" 48 55 34\n"insurance.agent" "wc" 55 71 41\n"store.clerk" "wc" 29 50 16\n"carpenter" "bc" 21 23 33\n"electrician" "bc" 47 39 53\n"RR.engineer" "bc" 81 28 67\n"machinist" "bc" 36 32 57\n"auto.repairman" "bc" 22 22 26\n"plumber" "bc" 44 25 29\n"gas.stn.attendant" "bc" 15 29 10\n"coal.miner" "bc" 7 7 15\n"streetcar.motorman" "bc" 42 26 19\n"taxi.driver" "bc" 9 19 10\n"truck.driver" "bc" 21 15 13\n"machine.operator" "bc" 21 20 24\n"barber" "bc" 16 26 20\n"bartender" "bc" 16 28 7\n"shoe.shiner" "bc" 9 17 3\n"cook" "bc" 14 22 16\n"soda.clerk" "bc" 12 30 6\n"watchman" "bc" 17 25 11\n"janitor" "bc" 7 20 8\n"policeman" "bc" 34 47 41\n"waiter" "bc" 8 32 10""")
    from pandas import read_table
    dta = read_table(dta, sep=" ")
    model = ols("prestige ~ income + education", dta).fit()
    assert_equal(model.fittedvalues.index, dta.index)

def test_formula_predict():
    from numpy import log
    formula = """TOTEMP ~ log(GNPDEFL) + log(GNP) + UNEMP + ARMED +
                    POP + YEAR"""
    data = load_pandas()
    dta = load_pandas().data
    results = ols(formula, dta).fit()
    npt.assert_almost_equal(results.fittedvalues.values,
                            results.predict(data.exog), 8)


""" Basic functions for manipulating 2d arrays

"""
from __future__ import division, absolute_import, print_function

import functools

from numpy.core.numeric import (
    absolute, asanyarray, arange, zeros, greater_equal, multiply, ones,
    asarray, where, int8, int16, int32, int64, empty, promote_types, diagonal,
    nonzero
    )
from numpy.core.overrides import set_module
from numpy.core import overrides
from numpy.core import iinfo, transpose

__all__ = [
    'diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'tri', 'triu',
    'tril', 'vander', 'histogram2d', 'mask_indices', 'tril_indices',
    'tril_indices_from', 'triu_indices', 'triu_indices_from', ]

array_function_dispatch = functools.partial(
    overrides.array_function_dispatch, module='numpy')

i1 = iinfo(int8)
i2 = iinfo(int16)
i4 = iinfo(int32)

def _min_int(low, high):
    """ get small int that fits the range """
    if high <= i1.max and low >= i1.min:
        return int8
    if high <= i2.max and low >= i2.min:
        return int16
    if high <= i4.max and low >= i4.min:
        return int32
    return int64

def _flip_dispatcher(m):
    return (m,)

@array_function_dispatch(_flip_dispatcher)
def fliplr(m):
    """
    Flip array in the left/right direction.

    Flip the entries in each row in the left/right direction.
    Columns are preserved, but appear in a different order than before.

    Parameters
    ----------
    m : array_like
        Input array, must be at least 2-D.

    Returns
    -------
    f : ndarray
        A view of `m` with the columns reversed.  Since a view
        is returned, this operation is :math:`\\mathcal O(1)`.

    See Also
    --------
    flipud : Flip array in the up/down direction.
    rot90 : Rotate array counterclockwise.

    Notes
    -----
    Equivalent to m[:,::-1]. Requires the array to be at least 2-D.

    Examples
    --------
    >>> A = np.diag([1.,2.,3.])
    >>> A
    array([[ 1.,  0.,  0.],
           [ 0.,  2.,  0.],
           [ 0.,  0.,  3.]])
    >>> np.fliplr(A)
    array([[ 0.,  0.,  1.],
           [ 0.,  2.,  0.],
           [ 3.,  0.,  0.]])

    >>> A = np.random.randn(2,3,5)
    >>> np.all(np.fliplr(A) == A[:,::-1,...])
    True

    """
    m = asanyarray(m)
    if m.ndim < 2:
        raise ValueError("Input must be >= 2-d.")
    return m[:, ::-1]

@array_function_dispatch(_flip_dispatcher)
def flipud(m):
    """
    Flip array in the up/down direction.

    Flip the entries in each column in the up/down direction.
    Rows are preserved, but appear in a different order than before.

    Parameters
    ----------
    m : array_like
        Input array.

    Returns
    -------
    out : array_like
        A view of `m` with the rows reversed.  Since a view is
        returned, this operation is :math:`\\mathcal O(1)`.

    See Also
    --------
    fliplr : Flip array in the left/right direction.
    rot90 : Rotate array counterclockwise.

    Notes
    -----
    Equivalent to ``m[::-1,...]``.
    Does not require the array to be two-dimensional.

    Examples
    --------
    >>> A = np.diag([1.0, 2, 3])
    >>> A
    array([[ 1.,  0.,  0.],
           [ 0.,  2.,  0.],
           [ 0.,  0.,  3.]])
    >>> np.flipud(A)
    array([[ 0.,  0.,  3.],
           [ 0.,  2.,  0.],
           [ 1.,  0.,  0.]])

    >>> A = np.random.randn(2,3,5)
    >>> np.all(np.flipud(A) == A[::-1,...])
    True

    >>> np.flipud([1,2])
    array([2, 1])

    """
    m = asanyarray(m)
    if m.ndim < 1:
        raise ValueError("Input must be >= 1-d.")
    return m[::-1, ...]

@set_module('numpy')
def eye(N, M=None, k=0, dtype=float, order='C'):
    """
    Return a 2-D array with ones on the diagonal and zeros elsewhere.

    Parameters
    ----------
    N : int
      Number of rows in the output.
    M : int, optional
      Number of columns in the output. If None, defaults to `N`.
    k : int, optional
      Index of the diagonal: 0 (the default) refers to the main diagonal,
      a positive value refers to an upper diagonal, and a negative value
      to a lower diagonal.
    dtype : data-type, optional
      Data-type of the returned array.
    order : {'C', 'F'}, optional
        Whether the output should be stored in row-major (C-style) or
        column-major (Fortran-style) order in memory.

        .. versionadded:: 1.14.0

    Returns
    -------
    I : ndarray of shape (N,M)
      An array where all elements are equal to zero, except for the `k`-th
      diagonal, whose values are equal to one.

    See Also
    --------
    identity : (almost) equivalent function
    diag : diagonal 2-D array from a 1-D array specified by the user.

    Examples
    --------
    >>> np.eye(2, dtype=int)
    array([[1, 0],
           [0, 1]])
    >>> np.eye(3, k=1)
    array([[ 0.,  1.,  0.],
           [ 0.,  0.,  1.],
           [ 0.,  0.,  0.]])

    """
    if M is None:
        M = N
    m = zeros((N, M), dtype=dtype, order=order)
    if k >= M:
        return m
    if k >= 0:
        i = k
    else:
        i = (-k) * M
    m[:M-k].flat[i::M+1] = 1
    return m

def _diag_dispatcher(v, k=None):
    return (v,)

@array_function_dispatch(_diag_dispatcher)
def diag(v, k=0):
    """
    Extract a diagonal or construct a diagonal array.

    See the more detailed documentation for ``numpy.diagonal`` if you use this
    function to extract a diagonal and wish to write to the resulting array;
    whether it returns a copy or a view depends on what version of numpy you
    are using.

    Parameters
    ----------
    v : array_like
        If `v` is a 2-D array, return a copy of its `k`-th diagonal.
        If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th
        diagonal.
    k : int, optional
        Diagonal in question. The default is 0. Use `k>0` for diagonals
        above the main diagonal, and `k<0` for diagonals below the main
        diagonal.

    Returns
    -------
    out : ndarray
        The extracted diagonal or constructed diagonal array.

    See Also
    --------
    diagonal : Return specified diagonals.
    diagflat : Create a 2-D array with the flattened input as a diagonal.
    trace : Sum along diagonals.
    triu : Upper triangle of an array.
    tril : Lower triangle of an array.

    Examples
    --------
    >>> x = np.arange(9).reshape((3,3))
    >>> x
    array([[0, 1, 2],
           [3, 4, 5],
           [6, 7, 8]])

    >>> np.diag(x)
    array([0, 4, 8])
    >>> np.diag(x, k=1)
    array([1, 5])
    >>> np.diag(x, k=-1)
    array([3, 7])

    >>> np.diag(np.diag(x))
    array([[0, 0, 0],
           [0, 4, 0],
           [0, 0, 8]])

    """
    v = asanyarray(v)
    s = v.shape
    if len(s) == 1:
        n = s[0]+abs(k)
        res = zeros((n, n), v.dtype)
        if k >= 0:
            i = k
        else:
            i = (-k) * n
        res[:n-k].flat[i::n+1] = v
        return res
    elif len(s) == 2:
        return diagonal(v, k)
    else:
        raise ValueError("Input must be 1- or 2-d.")

@array_function_dispatch(_diag_dispatcher)
def diagflat(v, k=0):
    """
    Create a two-dimensional array with the flattened input as a diagonal.

    Parameters
    ----------
    v : array_like
        Input data, which is flattened and set as the `k`-th
        diagonal of the output.
    k : int, optional
        Diagonal to set; 0, the default, corresponds to the "main" diagonal,
        a positive (negative) `k` giving the number of the diagonal above
        (below) the main.

    Returns
    -------
    out : ndarray
        The 2-D output array.

    See Also
    --------
    diag : MATLAB work-alike for 1-D and 2-D arrays.
    diagonal : Return specified diagonals.
    trace : Sum along diagonals.

    Examples
    --------
    >>> np.diagflat([[1,2], [3,4]])
    array([[1, 0, 0, 0],
           [0, 2, 0, 0],
           [0, 0, 3, 0],
           [0, 0, 0, 4]])

    >>> np.diagflat([1,2], 1)
    array([[0, 1, 0],
           [0, 0, 2],
           [0, 0, 0]])

    """
    try:
        wrap = v.__array_wrap__
    except AttributeError:
        wrap = None
    v = asarray(v).ravel()
    s = len(v)
    n = s + abs(k)
    res = zeros((n, n), v.dtype)
    if (k >= 0):
        i = arange(0, n-k)
        fi = i+k+i*n
    else:
        i = arange(0, n+k)
        fi = i+(i-k)*n
    res.flat[fi] = v
    if not wrap:
        return res
    return wrap(res)

@set_module('numpy')
def tri(N, M=None, k=0, dtype=float):
    """
    An array with ones at and below the given diagonal and zeros elsewhere.

    Parameters
    ----------
    N : int
        Number of rows in the array.
    M : int, optional
        Number of columns in the array.
        By default, `M` is taken equal to `N`.
    k : int, optional
        The sub-diagonal at and below which the array is filled.
        `k` = 0 is the main diagonal, while `k` < 0 is below it,
        and `k` > 0 is above.  The default is 0.
    dtype : dtype, optional
        Data type of the returned array.  The default is float.

    Returns
    -------
    tri : ndarray of shape (N, M)
        Array with its lower triangle filled with ones and zero elsewhere;
        in other words ``T[i,j] == 1`` for ``i <= j + k``, 0 otherwise.

    Examples
    --------
    >>> np.tri(3, 5, 2, dtype=int)
    array([[1, 1, 1, 0, 0],
           [1, 1, 1, 1, 0],
           [1, 1, 1, 1, 1]])

    >>> np.tri(3, 5, -1)
    array([[ 0.,  0.,  0.,  0.,  0.],
           [ 1.,  0.,  0.,  0.,  0.],
           [ 1.,  1.,  0.,  0.,  0.]])

    """
    if M is None:
        M = N

    m = greater_equal.outer(arange(N, dtype=_min_int(0, N)),
                            arange(-k, M-k, dtype=_min_int(-k, M - k)))

    # Avoid making a copy if the requested type is already bool
    m = m.astype(dtype, copy=False)

    return m

def _trilu_dispatcher(m, k=None):
    return (m,)

@array_function_dispatch(_trilu_dispatcher)
def tril(m, k=0):
    """
    Lower triangle of an array.

    Return a copy of an array with elements above the `k`-th diagonal zeroed.

    Parameters
    ----------
    m : array_like, shape (M, N)
        Input array.
    k : int, optional
        Diagonal above which to zero elements.  `k = 0` (the default) is the
        main diagonal, `k < 0` is below it and `k > 0` is above.

    Returns
    -------
    tril : ndarray, shape (M, N)
        Lower triangle of `m`, of same shape and data-type as `m`.

    See Also
    --------
    triu : same thing, only for the upper triangle

    Examples
    --------
    >>> np.tril([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)
    array([[ 0,  0,  0],
           [ 4,  0,  0],
           [ 7,  8,  0],
           [10, 11, 12]])

    """
    m = asanyarray(m)
    mask = tri(*m.shape[-2:], k=k, dtype=bool)

    return where(mask, m, zeros(1, m.dtype))

@array_function_dispatch(_trilu_dispatcher)
def triu(m, k=0):
    """
    Upper triangle of an array.

    Return a copy of a matrix with the elements below the `k`-th diagonal
    zeroed.

    Please refer to the documentation for `tril` for further details.

    See Also
    --------
    tril : lower triangle of an array

    Examples
    --------
    >>> np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)
    array([[ 1,  2,  3],
           [ 4,  5,  6],
           [ 0,  8,  9],
           [ 0,  0, 12]])

    """
    m = asanyarray(m)
    mask = tri(*m.shape[-2:], k=k-1, dtype=bool)

    return where(mask, zeros(1, m.dtype), m)

def _vander_dispatcher(x, N=None, increasing=None):
    return (x,)

# Originally borrowed from John Hunter and matplotlib
@array_function_dispatch(_vander_dispatcher)
def vander(x, N=None, increasing=False):
    """
    Generate a Vandermonde matrix.

    The columns of the output matrix are powers of the input vector. The
    order of the powers is determined by the `increasing` boolean argument.
    Specifically, when `increasing` is False, the `i`-th output column is
    the input vector raised element-wise to the power of ``N - i - 1``. Such
    a matrix with a geometric progression in each row is named for Alexandre-
    Theophile Vandermonde.

    Parameters
    ----------
    x : array_like
        1-D input array.
    N : int, optional
        Number of columns in the output.  If `N` is not specified, a square
        array is returned (``N = len(x)``).
    increasing : bool, optional
        Order of the powers of the columns.  If True, the powers increase
        from left to right, if False (the default) they are reversed.

        .. versionadded:: 1.9.0

    Returns
    -------
    out : ndarray
        Vandermonde matrix.  If `increasing` is False, the first column is
        ``x^(N-1)``, the second ``x^(N-2)`` and so forth. If `increasing` is
        True, the columns are ``x^0, x^1, ..., x^(N-1)``.

    See Also
    --------
    polynomial.polynomial.polyvander

    Examples
    --------
    >>> x = np.array([1, 2, 3, 5])
    >>> N = 3
    >>> np.vander(x, N)
    array([[ 1,  1,  1],
           [ 4,  2,  1],
           [ 9,  3,  1],
           [25,  5,  1]])

    >>> np.column_stack([x**(N-1-i) for i in range(N)])
    array([[ 1,  1,  1],
           [ 4,  2,  1],
           [ 9,  3,  1],
           [25,  5,  1]])

    >>> x = np.array([1, 2, 3, 5])
    >>> np.vander(x)
    array([[  1,   1,   1,   1],
           [  8,   4,   2,   1],
           [ 27,   9,   3,   1],
           [125,  25,   5,   1]])
    >>> np.vander(x, increasing=True)
    array([[  1,   1,   1,   1],
           [  1,   2,   4,   8],
           [  1,   3,   9,  27],
           [  1,   5,  25, 125]])

    The determinant of a square Vandermonde matrix is the product
    of the differences between the values of the input vector:

    >>> np.linalg.det(np.vander(x))
    48.000000000000043
    >>> (5-3)*(5-2)*(5-1)*(3-2)*(3-1)*(2-1)
    48

    """
    x = asarray(x)
    if x.ndim != 1:
        raise ValueError("x must be a one-dimensional array or sequence.")
    if N is None:
        N = len(x)

    v = empty((len(x), N), dtype=promote_types(x.dtype, int))
    tmp = v[:, ::-1] if not increasing else v

    if N > 0:
        tmp[:, 0] = 1
    if N > 1:
        tmp[:, 1:] = x[:, None]
        multiply.accumulate(tmp[:, 1:], out=tmp[:, 1:], axis=1)

    return v

def _histogram2d_dispatcher(x, y, bins=None, range=None, normed=None,
                            weights=None, density=None):
    return (x, y, bins, weights)

@array_function_dispatch(_histogram2d_dispatcher)
def histogram2d(x, y, bins=10, range=None, normed=None, weights=None,
                density=None):
    """
    Compute the bi-dimensional histogram of two data samples.

    Parameters
    ----------
    x : array_like, shape (N,)
        An array containing the x coordinates of the points to be
        histogrammed.
    y : array_like, shape (N,)
        An array containing the y coordinates of the points to be
        histogrammed.
    bins : int or array_like or [int, int] or [array, array], optional
        The bin specification:

          * If int, the number of bins for the two dimensions (nx=ny=bins).
          * If array_like, the bin edges for the two dimensions
            (x_edges=y_edges=bins).
          * If [int, int], the number of bins in each dimension
            (nx, ny = bins).
          * If [array, array], the bin edges in each dimension
            (x_edges, y_edges = bins).
          * A combination [int, array] or [array, int], where int
            is the number of bins and array is the bin edges.

    range : array_like, shape(2,2), optional
        The leftmost and rightmost edges of the bins along each dimension
        (if not specified explicitly in the `bins` parameters):
        ``[[xmin, xmax], [ymin, ymax]]``. All values outside of this range
        will be considered outliers and not tallied in the histogram.
    density : bool, optional
        If False, the default, returns the number of samples in each bin.
        If True, returns the probability *density* function at the bin,
        ``bin_count / sample_count / bin_area``.
    normed : bool, optional
        An alias for the density argument that behaves identically. To avoid
        confusion with the broken normed argument to `histogram`, `density`
        should be preferred.
    weights : array_like, shape(N,), optional
        An array of values ``w_i`` weighing each sample ``(x_i, y_i)``.
        Weights are normalized to 1 if `normed` is True. If `normed` is
        False, the values of the returned histogram are equal to the sum of
        the weights belonging to the samples falling into each bin.

    Returns
    -------
    H : ndarray, shape(nx, ny)
        The bi-dimensional histogram of samples `x` and `y`. Values in `x`
        are histogrammed along the first dimension and values in `y` are
        histogrammed along the second dimension.
    xedges : ndarray, shape(nx+1,)
        The bin edges along the first dimension.
    yedges : ndarray, shape(ny+1,)
        The bin edges along the second dimension.

    See Also
    --------
    histogram : 1D histogram
    histogramdd : Multidimensional histogram

    Notes
    -----
    When `normed` is True, then the returned histogram is the sample
    density, defined such that the sum over bins of the product
    ``bin_value * bin_area`` is 1.

    Please note that the histogram does not follow the Cartesian convention
    where `x` values are on the abscissa and `y` values on the ordinate
    axis.  Rather, `x` is histogrammed along the first dimension of the
    array (vertical), and `y` along the second dimension of the array
    (horizontal).  This ensures compatibility with `histogramdd`.

    Examples
    --------
    >>> from matplotlib.image import NonUniformImage
    >>> import matplotlib.pyplot as plt

    Construct a 2-D histogram with variable bin width. First define the bin
    edges:

    >>> xedges = [0, 1, 3, 5]
    >>> yedges = [0, 2, 3, 4, 6]

    Next we create a histogram H with random bin content:

    >>> x = np.random.normal(2, 1, 100)
    >>> y = np.random.normal(1, 1, 100)
    >>> H, xedges, yedges = np.histogram2d(x, y, bins=(xedges, yedges))
    >>> H = H.T  # Let each row list bins with common y range.

    :func:`imshow <matplotlib.pyplot.imshow>` can only display square bins:

    >>> fig = plt.figure(figsize=(7, 3))
    >>> ax = fig.add_subplot(131, title='imshow: square bins')
    >>> plt.imshow(H, interpolation='nearest', origin='low',
    ...         extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]])

    :func:`pcolormesh <matplotlib.pyplot.pcolormesh>` can display actual edges:

    >>> ax = fig.add_subplot(132, title='pcolormesh: actual edges',
    ...         aspect='equal')
    >>> X, Y = np.meshgrid(xedges, yedges)
    >>> ax.pcolormesh(X, Y, H)

    :class:`NonUniformImage <matplotlib.image.NonUniformImage>` can be used to
    display actual bin edges with interpolation:

    >>> ax = fig.add_subplot(133, title='NonUniformImage: interpolated',
    ...         aspect='equal', xlim=xedges[[0, -1]], ylim=yedges[[0, -1]])
    >>> im = NonUniformImage(ax, interpolation='bilinear')
    >>> xcenters = (xedges[:-1] + xedges[1:]) / 2
    >>> ycenters = (yedges[:-1] + yedges[1:]) / 2
    >>> im.set_data(xcenters, ycenters, H)
    >>> ax.images.append(im)
    >>> plt.show()

    """
    from numpy import histogramdd

    try:
        N = len(bins)
    except TypeError:
        N = 1

    if N != 1 and N != 2:
        xedges = yedges = asarray(bins)
        bins = [xedges, yedges]
    hist, edges = histogramdd([x, y], bins, range, normed, weights, density)
    return hist, edges[0], edges[1]

@set_module('numpy')
def mask_indices(n, mask_func, k=0):
    """
    Return the indices to access (n, n) arrays, given a masking function.

    Assume `mask_func` is a function that, for a square array a of size
    ``(n, n)`` with a possible offset argument `k`, when called as
    ``mask_func(a, k)`` returns a new array with zeros in certain locations
    (functions like `triu` or `tril` do precisely this). Then this function
    returns the indices where the non-zero values would be located.

    Parameters
    ----------
    n : int
        The returned indices will be valid to access arrays of shape (n, n).
    mask_func : callable
        A function whose call signature is similar to that of `triu`, `tril`.
        That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`.
        `k` is an optional argument to the function.
    k : scalar
        An optional argument which is passed through to `mask_func`. Functions
        like `triu`, `tril` take a second argument that is interpreted as an
        offset.

    Returns
    -------
    indices : tuple of arrays.
        The `n` arrays of indices corresponding to the locations where
        ``mask_func(np.ones((n, n)), k)`` is True.

    See Also
    --------
    triu, tril, triu_indices, tril_indices

    Notes
    -----
    .. versionadded:: 1.4.0

    Examples
    --------
    These are the indices that would allow you to access the upper triangular
    part of any 3x3 array:

    >>> iu = np.mask_indices(3, np.triu)

    For example, if `a` is a 3x3 array:

    >>> a = np.arange(9).reshape(3, 3)
    >>> a
    array([[0, 1, 2],
           [3, 4, 5],
           [6, 7, 8]])
    >>> a[iu]
    array([0, 1, 2, 4, 5, 8])

    An offset can be passed also to the masking function.  This gets us the
    indices starting on the first diagonal right of the main one:

    >>> iu1 = np.mask_indices(3, np.triu, 1)

    with which we now extract only three elements:

    >>> a[iu1]
    array([1, 2, 5])

    """
    m = ones((n, n), int)
    a = mask_func(m, k)
    return nonzero(a != 0)

@set_module('numpy')
def tril_indices(n, k=0, m=None):
    """
    Return the indices for the lower-triangle of an (n, m) array.

    Parameters
    ----------
    n : int
        The row dimension of the arrays for which the returned
        indices will be valid.
    k : int, optional
        Diagonal offset (see `tril` for details).
    m : int, optional
        .. versionadded:: 1.9.0

        The column dimension of the arrays for which the returned
        arrays will be valid.
        By default `m` is taken equal to `n`.

    Returns
    -------
    inds : tuple of arrays
        The indices for the triangle. The returned tuple contains two arrays,
        each with the indices along one dimension of the array.

    See also
    --------
    triu_indices : similar function, for upper-triangular.
    mask_indices : generic function accepting an arbitrary mask function.
    tril, triu

    Notes
    -----
    .. versionadded:: 1.4.0

    Examples
    --------
    Compute two different sets of indices to access 4x4 arrays, one for the
    lower triangular part starting at the main diagonal, and one starting two
    diagonals further right:

    >>> il1 = np.tril_indices(4)
    >>> il2 = np.tril_indices(4, 2)

    Here is how they can be used with a sample array:

    >>> a = np.arange(16).reshape(4, 4)
    >>> a
    array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11],
           [12, 13, 14, 15]])

    Both for indexing:

    >>> a[il1]
    array([ 0,  4,  5,  8,  9, 10, 12, 13, 14, 15])

    And for assigning values:

    >>> a[il1] = -1
    >>> a
    array([[-1,  1,  2,  3],
           [-1, -1,  6,  7],
           [-1, -1, -1, 11],
           [-1, -1, -1, -1]])

    These cover almost the whole array (two diagonals right of the main one):

    >>> a[il2] = -10
    >>> a
    array([[-10, -10, -10,   3],
           [-10, -10, -10, -10],
           [-10, -10, -10, -10],
           [-10, -10, -10, -10]])

    """
    return nonzero(tri(n, m, k=k, dtype=bool))

def _trilu_indices_form_dispatcher(arr, k=None):
    return (arr,)

@array_function_dispatch(_trilu_indices_form_dispatcher)
def tril_indices_from(arr, k=0):
    """
    Return the indices for the lower-triangle of arr.

    See `tril_indices` for full details.

    Parameters
    ----------
    arr : array_like
        The indices will be valid for square arrays whose dimensions are
        the same as arr.
    k : int, optional
        Diagonal offset (see `tril` for details).

    See Also
    --------
    tril_indices, tril

    Notes
    -----
    .. versionadded:: 1.4.0

    """
    if arr.ndim != 2:
        raise ValueError("input array must be 2-d")
    return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1])

@set_module('numpy')
def triu_indices(n, k=0, m=None):
    """
    Return the indices for the upper-triangle of an (n, m) array.

    Parameters
    ----------
    n : int
        The size of the arrays for which the returned indices will
        be valid.
    k : int, optional
        Diagonal offset (see `triu` for details).
    m : int, optional
        .. versionadded:: 1.9.0

        The column dimension of the arrays for which the returned
        arrays will be valid.
        By default `m` is taken equal to `n`.

    Returns
    -------
    inds : tuple, shape(2) of ndarrays, shape(`n`)
        The indices for the triangle. The returned tuple contains two arrays,
        each with the indices along one dimension of the array.  Can be used
        to slice a ndarray of shape(`n`, `n`).

    See also
    --------
    tril_indices : similar function, for lower-triangular.
    mask_indices : generic function accepting an arbitrary mask function.
    triu, tril

    Notes
    -----
    .. versionadded:: 1.4.0

    Examples
    --------
    Compute two different sets of indices to access 4x4 arrays, one for the
    upper triangular part starting at the main diagonal, and one starting two
    diagonals further right:

    >>> iu1 = np.triu_indices(4)
    >>> iu2 = np.triu_indices(4, 2)

    Here is how they can be used with a sample array:

    >>> a = np.arange(16).reshape(4, 4)
    >>> a
    array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11],
           [12, 13, 14, 15]])

    Both for indexing:

    >>> a[iu1]
    array([ 0,  1,  2,  3,  5,  6,  7, 10, 11, 15])

    And for assigning values:

    >>> a[iu1] = -1
    >>> a
    array([[-1, -1, -1, -1],
           [ 4, -1, -1, -1],
           [ 8,  9, -1, -1],
           [12, 13, 14, -1]])

    These cover only a small part of the whole array (two diagonals right
    of the main one):

    >>> a[iu2] = -10
    >>> a
    array([[ -1,  -1, -10, -10],
           [  4,  -1,  -1, -10],
           [  8,   9,  -1,  -1],
           [ 12,  13,  14,  -1]])

    """
    return nonzero(~tri(n, m, k=k-1, dtype=bool))

@array_function_dispatch(_trilu_indices_form_dispatcher)
def triu_indices_from(arr, k=0):
    """
    Return the indices for the upper-triangle of arr.

    See `triu_indices` for full details.

    Parameters
    ----------
    arr : ndarray, shape(N, N)
        The indices will be valid for square arrays.
    k : int, optional
        Diagonal offset (see `triu` for details).

    Returns
    -------
    triu_indices_from : tuple, shape(2) of ndarray, shape(N)
        Indices for the upper-triangle of `arr`.

    See Also
    --------
    triu_indices, triu

    Notes
    -----
    .. versionadded:: 1.4.0

    """
    if arr.ndim != 2:
        raise ValueError("input array must be 2-d")
    return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1])


"""
========================================
Plot multi-class SGD on the iris dataset
========================================

Plot decision surface of multi-class SGD on iris dataset.
The hyperplanes corresponding to the three one-versus-all (OVA) classifiers
are represented by the dashed lines.

"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn.linear_model import SGDClassifier

# 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
colors = "bry"

# shuffle
idx = np.arange(X.shape[0])
np.random.seed(13)
np.random.shuffle(idx)
X = X[idx]
y = y[idx]

# standardize
mean = X.mean(axis=0)
std = X.std(axis=0)
X = (X - mean) / std

h = .02  # step size in the mesh

clf = SGDClassifier(alpha=0.001, n_iter=100).fit(X, y)

# create a mesh to plot in
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))

# 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].
Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])
# Put the result into a color plot
Z = Z.reshape(xx.shape)
cs = plt.contourf(xx, yy, Z, cmap=plt.cm.Paired)
plt.axis('tight')

# Plot also the training points
for i, color in zip(clf.classes_, colors):
    idx = np.where(y == i)
    plt.scatter(X[idx, 0], X[idx, 1], c=color, label=iris.target_names[i],
                cmap=plt.cm.Paired)
plt.title("Decision surface of multi-class SGD")
plt.axis('tight')

# Plot the three one-against-all classifiers
xmin, xmax = plt.xlim()
ymin, ymax = plt.ylim()
coef = clf.coef_
intercept = clf.intercept_

def plot_hyperplane(c, color):
    def line(x0):
        return (-(x0 * coef[c, 0]) - intercept[c]) / coef[c, 1]

    plt.plot([xmin, xmax], [line(xmin), line(xmax)],
             ls="--", color=color)

for i, color in zip(clf.classes_, colors):
    plot_hyperplane(i, color)
plt.legend()
plt.show()


import numpy as np
import pandas as pd

arts = pd.DataFrame()

# Clean the dates so you only see numbers.
def clean_years(value):
    result = value
    chars_to_replace = ["c.", "©", ", CARCC", "no date", "n.d.", " SODRAC", ", CA", " CARCC", ""]
    chars_to_split = ["-", "/"]
    if isinstance(result, str):
        for char in chars_to_split:
            if char in result:
                result = result.split(char)[1].strip()
        for char in chars_to_replace:
            result = result.replace(char, "")
        if result == "":
            return np.nan
        else:
            return int(result)
    else:
        return result

arts['execution_date'] = arts['execution_date'].apply(clean_years)
arts.head()

# If a year is lower than 100, then is referred to 1900. For example, 78 is actually 1978, and that needs to be fixed too.
def clean_year_99(value):
    if value < 100:
        return value + 1900
    else:
        return value

arts["execution_date"] = arts["execution_date"].apply(clean_year_99)
arts.head()


from __future__ import print_function
from builtins import zip
from past.builtins import basestring
import csv
import logging
import subprocess
from tempfile import NamedTemporaryFile

from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol
from hive_service import ThriftHive
import pyhs2

from airflow.utils import AirflowException
from airflow.hooks.base_hook import BaseHook
from airflow.utils import TemporaryDirectory

class HiveCliHook(BaseHook):
    """
    Simple wrapper around the hive CLI.

    It also supports the ``beeline``
    a lighter CLI that runs JDBC and is replacing the heavier
    traditional CLI. To enable ``beeline``, set the use_beeline param in the
    extra field of your connection as in ``{ "use_beeline": true }``

    Note that you can also set default hive CLI parameters using the
    ``hive_cli_params`` to be used in your connection as in
    ``{"hive_cli_params": "-hiveconf mapred.job.tracker=some.jobtracker:444"}``
    """

    def __init__(
            self,
            hive_cli_conn_id="hive_cli_default"):
        conn = self.get_connection(hive_cli_conn_id)
        self.hive_cli_params = conn.extra_dejson.get('hive_cli_params', '')
        self.use_beeline = conn.extra_dejson.get('use_beeline', False)
        self.conn = conn

    def run_cli(self, hql, schema=None):
        """
        Run an hql statement using the hive cli

        >>> hh = HiveCliHook()
        >>> result = hh.run_cli("USE airflow;")
        >>> ("OK" in result)
        True
        """
        conn = self.conn
        schema = schema or conn.schema
        if schema:
            hql = "USE {schema};\n{hql}".format(**locals())

        with TemporaryDirectory(prefix='airflow_hiveop_') as tmp_dir:
            with NamedTemporaryFile(dir=tmp_dir) as f:
                f.write(hql)
                f.flush()
                fname = f.name
                hive_bin = 'hive'
                cmd_extra = []
                if self.use_beeline:
                    hive_bin = 'beeline'
                    jdbc_url = (
                        "jdbc:hive2://"
                        "{0}:{1}/{2}"
                        ";auth=noSasl"
                    ).format(conn.host, conn.port, conn.schema)
                    cmd_extra += ['-u', jdbc_url]
                    if conn.login:
                        cmd_extra += ['-n', conn.login]
                    if conn.password:
                        cmd_extra += ['-p', conn.password]
                    cmd_extra += ['-p', conn.login]
                hive_cmd = [hive_bin, '-f', fname] + cmd_extra
                if self.hive_cli_params:
                    hive_params_list = self.hive_cli_params.split()
                    hive_cmd.extend(hive_params_list)
                logging.info(" ".join(hive_cmd))
                sp = subprocess.Popen(
                    hive_cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    cwd=tmp_dir)
                all_err = ''
                self.sp = sp
                stdout = ''
                for line in iter(sp.stdout.readline, ''):
                    stdout += line
                    logging.info(line.strip())
                sp.wait()

                if sp.returncode:
                    raise AirflowException(all_err)

                return stdout

    def load_file(
            self,
            filepath,
            table,
            delimiter=",",
            field_dict=None,
            create=True,
            overwrite=True,
            partition=None,
            recreate=False):
        """
        Loads a local file into Hive

        Note that the table generated in Hive uses ``STORED AS textfile``
        which isn't the most efficient serialization format. If a
        large amount of data is loaded and/or if the tables gets
        queried considerably, you may want to use this operator only to
        stage the data into a temporary table before loading it into its
        final destination using a ``HiveOperator``.

        :param table: target Hive table, use dot notation to target a
            specific database
        :type table: str
        :param create: whether to create the table if it doesn't exist
        :type create: bool
        :param recreate: whether to drop and recreate the table at every
            execution
        :type recreate: bool
        :param partition: target partition as a dict of partition columns
            and values
        :type partition: dict
        :param delimiter: field delimiter in the file
        :type delimiter: str
        """
        hql = ''
        if recreate:
            hql += "DROP TABLE IF EXISTS {table};\n"
        if create or recreate:
            fields = ",\n    ".join(
                [k + ' ' + v for k, v in field_dict.items()])
            hql += "CREATE TABLE IF NOT EXISTS {table} (\n{fields})\n"
            if partition:
                pfields = ",\n    ".join(
                    [p + " STRING" for p in partition])
                hql += "PARTITIONED BY ({pfields})\n"
            hql += "ROW FORMAT DELIMITED\n"
            hql += "FIELDS TERMINATED BY '{delimiter}'\n"
            hql += "STORED AS textfile;"
        hql = hql.format(**locals())
        logging.info(hql)
        self.run_cli(hql)
        hql = "LOAD DATA LOCAL INPATH '{filepath}' "
        if overwrite:
            hql += "OVERWRITE "
        hql += "INTO TABLE {table} "
        if partition:
            pvals = ", ".join(
                ["{0}='{1}'".format(k, v) for k, v in partition.items()])
            hql += "PARTITION ({pvals});"
        hql = hql.format(**locals())
        logging.info(hql)
        self.run_cli(hql)

    def kill(self):
        if hasattr(self, 'sp'):
            if self.sp.poll() is None:
                print("Killing the Hive job")
                self.sp.kill()

class HiveMetastoreHook(BaseHook):
    '''
    Wrapper to interact with the Hive Metastore
    '''
    def __init__(self, metastore_conn_id='metastore_default'):
        self.metastore_conn = self.get_connection(metastore_conn_id)
        self.metastore = self.get_metastore_client()

    def __getstate__(self):
        # This is for pickling to work despite the thirft hive client not
        # being pickable
        d = dict(self.__dict__)
        del d['metastore']
        return d

    def __setstate__(self, d):
        self.__dict__.update(d)
        self.__dict__['metastore'] = self.get_metastore_client()

    def get_metastore_client(self):
        '''
        Returns a Hive thrift client.
        '''
        ms = self.metastore_conn
        transport = TSocket.TSocket(ms.host, ms.port)
        transport = TTransport.TBufferedTransport(transport)
        protocol = TBinaryProtocol.TBinaryProtocol(transport)
        return ThriftHive.Client(protocol)

    def get_conn(self):
        return self.metastore

    def check_for_partition(self, schema, table, partition):
        '''
        Checks whether a partition exists

        >>> hh = HiveMetastoreHook()
        >>> t = 'static_babynames_partitioned'
        >>> hh.check_for_partition('airflow', t, "ds='2015-01-01'")
        True
        '''
        self.metastore._oprot.trans.open()
        partitions = self.metastore.get_partitions_by_filter(
            schema, table, partition, 1)
        self.metastore._oprot.trans.close()
        if partitions:
            return True
        else:
            return False

    def get_table(self, table_name, db='default'):
        '''
        Get a metastore table object

        >>> hh = HiveMetastoreHook()
        >>> t = hh.get_table(db='airflow', table_name='static_babynames')
        >>> t.tableName
        'static_babynames'
        >>> [col.name for col in t.sd.cols]
        ['state', 'year', 'name', 'gender', 'num']
        '''
        self.metastore._oprot.trans.open()
        if db == 'default' and '.' in table_name:
            db, table_name = table_name.split('.')[:2]
        table = self.metastore.get_table(dbname=db, tbl_name=table_name)
        self.metastore._oprot.trans.close()
        return table

    def get_tables(self, db, pattern='*'):
        '''
        Get a metastore table object
        '''
        self.metastore._oprot.trans.open()
        tables = self.metastore.get_tables(db_name=db, pattern=pattern)
        objs = self.metastore.get_table_objects_by_name(db, tables)
        self.metastore._oprot.trans.close()
        return objs

    def get_databases(self, pattern='*'):
        '''
        Get a metastore table object
        '''
        self.metastore._oprot.trans.open()
        dbs = self.metastore.get_databases(pattern)
        self.metastore._oprot.trans.close()
        return dbs

    def get_partitions(
            self, schema, table_name, filter=None):
        '''
        Returns a list of all partitions in a table. Works only
        for tables with less than 32767 (java short max val).
        For subpartitioned table, the number might easily exceed this.

        >>> hh = HiveMetastoreHook()
        >>> t = 'static_babynames_partitioned'
        >>> parts = hh.get_partitions(schema='airflow', table_name=t)
        >>> len(parts)
        1
        >>> parts
        [{'ds': '2015-01-01'}]
        '''
        self.metastore._oprot.trans.open()
        table = self.metastore.get_table(dbname=schema, tbl_name=table_name)
        if len(table.partitionKeys) == 0:
            raise AirflowException("The table isn't partitioned")
        else:
            if filter:
                parts = self.metastore.get_partitions_by_filter(
                    db_name=schema, tbl_name=table_name,
                    filter=filter, max_parts=32767)
            else:
                parts = self.metastore.get_partitions(
                    db_name=schema, tbl_name=table_name, max_parts=32767)

            self.metastore._oprot.trans.close()
            pnames = [p.name for p in table.partitionKeys]
            return [dict(zip(pnames, p.values)) for p in parts]

    def max_partition(self, schema, table_name, field=None, filter=None):
        '''
        Returns the maximum value for all partitions in a table. Works only
        for tables that have a single partition key. For subpartitioned
        table, we recommend using signal tables.

        >>> hh = HiveMetastoreHook()
        >>> t = 'static_babynames_partitioned'
        >>> hh.max_partition(schema='airflow', table_name=t)
        '2015-01-01'
        '''
        parts = self.get_partitions(schema, table_name, filter)
        if not parts:
            return None
        elif len(parts[0]) == 1:
            field = list(parts[0].keys())[0]
        elif not field:
            raise AirflowException(
                "Please specify the field you want the max "
                "value for")

        return max([p[field] for p in parts])

class HiveServer2Hook(BaseHook):
    '''
    Wrapper around the pyhs2 library

    Note that the default authMechanism is NOSASL, to override it you
    can specify it in the ``extra`` of your connection in the UI as in
    ``{"authMechanism": "PLAIN"}``. Refer to the pyhs2 for more details.
    '''
    def __init__(self, hiveserver2_conn_id='hiveserver2_default'):
        self.hiveserver2_conn_id = hiveserver2_conn_id

    def get_conn(self):
        db = self.get_connection(self.hiveserver2_conn_id)
        return pyhs2.connect(
            host=db.host,
            port=db.port,
            authMechanism=db.extra_dejson.get('authMechanism', 'NOSASL'),
            user=db.login,
            database=db.schema or 'default')

    def get_results(self, hql, schema='default', arraysize=1000):
        with self.get_conn() as conn:
            if isinstance(hql, basestring):
                hql = [hql]
            results = {
                'data': [],
                'header': [],
            }
            for statement in hql:
                with conn.cursor() as cur:
                    cur.execute(statement)
                    records = cur.fetchall()
                    if records:
                        results = {
                            'data': records,
                            'header': cur.getSchema(),
                        }
            return results

    def to_csv(self, hql, csv_filepath, schema='default'):
        schema = schema or 'default'
        with self.get_conn() as conn:
            with conn.cursor() as cur:
                logging.info("Running query: " + hql)
                cur.execute(hql)
                schema = cur.getSchema()
                with open(csv_filepath, 'w') as f:
                    writer = csv.writer(f)
                    writer.writerow([c['columnName'] for c in cur.getSchema()])
                    i = 0
                    while cur.hasMoreRows:
                        rows = [row for row in cur.fetchmany() if row]
                        writer.writerows(rows)
                        i += len(rows)
                        logging.info("Written {0} rows so far.".format(i))
                    logging.info("Done. Loaded a total of {0} rows.".format(i))

    def get_records(self, hql, schema='default'):
        '''
        Get a set of records from a Hive query.

        >>> hh = HiveServer2Hook()
        >>> sql = "SELECT * FROM airflow.static_babynames LIMIT 100"
        >>> len(hh.get_records(sql))
        100
        '''
        return self.get_results(hql, schema=schema)['data']

    def get_pandas_df(self, hql, schema='default'):
        '''
        Get a pandas dataframe from a Hive query

        >>> hh = HiveServer2Hook()
        >>> sql = "SELECT * FROM airflow.static_babynames LIMIT 100"
        >>> df = hh.get_pandas_df(sql)
        >>> len(df.index)
        100
        '''
        import pandas as pd
        res = self.get_results(hql, schema=schema)
        df = pd.DataFrame(res['data'])
        df.columns = [c['columnName'] for c in res['header']]
        return df


"""Testing for Spectral Clustering methods"""

from sklearn.externals.six.moves import cPickle

dumps, loads = cPickle.dumps, cPickle.loads

import numpy as np
from scipy import sparse

from sklearn.utils import check_random_state
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_greater
from sklearn.utils.testing import assert_warns_message

from sklearn.cluster import SpectralClustering, spectral_clustering
from sklearn.cluster.spectral import spectral_embedding
from sklearn.cluster.spectral import discretize
from sklearn.metrics import pairwise_distances
from sklearn.metrics import adjusted_rand_score
from sklearn.metrics.pairwise import kernel_metrics, rbf_kernel
from sklearn.datasets.samples_generator import make_blobs

def test_spectral_clustering():
    S = np.array([[1.0, 1.0, 1.0, 0.2, 0.0, 0.0, 0.0],
                  [1.0, 1.0, 1.0, 0.2, 0.0, 0.0, 0.0],
                  [1.0, 1.0, 1.0, 0.2, 0.0, 0.0, 0.0],
                  [0.2, 0.2, 0.2, 1.0, 1.0, 1.0, 1.0],
                  [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0],
                  [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0],
                  [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0]])

    for eigen_solver in ('arpack', 'lobpcg'):
        for assign_labels in ('kmeans', 'discretize'):
            for mat in (S, sparse.csr_matrix(S)):
                model = SpectralClustering(random_state=0, n_clusters=2,
                                           affinity='precomputed',
                                           eigen_solver=eigen_solver,
                                           assign_labels=assign_labels
                                          ).fit(mat)
                labels = model.labels_
                if labels[0] == 0:
                    labels = 1 - labels

                assert_array_equal(labels, [1, 1, 1, 0, 0, 0, 0])

                model_copy = loads(dumps(model))
                assert_equal(model_copy.n_clusters, model.n_clusters)
                assert_equal(model_copy.eigen_solver, model.eigen_solver)
                assert_array_equal(model_copy.labels_, model.labels_)

def test_spectral_amg_mode():
    # Test the amg mode of SpectralClustering
    centers = np.array([
        [0., 0., 0.],
        [10., 10., 10.],
        [20., 20., 20.],
    ])
    X, true_labels = make_blobs(n_samples=100, centers=centers,
                                cluster_std=1., random_state=42)
    D = pairwise_distances(X)  # Distance matrix
    S = np.max(D) - D  # Similarity matrix
    S = sparse.coo_matrix(S)
    try:
        from pyamg import smoothed_aggregation_solver

        amg_loaded = True
    except ImportError:
        amg_loaded = False
    if amg_loaded:
        labels = spectral_clustering(S, n_clusters=len(centers),
                                     random_state=0, eigen_solver="amg")
        # We don't care too much that it's good, just that it *worked*.
        # There does have to be some lower limit on the performance though.
        assert_greater(np.mean(labels == true_labels), .3)
    else:
        assert_raises(ValueError, spectral_embedding, S,
                      n_components=len(centers),
                      random_state=0, eigen_solver="amg")

def test_spectral_unknown_mode():
    # Test that SpectralClustering fails with an unknown mode set.
    centers = np.array([
        [0., 0., 0.],
        [10., 10., 10.],
        [20., 20., 20.],
    ])
    X, true_labels = make_blobs(n_samples=100, centers=centers,
                                cluster_std=1., random_state=42)
    D = pairwise_distances(X)  # Distance matrix
    S = np.max(D) - D  # Similarity matrix
    S = sparse.coo_matrix(S)
    assert_raises(ValueError, spectral_clustering, S, n_clusters=2,
                  random_state=0, eigen_solver="<unknown>")

def test_spectral_unknown_assign_labels():
    # Test that SpectralClustering fails with an unknown assign_labels set.
    centers = np.array([
        [0., 0., 0.],
        [10., 10., 10.],
        [20., 20., 20.],
    ])
    X, true_labels = make_blobs(n_samples=100, centers=centers,
                                cluster_std=1., random_state=42)
    D = pairwise_distances(X)  # Distance matrix
    S = np.max(D) - D  # Similarity matrix
    S = sparse.coo_matrix(S)
    assert_raises(ValueError, spectral_clustering, S, n_clusters=2,
                  random_state=0, assign_labels="<unknown>")

def test_spectral_clustering_sparse():
    X, y = make_blobs(n_samples=20, random_state=0,
                      centers=[[1, 1], [-1, -1]], cluster_std=0.01)

    S = rbf_kernel(X, gamma=1)
    S = np.maximum(S - 1e-4, 0)
    S = sparse.coo_matrix(S)

    labels = SpectralClustering(random_state=0, n_clusters=2,
                                affinity='precomputed').fit(S).labels_
    assert_equal(adjusted_rand_score(y, labels), 1)

def test_affinities():
    # Note: in the following, random_state has been selected to have
    # a dataset that yields a stable eigen decomposition both when built
    # on OSX and Linux
    X, y = make_blobs(n_samples=20, random_state=0,
                      centers=[[1, 1], [-1, -1]], cluster_std=0.01
                     )
    # nearest neighbors affinity
    sp = SpectralClustering(n_clusters=2, affinity='nearest_neighbors',
                            random_state=0)
    assert_warns_message(UserWarning, 'not fully connected', sp.fit, X)
    assert_equal(adjusted_rand_score(y, sp.labels_), 1)

    sp = SpectralClustering(n_clusters=2, gamma=2, random_state=0)
    labels = sp.fit(X).labels_
    assert_equal(adjusted_rand_score(y, labels), 1)

    X = check_random_state(10).rand(10, 5) * 10

    kernels_available = kernel_metrics()
    for kern in kernels_available:
        # Additive chi^2 gives a negative similarity matrix which
        # doesn't make sense for spectral clustering
        if kern != 'additive_chi2':
            sp = SpectralClustering(n_clusters=2, affinity=kern,
                                    random_state=0)
            labels = sp.fit(X).labels_
            assert_equal((X.shape[0],), labels.shape)

    sp = SpectralClustering(n_clusters=2, affinity=lambda x, y: 1,
                            random_state=0)
    labels = sp.fit(X).labels_
    assert_equal((X.shape[0],), labels.shape)

    def histogram(x, y, **kwargs):
        # Histogram kernel implemented as a callable.
        assert_equal(kwargs, {})    # no kernel_params that we didn't ask for
        return np.minimum(x, y).sum()

    sp = SpectralClustering(n_clusters=2, affinity=histogram, random_state=0)
    labels = sp.fit(X).labels_
    assert_equal((X.shape[0],), labels.shape)

    # raise error on unknown affinity
    sp = SpectralClustering(n_clusters=2, affinity='<unknown>')
    assert_raises(ValueError, sp.fit, X)

def test_discretize(seed=8):
    # Test the discretize using a noise assignment matrix
    random_state = np.random.RandomState(seed)
    for n_samples in [50, 100, 150, 500]:
        for n_class in range(2, 10):
            # random class labels
            y_true = random_state.random_integers(0, n_class, n_samples)
            y_true = np.array(y_true, np.float)
            # noise class assignment matrix
            y_indicator = sparse.coo_matrix((np.ones(n_samples),
                                             (np.arange(n_samples),
                                              y_true)),
                                            shape=(n_samples,
                                                   n_class + 1))
            y_true_noisy = (y_indicator.toarray()
                            + 0.1 * random_state.randn(n_samples,
                                                       n_class + 1))
            y_pred = discretize(y_true_noisy, random_state)
            assert_greater(adjusted_rand_score(y_true, y_pred), 0.8)


from sklearn import datasets
from sklearn.neural_network import MLPClassifier
import traceback
from submissions.aartiste import election
from submissions.aartiste import county_demographics

class DataFrame:
    data = []
    feature_names = []
    target = []
    target_names = []

trumpECHP = DataFrame()

'''
Extract data from the CORGIS elections, and merge it with the
CORGIS demographics.  Both data sets are organized by county and state.
'''
joint = {}

elections = election.get_results()
for county in elections:
    try:
        st = county['Location']['State Abbreviation']
        countyST = county['Location']['County'] + st
        trump = county['Vote Data']['Donald Trump']['Percent of Votes']
        joint[countyST] = {}
        joint[countyST]['ST']= st
        joint[countyST]['Trump'] = trump
    except:
        traceback.print_exc()

demographics = county_demographics.get_all_counties()
for county in demographics:
    try:
        countyNames = county['County'].split()
        cName = ' '.join(countyNames[:-1])
        st = county['State']
        countyST = cName + st
        # elderly =
        # college =
        # home =
        # poverty =
        if countyST in joint:
            joint[countyST]['Elderly'] = county['Age']["Percent 65 and Older"]
            joint[countyST]['HighSchool'] = county['Education']["High School or Higher"]
            joint[countyST]['College'] = county['Education']["Bachelor's Degree or Higher"]
            joint[countyST]['White'] = county['Ethnicities']["White Alone, not Hispanic or Latino"]
            joint[countyST]['Persons'] = county['Housing']["Persons per Household"]
            joint[countyST]['Home'] = county['Housing']["Homeownership Rate"]
            joint[countyST]['Income'] = county['Income']["Median Houseold Income"]
            joint[countyST]['Poverty'] = county['Income']["Persons Below Poverty Level"]
            joint[countyST]['Sales'] = county['Sales']["Retail Sales per Capita"]
    except:
        traceback.print_exc()

'''
Remove the counties that did not appear in both samples.
'''
intersection = {}
for countyST in joint:
    if 'College' in joint[countyST]:
        intersection[countyST] = joint[countyST]

trumpECHP.data = []

'''
Build the input frame, row by row.
'''
for countyST in intersection:
    # choose the input values
    row = []
    for key in intersection[countyST]:
        if key in ['ST', 'Trump']:
            continue
        row.append(intersection[countyST][key])
    trumpECHP.data.append(row)

firstCounty = next(iter(intersection.keys()))
firstRow = intersection[firstCounty]
trumpECHP.feature_names = list(firstRow.keys())
trumpECHP.feature_names.remove('ST')
trumpECHP.feature_names.remove('Trump')

'''
Build the target list,
one entry for each row in the input frame.

The Naive Bayesian network is a classifier,
i.e. it sorts data points into bins.
The best it can do to estimate a continuous variable
is to break the domain into segments, and predict
the segment into which the variable's value will fall.
In this example, I'm breaking Trump's % into two
arbitrary segments.
'''
trumpECHP.target = []

def trumpTarget(percentage):
    if percentage > 45:
        return 1
    return 0

for countyST in intersection:
    # choose the target
    tt = trumpTarget(intersection[countyST]['Trump'])
    trumpECHP.target.append(tt)

trumpECHP.target_names = [
    'Trump <= 45%',
    'Trump >  45%',
]

mlpc = MLPClassifier(
    solver='sgd',
    learning_rate = 'adaptive',
)

Examples = {
    'TrumpDefault': {
        'frame': trumpECHP,
    },
    'TrumpSGD': {
        'frame': trumpECHP,
        'mlpc': mlpc
    },
}


# ----------------------------------------------------------------------
# Copyright (C) 2014-2015, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import datetime
import itertools
import numpy
import os
import pandas
try:
  import simplejson as json
except ImportError:
  import json

from nab.util import (absoluteFilePaths,
                      getProbationPeriod,
                      strf,
                      strp,
                      deepmap,
                      createPath,
                      writeJSON)

def bucket(rawTimes, buffer):
  """
  Buckets (groups) timestamps that are within the amount of time specified by
  buffer.
  """
  bucket = []
  rawBuckets = []

  current = None
  for t in rawTimes:
    if current is None:
      current = t
      bucket = [current]
      continue
    if (t - current) <= buffer:
      bucket.append(t)
    else:
      rawBuckets.append(bucket)
      current = t
      bucket = [current]
  if bucket:
    rawBuckets.append(bucket)

  return rawBuckets

def merge(rawBuckets, threshold):
  """
  Merges bucketed timestamps into one timestamp (most frequent, or earliest).
  """
  truths = []
  passed = []

  for bucket in rawBuckets:
    if len(bucket) >= threshold:
      truths.append(max(bucket, key=bucket.count))
    else:
      passed.append(bucket)

  return truths, passed

def checkForOverlap(labels, buffer, labelsFileName, dataFileName):
  """
  Raise a ValueError if the difference between any consecutive labels is smaller
  than the buffer.
  """
  for i in xrange(len(labels)-1):
    if labels[i+1] - labels[i] <= buffer:
      # import pdb; pdb.set_trace()
      raise ValueError("The labels {} and {} in \'{}\' labels for data file "
        "\'{}\' are too close to each other to be considered distinct "
        "anomalies. Please relabel."
        .format(labels[i], labels[i+1], labelsFileName, dataFileName))

class CorpusLabel(object):
  """
  Class to store and manipulate a single set of labels for the whole
  benchmark corpus.
  """

  def __init__(self, path, corpus):
    """
    Initializes a CorpusLabel object by getting the anomaly windows and labels.
    When this is done for combining raw user labels, we skip getLabels()
    because labels are not yet created.

    @param path    (string)      Name of file containing the set of labels.
    @param corpus  (nab.Corpus)  Corpus object.
    """
    self.path = path

    self.windows = None
    self.labels = None

    self.corpus = corpus
    self.getWindows()

    if "raw" not in self.path:
      # Do not get labels from files in the path nab/labels/raw
      self.getLabels()

  def getWindows(self):
    """
    Read JSON label file. Get timestamps as dictionaries with key:value pairs of
    a relative path and its corresponding list of windows.
    """
    def found(t, data):
      f = data["timestamp"][data["timestamp"] == pandas.tslib.Timestamp(t)]
      exists = (len(f) == 1)

      return exists

    with open(os.path.join(self.path)) as windowFile:
      windows = json.load(windowFile)

    self.windows = {}

    for relativePath in windows.keys():

      self.windows[relativePath] = deepmap(strp, windows[relativePath])

      if len(self.windows[relativePath]) == 0:
        continue

      data = self.corpus.dataFiles[relativePath].data
      if "raw" in self.path:
        timestamps = windows[relativePath]
      else:
        timestamps = list(itertools.chain.from_iterable(windows[relativePath]))

      # Check that timestamps are present in dataset
      if not all([found(t,data) for t in timestamps]):
        raise ValueError("In the label file %s, one of the timestamps used for "
                         "the datafile %s doesn't match; it does not exist in "
                         "the file. Timestamps in json label files have to "
                         "exactly match timestamps in corresponding datafiles."
                         % (self.path, relativePath))

  def validateLabels(self):
    """
    This is run at the end of the label combining process (see
    scripts/combine_labels.py) to validate the resulting ground truth windows,
    specifically that they are distinct (unique, non-overlapping).
    """
    with open(os.path.join(self.path)) as windowFile:
      windows = json.load(windowFile)

    self.windows = {}

    for relativePath in windows.keys():

      self.windows[relativePath] = deepmap(strp, windows[relativePath])

      if len(self.windows[relativePath]) == 0:
        continue

      num_windows = len(self.windows[relativePath])
      if num_windows > 1:
        if not all([(self.windows[relativePath][i+1][0]
                    - self.windows[relativePath][i][1]).total_seconds() >= 0
                    for i in xrange(num_windows-1)]):
          raise ValueError("In the label file %s, windows overlap." % self.path)

  def getLabels(self):
    """
    Get Labels as a dictionary of key-value pairs of a relative path and its
    corresponding binary vector of anomaly labels. Labels are simply a more
    verbose version of the windows.
    """
    self.labels = {}

    for relativePath, dataSet in self.corpus.dataFiles.iteritems():
      if self.windows.has_key(relativePath):
        windows = self.windows[relativePath]

        labels = pandas.DataFrame({"timestamp": dataSet.data["timestamp"]})
        labels['label'] = 0

        for t1, t2 in windows:
          moreThanT1 = labels[labels["timestamp"] >= t1]
          betweenT1AndT2 = moreThanT1[moreThanT1["timestamp"] <= t2]
          indices = betweenT1AndT2.loc[:,"label"].index
          labels["label"].values[indices.values] = 1

        self.labels[relativePath] = labels

      else:
        print "Warning: no label for datafile",relativePath

class LabelCombiner(object):
  """
  This class is used to combine labels from multiple human labelers, and the set
  of manual labels (known anomalies).
  The output is a single ground truth label file containing anomalies where
  there is enough human agreement. The class also computes the window around
  each anomaly.  The exact logic is described elsewhere in the NAB
  documentation.
  """

  def __init__(self, labelDir, corpus,
                     threshold, windowSize,
                     probationaryPercent, verbosity):
    """
    @param labelDir   (string)   A directory name containing user label files.
                                 This directory should contain one label file
                                 per human labeler.
    @param corpus     (Corpus)   Instance of Corpus class.
    @param threshold  (float)    A percentage between 0 and 1, specifying the
                                 agreement threshold.  It describes the level
                                 of agreement needed between individual
                                 labelers before a particular point in a
                                 data file is labeled as anomalous in the
                                 combined file.
    @param windowSize (float)    Estimated size of an anomaly window, as a
                                 ratio the dataset length.
    @param verbosity  (int)      0, 1, or 2 to print out select labeling
                                 metrics; 0 is none, 2 is the most.
    """
    self.labelDir = labelDir
    self.corpus = corpus
    self.threshold = threshold
    self.windowSize = windowSize
    self.probationaryPercent = probationaryPercent
    self.verbosity = verbosity

    self.userLabels = None
    self.nLabelers = None
    self.knownLabels = None

    self.combinedWindows = None

  def __str__(self):
    ans = ""
    ans += "labelDir:            %s\n" % self.labelDir
    ans += "corpus:              %s\n" % self.corpus
    ans += "number of labelers:  %d\n" % self.nLabelers
    ans += "agreement threshold: %d\n" % self.threshold
    return ans

  def write(self, labelsPath, windowsPath):
    """Write the combined labels and windows to destination directories."""
    if not os.path.isdir(labelsPath):
      createPath(labelsPath)
    if not os.path.isdir(windowsPath):
      createPath(windowsPath)

    writeJSON(labelsPath, self.labelTimestamps)
    writeJSON(windowsPath, self.combinedWindows)

  def combine(self):
    """Combine raw and known labels in anomaly windows."""
    self.getRawLabels()
    self.combineLabels()
    self.editPoorLabels()
    self.applyWindows()
    self.checkWindows()

  def getRawLabels(self):
    """Collect the raw user labels from specified directory."""
    labelPaths = absoluteFilePaths(self.labelDir)
    self.userLabels = []
    self.knownLabels = []
    for path in labelPaths:
      if "known" in path:
        self.knownLabels.append(CorpusLabel(path, self.corpus))
      else:
        self.userLabels.append(CorpusLabel(path, self.corpus))

    self.nLabelers = len(self.userLabels)
    if self.nLabelers == 0:
      raise ValueError("No users labels found")

  def combineLabels(self):
    """
    Combines raw user labels to create set of true anomaly labels.
    A buffer is used to bucket labels that identify the same anomaly. The buffer
    is half the estimated window size of an anomaly -- approximates an average
    of two anomalies per dataset, and no window can have > 1 anomaly.
    After bucketing, a label becomes a true anomaly if it was labeled by a
    proportion of the users greater than the defined threshold. Then the bucket
    is merged into one timestamp -- the ground truth label.
    The set of known anomaly labels are added as well. These have been manually
    labeled because we know the direct causes of the anomalies. They are added
    as if they are the result of the bucket-merge process.

    If verbosity > 0, the dictionary passedLabels -- the raw labels that did not
    pass the threshold qualification -- is printed to the console.
    """
    def setTruthLabels(dataSet, trueAnomalies):
      """Returns the indices of the ground truth anomalies for a data file."""
      timestamps = dataSet.data["timestamp"]
      labels = numpy.array(timestamps.isin(trueAnomalies), dtype=int)
      return [i for i in range(len(labels)) if labels[i]==1]

    self.labelTimestamps = {}
    self.labelIndices = {}
    for relativePath, dataSet in self.corpus.dataFiles.iteritems():
      if ("Known" in relativePath) or ("artificial" in relativePath):
        knownAnomalies = self.knownLabels[0].windows[relativePath]
        self.labelTimestamps[relativePath] = [str(t) for t in knownAnomalies]
        self.labelIndices[relativePath] = setTruthLabels(dataSet, knownAnomalies)
        continue

      # Calculate the window buffer -- used for bucketing labels identifying
      # the same anomaly.
      granularity = dataSet.data["timestamp"][1] - dataSet.data["timestamp"][0]
      buffer = datetime.timedelta(minutes=
        granularity.total_seconds()/60 * len(dataSet.data) * self.windowSize/10)

      rawTimesLists = []
      userCount = 0
      for user in self.userLabels:
        if relativePath in user.windows:
          # the user has labels for this file
          checkForOverlap(
            user.windows[relativePath], buffer, user.path, relativePath)
          rawTimesLists.append(user.windows[relativePath])
          userCount += 1
      if not rawTimesLists:
        # no labeled anomalies for this data file
        self.labelTimestamps[relativePath] = []
        self.labelIndices[relativePath] = setTruthLabels(dataSet, [])
        continue
      else:
        rawTimes = list(itertools.chain.from_iterable(rawTimesLists))
        rawTimes.sort()

      # Bucket and merge the anomaly timestamps.
      threshold = userCount * self.threshold
      trueAnomalies, passedAnomalies = merge(
        bucket(rawTimes, buffer), threshold)

      self.labelTimestamps[relativePath] = [str(t) for t in trueAnomalies]
      self.labelIndices[relativePath] = setTruthLabels(dataSet, trueAnomalies)

      if self.verbosity>0:
        print "----"
        print "For %s the passed raw labels and qualified true labels are,"\
              " respectively:" % relativePath
        print passedAnomalies
        print trueAnomalies

    return self.labelTimestamps, self.labelIndices

  def editPoorLabels(self):
    """
    This edits labels that have been flagged for manual revision. From
    inspecting the data and anomaly windows, we have determined some combined
    labels should be revised, or not included in the ground truth labels.
    """
    count = 0
    for relativePath, indices in self.labelIndices.iteritems():

      if "iio_us-east-1_i-a2eb1cd9_NetworkIn" in relativePath:
        self.labelIndices[relativePath] = [249, 339]

      count += len(indices)

    if self.verbosity > 0:
      print "============================================================="
      print "Total ground truth anomalies in benchmark dataset =", count

  def applyWindows(self):
    """
    This takes all the true anomalies, as calculated by combineLabels(), and
    adds a standard window. The window length is the class variable windowSize,
    and the location is centered on the anomaly timestamp.

    If verbosity = 2, the window metrics are printed to the console.
    """
    allWindows = {}
    for relativePath, anomalies in self.labelIndices.iteritems():
      data = self.corpus.dataFiles[relativePath].data
      length = len(data)
      num = len(anomalies)
      if num:
        windowLength = int(self.windowSize * length / len(anomalies))
      else:
        windowLength = int(self.windowSize * length)

      if self.verbosity==2:
        print "----"
        print "Window metrics for file", relativePath
        print "file length =", length, ";" \
              "number of windows =", num, ";" \
              "window length =", windowLength

      windows = []
      for a in anomalies:
        front = max(a - windowLength/2, 0)
        back = min(a + windowLength/2, length-1)

        windowLimit = [strf(data["timestamp"][front]),
                       strf(data["timestamp"][back])]

        windows.append(windowLimit)

      allWindows[relativePath] = windows

    self.combinedWindows = allWindows

  def checkWindows(self):
    """
    This takes the anomaly windows and checks for overlap with both each other
    and with the probationary period. Overlapping windows are merged into a
    single window. Windows overlapping with the probationary period are deleted.
    """
    for relativePath, windows in self.combinedWindows.iteritems():
      numWindows = len(windows)
      if numWindows > 0:

        fileLength = self.corpus.dataFiles[relativePath].data.shape[0]
        probationIndex = getProbationPeriod(
          self.probationaryPercent, fileLength)

        probationTimestamp = self.corpus.dataFiles[relativePath].data[
          "timestamp"][probationIndex]

        if (pandas.to_datetime(windows[0][0])
            -probationTimestamp).total_seconds() < 0:
          del windows[0]
          print ("The first window in {} overlaps with the probationary period "
                 ", so we're deleting it.".format(relativePath))

        i = 0
        while len(windows)-1 > i:
          if (pandas.to_datetime(windows[i+1][0])
              - pandas.to_datetime(windows[i][1])).total_seconds() <= 0:
            # merge windows
            windows[i] = [windows[i][0], windows[i+1][1]]
            del windows[i+1]
          i += 1


#!/usr/bin/env python
#
# Copyright 2013 Free Software Foundation, Inc.
#
# This file is part of GNU Radio
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
#

from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals
from gnuradio import gr, filter
from gnuradio import analog
from gnuradio import blocks
from gnuradio import eng_notation
from gnuradio.eng_arg import eng_float, intx
from argparse import ArgumentParser
import sys
import numpy

try:
    from matplotlib import pyplot
except ImportError:
    print("Error: could not from matplotlib import pyplot (http://matplotlib.sourceforge.net/)")
    sys.exit(1)

class example_fir_filter_fff(gr.top_block):
    def __init__(self, N, fs, bw, tw, atten, D):
        gr.top_block.__init__(self)

        self._nsamps = N
        self._fs = fs
        self._bw = bw
        self._tw = tw
        self._at = atten
        self._decim = D
        taps = filter.firdes.low_pass_2(1, self._fs, self._bw, self._tw, self._at)
        print("Num. Taps: ", len(taps))

        self.src  = analog.noise_source_f(analog.GR_GAUSSIAN, 1)
        self.head = blocks.head(gr.sizeof_float, self._nsamps)

        self.filt0 = filter.fir_filter_fff(self._decim, taps)

        self.vsnk_src = blocks.vector_sink_f()
        self.vsnk_out = blocks.vector_sink_f()

        self.connect(self.src, self.head, self.vsnk_src)
        self.connect(self.head, self.filt0, self.vsnk_out)

def main():
    parser = ArgumentParser(conflict_handler="resolve")
    parser.add_argument("-N", "--nsamples", type=int, default=10000,
                      help="Number of samples to process [default=%(default)r]")
    parser.add_argument("-s", "--samplerate", type=eng_float, default=8000,
                      help="System sample rate [default=%(default)r]")
    parser.add_argument("-B", "--bandwidth", type=eng_float, default=1000,
                      help="Filter bandwidth [default=%(default)r]")
    parser.add_argument("-T", "--transition", type=eng_float, default=100,
                      help="Transition band [default=%(default)r]")
    parser.add_argument("-A", "--attenuation", type=eng_float, default=80,
                      help="Stopband attenuation [default=%(default)r]")
    parser.add_argument("-D", "--decimation", type=int, default=1,
                      help="Decmation factor [default=%(default)r]")
    args = parser.parse_args()

    put = example_fir_filter_fff(args.nsamples,
                                 args.samplerate,
                                 args.bandwidth,
                                 args.transition,
                                 args.attenuation,
                                 args.decimation)
    put.run()

    data_src = numpy.array(put.vsnk_src.data())
    data_snk = numpy.array(put.vsnk_out.data())

    # Plot the signals PSDs
    nfft = 1024
    f1 = pyplot.figure(1, figsize=(12,10))
    s1 = f1.add_subplot(1,1,1)
    s1.psd(data_src, NFFT=nfft, noverlap=nfft / 4,
           Fs=args.samplerate)
    s1.psd(data_snk, NFFT=nfft, noverlap=nfft / 4,
           Fs=args.samplerate)

    f2 = pyplot.figure(2, figsize=(12,10))
    s2 = f2.add_subplot(1,1,1)
    s2.plot(data_src)
    s2.plot(data_snk.real, 'g')

    pyplot.show()

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass


import pandas as pd

from fireant.dataset.fields import Field
from fireant.utils import (
    alias_for_alias_selector,
    immutable,
)
from fireant.queries.builder.query_builder import QueryBuilder, QueryException, add_hints
from fireant.queries.execution import fetch_data

class DimensionLatestQueryBuilder(QueryBuilder):
    def __init__(self, dataset):
        super().__init__(dataset)

    @immutable
    def __call__(self, dimension: Field, *dimensions: Field):
        self._dimensions += [dimension] + list(dimensions)

    @property
    def sql(self):
        """
        Serializes this query builder as a set of SQL queries.  This method will always return a list of one
        query since only one query is required to retrieve dimension choices.

        This function only handles dimensions (select+group by) and filtering (where/having), which is everything
        needed for the query to fetch choices for dimensions.

        The dataset query extends this with metrics, references, and totals.
        """
        if not self.dimensions:
            raise QueryException("Must select at least one dimension to query latest values")

        query = self.dataset.database.make_latest_query(
            base_table=self.table,
            joins=self.dataset.joins,
            dimensions=self.dimensions,
        )
        return [query]

    def fetch(self, hint=None):
        queries = add_hints(self.sql, hint)
        max_rows_returned, data = fetch_data(self.dataset.database, queries, self.dimensions)
        data = self._get_latest_data_from_df(data)
        return self._transform_for_return(data, max_rows_returned=max_rows_returned)

    def _get_latest_data_from_df(self, df: pd.DataFrame) -> pd.Series:
        latest = df.reset_index().iloc[0]
        # Remove the row index as the name and trim the special dimension key characters from the dimension key
        latest.name = None
        latest.index = [alias_for_alias_selector(alias) for alias in latest.index]
        return latest