# -*- coding: utf-8 -*-
import warnings
from datetime import datetime, timedelta

import pytest

import numpy as np

import pandas as pd
import pandas.util.testing as tm
from pandas.errors import PerformanceWarning
from pandas import (Timestamp, Timedelta, Series,
                    DatetimeIndex, TimedeltaIndex,
                    date_range)

@pytest.fixture(params=[None, 'UTC', 'Asia/Tokyo',
                        'US/Eastern', 'dateutil/Asia/Singapore',
                        'dateutil/US/Pacific'])
def tz(request):
    return request.param

@pytest.fixture(params=[pd.offsets.Hour(2), timedelta(hours=2),
                        np.timedelta64(2, 'h'), Timedelta(hours=2)],
                ids=str)
def delta(request):
    # Several ways of representing two hours
    return request.param

@pytest.fixture(
    params=[
        datetime(2011, 1, 1),
        DatetimeIndex(['2011-01-01', '2011-01-02']),
        DatetimeIndex(['2011-01-01', '2011-01-02']).tz_localize('US/Eastern'),
        np.datetime64('2011-01-01'),
        Timestamp('2011-01-01')],
    ids=lambda x: type(x).__name__)
def addend(request):
    return request.param

class TestDatetimeIndexArithmetic(object):

    def test_dti_add_timestamp_raises(self):
        idx = DatetimeIndex(['2011-01-01', '2011-01-02'])
        msg = "cannot add DatetimeIndex and Timestamp"
        with tm.assert_raises_regex(TypeError, msg):
            idx + Timestamp('2011-01-01')

    def test_dti_radd_timestamp_raises(self):
        idx = DatetimeIndex(['2011-01-01', '2011-01-02'])
        msg = "cannot add DatetimeIndex and Timestamp"
        with tm.assert_raises_regex(TypeError, msg):
            Timestamp('2011-01-01') + idx

    # -------------------------------------------------------------
    # Binary operations DatetimeIndex and int

    def test_dti_add_int(self, tz, one):
        # Variants of `one` for #19012
        rng = pd.date_range('2000-01-01 09:00', freq='H',
                            periods=10, tz=tz)
        result = rng + one
        expected = pd.date_range('2000-01-01 10:00', freq='H',
                                 periods=10, tz=tz)
        tm.assert_index_equal(result, expected)

    def test_dti_iadd_int(self, tz, one):
        rng = pd.date_range('2000-01-01 09:00', freq='H',
                            periods=10, tz=tz)
        expected = pd.date_range('2000-01-01 10:00', freq='H',
                                 periods=10, tz=tz)
        rng += one
        tm.assert_index_equal(rng, expected)

    def test_dti_sub_int(self, tz, one):
        rng = pd.date_range('2000-01-01 09:00', freq='H',
                            periods=10, tz=tz)
        result = rng - one
        expected = pd.date_range('2000-01-01 08:00', freq='H',
                                 periods=10, tz=tz)
        tm.assert_index_equal(result, expected)

    def test_dti_isub_int(self, tz, one):
        rng = pd.date_range('2000-01-01 09:00', freq='H',
                            periods=10, tz=tz)
        expected = pd.date_range('2000-01-01 08:00', freq='H',
                                 periods=10, tz=tz)
        rng -= one
        tm.assert_index_equal(rng, expected)

    # -------------------------------------------------------------
    # Binary operations DatetimeIndex and timedelta-like

    def test_dti_add_timedeltalike(self, tz, delta):
        rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
        result = rng + delta
        expected = pd.date_range('2000-01-01 02:00',
                                 '2000-02-01 02:00', tz=tz)
        tm.assert_index_equal(result, expected)

    def test_dti_iadd_timedeltalike(self, tz, delta):
        rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
        expected = pd.date_range('2000-01-01 02:00',
                                 '2000-02-01 02:00', tz=tz)
        rng += delta
        tm.assert_index_equal(rng, expected)

    def test_dti_sub_timedeltalike(self, tz, delta):
        rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
        expected = pd.date_range('1999-12-31 22:00',
                                 '2000-01-31 22:00', tz=tz)
        result = rng - delta
        tm.assert_index_equal(result, expected)

    def test_dti_isub_timedeltalike(self, tz, delta):
        rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
        expected = pd.date_range('1999-12-31 22:00',
                                 '2000-01-31 22:00', tz=tz)
        rng -= delta
        tm.assert_index_equal(rng, expected)

    # -------------------------------------------------------------
    # Binary operations DatetimeIndex and TimedeltaIndex/array
    def test_dti_add_tdi(self, tz):
        # GH 17558
        dti = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
        tdi = pd.timedelta_range('0 days', periods=10)
        expected = pd.date_range('2017-01-01', periods=10, tz=tz)

        # add with TimdeltaIndex
        result = dti + tdi
        tm.assert_index_equal(result, expected)

        result = tdi + dti
        tm.assert_index_equal(result, expected)

        # add with timedelta64 array
        result = dti + tdi.values
        tm.assert_index_equal(result, expected)

        result = tdi.values + dti
        tm.assert_index_equal(result, expected)

    def test_dti_iadd_tdi(self, tz):
        # GH 17558
        dti = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
        tdi = pd.timedelta_range('0 days', periods=10)
        expected = pd.date_range('2017-01-01', periods=10, tz=tz)

        # iadd with TimdeltaIndex
        result = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
        result += tdi
        tm.assert_index_equal(result, expected)

        result = pd.timedelta_range('0 days', periods=10)
        result += dti
        tm.assert_index_equal(result, expected)

        # iadd with timedelta64 array
        result = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
        result += tdi.values
        tm.assert_index_equal(result, expected)

        result = pd.timedelta_range('0 days', periods=10)
        result += dti
        tm.assert_index_equal(result, expected)

    def test_dti_sub_tdi(self, tz):
        # GH 17558
        dti = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
        tdi = pd.timedelta_range('0 days', periods=10)
        expected = pd.date_range('2017-01-01', periods=10, tz=tz, freq='-1D')

        # sub with TimedeltaIndex
        result = dti - tdi
        tm.assert_index_equal(result, expected)

        msg = 'cannot subtract TimedeltaIndex and DatetimeIndex'
        with tm.assert_raises_regex(TypeError, msg):
            tdi - dti

        # sub with timedelta64 array
        result = dti - tdi.values
        tm.assert_index_equal(result, expected)

        msg = 'cannot perform __neg__ with this index type:'
        with tm.assert_raises_regex(TypeError, msg):
            tdi.values - dti

    def test_dti_isub_tdi(self, tz):
        # GH 17558
        dti = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
        tdi = pd.timedelta_range('0 days', periods=10)
        expected = pd.date_range('2017-01-01', periods=10, tz=tz, freq='-1D')

        # isub with TimedeltaIndex
        result = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
        result -= tdi
        tm.assert_index_equal(result, expected)

        msg = 'cannot subtract TimedeltaIndex and DatetimeIndex'
        with tm.assert_raises_regex(TypeError, msg):
            tdi -= dti

        # isub with timedelta64 array
        result = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
        result -= tdi.values
        tm.assert_index_equal(result, expected)

        msg = '|'.join(['cannot perform __neg__ with this index type:',
                        'ufunc subtract cannot use operands with types'])
        with tm.assert_raises_regex(TypeError, msg):
            tdi.values -= dti

    # -------------------------------------------------------------
    # Binary Operations DatetimeIndex and datetime-like
    # TODO: A couple other tests belong in this section.  Move them in
    # A PR where there isn't already a giant diff.

    def test_add_datetimelike_and_dti(self, addend):
        # GH#9631
        dti = DatetimeIndex(['2011-01-01', '2011-01-02'])
        msg = 'cannot add DatetimeIndex and {0}'.format(
            type(addend).__name__)
        with tm.assert_raises_regex(TypeError, msg):
            dti + addend
        with tm.assert_raises_regex(TypeError, msg):
            addend + dti

    def test_add_datetimelike_and_dti_tz(self, addend):
        # GH#9631
        dti_tz = DatetimeIndex(['2011-01-01',
                                '2011-01-02']).tz_localize('US/Eastern')
        msg = 'cannot add DatetimeIndex and {0}'.format(
            type(addend).__name__)
        with tm.assert_raises_regex(TypeError, msg):
            dti_tz + addend
        with tm.assert_raises_regex(TypeError, msg):
            addend + dti_tz

    # -------------------------------------------------------------

    def test_sub_dti_dti(self):
        # previously performed setop (deprecated in 0.16.0), now changed to
        # return subtraction -> TimeDeltaIndex (GH ...)

        dti = date_range('20130101', periods=3)
        dti_tz = date_range('20130101', periods=3).tz_localize('US/Eastern')
        dti_tz2 = date_range('20130101', periods=3).tz_localize('UTC')
        expected = TimedeltaIndex([0, 0, 0])

        result = dti - dti
        tm.assert_index_equal(result, expected)

        result = dti_tz - dti_tz
        tm.assert_index_equal(result, expected)

        with pytest.raises(TypeError):
            dti_tz - dti

        with pytest.raises(TypeError):
            dti - dti_tz

        with pytest.raises(TypeError):
            dti_tz - dti_tz2

        # isub
        dti -= dti
        tm.assert_index_equal(dti, expected)

        # different length raises ValueError
        dti1 = date_range('20130101', periods=3)
        dti2 = date_range('20130101', periods=4)
        with pytest.raises(ValueError):
            dti1 - dti2

        # NaN propagation
        dti1 = DatetimeIndex(['2012-01-01', np.nan, '2012-01-03'])
        dti2 = DatetimeIndex(['2012-01-02', '2012-01-03', np.nan])
        expected = TimedeltaIndex(['1 days', np.nan, np.nan])
        result = dti2 - dti1
        tm.assert_index_equal(result, expected)

    def test_sub_period(self):
        # GH 13078
        # not supported, check TypeError
        p = pd.Period('2011-01-01', freq='D')

        for freq in [None, 'D']:
            idx = pd.DatetimeIndex(['2011-01-01', '2011-01-02'], freq=freq)

            with pytest.raises(TypeError):
                idx - p

            with pytest.raises(TypeError):
                p - idx

    def test_ufunc_coercions(self):
        idx = date_range('2011-01-01', periods=3, freq='2D', name='x')

        delta = np.timedelta64(1, 'D')
        for result in [idx + delta, np.add(idx, delta)]:
            assert isinstance(result, DatetimeIndex)
            exp = date_range('2011-01-02', periods=3, freq='2D', name='x')
            tm.assert_index_equal(result, exp)
            assert result.freq == '2D'

        for result in [idx - delta, np.subtract(idx, delta)]:
            assert isinstance(result, DatetimeIndex)
            exp = date_range('2010-12-31', periods=3, freq='2D', name='x')
            tm.assert_index_equal(result, exp)
            assert result.freq == '2D'

        delta = np.array([np.timedelta64(1, 'D'), np.timedelta64(2, 'D'),
                          np.timedelta64(3, 'D')])
        for result in [idx + delta, np.add(idx, delta)]:
            assert isinstance(result, DatetimeIndex)
            exp = DatetimeIndex(['2011-01-02', '2011-01-05', '2011-01-08'],
                                freq='3D', name='x')
            tm.assert_index_equal(result, exp)
            assert result.freq == '3D'

        for result in [idx - delta, np.subtract(idx, delta)]:
            assert isinstance(result, DatetimeIndex)
            exp = DatetimeIndex(['2010-12-31', '2011-01-01', '2011-01-02'],
                                freq='D', name='x')
            tm.assert_index_equal(result, exp)
            assert result.freq == 'D'

    def test_datetimeindex_sub_timestamp_overflow(self):
        dtimax = pd.to_datetime(['now', pd.Timestamp.max])
        dtimin = pd.to_datetime(['now', pd.Timestamp.min])

        tsneg = Timestamp('1950-01-01')
        ts_neg_variants = [tsneg,
                           tsneg.to_pydatetime(),
                           tsneg.to_datetime64().astype('datetime64[ns]'),
                           tsneg.to_datetime64().astype('datetime64[D]')]

        tspos = Timestamp('1980-01-01')
        ts_pos_variants = [tspos,
                           tspos.to_pydatetime(),
                           tspos.to_datetime64().astype('datetime64[ns]'),
                           tspos.to_datetime64().astype('datetime64[D]')]

        for variant in ts_neg_variants:
            with pytest.raises(OverflowError):
                dtimax - variant

        expected = pd.Timestamp.max.value - tspos.value
        for variant in ts_pos_variants:
            res = dtimax - variant
            assert res[1].value == expected

        expected = pd.Timestamp.min.value - tsneg.value
        for variant in ts_neg_variants:
            res = dtimin - variant
            assert res[1].value == expected

        for variant in ts_pos_variants:
            with pytest.raises(OverflowError):
                dtimin - variant

    @pytest.mark.parametrize('box', [np.array, pd.Index])
    def test_dti_add_offset_array(self, tz, box):
        # GH#18849
        dti = pd.date_range('2017-01-01', periods=2, tz=tz)
        other = box([pd.offsets.MonthEnd(), pd.offsets.Day(n=2)])

        with tm.assert_produces_warning(PerformanceWarning):
            res = dti + other
        expected = DatetimeIndex([dti[n] + other[n] for n in range(len(dti))],
                                 name=dti.name, freq='infer')
        tm.assert_index_equal(res, expected)

        with tm.assert_produces_warning(PerformanceWarning):
            res2 = other + dti
        tm.assert_index_equal(res2, expected)

    @pytest.mark.parametrize('box', [np.array, pd.Index])
    def test_dti_sub_offset_array(self, tz, box):
        # GH#18824
        dti = pd.date_range('2017-01-01', periods=2, tz=tz)
        other = box([pd.offsets.MonthEnd(), pd.offsets.Day(n=2)])

        with tm.assert_produces_warning(PerformanceWarning):
            res = dti - other
        expected = DatetimeIndex([dti[n] - other[n] for n in range(len(dti))],
                                 name=dti.name, freq='infer')
        tm.assert_index_equal(res, expected)

    @pytest.mark.parametrize('names', [(None, None, None),
                                       ('foo', 'bar', None),
                                       ('foo', 'foo', 'foo')])
    def test_dti_with_offset_series(self, tz, names):
        # GH#18849
        dti = pd.date_range('2017-01-01', periods=2, tz=tz, name=names[0])
        other = Series([pd.offsets.MonthEnd(), pd.offsets.Day(n=2)],
                       name=names[1])

        expected_add = Series([dti[n] + other[n] for n in range(len(dti))],
                              name=names[2])

        with tm.assert_produces_warning(PerformanceWarning):
            res = dti + other
        tm.assert_series_equal(res, expected_add)

        with tm.assert_produces_warning(PerformanceWarning):
            res2 = other + dti
        tm.assert_series_equal(res2, expected_add)

        expected_sub = Series([dti[n] - other[n] for n in range(len(dti))],
                              name=names[2])

        with tm.assert_produces_warning(PerformanceWarning):
            res3 = dti - other
        tm.assert_series_equal(res3, expected_sub)

# GH 10699
@pytest.mark.parametrize('klass,assert_func', zip([Series, DatetimeIndex],
                                                  [tm.assert_series_equal,
                                                   tm.assert_index_equal]))
def test_datetime64_with_DateOffset(klass, assert_func):
    s = klass(date_range('2000-01-01', '2000-01-31'), name='a')
    result = s + pd.DateOffset(years=1)
    result2 = pd.DateOffset(years=1) + s
    exp = klass(date_range('2001-01-01', '2001-01-31'), name='a')
    assert_func(result, exp)
    assert_func(result2, exp)

    result = s - pd.DateOffset(years=1)
    exp = klass(date_range('1999-01-01', '1999-01-31'), name='a')
    assert_func(result, exp)

    s = klass([Timestamp('2000-01-15 00:15:00', tz='US/Central'),
               pd.Timestamp('2000-02-15', tz='US/Central')], name='a')
    result = s + pd.offsets.Day()
    result2 = pd.offsets.Day() + s
    exp = klass([Timestamp('2000-01-16 00:15:00', tz='US/Central'),
                 Timestamp('2000-02-16', tz='US/Central')], name='a')
    assert_func(result, exp)
    assert_func(result2, exp)

    s = klass([Timestamp('2000-01-15 00:15:00', tz='US/Central'),
               pd.Timestamp('2000-02-15', tz='US/Central')], name='a')
    result = s + pd.offsets.MonthEnd()
    result2 = pd.offsets.MonthEnd() + s
    exp = klass([Timestamp('2000-01-31 00:15:00', tz='US/Central'),
                 Timestamp('2000-02-29', tz='US/Central')], name='a')
    assert_func(result, exp)
    assert_func(result2, exp)

    # array of offsets - valid for Series only
    if klass is Series:
        with tm.assert_produces_warning(PerformanceWarning):
            s = klass([Timestamp('2000-1-1'), Timestamp('2000-2-1')])
            result = s + Series([pd.offsets.DateOffset(years=1),
                                 pd.offsets.MonthEnd()])
            exp = klass([Timestamp('2001-1-1'), Timestamp('2000-2-29')
                         ])
            assert_func(result, exp)

            # same offset
            result = s + Series([pd.offsets.DateOffset(years=1),
                                 pd.offsets.DateOffset(years=1)])
            exp = klass([Timestamp('2001-1-1'), Timestamp('2001-2-1')])
            assert_func(result, exp)

    s = klass([Timestamp('2000-01-05 00:15:00'),
               Timestamp('2000-01-31 00:23:00'),
               Timestamp('2000-01-01'),
               Timestamp('2000-03-31'),
               Timestamp('2000-02-29'),
               Timestamp('2000-12-31'),
               Timestamp('2000-05-15'),
               Timestamp('2001-06-15')])

    # DateOffset relativedelta fastpath
    relative_kwargs = [('years', 2), ('months', 5), ('days', 3),
                       ('hours', 5), ('minutes', 10), ('seconds', 2),
                       ('microseconds', 5)]
    for i, kwd in enumerate(relative_kwargs):
        op = pd.DateOffset(**dict([kwd]))
        assert_func(klass([x + op for x in s]), s + op)
        assert_func(klass([x - op for x in s]), s - op)
        op = pd.DateOffset(**dict(relative_kwargs[:i + 1]))
        assert_func(klass([x + op for x in s]), s + op)
        assert_func(klass([x - op for x in s]), s - op)

    # assert these are equal on a piecewise basis
    offsets = ['YearBegin', ('YearBegin', {'month': 5}),
               'YearEnd', ('YearEnd', {'month': 5}),
               'MonthBegin', 'MonthEnd',
               'SemiMonthEnd', 'SemiMonthBegin',
               'Week', ('Week', {'weekday': 3}),
               'BusinessDay', 'BDay', 'QuarterEnd', 'QuarterBegin',
               'CustomBusinessDay', 'CDay', 'CBMonthEnd',
               'CBMonthBegin', 'BMonthBegin', 'BMonthEnd',
               'BusinessHour', 'BYearBegin', 'BYearEnd',
               'BQuarterBegin', ('LastWeekOfMonth', {'weekday': 2}),
               ('FY5253Quarter', {'qtr_with_extra_week': 1,
                                  'startingMonth': 1,
                                  'weekday': 2,
                                  'variation': 'nearest'}),
               ('FY5253', {'weekday': 0,
                           'startingMonth': 2,
                           'variation':
                           'nearest'}),
               ('WeekOfMonth', {'weekday': 2,
                                'week': 2}),
               'Easter', ('DateOffset', {'day': 4}),
               ('DateOffset', {'month': 5})]

    with warnings.catch_warnings(record=True):
        for normalize in (True, False):
            for do in offsets:
                if isinstance(do, tuple):
                    do, kwargs = do
                else:
                    do = do
                    kwargs = {}

                    for n in [0, 5]:
                        if (do in ['WeekOfMonth', 'LastWeekOfMonth',
                                   'FY5253Quarter', 'FY5253'] and n == 0):
                            continue
                    op = getattr(pd.offsets, do)(n,
                                                 normalize=normalize,
                                                 **kwargs)
                    assert_func(klass([x + op for x in s]), s + op)
                    assert_func(klass([x - op for x in s]), s - op)
                    assert_func(klass([op + x for x in s]), op + s)


import numpy as np
import os.path as op

from numpy.testing import (assert_array_almost_equal, assert_array_equal,
                           assert_equal)
import pytest

import mne
from mne import Epochs, read_events, pick_types, create_info, EpochsArray
from mne.io import read_raw_fif
from mne.utils import _TempDir, run_tests_if_main, requires_h5py, grand_average
from mne.time_frequency.tfr import (morlet, tfr_morlet, _make_dpss,
                                    tfr_multitaper, AverageTFR, read_tfrs,
                                    write_tfrs, combine_tfr, cwt, _compute_tfr,
                                    EpochsTFR)
from mne.time_frequency import tfr_array_multitaper, tfr_array_morlet
from mne.viz.utils import _fake_click
from itertools import product
import matplotlib
matplotlib.use('Agg')  # for testing don't use X server

data_path = op.join(op.dirname(__file__), '..', '..', 'io', 'tests', 'data')
raw_fname = op.join(data_path, 'test_raw.fif')
event_fname = op.join(data_path, 'test-eve.fif')
raw_ctf_fname = op.join(data_path, 'test_ctf_raw.fif')

def test_tfr_ctf():
    """Test that TFRs can be calculated on CTF data."""
    raw = read_raw_fif(raw_ctf_fname).crop(0, 1)
    raw.apply_gradient_compensation(3)
    events = mne.make_fixed_length_events(raw, duration=0.5)
    epochs = mne.Epochs(raw, events)
    for method in (tfr_multitaper, tfr_morlet):
        method(epochs, [10], 1)  # smoke test

def test_morlet():
    """Test morlet with and without zero mean."""
    Wz = morlet(1000, [10], 2., zero_mean=True)
    W = morlet(1000, [10], 2., zero_mean=False)

    assert (np.abs(np.mean(np.real(Wz[0]))) < 1e-5)
    assert (np.abs(np.mean(np.real(W[0]))) > 1e-3)

def test_time_frequency():
    """Test time-frequency transform (PSD and ITC)."""
    # Set parameters
    event_id = 1
    tmin = -0.2
    tmax = 0.498  # Allows exhaustive decimation testing

    # Setup for reading the raw data
    raw = read_raw_fif(raw_fname)
    events = read_events(event_fname)

    include = []
    exclude = raw.info['bads'] + ['MEG 2443', 'EEG 053']  # bads + 2 more

    # picks MEG gradiometers
    picks = pick_types(raw.info, meg='grad', eeg=False,
                       stim=False, include=include, exclude=exclude)

    picks = picks[:2]
    epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks)
    data = epochs.get_data()
    times = epochs.times
    nave = len(data)

    epochs_nopicks = Epochs(raw, events, event_id, tmin, tmax)

    freqs = np.arange(6, 20, 5)  # define frequencies of interest
    n_cycles = freqs / 4.

    # Test first with a single epoch
    power, itc = tfr_morlet(epochs[0], freqs=freqs, n_cycles=n_cycles,
                            use_fft=True, return_itc=True)
    # Now compute evoked
    evoked = epochs.average()
    power_evoked = tfr_morlet(evoked, freqs, n_cycles, use_fft=True,
                              return_itc=False)
    pytest.raises(ValueError, tfr_morlet, evoked, freqs, 1., return_itc=True)
    power, itc = tfr_morlet(epochs, freqs=freqs, n_cycles=n_cycles,
                            use_fft=True, return_itc=True)
    power_, itc_ = tfr_morlet(epochs, freqs=freqs, n_cycles=n_cycles,
                              use_fft=True, return_itc=True, decim=slice(0, 2))
    # Test picks argument and average parameter
    pytest.raises(ValueError, tfr_morlet, epochs, freqs=freqs,
                  n_cycles=n_cycles, return_itc=True, average=False)

    power_picks, itc_picks = \
        tfr_morlet(epochs_nopicks,
                   freqs=freqs, n_cycles=n_cycles, use_fft=True,
                   return_itc=True, picks=picks, average=True)

    epochs_power_picks = \
        tfr_morlet(epochs_nopicks,
                   freqs=freqs, n_cycles=n_cycles, use_fft=True,
                   return_itc=False, picks=picks, average=False)
    power_picks_avg = epochs_power_picks.average()
    # the actual data arrays here are equivalent, too...
    assert_array_almost_equal(power.data, power_picks.data)
    assert_array_almost_equal(power.data, power_picks_avg.data)
    assert_array_almost_equal(itc.data, itc_picks.data)
    assert_array_almost_equal(power.data, power_evoked.data)
    # complex output
    pytest.raises(ValueError, tfr_morlet, epochs, freqs, n_cycles,
                  return_itc=False, average=True, output="complex")
    pytest.raises(ValueError, tfr_morlet, epochs, freqs, n_cycles,
                  output="complex", average=False, return_itc=True)
    epochs_power_complex = tfr_morlet(epochs, freqs, n_cycles,
                                      output="complex", average=False,
                                      return_itc=False)
    epochs_power_2 = abs(epochs_power_complex)
    epochs_power_3 = epochs_power_2.copy()
    epochs_power_3.data[:] = np.inf  # test that it's actually copied
    assert_array_almost_equal(epochs_power_2.data, epochs_power_picks.data)
    power_2 = epochs_power_2.average()
    assert_array_almost_equal(power_2.data, power.data)

    print(itc)  # test repr
    print(itc.ch_names)  # test property
    itc += power  # test add
    itc -= power  # test sub

    power = power.apply_baseline(baseline=(-0.1, 0), mode='logratio')

    assert 'meg' in power
    assert 'grad' in power
    assert 'mag' not in power
    assert 'eeg' not in power

    assert_equal(power.nave, nave)
    assert_equal(itc.nave, nave)
    assert (power.data.shape == (len(picks), len(freqs), len(times)))
    assert (power.data.shape == itc.data.shape)
    assert (power_.data.shape == (len(picks), len(freqs), 2))
    assert (power_.data.shape == itc_.data.shape)
    assert (np.sum(itc.data >= 1) == 0)
    assert (np.sum(itc.data <= 0) == 0)

    # grand average
    itc2 = itc.copy()
    itc2.info['bads'] = [itc2.ch_names[0]]  # test channel drop
    gave = grand_average([itc2, itc])
    assert_equal(gave.data.shape, (itc2.data.shape[0] - 1,
                                   itc2.data.shape[1],
                                   itc2.data.shape[2]))
    assert_equal(itc2.ch_names[1:], gave.ch_names)
    assert_equal(gave.nave, 2)
    itc2.drop_channels(itc2.info["bads"])
    assert_array_almost_equal(gave.data, itc2.data)
    itc2.data = np.ones(itc2.data.shape)
    itc.data = np.zeros(itc.data.shape)
    itc2.nave = 2
    itc.nave = 1
    itc.drop_channels([itc.ch_names[0]])
    combined_itc = combine_tfr([itc2, itc])
    assert_array_almost_equal(combined_itc.data,
                              np.ones(combined_itc.data.shape) * 2 / 3)

    # more tests
    power, itc = tfr_morlet(epochs, freqs=freqs, n_cycles=2, use_fft=False,
                            return_itc=True)

    assert (power.data.shape == (len(picks), len(freqs), len(times)))
    assert (power.data.shape == itc.data.shape)
    assert (np.sum(itc.data >= 1) == 0)
    assert (np.sum(itc.data <= 0) == 0)

    tfr = tfr_morlet(epochs[0], freqs, use_fft=True, n_cycles=2, average=False,
                     return_itc=False).data[0]
    assert (tfr.shape == (len(picks), len(freqs), len(times)))
    tfr2 = tfr_morlet(epochs[0], freqs, use_fft=True, n_cycles=2,
                      decim=slice(0, 2), average=False,
                      return_itc=False).data[0]
    assert (tfr2.shape == (len(picks), len(freqs), 2))

    single_power = tfr_morlet(epochs, freqs, 2, average=False,
                              return_itc=False).data
    single_power2 = tfr_morlet(epochs, freqs, 2, decim=slice(0, 2),
                               average=False, return_itc=False).data
    single_power3 = tfr_morlet(epochs, freqs, 2, decim=slice(1, 3),
                               average=False, return_itc=False).data
    single_power4 = tfr_morlet(epochs, freqs, 2, decim=slice(2, 4),
                               average=False, return_itc=False).data

    assert_array_almost_equal(np.mean(single_power, axis=0), power.data)
    assert_array_almost_equal(np.mean(single_power2, axis=0),
                              power.data[:, :, :2])
    assert_array_almost_equal(np.mean(single_power3, axis=0),
                              power.data[:, :, 1:3])
    assert_array_almost_equal(np.mean(single_power4, axis=0),
                              power.data[:, :, 2:4])

    power_pick = power.pick_channels(power.ch_names[:10:2])
    assert_equal(len(power_pick.ch_names), len(power.ch_names[:10:2]))
    assert_equal(power_pick.data.shape[0], len(power.ch_names[:10:2]))
    power_drop = power.drop_channels(power.ch_names[1:10:2])
    assert_equal(power_drop.ch_names, power_pick.ch_names)
    assert_equal(power_pick.data.shape[0], len(power_drop.ch_names))

    mne.equalize_channels([power_pick, power_drop])
    assert_equal(power_pick.ch_names, power_drop.ch_names)
    assert_equal(power_pick.data.shape, power_drop.data.shape)

    # Test decimation:
    # 2: multiple of len(times) even
    # 3: multiple odd
    # 8: not multiple, even
    # 9: not multiple, odd
    for decim in [2, 3, 8, 9]:
        for use_fft in [True, False]:
            power, itc = tfr_morlet(epochs, freqs=freqs, n_cycles=2,
                                    use_fft=use_fft, return_itc=True,
                                    decim=decim)
            assert_equal(power.data.shape[2],
                         np.ceil(float(len(times)) / decim))
    freqs = list(range(50, 55))
    decim = 2
    _, n_chan, n_time = data.shape
    tfr = tfr_morlet(epochs[0], freqs, 2., decim=decim, average=False,
                     return_itc=False).data[0]
    assert_equal(tfr.shape, (n_chan, len(freqs), n_time // decim))

    # Test cwt modes
    Ws = morlet(512, [10, 20], n_cycles=2)
    pytest.raises(ValueError, cwt, data[0, :, :], Ws, mode='foo')
    for use_fft in [True, False]:
        for mode in ['same', 'valid', 'full']:
            cwt(data[0], Ws, use_fft=use_fft, mode=mode)

    # Test decim parameter checks
    pytest.raises(TypeError, tfr_morlet, epochs, freqs=freqs,
                  n_cycles=n_cycles, use_fft=True, return_itc=True,
                  decim='decim')

    # When convolving in time, wavelets must not be longer than the data
    pytest.raises(ValueError, cwt, data[0, :, :Ws[0].size - 1], Ws,
                  use_fft=False)
    with pytest.warns(UserWarning, match='one of the wavelets is longer'):
        cwt(data[0, :, :Ws[0].size - 1], Ws, use_fft=True)

    # Check for off-by-one errors when using wavelets with an even number of
    # samples
    psd = cwt(data[0], [Ws[0][:-1]], use_fft=False, mode='full')
    assert_equal(psd.shape, (2, 1, 420))

def test_dpsswavelet():
    """Test DPSS tapers."""
    freqs = np.arange(5, 25, 3)
    Ws = _make_dpss(1000, freqs=freqs, n_cycles=freqs / 2., time_bandwidth=4.0,
                    zero_mean=True)

    assert (len(Ws) == 3)  # 3 tapers expected

    # Check that zero mean is true
    assert (np.abs(np.mean(np.real(Ws[0][0]))) < 1e-5)

    assert (len(Ws[0]) == len(freqs))  # As many wavelets as asked for

@pytest.mark.slowtest
def test_tfr_multitaper():
    """Test tfr_multitaper."""
    sfreq = 200.0
    ch_names = ['SIM0001', 'SIM0002']
    ch_types = ['grad', 'grad']
    info = create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)

    n_times = int(sfreq)  # Second long epochs
    n_epochs = 3
    seed = 42
    rng = np.random.RandomState(seed)
    noise = 0.1 * rng.randn(n_epochs, len(ch_names), n_times)
    t = np.arange(n_times, dtype=np.float) / sfreq
    signal = np.sin(np.pi * 2. * 50. * t)  # 50 Hz sinusoid signal
    signal[np.logical_or(t < 0.45, t > 0.55)] = 0.  # Hard windowing
    on_time = np.logical_and(t >= 0.45, t <= 0.55)
    signal[on_time] *= np.hanning(on_time.sum())  # Ramping
    dat = noise + signal

    reject = dict(grad=4000.)
    events = np.empty((n_epochs, 3), int)
    first_event_sample = 100
    event_id = dict(sin50hz=1)
    for k in range(n_epochs):
        events[k, :] = first_event_sample + k * n_times, 0, event_id['sin50hz']

    epochs = EpochsArray(data=dat, info=info, events=events, event_id=event_id,
                         reject=reject)

    freqs = np.arange(35, 70, 5, dtype=np.float)

    power, itc = tfr_multitaper(epochs, freqs=freqs, n_cycles=freqs / 2.,
                                time_bandwidth=4.0)
    power2, itc2 = tfr_multitaper(epochs, freqs=freqs, n_cycles=freqs / 2.,
                                  time_bandwidth=4.0, decim=slice(0, 2))
    picks = np.arange(len(ch_names))
    power_picks, itc_picks = tfr_multitaper(epochs, freqs=freqs,
                                            n_cycles=freqs / 2.,
                                            time_bandwidth=4.0, picks=picks)
    power_epochs = tfr_multitaper(epochs, freqs=freqs,
                                  n_cycles=freqs / 2., time_bandwidth=4.0,
                                  return_itc=False, average=False)
    power_averaged = power_epochs.average()
    power_evoked = tfr_multitaper(epochs.average(), freqs=freqs,
                                  n_cycles=freqs / 2., time_bandwidth=4.0,
                                  return_itc=False, average=False).average()

    print(power_evoked)  # test repr for EpochsTFR

    # Test channel picking
    power_epochs_picked = power_epochs.copy().drop_channels(['SIM0002'])
    assert_equal(power_epochs_picked.data.shape, (3, 1, 7, 200))
    assert_equal(power_epochs_picked.ch_names, ['SIM0001'])

    pytest.raises(ValueError, tfr_multitaper, epochs,
                  freqs=freqs, n_cycles=freqs / 2.,
                  return_itc=True, average=False)

    # test picks argument
    assert_array_almost_equal(power.data, power_picks.data)
    assert_array_almost_equal(power.data, power_averaged.data)
    assert_array_almost_equal(power.times, power_epochs.times)
    assert_array_almost_equal(power.times, power_averaged.times)
    assert_equal(power.nave, power_averaged.nave)
    assert_equal(power_epochs.data.shape, (3, 2, 7, 200))
    assert_array_almost_equal(itc.data, itc_picks.data)
    # one is squared magnitude of the average (evoked) and
    # the other is average of the squared magnitudes (epochs PSD)
    # so values shouldn't match, but shapes should
    assert_array_equal(power.data.shape, power_evoked.data.shape)
    pytest.raises(AssertionError, assert_array_almost_equal,
                  power.data, power_evoked.data)

    tmax = t[np.argmax(itc.data[0, freqs == 50, :])]
    fmax = freqs[np.argmax(power.data[1, :, t == 0.5])]
    assert (tmax > 0.3 and tmax < 0.7)
    assert not np.any(itc.data < 0.)
    assert (fmax > 40 and fmax < 60)
    assert (power2.data.shape == (len(picks), len(freqs), 2))
    assert (power2.data.shape == itc2.data.shape)

    # Test decim parameter checks and compatibility between wavelets length
    # and instance length in the time dimension.
    pytest.raises(TypeError, tfr_multitaper, epochs, freqs=freqs,
                  n_cycles=freqs / 2., time_bandwidth=4.0, decim=(1,))
    pytest.raises(ValueError, tfr_multitaper, epochs, freqs=freqs,
                  n_cycles=1000, time_bandwidth=4.0)

def test_crop():
    """Test TFR cropping."""
    data = np.zeros((3, 2, 3))
    times = np.array([.1, .2, .3])
    freqs = np.array([.10, .20])
    info = mne.create_info(['MEG 001', 'MEG 002', 'MEG 003'], 1000.,
                           ['mag', 'mag', 'mag'])
    tfr = AverageTFR(info, data=data, times=times, freqs=freqs,
                     nave=20, comment='test', method='crazy-tfr')
    tfr.crop(0.2, 0.3)
    assert_array_equal(tfr.times, [0.2, 0.3])
    assert_equal(tfr.data.shape[-1], 2)

@requires_h5py
def test_io():
    """Test TFR IO capacities."""
    tempdir = _TempDir()
    fname = op.join(tempdir, 'test-tfr.h5')
    data = np.zeros((3, 2, 3))
    times = np.array([.1, .2, .3])
    freqs = np.array([.10, .20])

    info = mne.create_info(['MEG 001', 'MEG 002', 'MEG 003'], 1000.,
                           ['mag', 'mag', 'mag'])
    tfr = AverageTFR(info, data=data, times=times, freqs=freqs,
                     nave=20, comment='test', method='crazy-tfr')
    tfr.save(fname)
    tfr2 = read_tfrs(fname, condition='test')

    assert_array_equal(tfr.data, tfr2.data)
    assert_array_equal(tfr.times, tfr2.times)
    assert_array_equal(tfr.freqs, tfr2.freqs)
    assert_equal(tfr.comment, tfr2.comment)
    assert_equal(tfr.nave, tfr2.nave)

    pytest.raises(IOError, tfr.save, fname)

    tfr.comment = None
    tfr.save(fname, overwrite=True)
    assert_equal(read_tfrs(fname, condition=0).comment, tfr.comment)
    tfr.comment = 'test-A'
    tfr2.comment = 'test-B'

    fname = op.join(tempdir, 'test2-tfr.h5')
    write_tfrs(fname, [tfr, tfr2])
    tfr3 = read_tfrs(fname, condition='test-A')
    assert_equal(tfr.comment, tfr3.comment)

    assert (isinstance(tfr.info, mne.Info))

    tfrs = read_tfrs(fname, condition=None)
    assert_equal(len(tfrs), 2)
    tfr4 = tfrs[1]
    assert_equal(tfr2.comment, tfr4.comment)

    pytest.raises(ValueError, read_tfrs, fname, condition='nonono')

    # Test save of EpochsTFR.
    data = np.zeros((5, 3, 2, 3))
    tfr = EpochsTFR(info, data=data, times=times, freqs=freqs,
                    comment='test', method='crazy-tfr')
    tfr.save(fname, True)
    read_tfr = read_tfrs(fname)[0]
    assert_array_equal(tfr.data, read_tfr.data)

def test_plot():
    """Test TFR plotting."""
    import matplotlib.pyplot as plt

    data = np.zeros((3, 2, 3))
    times = np.array([.1, .2, .3])
    freqs = np.array([.10, .20])
    info = mne.create_info(['MEG 001', 'MEG 002', 'MEG 003'], 1000.,
                           ['mag', 'mag', 'mag'])
    tfr = AverageTFR(info, data=data, times=times, freqs=freqs,
                     nave=20, comment='test', method='crazy-tfr')
    tfr.plot([1, 2], title='title', colorbar=False,
             mask=np.ones(tfr.data.shape[1:], bool))
    plt.close('all')
    ax = plt.subplot2grid((2, 2), (0, 0))
    ax2 = plt.subplot2grid((2, 2), (1, 1))
    ax3 = plt.subplot2grid((2, 2), (0, 1))
    tfr.plot(picks=[0, 1, 2], axes=[ax, ax2, ax3])
    plt.close('all')

    tfr.plot([1, 2], title='title', colorbar=False, exclude='bads')
    plt.close('all')

    tfr.plot_topo(picks=[1, 2])
    plt.close('all')

    fig = tfr.plot(picks=[1], cmap='RdBu_r')  # interactive mode on by default
    fig.canvas.key_press_event('up')
    fig.canvas.key_press_event(' ')
    fig.canvas.key_press_event('down')

    cbar = fig.get_axes()[0].CB  # Fake dragging with mouse.
    ax = cbar.cbar.ax
    _fake_click(fig, ax, (0.1, 0.1))
    _fake_click(fig, ax, (0.1, 0.2), kind='motion')
    _fake_click(fig, ax, (0.1, 0.3), kind='release')

    _fake_click(fig, ax, (0.1, 0.1), button=3)
    _fake_click(fig, ax, (0.1, 0.2), button=3, kind='motion')
    _fake_click(fig, ax, (0.1, 0.3), kind='release')

    fig.canvas.scroll_event(0.5, 0.5, -0.5)  # scroll down
    fig.canvas.scroll_event(0.5, 0.5, 0.5)  # scroll up

    plt.close('all')

def test_plot_joint():
    """Test TFR joint plotting."""
    import matplotlib.pyplot as plt

    raw = read_raw_fif(raw_fname)
    times = np.linspace(-0.1, 0.1, 200)
    n_freqs = 3
    nave = 1
    rng = np.random.RandomState(42)
    data = rng.randn(len(raw.ch_names), n_freqs, len(times))
    tfr = AverageTFR(raw.info, data, times, np.arange(n_freqs), nave)

    topomap_args = {'res': 8, 'contours': 0, 'sensors': False}

    for combine in ('mean', 'rms', None):
        tfr.plot_joint(title='auto', colorbar=True,
                       combine=combine, topomap_args=topomap_args)
        plt.close('all')

    # check various timefreqs
    for timefreqs in (
            {(tfr.times[0], tfr.freqs[1]): (0.1, 0.5),
             (tfr.times[-1], tfr.freqs[-1]): (0.2, 0.6)},
            [(tfr.times[1], tfr.freqs[1])]):
        tfr.plot_joint(timefreqs=timefreqs, topomap_args=topomap_args)
        plt.close('all')

    # test bad timefreqs
    timefreqs = ([(-100, 1)], tfr.times[1], [1],
                 [(tfr.times[1], tfr.freqs[1], tfr.freqs[1])])
    for these_timefreqs in timefreqs:
        pytest.raises(ValueError, tfr.plot_joint, these_timefreqs)

    # test that the object is not internally modified
    tfr_orig = tfr.copy()
    tfr.plot_joint(baseline=(0, None), exclude=[tfr.ch_names[0]],
                   topomap_args=topomap_args)
    plt.close('all')
    assert_array_equal(tfr.data, tfr_orig.data)
    assert (set(tfr.ch_names) == set(tfr_orig.ch_names))
    assert (set(tfr.times) == set(tfr_orig.times))

def test_add_channels():
    """Test tfr splitting / re-appending channel types."""
    data = np.zeros((6, 2, 3))
    times = np.array([.1, .2, .3])
    freqs = np.array([.10, .20])
    info = mne.create_info(
        ['MEG 001', 'MEG 002', 'MEG 003', 'EEG 001', 'EEG 002', 'STIM 001'],
        1000., ['mag', 'mag', 'mag', 'eeg', 'eeg', 'stim'])
    tfr = AverageTFR(info, data=data, times=times, freqs=freqs,
                     nave=20, comment='test', method='crazy-tfr')
    tfr_eeg = tfr.copy().pick_types(meg=False, eeg=True)
    tfr_meg = tfr.copy().pick_types(meg=True)
    tfr_stim = tfr.copy().pick_types(meg=False, stim=True)
    tfr_eeg_meg = tfr.copy().pick_types(meg=True, eeg=True)
    tfr_new = tfr_meg.copy().add_channels([tfr_eeg, tfr_stim])
    assert all(ch in tfr_new.ch_names
               for ch in tfr_stim.ch_names + tfr_meg.ch_names)
    tfr_new = tfr_meg.copy().add_channels([tfr_eeg])

    assert all(ch in tfr_new.ch_names
               for ch in tfr.ch_names if ch != 'STIM 001')
    assert_array_equal(tfr_new.data, tfr_eeg_meg.data)
    assert all(ch not in tfr_new.ch_names for ch in tfr_stim.ch_names)

    # Now test errors
    tfr_badsf = tfr_eeg.copy()
    tfr_badsf.info['sfreq'] = 3.1415927
    tfr_eeg = tfr_eeg.crop(-.1, .1)

    pytest.raises(RuntimeError, tfr_meg.add_channels, [tfr_badsf])
    pytest.raises(AssertionError, tfr_meg.add_channels, [tfr_eeg])
    pytest.raises(ValueError, tfr_meg.add_channels, [tfr_meg])
    pytest.raises(TypeError, tfr_meg.add_channels, tfr_badsf)

def test_compute_tfr():
    """Test _compute_tfr function."""
    # Set parameters
    event_id = 1
    tmin = -0.2
    tmax = 0.498  # Allows exhaustive decimation testing

    # Setup for reading the raw data
    raw = read_raw_fif(raw_fname)
    events = read_events(event_fname)

    exclude = raw.info['bads'] + ['MEG 2443', 'EEG 053']  # bads + 2 more

    # picks MEG gradiometers
    picks = pick_types(raw.info, meg='grad', eeg=False,
                       stim=False, include=[], exclude=exclude)

    picks = picks[:2]
    epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks)
    data = epochs.get_data()
    sfreq = epochs.info['sfreq']
    freqs = np.arange(10, 20, 3).astype(float)

    # Check all combination of options
    for func, use_fft, zero_mean, output in product(
        (tfr_array_multitaper, tfr_array_morlet), (False, True), (False, True),
        ('complex', 'power', 'phase',
         'avg_power_itc', 'avg_power', 'itc')):
        # Check exception
        if (func == tfr_array_multitaper) and (output == 'phase'):
            pytest.raises(NotImplementedError, func, data, sfreq=sfreq,
                          freqs=freqs, output=output)
            continue

        # Check runs
        out = func(data, sfreq=sfreq, freqs=freqs, use_fft=use_fft,
                   zero_mean=zero_mean, n_cycles=2., output=output)
        # Check shapes
        shape = np.r_[data.shape[:2], len(freqs), data.shape[2]]
        if ('avg' in output) or ('itc' in output):
            assert_array_equal(shape[1:], out.shape)
        else:
            assert_array_equal(shape, out.shape)

        # Check types
        if output in ('complex', 'avg_power_itc'):
            assert_equal(np.complex, out.dtype)
        else:
            assert_equal(np.float, out.dtype)
        assert (np.all(np.isfinite(out)))

    # Check errors params
    for _data in (None, 'foo', data[0]):
        pytest.raises(ValueError, _compute_tfr, _data, freqs, sfreq)
    for _freqs in (None, 'foo', [[0]]):
        pytest.raises(ValueError, _compute_tfr, data, _freqs, sfreq)
    for _sfreq in (None, 'foo'):
        pytest.raises(ValueError, _compute_tfr, data, freqs, _sfreq)
    for key in ('output', 'method', 'use_fft', 'decim', 'n_jobs'):
        for value in (None, 'foo'):
            kwargs = {key: value}  # FIXME pep8
            pytest.raises(ValueError, _compute_tfr, data, freqs, sfreq,
                          **kwargs)

    # No time_bandwidth param in morlet
    pytest.raises(ValueError, _compute_tfr, data, freqs, sfreq,
                  method='morlet', time_bandwidth=1)
    # No phase in multitaper XXX Check ?
    pytest.raises(NotImplementedError, _compute_tfr, data, freqs, sfreq,
                  method='multitaper', output='phase')

    # Inter-trial coherence tests
    out = _compute_tfr(data, freqs, sfreq, output='itc', n_cycles=2.)
    assert (np.sum(out >= 1) == 0)
    assert (np.sum(out <= 0) == 0)

    # Check decim shapes
    # 2: multiple of len(times) even
    # 3: multiple odd
    # 8: not multiple, even
    # 9: not multiple, odd
    for decim in (2, 3, 8, 9, slice(0, 2), slice(1, 3), slice(2, 4)):
        _decim = slice(None, None, decim) if isinstance(decim, int) else decim
        n_time = len(np.arange(data.shape[2])[_decim])
        shape = np.r_[data.shape[:2], len(freqs), n_time]
        for method in ('multitaper', 'morlet'):
            # Single trials
            out = _compute_tfr(data, freqs, sfreq, method=method, decim=decim,
                               n_cycles=2.)
            assert_array_equal(shape, out.shape)
            # Averages
            out = _compute_tfr(data, freqs, sfreq, method=method, decim=decim,
                               output='avg_power', n_cycles=2.)
            assert_array_equal(shape[1:], out.shape)

run_tests_if_main()


#!/usr/bin/env python

import sys
import operator
import pandas as pd
import numpy as np
from sklearn import cross_validation
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.cross_validation import train_test_split
from sklearn.preprocessing import label_binarize
from sklearn.metrics import roc_curve, auc
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
from scipy import interp
from dataset import load_data

# obtains the classifications from the final curated dataset
def get_targets():
	with open('/Users/schencro/Desktop/FINAL_DATASET/Curated_Dataset/FINAL_CURATED_TABLE.csv','r') as table:
		typed = {}
		for line in table:
			line = line.split(',')
			acc = line[1].rstrip(' ')
			typed.update({acc:line[2]})
	return typed

# obtain FINAL_DATASET for model (all data)
def get_data():
	with open('/Users/schencro/Desktop/FINAL_DATASET/Curated_Dataset/FINAL_CURATED_SCORES.csv', 'r') as scores:
		scores = scores.readlines()
	formatted = []
	for item in scores:
		item = item.rstrip('\n')
		item = item.split(',')
		sample = [item[0]]
		for i in range(1, len(item)):
			 ind = float(item[i])
			 sample.append(ind)
		formatted.append(sample)
	scores = None
	return formatted

# get arrays after fetching the proper classification and getting that classifications set of scores
def get_arrays(types, scores):
	order_types = []
	out_scores = []
	for item in scores:
		acc = item[0]
		ctype = types[acc]
		order_types.append(ctype)
		del item[0]
		out_scores.append(item)

	# the arrays needed for cross validation
	type_array = np.asarray(order_types)
	scores = np.asarray(out_scores)

	# cleanup
	item = None
	ourder_types = None
	out_scores = None

	return scores, type_array

# ExtraTreesClassifier model
def extratrees_model(x, y):
	clf = ExtraTreesClassifier(n_estimators=25, class_weight={"Type A":0.3,"Type B":0.5,"Neither":0.2}, bootstrap=False, max_features=125, criterion='gini', n_jobs=-1)
	clf = clf.fit(x, y)
	return clf

# Voting model
def results_vote(x, y):
	pass

# Section for running loops on different parameters
def tune_model_parameters(data, targets):
	# cross validate and tuning of the ExtraTreesClassifier parameters
	my_range = range(1,20)
	n_scores = []
	for n in my_range:
		clf = ExtraTreesClassifier(n_estimators=25, class_weight={"Type A":0.3,"Type B":0.5,"Neither":0.2}, bootstrap=False, max_features=125, criterion='gini', n_jobs=-1)
		scores = cross_validation.cross_val_score(clf, data, targets, cv=10, scoring='accuracy')
		n_scores.append(scores.mean())

	plt.plot(my_range,n_scores)
	plt.xlabel('Number of Trees in the Forest')
	plt.ylabel('Cross-Validated Accuracy (10-fold Mean)')
	plt.show()
	#plt.savefig('/Users/ryan/Desktop/FINAL_DATASET/Curated_Dataset/Capsid_Classifier/max_features_10_126.png', bbox_inches = 'tight')

	# get the parameter with the maximum mean output
	m = max(n_scores)
	mi = min(n_scores)
	print 'Max Accuracy: ' + repr(m)
	index = [i for i, j in enumerate(n_scores) if j == m]
	for i in index:
		print 'Parameter value max: ' + repr(my_range[i])
	indexmi = [i for i, j in enumerate(n_scores) if j == mi]
	print 'Min Accuracy: ' + repr(mi)
	for i in indexmi:
		print 'Parameter value min: ' + repr(my_range[i])

# get ROC curves for the predictions
def get_roc(data, targets):
	# binarize the classifactions
	bi_targets = label_binarize(targets, classes=['Type A', 'Type B', 'Neither'])
	#print bi_targets
	#print targets
	n_classes = bi_targets.shape[1]
	#print n_classes

	# shuffle and split training and test sets
	X_train, X_test, y_train, y_test = train_test_split(data, bi_targets, train_size=.8)

	# convert array to array of strings instead of arrays of arrays for the classifier (for the weights)
	string_test = []
	for i in range(0, len(y_train)):
		string_test.append(str(y_train[i]))
	string_test = np.asarray(string_test)

	clf = ExtraTreesClassifier(n_estimators=25, class_weight={"[1 0 0]":0.4,"[0 1 0]":0.5,"[0 1 0]":0.1}, bootstrap=False, max_features=125, criterion='gini', n_jobs = -1)
	model = clf.fit(X_train, string_test)

	y_score = model.predict(X_test)

	# get output of scores from string list into a np array
	array_scores = []
	for item in y_score:
		ind = item.split(' ')
		ind0 = ind[0].lstrip('[')
		ind1 = ind[1]
		ind2 = ind[2].rstrip(']')
		ind = [int(ind0),int(ind1), int(ind2)]
		array_scores.append(ind)
	array_scores = np.asarray(array_scores)
	print array_scores

	# Compute ROC curve and ROC area for each class
	fpr = dict()
	tpr = dict()
	roc_auc = dict()
	for i in range(n_classes):
		fpr[i], tpr[i], _ = roc_curve(y_test[:, i], array_scores[:, i])
		roc_auc[i] = auc(fpr[i], tpr[i])

	# Compute micro-average ROC curve and ROC area
	fpr["micro"], tpr["micro"], _ = roc_curve(y_test.ravel(), array_scores.ravel())
	roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])
	'''
	plt.figure()
	plt.plot(fpr[2], tpr[2], label='ROC curve (area = %0.2f)' % roc_auc[2])
	plt.plot([0, 1], [0, 1], 'k--')
	plt.xlim([0.0, 1.0])
	plt.ylim([0.0, 1.05])
	plt.xlabel('False Positive Rate')
	plt.ylabel('True Positive Rate')
	plt.title('Receiver operating characteristic example')
	plt.legend(loc="lower right")
	plt.show()
	'''
	# Plot ROC curves for the multiclass problem

	# Compute macro-average ROC curve and ROC area

	# First aggregate all false positive rates
	all_fpr = np.unique(np.concatenate([fpr[i] for i in range(n_classes)]))

	# Then interpolate all ROC curves at this points
	mean_tpr = np.zeros_like(all_fpr)
	for i in range(n_classes):
	    mean_tpr += interp(all_fpr, fpr[i], tpr[i])

	# Finally average it and compute AUC
	mean_tpr /= n_classes

	fpr["macro"] = all_fpr
	tpr["macro"] = mean_tpr
	roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])

	# Plot all ROC curves
	plt.figure()
	plt.plot(fpr["micro"], tpr["micro"],
	         label='micro-average ROC curve (area = {0:0.2f})'
	               ''.format(roc_auc["micro"]),
	         linewidth=2)

	plt.plot(fpr["macro"], tpr["macro"],
	         label='macro-average ROC curve (area = {0:0.2f})'
	               ''.format(roc_auc["macro"]),
	         linewidth=2)

	for i in range(n_classes):
	    plt.plot(fpr[i], tpr[i], label='ROC curve of class {0} (area = {1:0.2f})'
	                                   ''.format(i, roc_auc[i]))

	plt.plot([0, 1], [0, 1], 'k--')
	plt.xlim([0.0, 1.0])
	plt.ylim([0.0, 1.05])
	plt.xlabel('False Positive Rate')
	plt.ylabel('True Positive Rate')
	plt.title('Receiver operating characteristics')
	plt.legend(loc="lower right")
	plt.savefig('/Users/schencro/Desktop/FINAL_DATASET/Curated_Dataset/Capsid_Classifier/ROC_curves.eps', bbox_inches = 'tight')

# plot confusion matrices
def plot_confusion_matrix(cm, labels, title='Confusion matrix', cmap=plt.cm.Greens):
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(labels))
    plt.xticks(tick_marks, labels, rotation=45)
    plt.yticks(tick_marks, labels)
    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')

def cm_model_p1(X_train, y_train):
	clf = ExtraTreesClassifier(n_estimators=25, class_weight={"Type A":0.3,"Type B":0.5,"Neither":0.2}, bootstrap=False, max_features=125, criterion='gini', n_jobs=-1)
	model = clf.fit(X_train, y_train)
	return model

def cm_model_p2(model, X_test):
	# generate 100 predictions and vote for the majority for final prediction
	hundred_pred = []
	for i in range(0,100):
		y_pred = model.predict(X_test)
		hundred_pred.append(y_pred)
	final_pred = []
	for i in range(0, len(hundred_pred[0])):
		types = []
		for k,t in enumerate(hundred_pred):
			types.append(hundred_pred[k][i])
		counts = [types.count('Type A'),types.count('Type B'),types.count('Neither')]
		index, value = max(enumerate(counts), key=operator.itemgetter(1))
		if index == 0:
			final_pred.append('Type A')
		elif index == 1:
			final_pred.append('Type B')
		elif index == 2:
			final_pred.append('Neither')
		else:
			pass
	y_pred = np.asarray(final_pred)
	return y_pred

# Generate confusion matrix
def get_conf_matrix(data, targets):
	# shuffle and split training and test sets
	X_train, X_test, y_train, y_test = train_test_split(data, targets, train_size=.8)

	# gets the model for predictions
	model = cm_model_p1(X_train, y_train)

	# generate 100 confusion matrices, get mean value for each
	out_cm = np.zeros((3,3))
	for i in range(0,100):
		y_pred = cm_model_p2(model, X_test)
		# Compute confusion matrix
		labels = ['Type A', 'Type B', 'Neither']
		cm = confusion_matrix(y_test, y_pred, labels=labels)
		np.set_printoptions(precision=2)
		# Normalize the confusion matrix by row (i.e by the number of samples
		# in each class)
		cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
		out_cm += cm_normalized
	print out_cm
	cm_normalized = np.divide(out_cm, 100.0)

	print('Normalized confusion matrix (Mean of 100 predictions)')
	print(cm_normalized)
	plt.figure()
	plot_confusion_matrix(cm_normalized, labels, title='Normalized confusion matrix')
	# plt.show()
	plt.savefig('/Users/schencro/Desktop/FINAL_DATASET/Curated_Dataset/Capsid_Classifier/confusion_matrix_RYANFINAL_100mean.eps', bbox_inches = 'tight')

def main():
	'''
	# Use these three to get the data loaded, targets loaded, and the accessions stripped (Otherwise use dataset.py load_data())
	# get classifications
	type_dict = get_targets()

	# load data
	scores = get_data()

	# get arrays of scores and targets
	data, targets = get_arrays(type_dict, scores)
	'''

	data, targets = load_data()

	# tune model parameters
	#tune_model_parameters(data,targets)

	# get ROC curves
	#get_roc(data, targets)

	# get confusion matrix
	get_conf_matrix(data, targets)

	'''I WANT TO RE-RUN the ROC curves and the Confusion matrix data using predictions from a cross-validation rather than train/test_split'''

if __name__ == "__main__":
	main()


import os
import re

from pandas import DataFrame

class WaferRun:

    def __init__(self, run_id, wafer_id, label, measurements):
        self.run_id = int(run_id)
        self.wafer_id = int(wafer_id)
        self.label = int(label)
        self.measurements = DataFrame(measurements)
        self.measurements.sort(axis=1, inplace=True)
        self.measurements.sort_index(inplace=True)

    @staticmethod
    def from_files(path, run_id, wafer_id):
        fn_base = os.path.join(path, '{0}_{1:02}'.format(run_id, wafer_id))

        try:
            df = DataFrame({11: DataFrame.from_csv(fn_base + '.11', header=None, sep='\t', index_col=None, parse_dates=False)[1],
                            12: DataFrame.from_csv(fn_base + '.12', header=None, sep='\t', index_col=None, parse_dates=False)[1],
                            15: DataFrame.from_csv(fn_base + '.15', header=None, sep='\t', index_col=None, parse_dates=False)[1],
                            6: DataFrame.from_csv(fn_base + '.6', header=None, sep='\t', index_col=None, parse_dates=False)[1],
                            7: DataFrame.from_csv(fn_base + '.7', header=None, sep='\t', index_col=None, parse_dates=False)[1],
                            8: DataFrame.from_csv(fn_base + '.8', header=None, sep='\t', index_col=None, parse_dates=False)[1]})
        except:
            return None

        m = re.search('/(normal|abnormal)', path)
        if m is None:
            return None

        label = 1 if m.group(1) == 'abnormal' else -1

        return WaferRun(run_id, wafer_id, label, df)

    def as_nparray(self):
        """Spits out data as a T x D numpy.array (T=# samples, D=# variables)

        Notes:
        Notice what we do here: we start with a pandas.DataFrame where each channel
        is a column (so you can think of it as a T x D matrix). We first rename the
        columns to channel numbers,then sort the columns, then sort the index, then
        transform to numpy.array.
        """
        return self.measurements.sort(axis=1).sort_index().reset_index().as_matrix().astype(float)


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

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

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

n_neighbors = 15

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

h = .02  # step size in the mesh

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

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

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

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

plt.show()


from scipy.linalg import expm, norm
import numpy as np

def rot_mat(axis, theta):
    return expm(np.cross(np.eye(3), axis/norm(axis)*theta))

def rotate_vector(v, axis, theta):
    M = rot_mat(axis, theta)
    return np.tensordot(M,v,axes=([0],[1])).T #np.dot(M, v)

def rotate_around_z(v, theta):
    return rotate_vector(v, np.array([0.,0.,1.]), theta)

def is_odd(num):
    return num & 0x1

def is_inside_hexagon(x, y, d=None, x0=0., y0=0.):
    p_eps = 10.*np.finfo(float).eps
    if d is None:
        d = y.max() - y.min() + p_eps
    dx = np.abs(x - x0)/d
    dy = np.abs(y - y0)/d
    a = 0.25 * np.sqrt(3.0)
    return np.logical_and(dx <= a, a*dy + 0.25*dx <= 0.5*a)

def get_hex_plane(plane_idx, inradius, z_height, z_center, np_xy,
                  np_z):

    # We use 10* float machine precision to correct the ccordinates
    # to avoid leaving the computational domain due to precision
    # problems
    p_eps = 10.*np.finfo(float).eps

    ri = inradius # short for inradius
    rc = inradius/np.sqrt(3.)*2. # short for circumradius

    if np_z == 'auto':
        np_z = int(np.round(float(np_xy)/2./rc*z_height))

    # XY-plane (no hexagonal shape!)
    if plane_idx == 6:
        X = np.linspace(-ri+p_eps, ri-p_eps, np_xy)
        Y = np.linspace(-rc+p_eps, rc-p_eps, np_xy)
        XY = np.meshgrid(X,Y)
        XYrs = np.concatenate((XY[0][..., np.newaxis],
                               XY[1][..., np.newaxis]),
                              axis=2)
        Z = np.ones((np_xy, np_xy, 1))*z_center
        pl = np.concatenate((XYrs, Z), axis=2)
        pl = pl.reshape(-1, pl.shape[-1])

        # Restrict to hexagon
        idx_hex = is_inside_hexagon(pl[:,0], pl[:,1])
        return pl[idx_hex]

    # Vertical planes
    elif plane_idx < 6:
        r = rc if is_odd(plane_idx) else ri
        r = r-p_eps
        xy_line = np.empty((np_xy,2))
        xy_line[:,0] = np.linspace(-r, r, np_xy)
        xy_line[:,1] = 0.
        z_points = np.linspace(0.+p_eps, z_height-p_eps, np_z)

        # Construct the plane
        plane = np.empty((np_xy*np_z, 3))
        for i, xy in enumerate(xy_line):
            for j, z in enumerate(z_points):
                idx = i*np_z + j
                plane[idx, :2] = xy
                plane[idx, 2] = z

        # Rotate the plane
        return rotate_around_z(plane, plane_idx*np.pi/6.)
    else:
        raise ValueError('`plane_idx` must be in [0...6].')

def get_hex_planes_point_list(inradius, z_height, z_center, np_xy, np_z,
                              plane_indices=[0,1,2,3,6]):
    # Construct the desired planes
    planes = []
    for i in plane_indices:
        planes.append(get_hex_plane(i, inradius, z_height, z_center,
                                    np_xy, np_z))

    # Flatten and save lengths
    lengths = [len(p) for p in planes]
    return np.vstack(planes), np.array(lengths)

def hex_planes_point_list_for_keys(keys, plane_indices=[0,1,2,3,6]):
    if not 'uol' in keys:
        keys['uol'] = 1.e-9
    inradius = keys['p'] * keys['uol'] /2.
    z_height = (keys['h'] + keys['h_sub'] + keys['h_sup']) * keys['uol']
    z_center = (keys['h_sub']+keys['h']/2.) * keys['uol']
    np_xy = keys['hex_np_xy']
    if not 'hex_np_z' in keys:
        np_z = 'auto'
    return get_hex_planes_point_list(inradius, z_height, z_center, np_xy,
                                     np_z)

def plane_idx_iter(lengths_):
    """Yields the plane index plus lower index `idx_i` and upper index
    `idx_f` of the point list representing this plane
    (i.e. pointlist[idx_i:idx_f]).

    """
    i = 0
    while i < len(lengths_):
        yield i, lengths_[:i].sum(), lengths_[:(i+1)].sum()
        i += 1

def plot_planes(pointlist, lengths):
    import matplotlib.pyplot as plt
    import seaborn as sns
    from mpl_toolkits.mplot3d import Axes3D

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    colors = sns.color_palette('husl', len(lengths))
    for i, idx_i, idx_f in plane_idx_iter(lengths):
        pl = pointlist[idx_i:idx_f]
        ax.scatter(pl[:,0], pl[:,1], pl[:,2], s=10., c=colors[i],
                   label='plane {}'.format(i+1), linewidth=0.)
    _ = plt.legend(loc='upper left')


import unittest
import sys

import numpy as np

from sklearn.mixture import DPGMM, VBGMM
from sklearn.mixture.dpgmm import log_normalize
from sklearn.datasets import make_blobs
from sklearn.utils.testing import assert_array_less, assert_equal
from sklearn.mixture.tests.test_gmm import GMMTester
from sklearn.externals.six.moves import cStringIO as StringIO

np.seterr(all='warn')

def test_class_weights():
    # check that the class weights are updated
    # simple 3 cluster dataset
    X, y = make_blobs(random_state=1)
    for Model in [DPGMM, VBGMM]:
        dpgmm = Model(n_components=10, random_state=1, alpha=20, n_iter=50)
        dpgmm.fit(X)
        # get indices of components that are used:
        indices = np.unique(dpgmm.predict(X))
        active = np.zeros(10, dtype=np.bool)
        active[indices] = True
        # used components are important
        assert_array_less(.1, dpgmm.weights_[active])
        # others are not
        assert_array_less(dpgmm.weights_[~active], .05)

def test_verbose_boolean():
    # checks that the output for the verbose output is the same
    # for the flag values '1' and 'True'
    # simple 3 cluster dataset
    X, y = make_blobs(random_state=1)
    for Model in [DPGMM, VBGMM]:
        dpgmm_bool = Model(n_components=10, random_state=1, alpha=20,
                           n_iter=50, verbose=True)
        dpgmm_int = Model(n_components=10, random_state=1, alpha=20,
                          n_iter=50, verbose=1)

        old_stdout = sys.stdout
        sys.stdout = StringIO()
        try:
            # generate output with the boolean flag
            dpgmm_bool.fit(X)
            verbose_output = sys.stdout
            verbose_output.seek(0)
            bool_output = verbose_output.readline()
            # generate output with the int flag
            dpgmm_int.fit(X)
            verbose_output = sys.stdout
            verbose_output.seek(0)
            int_output = verbose_output.readline()
            assert_equal(bool_output, int_output)
        finally:
            sys.stdout = old_stdout

def test_verbose_first_level():
    # simple 3 cluster dataset
    X, y = make_blobs(random_state=1)
    for Model in [DPGMM, VBGMM]:
        dpgmm = Model(n_components=10, random_state=1, alpha=20, n_iter=50,
                      verbose=1)

        old_stdout = sys.stdout
        sys.stdout = StringIO()
        try:
            dpgmm.fit(X)
        finally:
            sys.stdout = old_stdout

def test_verbose_second_level():
    # simple 3 cluster dataset
    X, y = make_blobs(random_state=1)
    for Model in [DPGMM, VBGMM]:
        dpgmm = Model(n_components=10, random_state=1, alpha=20, n_iter=50,
                      verbose=2)

        old_stdout = sys.stdout
        sys.stdout = StringIO()
        try:
            dpgmm.fit(X)
        finally:
            sys.stdout = old_stdout

def test_log_normalize():
    v = np.array([0.1, 0.8, 0.01, 0.09])
    a = np.log(2 * v)
    assert np.allclose(v, log_normalize(a), rtol=0.01)

def do_model(self, **kwds):
    return VBGMM(verbose=False, **kwds)

class DPGMMTester(GMMTester):
    model = DPGMM
    do_test_eval = False

    def score(self, g, train_obs):
        _, z = g.score_samples(train_obs)
        return g.lower_bound(train_obs, z)

class TestDPGMMWithSphericalCovars(unittest.TestCase, DPGMMTester):
    covariance_type = 'spherical'
    setUp = GMMTester._setUp

class TestDPGMMWithDiagCovars(unittest.TestCase, DPGMMTester):
    covariance_type = 'diag'
    setUp = GMMTester._setUp

class TestDPGMMWithTiedCovars(unittest.TestCase, DPGMMTester):
    covariance_type = 'tied'
    setUp = GMMTester._setUp

class TestDPGMMWithFullCovars(unittest.TestCase, DPGMMTester):
    covariance_type = 'full'
    setUp = GMMTester._setUp

class VBGMMTester(GMMTester):
    model = do_model
    do_test_eval = False

    def score(self, g, train_obs):
        _, z = g.score_samples(train_obs)
        return g.lower_bound(train_obs, z)

class TestVBGMMWithSphericalCovars(unittest.TestCase, VBGMMTester):
    covariance_type = 'spherical'
    setUp = GMMTester._setUp

class TestVBGMMWithDiagCovars(unittest.TestCase, VBGMMTester):
    covariance_type = 'diag'
    setUp = GMMTester._setUp

class TestVBGMMWithTiedCovars(unittest.TestCase, VBGMMTester):
    covariance_type = 'tied'
    setUp = GMMTester._setUp

class TestVBGMMWithFullCovars(unittest.TestCase, VBGMMTester):
    covariance_type = 'full'
    setUp = GMMTester._setUp


""" functions for developing

Author: Yi Zhang <beingzy@gmail.com>
Date: 2016/03/10
"""
import os
import os.path
from os.path import dirname, abspath, join
import pandas as pd

def get_file_parent_dir_path(level=1):
    """ return the path of the parent directory of current file """
    current_dir_path = dirname(abspath(__file__))
    path_sep = os.path.sep
    components = current_dir_path.split(path_sep)
    return path_sep.join(components[:-level])

def load_sample_test_data():
    """ load small test data """
    _root_dir = get_file_parent_dir_path(level=2)
    _data_dir = join(_root_dir, 'data', 'small_test')

    user_profile_fpath = join(_data_dir, "user_profile.csv")
    user_connections_fpath = join(_data_dir, "connections.csv")

    int_user_profile_df = pd.read_csv(user_profile_fpath, header=0, sep=',')
    user_connections_df = pd.read_csv(user_connections_fpath, header=0, sep=',')

    user_ids = int_user_profile_df.id.tolist()
    # remove id columns and cetegorical feature column
    user_profile_df = int_user_profile_df.drop(["id", "feat_3"], axis=1, inplace=False).as_matrix()
    user_connections_df = user_connections_df.as_matrix()

    return user_ids, user_profile_df, user_connections_df

def load_simulated_test_data():
    """ load simulationd data  with defined two groups """
    _root_dir = get_file_parent_dir_path(level=2)
    _data_dir = join(_root_dir, 'data', 'sim_two_groups')

    user_profile_fpath = join(_data_dir, "user_profiles.csv")
    user_connections_fpath = join(_data_dir, "friendships.csv")

    # prepare user profile information
    user_profile_df = pd.read_csv(user_profile_fpath, header=0, sep=",")
    # unpack data
    user_ids = user_profile_df.ID.tolist()
    user_true_groups = user_profile_df.decision_style.tolist()
    user_profile_df = user_profile_df.drop(["ID", "decision_style"], axis=1, inplace=False).as_matrix()

    user_connections_df = pd.read_csv(user_connections_fpath, header=0, sep=",")
    user_connections_df = (user_connections_df[user_connections_df.isFriend==1]
                       .drop('isFriend', axis=1, inplace=False).astype(int).as_matrix())

    return user_ids, user_profile_df, user_connections_df, user_true_groups


import wx, numpy

from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg
from matplotlib.figure import Figure

def supershape_radius(phi, a, b, m, n1, n2, n3):
	theta = .25 * m * phi
	cos = numpy.fabs(numpy.cos(theta) / a) ** n2
	sin = numpy.fabs(numpy.sin(theta) / b) ** n3
	r = (cos + sin) ** (-1. / n1)
	r /= numpy.max(r)
	return r

class SuperShapeFrame(wx.Frame):
	def __init__(self, parent, id, title):
		wx.Frame.__init__(self, parent, id, title,
		                  style = wx.DEFAULT_FRAME_STYLE ^ wx.RESIZE_BORDER,
		                  size = (480, 480))
		self.fig = Figure((6, 6), dpi = 80)

		self.panel = wx.Panel(self, -1)
		sizer = wx.BoxSizer(wx.VERTICAL)
		sizer.Add(FigureCanvasWxAgg(self.panel, -1, self.fig), 1)
		self.panel.SetSizer(sizer)

		self.draw_figure()

	def draw_figure(self):
		phi = numpy.linspace(0, 2 * numpy.pi, 1024)
		r = supershape_radius(phi, 1, 1, 3, 2, 18, 18)
		ax = self.fig.add_subplot(111, polar = True)
		ax.plot(phi, r, lw = 3.)

		self.fig.canvas.draw()

app = wx.App(redirect = True)
top = SuperShapeFrame(None, -1, 'SuperShape')
top.Show()
app.MainLoop()


'''
Copyright (C) 2015 Dato, Inc.
All rights reserved.

This software may be modified and distributed under the terms
of the BSD license. See the DATO-PYTHON-LICENSE file for details.
'''
from distutils.version import StrictVersion
import logging

def __get_version(version):
    if 'dev' in str(version):
        version = version[:version.find('.dev')]

    return StrictVersion(version)

HAS_PANDAS = True
PANDAS_MIN_VERSION = '0.13.0'
try:
    import pandas
    if __get_version(pandas.__version__) < StrictVersion(PANDAS_MIN_VERSION):
        HAS_PANDAS = False
        logging.warn(('Pandas version %s is not supported. Minimum required version: %s. '
                      'Pandas support will be disabled.')
                      % (pandas.__version__, PANDAS_MIN_VERSION) )
except:
    HAS_PANDAS = False
    import pandas_mock as pandas

HAS_NUMPY = True
NUMPY_MIN_VERSION = '1.8.0'
try:
    import numpy

    if __get_version(numpy.__version__) < StrictVersion(NUMPY_MIN_VERSION):
        HAS_NUMPY = False
        logging.warn(('Numpy version %s is not supported. Minimum required version: %s. '
                      'Numpy support will be disabled.')
                      % (numpy.__version__, NUMPY_MIN_VERSION) )
except:
    HAS_NUMPY = False
    import numpy_mock as numpy


"""
Constant-pressure, adiabatic kinetics simulation with sensitivity analysis
"""

import sys
import numpy as np

import cantera as ct

gri3 = ct.Solution('gri30.xml')
temp = 1500.0
pres = ct.one_atm

gri3.TPX = temp, pres, 'CH4:0.1, O2:2, N2:7.52'
r = ct.IdealGasConstPressureReactor(gri3, name='R1')
sim = ct.ReactorNet([r])

# enable sensitivity with respect to the rates of the first 10
# reactions (reactions 0 through 9)
for i in range(10):
    r.add_sensitivity_reaction(i)

# set the tolerances for the solution and for the sensitivity coefficients
sim.rtol = 1.0e-6
sim.atol = 1.0e-15
sim.rtol_sensitivity = 1.0e-6
sim.atol_sensitivity = 1.0e-6

n_times = 400
tim = np.zeros(n_times)
data = np.zeros((n_times,6))

time = 0.0
for n in range(n_times):
    time += 5.0e-6
    sim.advance(time)
    tim[n] = 1000 * time
    data[n,0] = r.T
    data[n,1:4] = r.thermo['OH','H','CH4'].X

    # sensitivity of OH to reaction 2
    data[n,4] = sim.sensitivity('OH',2)

    # sensitivity of OH to reaction 3
    data[n,5] = sim.sensitivity('OH',3)

    print('%10.3e %10.3f %10.3f %14.6e %10.3f %10.3f' %
          (sim.time, r.T, r.thermo.P, r.thermo.u,  data[n,4],  data[n,5]))

# plot the results if matplotlib is installed.
# see http://matplotlib.org/ to get it
if '--plot' in sys.argv:
    import matplotlib.pyplot as plt
    plt.subplot(2,2,1)
    plt.plot(tim,data[:,0])
    plt.xlabel('Time (ms)')
    plt.ylabel('Temperature (K)')
    plt.subplot(2,2,2)
    plt.plot(tim,data[:,1])
    plt.xlabel('Time (ms)')
    plt.ylabel('OH Mole Fraction')
    plt.subplot(2,2,3)
    plt.plot(tim,data[:,2])
    plt.xlabel('Time (ms)')
    plt.ylabel('H Mole Fraction')
    plt.subplot(2,2,4)
    plt.plot(tim,data[:,3])
    plt.xlabel('Time (ms)')
    plt.ylabel('H2 Mole Fraction')
    plt.tight_layout()

    plt.figure(2)
    plt.plot(tim,data[:,4],'-',tim,data[:,5],'-g')
    plt.legend([sim.sensitivity_parameter_name(2),sim.sensitivity_parameter_name(3)],'best')
    plt.xlabel('Time (ms)')
    plt.ylabel('OH Sensitivity')
    plt.tight_layout()
    plt.show()
else:
    print("""To view a plot of these results, run this script with the option '--plot""")


"""
===================================
Compare cross decomposition methods
===================================

Simple usage of various cross decomposition algorithms:
- PLSCanonical
- PLSRegression, with multivariate response, a.k.a. PLS2
- PLSRegression, with univariate response, a.k.a. PLS1
- CCA

Given 2 multivariate covarying two-dimensional datasets, X, and Y,
PLS extracts the 'directions of covariance', i.e. the components of each
datasets that explain the most shared variance between both datasets.
This is apparent on the **scatterplot matrix** display: components 1 in
dataset X and dataset Y are maximally correlated (points lie around the
first diagonal). This is also true for components 2 in both dataset,
however, the correlation across datasets for different components is
weak: the point cloud is very spherical.
"""
print(__doc__)

import numpy as np
import pylab as pl
from sklearn.cross_decomposition import PLSCanonical, PLSRegression, CCA

###############################################################################
# Dataset based latent variables model

n = 500
# 2 latents vars:
l1 = np.random.normal(size=n)
l2 = np.random.normal(size=n)

latents = np.array([l1, l1, l2, l2]).T
X = latents + np.random.normal(size=4 * n).reshape((n, 4))
Y = latents + np.random.normal(size=4 * n).reshape((n, 4))

X_train = X[:n / 2]
Y_train = Y[:n / 2]
X_test = X[n / 2:]
Y_test = Y[n / 2:]

print("Corr(X)")
print(np.round(np.corrcoef(X.T), 2))
print("Corr(Y)")
print(np.round(np.corrcoef(Y.T), 2))

###############################################################################
# Canonical (symmetric) PLS

# Transform data
# ~~~~~~~~~~~~~~
plsca = PLSCanonical(n_components=2)
plsca.fit(X_train, Y_train)
X_train_r, Y_train_r = plsca.transform(X_train, Y_train)
X_test_r, Y_test_r = plsca.transform(X_test, Y_test)

# Scatter plot of scores
# ~~~~~~~~~~~~~~~~~~~~~~
# 1) On diagonal plot X vs Y scores on each components
pl.figure(figsize=(12, 8))
pl.subplot(221)
pl.plot(X_train_r[:, 0], Y_train_r[:, 0], "ob", label="train")
pl.plot(X_test_r[:, 0], Y_test_r[:, 0], "or", label="test")
pl.xlabel("x scores")
pl.ylabel("y scores")
pl.title('Comp. 1: X vs Y (test corr = %.2f)' %
         np.corrcoef(X_test_r[:, 0], Y_test_r[:, 0])[0, 1])
pl.xticks(())
pl.yticks(())
pl.legend(loc="best")

pl.subplot(224)
pl.plot(X_train_r[:, 1], Y_train_r[:, 1], "ob", label="train")
pl.plot(X_test_r[:, 1], Y_test_r[:, 1], "or", label="test")
pl.xlabel("x scores")
pl.ylabel("y scores")
pl.title('Comp. 2: X vs Y (test corr = %.2f)' %
         np.corrcoef(X_test_r[:, 1], Y_test_r[:, 1])[0, 1])
pl.xticks(())
pl.yticks(())
pl.legend(loc="best")

# 2) Off diagonal plot components 1 vs 2 for X and Y
pl.subplot(222)
pl.plot(X_train_r[:, 0], X_train_r[:, 1], "*b", label="train")
pl.plot(X_test_r[:, 0], X_test_r[:, 1], "*r", label="test")
pl.xlabel("X comp. 1")
pl.ylabel("X comp. 2")
pl.title('X comp. 1 vs X comp. 2 (test corr = %.2f)'
         % np.corrcoef(X_test_r[:, 0], X_test_r[:, 1])[0, 1])
pl.legend(loc="best")
pl.xticks(())
pl.yticks(())

pl.subplot(223)
pl.plot(Y_train_r[:, 0], Y_train_r[:, 1], "*b", label="train")
pl.plot(Y_test_r[:, 0], Y_test_r[:, 1], "*r", label="test")
pl.xlabel("Y comp. 1")
pl.ylabel("Y comp. 2")
pl.title('Y comp. 1 vs Y comp. 2 , (test corr = %.2f)'
         % np.corrcoef(Y_test_r[:, 0], Y_test_r[:, 1])[0, 1])
pl.legend(loc="best")
pl.xticks(())
pl.yticks(())
pl.show()

###############################################################################
# PLS regression, with multivariate response, a.k.a. PLS2

n = 1000
q = 3
p = 10
X = np.random.normal(size=n * p).reshape((n, p))
B = np.array([[1, 2] + [0] * (p - 2)] * q).T
# each Yj = 1*X1 + 2*X2 + noize
Y = np.dot(X, B) + np.random.normal(size=n * q).reshape((n, q)) + 5

pls2 = PLSRegression(n_components=3)
pls2.fit(X, Y)
print("True B (such that: Y = XB + Err)")
print(B)
# compare pls2.coefs with B
print("Estimated B")
print(np.round(pls2.coefs, 1))
pls2.predict(X)

###############################################################################
# PLS regression, with univariate response, a.k.a. PLS1

n = 1000
p = 10
X = np.random.normal(size=n * p).reshape((n, p))
y = X[:, 0] + 2 * X[:, 1] + np.random.normal(size=n * 1) + 5
pls1 = PLSRegression(n_components=3)
pls1.fit(X, y)
# note that the number of compements exceeds 1 (the dimension of y)
print("Estimated betas")
print(np.round(pls1.coefs, 1))

###############################################################################
# CCA (PLS mode B with symmetric deflation)

cca = CCA(n_components=2)
cca.fit(X_train, Y_train)
X_train_r, Y_train_r = plsca.transform(X_train, Y_train)
X_test_r, Y_test_r = plsca.transform(X_test, Y_test)


# -*- coding: utf-8 -*-
"""
    femagtools.forcedens
    ~~~~~~~~~~~~~~~~~~~~

    Read Force Density Plot Files

"""
import os
import re
import glob
import numpy as np
import logging

logger = logging.getLogger('femagtools.forcedens')

filename_pat = re.compile(r'^(\w+)_(\d{3}).PLT(\d+)')
num_pat = re.compile(r'([+-]?\d+(?:\.\d+)?(?:[eE][+-]\d+)?)\s*')
pos_pat = re.compile(r'^\s*POSITION\s*\[(\w+)\]')
unit_pat = re.compile(r'\[([^\]]+)')

def _readSections(f):
    """return list of PLT sections

    sections are surrounded by lines starting with '[***'
    or 2d arrays with 7 columns

    Args:
      param f (file) PLT file to be read

    Returns:
      list of sections
    """

    section = []
    for line in f:
        if line.startswith('[****') or pos_pat.match(line):
            if section:
                if len(section) > 2 and section[1].startswith('Date'):
                    yield section[1:]
                else:
                    yield section
                if line.startswith('[****'):
                    section = []
                else:
                    section = [line.strip()]
        else:
            section.append(line.strip())
    yield section

class ForceDensity(object):

    def __init__(self):
        self.type = ''
        self.positions = []
        pass

    def __read_version(self, content):
        rec = content[0].split(' ')
        if len(rec) > 3:
            self.version = rec[3]
        else:
            self.version = rec[-1]

    def __read_project_filename(self, content):
        self.project = content[1].strip()

    def __read_nodes_and_mesh(self, content):
        self.nodes, self.elements, self.quality = \
            [float(r[0]) for r in [num_pat.findall(l)
                                   for l in content[:3]]]
        for l in content[3:]:
            m = re.match(r'\*+([^\*]+)\*+', l)
            if m:
                self.type = m.group(1).strip()
                return

    def __read_date_and_title(self, content):
        d = content[0].split(':')[1].strip().split()
        dd, MM, yy = d[0].split('.')
        hh, mm = ''.join(d[1:-1]).split('.')
        self.date = '{}-{}-{}T{:02}:{:02}'.format(
            yy, MM, dd, int(hh), int(mm))

        if len(content) > 6:
            self.title = content[2].strip() + ', ' + content[6].strip()
        else:
            self.title = content[2].strip()

        self.current = float(num_pat.findall(content[4])[0])

    def __read_filename(self, content):
        self.filename = content[0].split(':')[1].strip()

    def __read_position(self, content):
        d = dict(position=float(num_pat.findall(content[0])[-1]),
                 unit=unit_pat.findall(content[0].split()[1])[0])
        cols = content[2].split()
        labels = cols[::2] # either X, FN, FT, B_N, B_T
                           # or X FX FY B_X B_Y
        d['column_units'] = {k: u for k, u in zip(labels,
                                                  [unit_pat.findall(u)[0]
                                                   for u in cols[1::2]])}
        m = []
        for l in content[4:]:
            rec = l.split()[1:]
            if len(rec) == len(labels):
                m.append([float(x) for x in rec])
        d.update({k: v for k, v in zip(labels, list(zip(*m)))})

        self.positions.append(d)

    def read(self, filename):
        with open(filename) as f:
            for s in _readSections(f.readlines()):
                logger.debug('Section %s' % s[0:2])
                if s[0].startswith('FEMAG'):
                    self.__read_version(s)
                elif s[0].startswith('Project'):
                    self.__read_project_filename(s)
                elif s[0].startswith('Number'):
                    self.__read_nodes_and_mesh(s)
                elif s[0].startswith('File'):
                    self.__read_filename(s)
                elif s[0].startswith('Date'):
                    self.__read_date_and_title(s)
                elif s[0].startswith('POSITION'):
                    self.__read_position(s)

    def fft(self):
        """return FFT of FN"""
        import scipy.fftpack
        try:
            ntiles = int(360/self.positions[0]['X'][-1])
            FN = np.tile(
                np.array([p['FN'][:-1] for p in self.positions[:-1]]),
                (ntiles, ntiles))
        except AttributeError:
            return []

        N = FN.shape[0]
        fdn = scipy.fftpack.fft2(FN)
        dim = N//ntiles//2
        return np.abs(fdn)[1:dim, 1:dim]/N

    def items(self):
        return [(k, getattr(self, k)) for k in ('version',
                                                'type',
                                                'title',
                                                'current',
                                                'filename',
                                                'date',
                                                'positions')]

    def __str__(self):
        "return string format of this object"
        if self.type:
            return "\n".join([
                'FEMAG {}: {}'.format(self.version, self.type),
                'File: {}  {}'.format(self.filename, self.date)] +
                             ['{}: {}'.format(k, v)
                              for k, v in self.items()])
        return "{}"

def read(filename):
    f = ForceDensity()
    f.read(filename)
    return f

def readall(workdir='.'):
    """collect all recent PLT files
    returns list of ForceDensity objects
    """
    plt = dict()
    pltfiles = sorted(glob.glob(os.path.join(workdir, '*_*.PLT*')))
    base = os.path.basename(pltfiles[-1])
    lastserie = filename_pat.match(base).groups()[1]
    for p in pltfiles:
        base = os.path.basename(p)
        m = filename_pat.match(base)
        if m and lastserie == m.groups()[1]:
            model, i, k = m.groups()
            fdens = ForceDensity()
            fdens.read(p)
            logging.info("%s: %s", p, fdens.title)
            if model in plt:
                plt[model].append(fdens)
            else:
                plt[model] = [fdens]
    return plt

if __name__ == "__main__":
    import matplotlib.pyplot as pl
    import femagtools.plot
    import sys
    if len(sys.argv) == 2:
        filename = sys.argv[1]
    else:
        filename = sys.stdin.readline().strip()

    fdens = read(filename)

    # Show the results

    title = '{}, Rotor position {}'.format(
        fdens.title, fdens.positions[0]['position'])
    pos = fdens.positions[0]['X']
    FT_FN = (fdens.positions[0]['FT'],
             fdens.positions[0]['FN'])
    femagtools.plot.forcedens(title, pos, FT_FN)
    pl.show()

    title = 'Force Density Harmonics'
    femagtools.plot.forcedens_fft(title, fdens)
    pl.show()


import numpy as np
import scipy.sparse as sp

from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_almost_equal
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_less
from sklearn.utils.testing import assert_true

from sklearn.utils.testing import assert_greater
from sklearn.utils.testing import ignore_warnings

from sklearn.linear_model.coordinate_descent import (Lasso, ElasticNet,
                                                     LassoCV, ElasticNetCV)

def test_sparse_coef():
    # Check that the sparse_coef propery works
    clf = ElasticNet()
    clf.coef_ = [1, 2, 3]

    assert_true(sp.isspmatrix(clf.sparse_coef_))
    assert_equal(clf.sparse_coef_.toarray().tolist()[0], clf.coef_)

def test_normalize_option():
    # Check that the normalize option in enet works
    X = sp.csc_matrix([[-1], [0], [1]])
    y = [-1, 0, 1]
    clf_dense = ElasticNet(fit_intercept=True, normalize=True)
    clf_sparse = ElasticNet(fit_intercept=True, normalize=True)
    clf_dense.fit(X, y)
    X = sp.csc_matrix(X)
    clf_sparse.fit(X, y)
    assert_almost_equal(clf_dense.dual_gap_, 0)
    assert_array_almost_equal(clf_dense.coef_, clf_sparse.coef_)

def test_lasso_zero():
    # Check that the sparse lasso can handle zero data without crashing
    X = sp.csc_matrix((3, 1))
    y = [0, 0, 0]
    T = np.array([[1], [2], [3]])
    clf = Lasso().fit(X, y)
    pred = clf.predict(T)
    assert_array_almost_equal(clf.coef_, [0])
    assert_array_almost_equal(pred, [0, 0, 0])
    assert_almost_equal(clf.dual_gap_,  0)

def test_enet_toy_list_input():
    # Test ElasticNet for various values of alpha and l1_ratio with list X

    X = np.array([[-1], [0], [1]])
    X = sp.csc_matrix(X)
    Y = [-1, 0, 1]       # just a straight line
    T = np.array([[2], [3], [4]])  # test sample

    # this should be the same as unregularized least squares
    clf = ElasticNet(alpha=0, l1_ratio=1.0)
    # catch warning about alpha=0.
    # this is discouraged but should work.
    ignore_warnings(clf.fit)(X, Y)
    pred = clf.predict(T)
    assert_array_almost_equal(clf.coef_, [1])
    assert_array_almost_equal(pred, [2, 3, 4])
    assert_almost_equal(clf.dual_gap_, 0)

    clf = ElasticNet(alpha=0.5, l1_ratio=0.3, max_iter=1000)
    clf.fit(X, Y)
    pred = clf.predict(T)
    assert_array_almost_equal(clf.coef_, [0.50819], decimal=3)
    assert_array_almost_equal(pred, [1.0163,  1.5245,  2.0327], decimal=3)
    assert_almost_equal(clf.dual_gap_, 0)

    clf = ElasticNet(alpha=0.5, l1_ratio=0.5)
    clf.fit(X, Y)
    pred = clf.predict(T)
    assert_array_almost_equal(clf.coef_, [0.45454], 3)
    assert_array_almost_equal(pred, [0.9090,  1.3636,  1.8181], 3)
    assert_almost_equal(clf.dual_gap_, 0)

def test_enet_toy_explicit_sparse_input():
    # Test ElasticNet for various values of alpha and l1_ratio with sparse X
    f = ignore_warnings
    # training samples
    X = sp.lil_matrix((3, 1))
    X[0, 0] = -1
    # X[1, 0] = 0
    X[2, 0] = 1
    Y = [-1, 0, 1]       # just a straight line (the identity function)

    # test samples
    T = sp.lil_matrix((3, 1))
    T[0, 0] = 2
    T[1, 0] = 3
    T[2, 0] = 4

    # this should be the same as lasso
    clf = ElasticNet(alpha=0, l1_ratio=1.0)
    f(clf.fit)(X, Y)
    pred = clf.predict(T)
    assert_array_almost_equal(clf.coef_, [1])
    assert_array_almost_equal(pred, [2, 3, 4])
    assert_almost_equal(clf.dual_gap_, 0)

    clf = ElasticNet(alpha=0.5, l1_ratio=0.3, max_iter=1000)
    clf.fit(X, Y)
    pred = clf.predict(T)
    assert_array_almost_equal(clf.coef_, [0.50819], decimal=3)
    assert_array_almost_equal(pred, [1.0163,  1.5245,  2.0327], decimal=3)
    assert_almost_equal(clf.dual_gap_, 0)

    clf = ElasticNet(alpha=0.5, l1_ratio=0.5)
    clf.fit(X, Y)
    pred = clf.predict(T)
    assert_array_almost_equal(clf.coef_, [0.45454], 3)
    assert_array_almost_equal(pred, [0.9090,  1.3636,  1.8181], 3)
    assert_almost_equal(clf.dual_gap_, 0)

def make_sparse_data(n_samples=100, n_features=100, n_informative=10, seed=42,
                     positive=False, n_targets=1):
    random_state = np.random.RandomState(seed)

    # build an ill-posed linear regression problem with many noisy features and
    # comparatively few samples

    # generate a ground truth model
    w = random_state.randn(n_features, n_targets)
    w[n_informative:] = 0.0  # only the top features are impacting the model
    if positive:
        w = np.abs(w)

    X = random_state.randn(n_samples, n_features)
    rnd = random_state.uniform(size=(n_samples, n_features))
    X[rnd > 0.5] = 0.0  # 50% of zeros in input signal

    # generate training ground truth labels
    y = np.dot(X, w)
    X = sp.csc_matrix(X)
    if n_targets == 1:
        y = np.ravel(y)
    return X, y

def _test_sparse_enet_not_as_toy_dataset(alpha, fit_intercept, positive):
    n_samples, n_features, max_iter = 100, 100, 1000
    n_informative = 10

    X, y = make_sparse_data(n_samples, n_features, n_informative,
                            positive=positive)

    X_train, X_test = X[n_samples // 2:], X[:n_samples // 2]
    y_train, y_test = y[n_samples // 2:], y[:n_samples // 2]

    s_clf = ElasticNet(alpha=alpha, l1_ratio=0.8, fit_intercept=fit_intercept,
                       max_iter=max_iter, tol=1e-7, positive=positive,
                       warm_start=True)
    s_clf.fit(X_train, y_train)

    assert_almost_equal(s_clf.dual_gap_, 0, 4)
    assert_greater(s_clf.score(X_test, y_test), 0.85)

    # check the convergence is the same as the dense version
    d_clf = ElasticNet(alpha=alpha, l1_ratio=0.8, fit_intercept=fit_intercept,
                       max_iter=max_iter, tol=1e-7, positive=positive,
                       warm_start=True)
    d_clf.fit(X_train.toarray(), y_train)

    assert_almost_equal(d_clf.dual_gap_, 0, 4)
    assert_greater(d_clf.score(X_test, y_test), 0.85)

    assert_almost_equal(s_clf.coef_, d_clf.coef_, 5)
    assert_almost_equal(s_clf.intercept_, d_clf.intercept_, 5)

    # check that the coefs are sparse
    assert_less(np.sum(s_clf.coef_ != 0.0), 2 * n_informative)

def test_sparse_enet_not_as_toy_dataset():
    _test_sparse_enet_not_as_toy_dataset(alpha=0.1, fit_intercept=False,
                                         positive=False)
    _test_sparse_enet_not_as_toy_dataset(alpha=0.1, fit_intercept=True,
                                         positive=False)
    _test_sparse_enet_not_as_toy_dataset(alpha=1e-3, fit_intercept=False,
                                         positive=True)
    _test_sparse_enet_not_as_toy_dataset(alpha=1e-3, fit_intercept=True,
                                         positive=True)

def test_sparse_lasso_not_as_toy_dataset():
    n_samples = 100
    max_iter = 1000
    n_informative = 10
    X, y = make_sparse_data(n_samples=n_samples, n_informative=n_informative)

    X_train, X_test = X[n_samples // 2:], X[:n_samples // 2]
    y_train, y_test = y[n_samples // 2:], y[:n_samples // 2]

    s_clf = Lasso(alpha=0.1, fit_intercept=False, max_iter=max_iter, tol=1e-7)
    s_clf.fit(X_train, y_train)
    assert_almost_equal(s_clf.dual_gap_, 0, 4)
    assert_greater(s_clf.score(X_test, y_test), 0.85)

    # check the convergence is the same as the dense version
    d_clf = Lasso(alpha=0.1, fit_intercept=False, max_iter=max_iter, tol=1e-7)
    d_clf.fit(X_train.toarray(), y_train)
    assert_almost_equal(d_clf.dual_gap_, 0, 4)
    assert_greater(d_clf.score(X_test, y_test), 0.85)

    # check that the coefs are sparse
    assert_equal(np.sum(s_clf.coef_ != 0.0), n_informative)

def test_enet_multitarget():
    n_targets = 3
    X, y = make_sparse_data(n_targets=n_targets)

    estimator = ElasticNet(alpha=0.01, fit_intercept=True, precompute=None)
    # XXX: There is a bug when precompute is not None!
    estimator.fit(X, y)
    coef, intercept, dual_gap = (estimator.coef_,
                                 estimator.intercept_,
                                 estimator.dual_gap_)

    for k in range(n_targets):
        estimator.fit(X, y[:, k])
        assert_array_almost_equal(coef[k, :], estimator.coef_)
        assert_array_almost_equal(intercept[k], estimator.intercept_)
        assert_array_almost_equal(dual_gap[k], estimator.dual_gap_)

def test_path_parameters():
    X, y = make_sparse_data()
    max_iter = 50
    n_alphas = 10
    clf = ElasticNetCV(n_alphas=n_alphas, eps=1e-3, max_iter=max_iter,
                       l1_ratio=0.5, fit_intercept=False)
    ignore_warnings(clf.fit)(X, y)  # new params
    assert_almost_equal(0.5, clf.l1_ratio)
    assert_equal(n_alphas, clf.n_alphas)
    assert_equal(n_alphas, len(clf.alphas_))
    sparse_mse_path = clf.mse_path_
    ignore_warnings(clf.fit)(X.toarray(), y)  # compare with dense data
    assert_almost_equal(clf.mse_path_, sparse_mse_path)

def test_same_output_sparse_dense_lasso_and_enet_cv():
    X, y = make_sparse_data(n_samples=40, n_features=10)
    for normalize in [True, False]:
        clfs = ElasticNetCV(max_iter=100, cv=5, normalize=normalize)
        ignore_warnings(clfs.fit)(X, y)
        clfd = ElasticNetCV(max_iter=100, cv=5, normalize=normalize)
        ignore_warnings(clfd.fit)(X.toarray(), y)
        assert_almost_equal(clfs.alpha_, clfd.alpha_, 7)
        assert_almost_equal(clfs.intercept_, clfd.intercept_, 7)
        assert_array_almost_equal(clfs.mse_path_, clfd.mse_path_)
        assert_array_almost_equal(clfs.alphas_, clfd.alphas_)

        clfs = LassoCV(max_iter=100, cv=4, normalize=normalize)
        ignore_warnings(clfs.fit)(X, y)
        clfd = LassoCV(max_iter=100, cv=4, normalize=normalize)
        ignore_warnings(clfd.fit)(X.toarray(), y)
        assert_almost_equal(clfs.alpha_, clfd.alpha_, 7)
        assert_almost_equal(clfs.intercept_, clfd.intercept_, 7)
        assert_array_almost_equal(clfs.mse_path_, clfd.mse_path_)
        assert_array_almost_equal(clfs.alphas_, clfd.alphas_)