# -*- coding: utf-8 -*-
"""
Functions to calculate spike-triggered average and spike-field coherence of
analog signals.

.. autosummary::
    :toctree: _toctree/sta

    spike_triggered_average
    spike_field_coherence

:copyright: Copyright 2014-2020 by the Elephant team, see `doc/authors.rst`.
:license: Modified BSD, see LICENSE.txt for details.
"""

from __future__ import division, print_function, unicode_literals

import warnings

import numpy as np
import quantities as pq
import scipy.signal
from neo.core import AnalogSignal, SpikeTrain

from .conversion import BinnedSpikeTrain

__all__ = [
    "spike_triggered_average",
    "spike_field_coherence"
]

def spike_triggered_average(signal, spiketrains, window):
    """
    Calculates the spike-triggered averages of analog signals in a time window
    relative to the spike times of a corresponding spiketrain for multiple
    signals each. The function receives n analog signals and either one or
    n spiketrains. In case it is one spiketrain this one is muliplied n-fold
    and used for each of the n analog signals.

    Parameters
    ----------
    signal : neo AnalogSignal object
        'signal' contains n analog signals.
    spiketrains : one SpikeTrain or one numpy ndarray or a list of n of either of these.
        'spiketrains' contains the times of the spikes in the spiketrains.
    window : tuple of 2 Quantity objects with dimensions of time.
        'window' is the start time and the stop time, relative to a spike, of
        the time interval for signal averaging.
        If the window size is not a multiple of the sampling interval of the
        signal the window will be extended to the next multiple.

    Returns
    -------
    result_sta : neo AnalogSignal object
        'result_sta' contains the spike-triggered averages of each of the
        analog signals with respect to the spikes in the corresponding
        spiketrains. The length of 'result_sta' is calculated as the number
        of bins from the given start and stop time of the averaging interval
        and the sampling rate of the analog signal. If for an analog signal
        no spike was either given or all given spikes had to be ignored
        because of a too large averaging interval, the corresponding returned
        analog signal has all entries as nan. The number of used spikes and
        unused spikes for each analog signal are returned as annotations to
        the returned AnalogSignal object.

    Examples
    --------

    >>> signal = neo.AnalogSignal(np.array([signal1, signal2]).T, units='mV',
    ...                                sampling_rate=10/ms)
    >>> stavg = spike_triggered_average(signal, [spiketrain1, spiketrain2],
    ...                                 (-5 * ms, 10 * ms))

    """

    # checking compatibility of data and data types
    # window_starttime: time to specify the start time of the averaging
    # interval relative to a spike
    # window_stoptime: time to specify the stop time of the averaging
    # interval relative to a spike
    window_starttime, window_stoptime = window
    if not (isinstance(window_starttime, pq.quantity.Quantity) and
            window_starttime.dimensionality.simplified ==
            pq.Quantity(1, "s").dimensionality):
        raise TypeError("The start time of the window (window[0]) "
                        "must be a time quantity.")
    if not (isinstance(window_stoptime, pq.quantity.Quantity) and
            window_stoptime.dimensionality.simplified ==
            pq.Quantity(1, "s").dimensionality):
        raise TypeError("The stop time of the window (window[1]) "
                        "must be a time quantity.")
    if window_stoptime <= window_starttime:
        raise ValueError("The start time of the window (window[0]) must be "
                         "earlier than the stop time of the window (window[1]).")

    # checks on signal
    if not isinstance(signal, AnalogSignal):
        raise TypeError(
            "Signal must be an AnalogSignal, not %s." % type(signal))
    if len(signal.shape) > 1:
        # num_signals: number of analog signals
        num_signals = signal.shape[1]
    else:
        raise ValueError("Empty analog signal, hence no averaging possible.")
    if window_stoptime - window_starttime > signal.t_stop - signal.t_start:
        raise ValueError("The chosen time window is larger than the "
                         "time duration of the signal.")

    # spiketrains type check
    if isinstance(spiketrains, (np.ndarray, SpikeTrain)):
        spiketrains = [spiketrains]
    elif isinstance(spiketrains, list):
        for st in spiketrains:
            if not isinstance(st, (np.ndarray, SpikeTrain)):
                raise TypeError(
                    "spiketrains must be a SpikeTrain, a numpy ndarray, or a "
                    "list of one of those, not %s." % type(spiketrains))
    else:
        raise TypeError(
            "spiketrains must be a SpikeTrain, a numpy ndarray, or a list of "
            "one of those, not %s." % type(spiketrains))

    # multiplying spiketrain in case only a single spiketrain is given
    if len(spiketrains) == 1 and num_signals != 1:
        template = spiketrains[0]
        spiketrains = []
        for i in range(num_signals):
            spiketrains.append(template)

    # checking for matching numbers of signals and spiketrains
    if num_signals != len(spiketrains):
        raise ValueError(
            "The number of signals and spiketrains has to be the same.")

    # checking the times of signal and spiketrains
    for i in range(num_signals):
        if spiketrains[i].t_start < signal.t_start:
            raise ValueError(
                "The spiketrain indexed by %i starts earlier than "
                "the analog signal." % i)
        if spiketrains[i].t_stop > signal.t_stop:
            raise ValueError(
                "The spiketrain indexed by %i stops later than "
                "the analog signal." % i)

    # *** Main algorithm: ***

    # window_bins: number of bins of the chosen averaging interval
    window_bins = int(np.ceil(((window_stoptime - window_starttime) *
        signal.sampling_rate).simplified))
    # result_sta: array containing finally the spike-triggered averaged signal
    result_sta = AnalogSignal(np.zeros((window_bins, num_signals)),
        sampling_rate=signal.sampling_rate, units=signal.units)
    # setting of correct times of the spike-triggered average
    # relative to the spike
    result_sta.t_start = window_starttime
    used_spikes = np.zeros(num_signals, dtype=int)
    unused_spikes = np.zeros(num_signals, dtype=int)
    total_used_spikes = 0

    for i in range(num_signals):
        # summing over all respective signal intervals around spiketimes
        for spiketime in spiketrains[i]:
            # checks for sufficient signal data around spiketime
            if (spiketime + window_starttime >= signal.t_start and
                    spiketime + window_stoptime <= signal.t_stop):
                # calculating the startbin in the analog signal of the
                # averaging window for spike
                startbin = int(np.floor(((spiketime + window_starttime -
                    signal.t_start) * signal.sampling_rate).simplified))
                # adds the signal in selected interval relative to the spike
                result_sta[:, i] += signal[
                    startbin: startbin + window_bins, i]
                # counting of the used spikes
                used_spikes[i] += 1
            else:
                # counting of the unused spikes
                unused_spikes[i] += 1

        # normalization
        result_sta[:, i] = result_sta[:, i] / used_spikes[i]

        total_used_spikes += used_spikes[i]

    if total_used_spikes == 0:
        warnings.warn(
            "No spike at all was either found or used for averaging")
    result_sta.annotate(used_spikes=used_spikes, unused_spikes=unused_spikes)

    return result_sta

def spike_field_coherence(signal, spiketrain, **kwargs):
    """
    Calculates the spike-field coherence between a analog signal(s) and a
    (binned) spike train.

    The current implementation makes use of scipy.signal.coherence(). Additional
    kwargs will will be directly forwarded to scipy.signal.coherence(),
    except for the axis parameter and the sampling frequency, which will be
    extracted from the input signals.

    The spike_field_coherence function receives an analog signal array and
    either a binned spike train or a spike train containing the original spike
    times. In case of original spike times the spike train is binned according
    to the sampling rate of the analog signal array.

    The AnalogSignal object can contain one or multiple signal traces. In case
    of multiple signal traces, the spike field coherence is calculated
    individually for each signal trace and the spike train.

    Parameters
    ----------
    signal : neo AnalogSignal object
        'signal' contains n analog signals.
    spiketrain : SpikeTrain or BinnedSpikeTrain
        Single spike train to perform the analysis on. The bin_size of the
        binned spike train must match the sampling_rate of signal.
    **kwargs:
        All kwargs are passed to `scipy.signal.coherence()`.

    Returns
    -------
    coherence : complex Quantity array
        contains the coherence values calculated for each analog signal trace
        in combination with the spike train. The first dimension corresponds to
        the frequency, the second to the number of the signal trace.
    frequencies : Quantity array
        contains the frequency values corresponding to the first dimension of
        the 'coherence' array

    Examples
    --------

    Plot the SFC between a regular spike train at 20 Hz, and two sinusoidal
    time series at 20 Hz and 23 Hz, respectively.

    >>> import numpy as np
    >>> import matplotlib.pyplot as plt
    >>> from quantities import s, ms, mV, Hz, kHz
    >>> import neo, elephant

    >>> t = pq.Quantity(range(10000),units='ms')
    >>> f1, f2 = 20. * Hz, 23. * Hz
    >>> signal = neo.AnalogSignal(np.array([
            np.sin(f1 * 2. * np.pi * t.rescale(s)),
            np.sin(f2 * 2. * np.pi * t.rescale(s))]).T,
            units=pq.mV, sampling_rate=1. * kHz)
    >>> spiketrain = neo.SpikeTrain(
        range(t[0], t[-1], 50), units='ms',
        t_start=t[0], t_stop=t[-1])
    >>> sfc, freqs = elephant.sta.spike_field_coherence(
        signal, spiketrain, window='boxcar')

    >>> plt.plot(freqs, sfc[:,0])
    >>> plt.plot(freqs, sfc[:,1])
    >>> plt.xlabel('Frequency [Hz]')
    >>> plt.ylabel('SFC')
    >>> plt.xlim((0, 60))
    >>> plt.show()
    """

    if not hasattr(scipy.signal, 'coherence'):
        raise AttributeError('scipy.signal.coherence is not available. The sfc '
                             'function uses scipy.signal.coherence for '
                             'the coherence calculation. This function is '
                             'available for scipy version 0.16 or newer. '
                             'Please update you scipy version.')

    # spiketrains type check
    if not isinstance(spiketrain, (SpikeTrain, BinnedSpikeTrain)):
        raise TypeError(
            "spiketrain must be of type SpikeTrain or BinnedSpikeTrain, "
            "not %s." % type(spiketrain))

    # checks on analogsignal
    if not isinstance(signal, AnalogSignal):
        raise TypeError(
            "Signal must be an AnalogSignal, not %s." % type(signal))
    if len(signal.shape) > 1:
        # num_signals: number of individual traces in the analog signal
        num_signals = signal.shape[1]
    elif len(signal.shape) == 1:
        num_signals = 1
    else:
        raise ValueError("Empty analog signal.")
    len_signals = signal.shape[0]

    # bin spiketrain if necessary
    if isinstance(spiketrain, SpikeTrain):
        spiketrain = BinnedSpikeTrain(
            spiketrain, bin_size=signal.sampling_period)

    # check the start and stop times of signal and spike trains
    if spiketrain.t_start < signal.t_start:
        raise ValueError(
            "The spiketrain starts earlier than the analog signal.")
    if spiketrain.t_stop > signal.t_stop:
        raise ValueError(
            "The spiketrain stops later than the analog signal.")

    # check equal time resolution for both signals
    if spiketrain.bin_size != signal.sampling_period:
        raise ValueError(
            "The spiketrain and signal must have a "
            "common sampling frequency / bin_size")

    # calculate how many bins to add on the left of the binned spike train
    delta_t = spiketrain.t_start - signal.t_start
    if delta_t % spiketrain.bin_size == 0:
        left_edge = int((delta_t / spiketrain.bin_size).magnitude)
    else:
        raise ValueError("Incompatible binning of spike train and LFP")
    right_edge = int(left_edge + spiketrain.n_bins)

    # duplicate spike trains
    spiketrain_array = np.zeros((1, len_signals))
    spiketrain_array[0, left_edge:right_edge] = spiketrain.to_array()
    spiketrains_array = np.repeat(spiketrain_array, repeats=num_signals, axis=0).transpose()

    # calculate coherence
    frequencies, sfc = scipy.signal.coherence(
        spiketrains_array, signal.magnitude,
        fs=signal.sampling_rate.rescale('Hz').magnitude,
        axis=0, **kwargs)

    return (pq.Quantity(sfc, units=pq.dimensionless),
            pq.Quantity(frequencies, units=pq.Hz))


# Author: Peter Prettenhofer <peter.prettenhofer@gmail.com>
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt

import gc

from time import time

from sklearn.linear_model import Ridge, SGDRegressor, ElasticNet
from sklearn.metrics import mean_squared_error
from sklearn.datasets.samples_generator import make_regression

"""
Benchmark for SGD regression

Compares SGD regression against coordinate descent and Ridge
on synthetic data.
"""

print(__doc__)

if __name__ == "__main__":
    list_n_samples = np.linspace(100, 10000, 5).astype(np.int)
    list_n_features = [10, 100, 1000]
    n_test = 1000
    max_iter = 1000
    noise = 0.1
    alpha = 0.01
    sgd_results = np.zeros((len(list_n_samples), len(list_n_features), 2))
    elnet_results = np.zeros((len(list_n_samples), len(list_n_features), 2))
    ridge_results = np.zeros((len(list_n_samples), len(list_n_features), 2))
    asgd_results = np.zeros((len(list_n_samples), len(list_n_features), 2))
    for i, n_train in enumerate(list_n_samples):
        for j, n_features in enumerate(list_n_features):
            X, y, coef = make_regression(
                n_samples=n_train + n_test, n_features=n_features,
                noise=noise, coef=True)

            X_train = X[:n_train]
            y_train = y[:n_train]
            X_test = X[n_train:]
            y_test = y[n_train:]

            print("=======================")
            print("Round %d %d" % (i, j))
            print("n_features:", n_features)
            print("n_samples:", n_train)

            # Shuffle data
            idx = np.arange(n_train)
            np.random.seed(13)
            np.random.shuffle(idx)
            X_train = X_train[idx]
            y_train = y_train[idx]

            std = X_train.std(axis=0)
            mean = X_train.mean(axis=0)
            X_train = (X_train - mean) / std
            X_test = (X_test - mean) / std

            std = y_train.std(axis=0)
            mean = y_train.mean(axis=0)
            y_train = (y_train - mean) / std
            y_test = (y_test - mean) / std

            gc.collect()
            print("- benchmarking ElasticNet")
            clf = ElasticNet(alpha=alpha, l1_ratio=0.5, fit_intercept=False)
            tstart = time()
            clf.fit(X_train, y_train)
            elnet_results[i, j, 0] = mean_squared_error(clf.predict(X_test),
                                                        y_test)
            elnet_results[i, j, 1] = time() - tstart

            gc.collect()
            print("- benchmarking SGD")
            clf = SGDRegressor(alpha=alpha / n_train, fit_intercept=False,
                               max_iter=max_iter, learning_rate="invscaling",
                               eta0=.01, power_t=0.25, tol=1e-3)

            tstart = time()
            clf.fit(X_train, y_train)
            sgd_results[i, j, 0] = mean_squared_error(clf.predict(X_test),
                                                      y_test)
            sgd_results[i, j, 1] = time() - tstart

            gc.collect()
            print("max_iter", max_iter)
            print("- benchmarking A-SGD")
            clf = SGDRegressor(alpha=alpha / n_train, fit_intercept=False,
                               max_iter=max_iter, learning_rate="invscaling",
                               eta0=.002, power_t=0.05, tol=1e-3,
                               average=(max_iter * n_train // 2))

            tstart = time()
            clf.fit(X_train, y_train)
            asgd_results[i, j, 0] = mean_squared_error(clf.predict(X_test),
                                                       y_test)
            asgd_results[i, j, 1] = time() - tstart

            gc.collect()
            print("- benchmarking RidgeRegression")
            clf = Ridge(alpha=alpha, fit_intercept=False)
            tstart = time()
            clf.fit(X_train, y_train)
            ridge_results[i, j, 0] = mean_squared_error(clf.predict(X_test),
                                                        y_test)
            ridge_results[i, j, 1] = time() - tstart

    # Plot results
    i = 0
    m = len(list_n_features)
    plt.figure('scikit-learn SGD regression benchmark results',
               figsize=(5 * 2, 4 * m))
    for j in range(m):
        plt.subplot(m, 2, i + 1)
        plt.plot(list_n_samples, np.sqrt(elnet_results[:, j, 0]),
                 label="ElasticNet")
        plt.plot(list_n_samples, np.sqrt(sgd_results[:, j, 0]),
                 label="SGDRegressor")
        plt.plot(list_n_samples, np.sqrt(asgd_results[:, j, 0]),
                 label="A-SGDRegressor")
        plt.plot(list_n_samples, np.sqrt(ridge_results[:, j, 0]),
                 label="Ridge")
        plt.legend(prop={"size": 10})
        plt.xlabel("n_train")
        plt.ylabel("RMSE")
        plt.title("Test error - %d features" % list_n_features[j])
        i += 1

        plt.subplot(m, 2, i + 1)
        plt.plot(list_n_samples, np.sqrt(elnet_results[:, j, 1]),
                 label="ElasticNet")
        plt.plot(list_n_samples, np.sqrt(sgd_results[:, j, 1]),
                 label="SGDRegressor")
        plt.plot(list_n_samples, np.sqrt(asgd_results[:, j, 1]),
                 label="A-SGDRegressor")
        plt.plot(list_n_samples, np.sqrt(ridge_results[:, j, 1]),
                 label="Ridge")
        plt.legend(prop={"size": 10})
        plt.xlabel("n_train")
        plt.ylabel("Time [sec]")
        plt.title("Training time - %d features" % list_n_features[j])
        i += 1

    plt.subplots_adjust(hspace=.30)

    plt.show()


#!/usr/bin/python
# -*- coding: utf-8 -*-
import matplotlib
matplotlib.use('TkAgg')

O.engines=[
	ForceResetter(),
	InsertionSortCollider([Bo1_Sphere_Aabb(),Bo1_Box_Aabb()]),
	InteractionLoop(
		[Ig2_Sphere_Sphere_ScGeom(),Ig2_Box_Sphere_ScGeom()],
		[Ip2_FrictMat_FrictMat_FrictPhys()],
		[Law2_ScGeom_FrictPhys_CundallStrack()]
	),
	NewtonIntegrator(damping=0.0,gravity=(0,0,-9.81)),
	###
	### NOTE this extra engine:
	###
	### You want snapshot to be taken every 1 sec (realTimeLim) or every 50 iterations (iterLim),
	### whichever comes soones. virtTimeLim attribute is unset, hence virtual time period is not taken into account.
	PyRunner(iterPeriod=20,command='myAddPlotData()')
]
O.bodies.append(box(center=[0,0,0],extents=[.5,.5,.5],fixed=True,color=[1,0,0]))
O.bodies.append(sphere([0,0,2],1,color=[0,1,0]))
O.dt=.002*PWaveTimeStep()

############################################
##### now the part pertaining to plots #####
############################################

from yade import plot
## we will have 2 plots:
## 1. t as function of i (joke test function)
## 2. i as function of t on left y-axis ('|||' makes the separation) and z_sph, v_sph (as green circles connected with line) and z_sph_half again as function of t
plot.plots={'i':('t'),'t':('z_sph',None,('v_sph','go-'),'z_sph_half')}

## this function is called by plotDataCollector
## it should add data with the labels that we will plot
## if a datum is not specified (but exists), it will be NaN and will not be plotted
def myAddPlotData():
	sph=O.bodies[1]
	## store some numbers under some labels
	plot.addData(t=O.time,i=O.iter,z_sph=sph.state.pos[2],z_sph_half=.5*sph.state.pos[2],v_sph=sph.state.vel.norm())
print "Now calling plot.plot() to show the figures. The timestep is artificially low so that you can watch graphs being updated live."
plot.liveInterval=.2
plot.plot(subPlots=False)
O.run(int(5./O.dt));
#plot.saveGnuplot('/tmp/a')
## you can also access the data in plot.data['i'], plot.data['t'] etc, under the labels they were saved.


#!/usr/bin/env python

"""
    --------------------------------------------------------------------------
    Routine to perform aperture photometry on CHIMERA science frames.

    Usage: python fastphot.py [options] image coords

    Authors:
        Navtej Saini, Lee Rosenthal

    Organization:
        Caltech, Pasadena, CA, USA

    Version:
        7 January 2016     0.1     Initial implementation
        9 February 2016    0.2     User input for photometric zero point
        28 July 2017       0.3     Allow processing of multiple stars.
    --------------------------------------------------------------------------
"""

import os, sys
import numpy as np, warnings
from StringIO import StringIO
from optparse import OptionParser

try:
    import matplotlib.pylab as plt
except ImportError:
    plot_flag = False
else:
    try:
        import seaborn
    except ImportError:
        pass
    plot_flag = True

import chimera

def plotter(phot_data, nframes, exptime, outfile):
    """
    Plot light curve.

    Parameters
    ----------
    phot_data : numpy array
        Photometry array

    nframes : int
        Number of image cube frames

    exptime : float
        Kinetic or accumulation time

    outfile : string
        Name of the out png image

    Returns
    -------
    None
    """
    params = {'backend': 'ps',
	      'font.size': 10,
              'axes.labelweight': 'medium',
	      'figure.dpi' : 300,
              'savefig.dpi': 300,
              'savefig.jpeg_quality': 100
              }
    plt.rcParams.update(params)

    ts = np.linspace(0, nframes*exptime, nframes)
    plt.figure(figsize=(6,4))
    plt.title("Normalized Light Curve : %s" %phot_data[0]['DATETIME'].split('T')[0])
    plt.xlabel("Time (secs)")
    plt.ylabel("Normalized Flux")
    plt.plot(ts, phot_data['FLUX_ADU']/np.mean(phot_data['FLUX_ADU']), "r-")
    plt.savefig(outfile, dpi = 300, bbox_inches = "tight")

    return

def process(infile, coords, method, inner_radius, outer_radius, cen_method, window_size, output, zmag):
    """
    Entry point function to process science image.

    Parameters
    ----------
    infile : string
        Science image or list of science images

    coords : string
        Input text file with coordinates of stars

    method : string
        FWHM of the stelar psf in pixels

    inner_radius : float
        Sky background sigma

    outer_radius : int
        Inner sky annulus radius in pixels

    cen_method : string
        Centroid method

    window_size : int
        Centroid finding window size in pixels

    output : string
        Output file name

    zmag : float
        Photometric zero point

    Returns
    -------
    None
    """
    print "FASTPHOT: CHIMERA Fast Aperture Photometry Routine"

    inner_radius = float(inner_radius)
    outer_radius = float(outer_radius)

    # Check if input is a string of FITS images or a text file with file names
    if infile[0] == "@":
        infile = infile[1:]

        if not os.path.exists(infile):
            print "REGISTER: Not able to locate file %s" %infile

        image_cubes = []
        with open(infile, "r") as fd:
            for line in fd.readlines():
                if len(line) > 1:
                    image_cubes.append(line.replace("\n", ""))
    else:
        image_cubes = infile.split(",")

    # Number of images
    ncubes = len(image_cubes)
    pos = np.loadtxt(coords, ndmin = 2)
    nstars = len(pos)

    total_phot_data = []

    for i in range(ncubes):
        sci_file = image_cubes[i]
        print "  Processing science image %s" %sci_file

        # Read FITS image and star coordinate
        image = chimera.fitsread(sci_file)

        # Instantiate an Aperphot object
        ap = chimera.Aperphot(sci_file, coords)

        # Set fwhmpsf, sigma, annulus, dannulus and zmag
        ap.method = method
        ap.inner_radius = inner_radius
        ap.outer_radius = outer_radius

        if zmag != "":
            ap.zmag = float(zmag)

        # Determine nominal aperture radius for photometry
        if i == 0:
            nom_aper = ap.cog(window_size, cen_method)

        print "  Nominal aperture radius : %4.1f pixels" %nom_aper

        # Perform aperture photometry on all the frames
        dtype = [("DATETIME", "S25"),("XCEN", "f4"),("YCEN", "f4"),("MSKY", "f8"),("NSKY", "f8"),("AREA", "f8"),("FLUX_ADU", "f8"),("FLUX_ELEC", "f8"),("FERR", "f8"),("MAG", "f8")]
        phot_data = np.zeros([nstars, ap.nframes], dtype = dtype)
        for j in range(ap.nframes):
            print "    Processing frame number : %d" %(j+1)

            objpos = chimera.recenter(image[j,:,:], pos, window_size, cen_method)
            aperphot_data = ap.phot(image[j,:,:], objpos, nom_aper)
            pos = np.copy(objpos)

            phot_data[:,j]['DATETIME'] = ap.addtime(j * ap.kintime).isoformat()
            phot_data[:,j]['XCEN'] = aperphot_data["xcenter_raw"]
            phot_data[:,j]['YCEN'] = aperphot_data["ycenter_raw"]
            phot_data[:,j]['MSKY'] = aperphot_data["msky"]
            phot_data[:,j]['NSKY'] = aperphot_data["nsky"]
            phot_data[:,j]['AREA'] = aperphot_data["area"]
            phot_data[:,j]['FLUX_ADU'] = aperphot_data["flux"]
            phot_data[:,j]['FLUX_ELEC'] = phot_data[:,j]['FLUX_ADU'] * ap.epadu
            phot_data[:,j]['MAG'] = ap.zmag - 2.5 * np.log10(phot_data[:,j]['FLUX_ELEC']/ap.exptime)

            # Calculate error in flux - using the formula
            # err = sqrt(flux * gain + npix * (1 + (npix/nsky)) * (flux_sky * gain + R**2))
            phot_data[:,j]['FERR'] = np.sqrt(phot_data[:,j]['FLUX_ELEC'] + phot_data[:,j]['AREA'] * (1 + phot_data[j]['AREA']/phot_data[j]['NSKY']) * (phot_data[j]['MSKY'] * ap.epadu + ap.readnoise**2))

        total_phot_data.append(phot_data)

        # Save photometry data in numpy binary format
        print "  Saving photometry data as numpy binary"
        if output != "":
            npy_outfile = output + ".npy"
        else:
            npy_outfile = sci_file.replace(".fits", ".phot.npy")
        if os.path.exists(npy_outfile):
            os.remove(npy_outfile)

        #np.save(npy_outfile, phot_data)

        # Plot first pass light curve
        if plot_flag:
            print "  Plotting normalized light curve"
            if output != "":
                plt_outfile = output + ".png"
            else:
                plt_outfile = sci_file.replace(".fits", ".lc.png")
            plotter(phot_data, ap.nframes, ap.kintime, plt_outfile)

    # Convert the total_phot_data to array and reshape it
    print '  Saving consolidated photometry data...'
    total_phot_data_arr = np.concatenate(total_phot_data, axis=1)

    # Save the array as npy file
    if output != "":
        np.save(output+"phot_total.npy", total_phot_data_arr)
    else: np.save("phot_total.npy", total_phot_data_arr)

    return

if __name__ == "__main__":
    usage = "Usage: python %prog [options] sci_image coords"
    description = "Description. Utility to perform fast aperture photometry in CHIMERA science images."
    parser = OptionParser(usage = usage, version = "%prog 0.2", description = description)
    parser.add_option("-v", "--verbose",
                      action="store_true", dest="verbose", default = False,
                      help = "print result messages to stdout"
                      )
    parser.add_option("-q", "--quiet",
                    action="store_false", dest="verbose", default = True,
                    help = "don't print result messages to stdout"
                    )
    parser.add_option("-m", "--method", dest = "method",
                    action="store", metavar="METHOD", help = "Method to use for determining overlap between aperture and pixels (default is exact)",
                    default = "exact"
                    )
    parser.add_option("-i", "--inner_radius", dest = "inner_radius",
                    action="store", metavar="INNER_RADIUS", help = "Inner radius of sky annlus in pixels (default is 14)",
                    default = 14
                    )
    parser.add_option("-d", "--outer_radius", dest = "outer_radius",
                    action="store", metavar="OUTER_RADIUS", help = "Radius of sky annulus in pixels (default is 16)",
                    default = 16
                    )
    parser.add_option("-c", "--cen_method", dest = "cen_method",
                    action="store", metavar="CEN_METHOD", help = "Centroid method (default is 2dg)",
                    default = "2dg"
                    )
    parser.add_option("-w", "--window_size", dest = "window_size",
                    action="store", metavar="WINDOW_SIZE", help = "Window size for centroid (default is 35)",
                    default = 35
                    )
    parser.add_option("-o", "--output", dest = "output",
                    action="store", metavar="OUTPUT", help = "Output file name",
                    default = ""
                    )
    parser.add_option("-z", "--zmag", dest = "zmag",
                    action="store", metavar="ZMAG", help = "Photometric zeroo point",
                    default = ""
                    )

    (options, args) = parser.parse_args()
    if len(args) != 2:
        parser.error("FASTPHOT: Incorrect number of arguments")

    # Check verbosity
    if not options.verbose:
        output = StringIO()
        old_stdout = sys.stdout
        sys.stdout = output

    # Switch off warnings
    warnings.filterwarnings('ignore')

    process(args[0], args[1], options.method, options.inner_radius, options.outer_radius, options.cen_method, options.window_size, options.output, options.zmag)

    # Reset verbosity
    if not options.verbose:
        sys.stdout = old_stdout


"""
=========================================================
Comparing different clustering algorithms on toy datasets
=========================================================

This example shows characteristics of different
clustering algorithms on datasets that are "interesting"
but still in 2D. With the exception of the last dataset,
the parameters of each of these dataset-algorithm pairs
has been tuned to produce good clustering results. Some
algorithms are more sensitive to parameter values than
others.

The last dataset is an example of a 'null' situation for
clustering: the data is homogeneous, and there is no good
clustering. For this example, the null dataset uses the
same parameters as the dataset in the row above it, which
represents a mismatch in the parameter values and the
data structure.

While these examples give some intuition about the
algorithms, this intuition might not apply to very high
dimensional data.
"""
print(__doc__)

import time
import warnings

import numpy as np
import matplotlib.pyplot as plt

from sklearn import cluster, datasets, mixture
from sklearn.neighbors import kneighbors_graph
from sklearn.preprocessing import StandardScaler
from itertools import cycle, islice

np.random.seed(0)

# ============
# Generate datasets. We choose the size big enough to see the scalability
# of the algorithms, but not too big to avoid too long running times
# ============
n_samples = 1500
noisy_circles = datasets.make_circles(n_samples=n_samples, factor=.5,
                                      noise=.05)
noisy_moons = datasets.make_moons(n_samples=n_samples, noise=.05)
blobs = datasets.make_blobs(n_samples=n_samples, random_state=8)
no_structure = np.random.rand(n_samples, 2), None

# Anisotropicly distributed data
random_state = 170
X, y = datasets.make_blobs(n_samples=n_samples, random_state=random_state)
transformation = [[0.6, -0.6], [-0.4, 0.8]]
X_aniso = np.dot(X, transformation)
aniso = (X_aniso, y)

# blobs with varied variances
varied = datasets.make_blobs(n_samples=n_samples,
                             cluster_std=[1.0, 2.5, 0.5],
                             random_state=random_state)

# ============
# Set up cluster parameters
# ============
plt.figure(figsize=(9 * 2 + 3, 12.5))
plt.subplots_adjust(left=.02, right=.98, bottom=.001, top=.96, wspace=.05,
                    hspace=.01)

plot_num = 1

default_base = {'quantile': .3,
                'eps': .3,
                'damping': .9,
                'preference': -200,
                'n_neighbors': 10,
                'n_clusters': 3}

datasets = [
    (noisy_circles, {'damping': .77, 'preference': -240,
                     'quantile': .2, 'n_clusters': 2}),
    (noisy_moons, {'damping': .75, 'preference': -220, 'n_clusters': 2}),
    (varied, {'eps': .18, 'n_neighbors': 2}),
    (aniso, {'eps': .15, 'n_neighbors': 2}),
    (blobs, {}),
    (no_structure, {})]

for i_dataset, (dataset, algo_params) in enumerate(datasets):
    # update parameters with dataset-specific values
    params = default_base.copy()
    params.update(algo_params)

    X, y = dataset

    # normalize dataset for easier parameter selection
    X = StandardScaler().fit_transform(X)

    # estimate bandwidth for mean shift
    bandwidth = cluster.estimate_bandwidth(X, quantile=params['quantile'])

    # connectivity matrix for structured Ward
    connectivity = kneighbors_graph(
        X, n_neighbors=params['n_neighbors'], include_self=False)
    # make connectivity symmetric
    connectivity = 0.5 * (connectivity + connectivity.T)

    # ============
    # Create cluster objects
    # ============
    ms = cluster.MeanShift(bandwidth=bandwidth, bin_seeding=True)
    two_means = cluster.MiniBatchKMeans(n_clusters=params['n_clusters'])
    ward = cluster.AgglomerativeClustering(
        n_clusters=params['n_clusters'], linkage='ward',
        connectivity=connectivity)
    spectral = cluster.SpectralClustering(
        n_clusters=params['n_clusters'], eigen_solver='arpack',
        affinity="nearest_neighbors")
    dbscan = cluster.DBSCAN(eps=params['eps'])
    affinity_propagation = cluster.AffinityPropagation(
        damping=params['damping'], preference=params['preference'])
    average_linkage = cluster.AgglomerativeClustering(
        linkage="average", affinity="cityblock",
        n_clusters=params['n_clusters'], connectivity=connectivity)
    birch = cluster.Birch(n_clusters=params['n_clusters'])
    gmm = mixture.GaussianMixture(
        n_components=params['n_clusters'], covariance_type='full')

    clustering_algorithms = (
        ('MiniBatchKMeans', two_means),
        ('AffinityPropagation', affinity_propagation),
        ('MeanShift', ms),
        ('SpectralClustering', spectral),
        ('Ward', ward),
        ('AgglomerativeClustering', average_linkage),
        ('DBSCAN', dbscan),
        ('Birch', birch),
        ('GaussianMixture', gmm)
    )

    for name, algorithm in clustering_algorithms:
        t0 = time.time()

        # catch warnings related to kneighbors_graph
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message="the number of connected components of the " +
                "connectivity matrix is [0-9]{1,2}" +
                " > 1. Completing it to avoid stopping the tree early.",
                category=UserWarning)
            warnings.filterwarnings(
                "ignore",
                message="Graph is not fully connected, spectral embedding" +
                " may not work as expected.",
                category=UserWarning)
            algorithm.fit(X)

        t1 = time.time()
        if hasattr(algorithm, 'labels_'):
            y_pred = algorithm.labels_.astype(np.int)
        else:
            y_pred = algorithm.predict(X)

        plt.subplot(len(datasets), len(clustering_algorithms), plot_num)
        if i_dataset == 0:
            plt.title(name, size=18)

        colors = np.array(list(islice(cycle(['#377eb8', '#ff7f00', '#4daf4a',
                                             '#f781bf', '#a65628', '#984ea3',
                                             '#999999', '#e41a1c', '#dede00']),
                                      int(max(y_pred) + 1))))
        plt.scatter(X[:, 0], X[:, 1], s=10, color=colors[y_pred])

        plt.xlim(-2.5, 2.5)
        plt.ylim(-2.5, 2.5)
        plt.xticks(())
        plt.yticks(())
        plt.text(.99, .01, ('%.2fs' % (t1 - t0)).lstrip('0'),
                 transform=plt.gca().transAxes, size=15,
                 horizontalalignment='right')
        plot_num += 1

plt.show()


import numpy as np
from scipy.optimize import differential_evolution as DE
import time as t
from .abcpmc import sampler, threshold
from ..synth_clust import synth_cluster
from . import likelihood
from .emcee_algor import varPars, closeSol, discreteParams, convergenceVals

def main(
    lkl_method, e_max, err_lst, completeness, max_mag_syn,
        fundam_params, obs_clust, theor_tracks, R_V, ext_coefs, st_dist_mass,
        N_fc, cmpl_rnd, err_rnd, nwalkers_abc, nsteps_abc, nburn_abc,
        priors_abc):

    varIdxs, ndim, ranges = varPars(fundam_params)

    def dist(synth_clust, obs_clust):
        lkl = np.inf
        if synth_clust:
            lkl = likelihood.main(lkl_method, synth_clust, obs_clust)
        return lkl

    def postfn(model):
        # Re-scale z and M
        model_scale = [
            model[0] / 100., model[1], model[2], model[3] * 10.,
            model[4] * 1000., model[5]]

        check_ranges = [
            r[0] <= p <= r[1] for p, r in zip(*[model_scale, ranges[varIdxs]])]

        synth_clust = []
        # If some parameter is outside of the given ranges, don't bother
        # obtaining the proper model.
        if all(check_ranges):
            model_proper = closeSol(fundam_params, varIdxs, model_scale)
            # Metallicity and age indexes to identify isochrone.
            m_i = fundam_params[0].index(model_proper[0])
            a_i = fundam_params[1].index(model_proper[1])
            isochrone = theor_tracks[m_i][a_i]

            # Generate synthetic cluster.
            synth_clust = synth_cluster.main(
                e_max, err_lst, completeness, max_mag_syn, st_dist_mass,
                isochrone, R_V, ext_coefs, N_fc, cmpl_rnd, err_rnd,
                model_proper)

        return synth_clust

    # TODO add these parameters to the input params file
    alpha, init_eps = 95, None
    N_conv, tol_conv = 50., 0.01
    max_secs = 22. * 60. * 60.
    # Break out when AF is low.
    # af_low, af_min_steps = 0.001, .1
    max_t_walker = 30.
    # eps_stuck_perc, N_eps_stuck_max = .005, 100

    # Start timing.
    elapsed = 0.
    available_secs = max(30, max_secs)
    start_t = t.time()

    abcsampler = sampler.Sampler(
        N=nwalkers_abc, Y=obs_clust, postfn=postfn, dist=dist)
    # Set proposal
    # sampler.particle_proposal_cls = sampler.OLCMParticleProposal

    if init_eps is None:
        # Estimate initial threshold value using DE.
        def lnprob(model):
            synth_clust = postfn(model)
            return dist(synth_clust, obs_clust)
        # Scale parameters bounds.
        bounds = [
            ranges[0] * 100., ranges[1], ranges[2], ranges[3] / 10.,
            ranges[4] / 1000., ranges[5]]
        result = DE(lnprob, bounds, maxiter=20)
        init_eps = 4. * result.fun
    print(" Initial threshold value: {:.2f}".format(init_eps))

    # old_eps = init_eps
    # TODO pass type of threshold from params file
    # eps = threshold.LinearEps(T, 5000, init_eps)
    eps = threshold.ConstEps(nsteps_abc, init_eps)

    # Stddev values as full range.
    std = np.eye(ndim) * (ranges.max(axis=1) - ranges.min(axis=1))
    # Means as middle points in ranges.
    means = (ranges.max(axis=1) + ranges.min(axis=1)) / 2.
    # Scale values.
    std[0], means[0] = std[0] * 100, means[0] * 100
    std[3], means[3] = std[3] / 10, means[3] / 10
    std[4], means[4] = std[4] / 1000., means[4] / 1000.
    # Gaussian prior.
    print(means)
    print(std)
    prior = sampler.GaussianPrior(mu=means, sigma=std)

    # # We'll track how the average autocorrelation time estimate changes
    # tau_index, autocorr_vals = 0, np.empty(nsteps_abc)
    # # This will be useful to testing convergence
    # old_tau = np.inf

    # Check for convergence every 2% of steps or 100, whichever value
    # is lower.
    # N_steps_conv = min(int(nsteps_abc * 0.02), 100)

    map_sol_old, N_models, prob_mean = [[], np.inf], 0, []
    # N_eps_stuck = 0
    chains_nruns, maf_steps, map_lkl = [], [], []
    milestones = list(range(5, 101, 5))
    for pool in abcsampler.sample(prior, eps):

        print(
            pool.t, pool.eps, pool.ratio, np.min(pool.dists),
            np.mean(pool.dists))

        chains_nruns.append(pool.thetas)
        maf = pool.ratio
        maf_steps.append([pool.t, maf])
        N_models += nwalkers_abc / maf

        # reduce eps value
        # old_eps = eps.eps
        eps.eps = np.percentile(pool.dists, alpha)

        # # Check if threshold is stuck.
        # if abs(eps.eps - old_eps) < eps_stuck_perc * eps.eps:
        #     N_eps_stuck += 1
        # else:
        #     N_eps_stuck = 0
        # if N_eps_stuck > N_eps_stuck_max:
        #     print("  Threshold is stuck (runs={}).".format(pool.t + 1))
        #     break

        # if maf < af_low and pool.t > int(af_min_steps * nsteps_abc):
        #     print("  AF<{} (runs={})".format(af_low, pool.t + 1))
        #     break

        if t.time() - start_t > (max_t_walker * nwalkers_abc):
            print("  Sampler is stuck (runs={})".format(pool.t + 1))
            break

        elapsed += t.time() - start_t
        if elapsed >= available_secs:
            print("  Time consumed (runs={})".format(pool.t + 1))
            break
        start_t = t.time()

        # # Only check convergence every 'N_steps_conv' steps
        # if (pool.t + 1) % N_steps_conv:
        #     continue

        # # Compute the autocorrelation time so far. Using tol=0 means that
        # # we'll always get an estimate even if it isn't trustworthy.
        # try:
        #     tau = autocorr.integrated_time(np.array(chains_nruns), tol=0)
        #     autocorr_vals[tau_index] = np.nanmean(tau)
        #     tau_index += 1

        #     # Check convergence
        #     converged = np.all(tau * N_conv < (pool.t + 1))
        #     converged &= np.all(np.abs(old_tau - tau) / tau < tol_conv)
        #     if converged:
        #         print("  Convergence achieved (runs={}).".format(pool.t + 1))
        #         break
        #     old_tau = tau
        # except FloatingPointError:
        #     pass

        # Store MAP solution in this iteration.
        prob_mean.append([pool.t, np.mean(pool.dists)])
        idx_best = np.argmin(pool.dists)
        # Update if a new optimal solution was found.
        if pool.dists[idx_best] < map_sol_old[1]:
            pars = pool.thetas[idx_best]
            # pars = scaleParams(model)
            pars = [pars[0] / 100., pars[1], pars[2], pars[3] * 10.,
                    pars[4] * 1000., pars[5]]
            map_sol_old = [
                closeSol(fundam_params, varIdxs, pars),
                pool.dists[idx_best]]
        map_lkl.append([pool.t, map_sol_old[1]])

        # Print progress.
        percentage_complete = (100. * (pool.t + 1) / nsteps_abc)
        if len(milestones) > 0 and percentage_complete >= milestones[0]:
            map_sol, logprob = map_sol_old
            print("{:>3}% ({:.3f}) LP={:.1f} ({:g}, {:g}, {:.3f}, {:.2f}"
                  ", {:g}, {:.2f})".format(
                      milestones[0], maf, logprob, *map_sol) +
                  " [{:.0f} m/s]".format(N_models / elapsed))
            milestones = milestones[1:]

    runs = pool.t + 1

    # Evolution of the mean autocorrelation time.
    tau_autocorr = np.array([np.nan] * 10)  # autocorr_vals[:tau_index]
    tau_index = np.nan
    N_steps_conv = runs

    # Final MAP fit.
    idx_best = np.argmin(pool.dists)
    pars = pool.thetas[idx_best]
    # pars = scaleParams(model)
    pars = [
        pars[0] / 100., pars[1], pars[2], pars[3] * 10., pars[4] * 1000.,
        pars[5]]
    map_sol = closeSol(fundam_params, varIdxs, pars)
    map_lkl_final = pool.dists[idx_best]

    abcsampler.close()

    # Shape: (runs, nwalkers, ndim)
    chains_nruns = np.array(chains_nruns)
    # De-scale parameters.
    chains_nruns[:, :, 0] = chains_nruns[:, :, 0] / 100.
    chains_nruns[:, :, 3] = chains_nruns[:, :, 3] * 10.
    chains_nruns[:, :, 4] = chains_nruns[:, :, 4] * 1000.

    # Burn-in range.
    Nb = int(runs * nburn_abc)

    # Burn-in. Shape: (ndim, nwalkers, runs)
    pars_chains_bi = discreteParams(
        fundam_params, varIdxs, chains_nruns[:Nb, :, :]).T

    # Change values for the discrete parameters with the closest valid values.
    chains_nruns = discreteParams(
        fundam_params, varIdxs, chains_nruns[Nb:, :, :])

    mcmc_trace = chains_nruns.reshape(-1, ndim).T

    # import matplotlib.pyplot as plt
    # import corner
    # corner.corner(
    #     mcmc_trace.T, quantiles=[0.16, 0.5, 0.84], show_titles=True)
    #     # levels=(1 - np.exp(-0.5),))
    # plt.savefig("corner.png", dpi=300)

    # Convergence parameters.
    acorr_t, max_at_c, min_at_c, geweke_z, emcee_acorf, mcmc_ess, minESS,\
        mESS, mESS_epsilon = convergenceVals(
            'abc', ndim, varIdxs, N_conv, chains_nruns, mcmc_trace)

    # Store mean solution.
    mean_sol = closeSol(fundam_params, varIdxs, np.mean(mcmc_trace, axis=1))

    isoch_fit_params = {
        'varIdxs': varIdxs, 'nsteps_abc': runs, 'mean_sol': mean_sol,
        'nburn_abc': Nb, 'map_sol': map_sol, 'map_lkl': map_lkl,
        'map_lkl_final': map_lkl_final, 'prob_mean': prob_mean,
        'mcmc_elapsed': elapsed, 'mcmc_trace': mcmc_trace,
        'pars_chains_bi': pars_chains_bi, 'pars_chains': chains_nruns.T,
        'maf_steps': maf_steps, 'autocorr_time': acorr_t,
        'max_at_c': max_at_c, 'min_at_c': min_at_c,
        'minESS': minESS, 'mESS': mESS, 'mESS_epsilon': mESS_epsilon,
        'emcee_acorf': emcee_acorf, 'geweke_z': geweke_z,
        'mcmc_ess': mcmc_ess,
        'N_steps_conv': N_steps_conv, 'N_conv': N_conv, 'tol_conv': tol_conv,
        'tau_index': tau_index, 'tau_autocorr': tau_autocorr
    }

    return isoch_fit_params


#K-Means++ Clustering with Gap Statistic to determine the optimal number of clusters
import sys
import numpy as np
import scipy.io as sio
#import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.svm import SVC

filename = sys.argv[1]
datafile = sio.loadmat(filename)
data = datafile['bow']
sizedata=[len(data), len(data[0])]
disp = []
optimal_ks = []
#Determining the optimal number of k with gap statistic method
def gap_statistic(data):
	sizedata = [len(data),len(data[0])]
	SD = []
	gap = []
	for knum in xrange(1,20):
		#I assumed that the number of clusters in my data won't be more than 20, this can be changed accordingly
		print knum
		#Clustering original Data
		kmeanspp = KMeans(n_clusters=knum,init = 'k-means++',max_iter = 100,n_jobs = 1)
		kmeanspp.fit(data)
		dispersion = kmeanspp.inertia_
		#Clustering Reference Data
		nrefs = 10
		refDisp = np.zeros(nrefs)
		for nref in xrange(nrefs):
			refdata = np.random.random_sample(tuple(sizedata))
			refkmeans = KMeans(n_clusters=knum,init='k-means++',max_iter=100,n_jobs=1)
			refkmeans.fit(refdata)
			refdisp = refkmeans.inertia_
			refDisp[nref]=np.log(refdisp)
		mean_log_refdisp = np.mean(refDisp)
		gap.append(mean_log_refdisp-np.log(dispersion))
		sd = (sum([(r-m)**2 for r,m in zip(refDisp,[mean_log_refdisp]*nrefs)])/nrefs)**0.5
		SD.append(sd)
	SD = [sd*((1+(1/nrefs))**0.5) for sd in SD]
	opt_k = None
	diff = []
	for i in xrange(len(gap)-1):
		diff = (SD[i+1]-(gap[i+1]-gap[i]))
		if diff>0:
			opt_k = i+10
			break
	if opt_k < 20:
		#print opt_k
		return opt_k
	else:
		return 20
		#Returning 20 if opt_k is more than 20 in my case, as I wanted not to search more than 20.
		# Not required if range is larger.

ntrials = 50
for ntrial in xrange(ntrials):
	print 'ntrial: ',ntrial
	optimal_ks.append(gap_statistic(data))
#For plotting the gap statistic measure
#plt.plot(np.linspace(10,19,10,True),gap)
#plt.show()
unique_opt_k = list(set(optimal_ks))
k_count = {}
count_opt_k = 0
second_opt_k = 0
opt_k = 0
for u_o_k in unique_opt_k:
	count = optimal_ks.count(u_o_k)
	k_count[u_o_k]=count
	if count>count_opt_k:
		count_opt_k = count
		opt_k = u_o_k
	elif count==count_opt_k:
		second_opt_k = u_o_k
print opt_k
print k_count
#Clusterin with optimal number of k
kmeanspp = KMeans(n_clusters = opt_k,init='k-means++',max_iter=100,n_jobs=1)
kmeanspp.fit(data)
centers = kmeanspp.cluster_centers_
clusterlabels = kmeanspp.labels_
print clusterlabels
mdict = {}
mdict['clusterlabels'] = clusterlabels

sio.savemat('clusterlabels.mat',mdict,format = '4',oned_as = 'column')
print 'dan dana dan done...'


"""
Tests for DatetimeIndex timezone-related methods
"""
from datetime import date, datetime, time, timedelta, tzinfo

import dateutil
from dateutil.tz import gettz, tzlocal
import numpy as np
import pytest
import pytz

from pandas._libs.tslibs import conversion, timezones
import pandas.util._test_decorators as td

import pandas as pd
from pandas import (
    DatetimeIndex,
    Index,
    Timestamp,
    bdate_range,
    date_range,
    isna,
    to_datetime,
)
import pandas.util.testing as tm

class FixedOffset(tzinfo):
    """Fixed offset in minutes east from UTC."""

    def __init__(self, offset, name):
        self.__offset = timedelta(minutes=offset)
        self.__name = name

    def utcoffset(self, dt):
        return self.__offset

    def tzname(self, dt):
        return self.__name

    def dst(self, dt):
        return timedelta(0)

fixed_off = FixedOffset(-420, "-07:00")
fixed_off_no_name = FixedOffset(-330, None)

class TestDatetimeIndexTimezones:
    # -------------------------------------------------------------
    # DatetimeIndex.tz_convert
    def test_tz_convert_nat(self):
        # GH#5546
        dates = [pd.NaT]
        idx = DatetimeIndex(dates)
        idx = idx.tz_localize("US/Pacific")
        tm.assert_index_equal(idx, DatetimeIndex(dates, tz="US/Pacific"))
        idx = idx.tz_convert("US/Eastern")
        tm.assert_index_equal(idx, DatetimeIndex(dates, tz="US/Eastern"))
        idx = idx.tz_convert("UTC")
        tm.assert_index_equal(idx, DatetimeIndex(dates, tz="UTC"))

        dates = ["2010-12-01 00:00", "2010-12-02 00:00", pd.NaT]
        idx = DatetimeIndex(dates)
        idx = idx.tz_localize("US/Pacific")
        tm.assert_index_equal(idx, DatetimeIndex(dates, tz="US/Pacific"))
        idx = idx.tz_convert("US/Eastern")
        expected = ["2010-12-01 03:00", "2010-12-02 03:00", pd.NaT]
        tm.assert_index_equal(idx, DatetimeIndex(expected, tz="US/Eastern"))

        idx = idx + pd.offsets.Hour(5)
        expected = ["2010-12-01 08:00", "2010-12-02 08:00", pd.NaT]
        tm.assert_index_equal(idx, DatetimeIndex(expected, tz="US/Eastern"))
        idx = idx.tz_convert("US/Pacific")
        expected = ["2010-12-01 05:00", "2010-12-02 05:00", pd.NaT]
        tm.assert_index_equal(idx, DatetimeIndex(expected, tz="US/Pacific"))

        idx = idx + np.timedelta64(3, "h")
        expected = ["2010-12-01 08:00", "2010-12-02 08:00", pd.NaT]
        tm.assert_index_equal(idx, DatetimeIndex(expected, tz="US/Pacific"))

        idx = idx.tz_convert("US/Eastern")
        expected = ["2010-12-01 11:00", "2010-12-02 11:00", pd.NaT]
        tm.assert_index_equal(idx, DatetimeIndex(expected, tz="US/Eastern"))

    @pytest.mark.parametrize("prefix", ["", "dateutil/"])
    def test_dti_tz_convert_compat_timestamp(self, prefix):
        strdates = ["1/1/2012", "3/1/2012", "4/1/2012"]
        idx = DatetimeIndex(strdates, tz=prefix + "US/Eastern")

        conv = idx[0].tz_convert(prefix + "US/Pacific")
        expected = idx.tz_convert(prefix + "US/Pacific")[0]

        assert conv == expected

    def test_dti_tz_convert_hour_overflow_dst(self):
        # Regression test for:
        # https://github.com/pandas-dev/pandas/issues/13306

        # sorted case US/Eastern -> UTC
        ts = ["2008-05-12 09:50:00", "2008-12-12 09:50:35", "2009-05-12 09:50:32"]
        tt = DatetimeIndex(ts).tz_localize("US/Eastern")
        ut = tt.tz_convert("UTC")
        expected = Index([13, 14, 13])
        tm.assert_index_equal(ut.hour, expected)

        # sorted case UTC -> US/Eastern
        ts = ["2008-05-12 13:50:00", "2008-12-12 14:50:35", "2009-05-12 13:50:32"]
        tt = DatetimeIndex(ts).tz_localize("UTC")
        ut = tt.tz_convert("US/Eastern")
        expected = Index([9, 9, 9])
        tm.assert_index_equal(ut.hour, expected)

        # unsorted case US/Eastern -> UTC
        ts = ["2008-05-12 09:50:00", "2008-12-12 09:50:35", "2008-05-12 09:50:32"]
        tt = DatetimeIndex(ts).tz_localize("US/Eastern")
        ut = tt.tz_convert("UTC")
        expected = Index([13, 14, 13])
        tm.assert_index_equal(ut.hour, expected)

        # unsorted case UTC -> US/Eastern
        ts = ["2008-05-12 13:50:00", "2008-12-12 14:50:35", "2008-05-12 13:50:32"]
        tt = DatetimeIndex(ts).tz_localize("UTC")
        ut = tt.tz_convert("US/Eastern")
        expected = Index([9, 9, 9])
        tm.assert_index_equal(ut.hour, expected)

    @pytest.mark.parametrize("tz", ["US/Eastern", "dateutil/US/Eastern"])
    def test_dti_tz_convert_hour_overflow_dst_timestamps(self, tz):
        # Regression test for GH#13306

        # sorted case US/Eastern -> UTC
        ts = [
            Timestamp("2008-05-12 09:50:00", tz=tz),
            Timestamp("2008-12-12 09:50:35", tz=tz),
            Timestamp("2009-05-12 09:50:32", tz=tz),
        ]
        tt = DatetimeIndex(ts)
        ut = tt.tz_convert("UTC")
        expected = Index([13, 14, 13])
        tm.assert_index_equal(ut.hour, expected)

        # sorted case UTC -> US/Eastern
        ts = [
            Timestamp("2008-05-12 13:50:00", tz="UTC"),
            Timestamp("2008-12-12 14:50:35", tz="UTC"),
            Timestamp("2009-05-12 13:50:32", tz="UTC"),
        ]
        tt = DatetimeIndex(ts)
        ut = tt.tz_convert("US/Eastern")
        expected = Index([9, 9, 9])
        tm.assert_index_equal(ut.hour, expected)

        # unsorted case US/Eastern -> UTC
        ts = [
            Timestamp("2008-05-12 09:50:00", tz=tz),
            Timestamp("2008-12-12 09:50:35", tz=tz),
            Timestamp("2008-05-12 09:50:32", tz=tz),
        ]
        tt = DatetimeIndex(ts)
        ut = tt.tz_convert("UTC")
        expected = Index([13, 14, 13])
        tm.assert_index_equal(ut.hour, expected)

        # unsorted case UTC -> US/Eastern
        ts = [
            Timestamp("2008-05-12 13:50:00", tz="UTC"),
            Timestamp("2008-12-12 14:50:35", tz="UTC"),
            Timestamp("2008-05-12 13:50:32", tz="UTC"),
        ]
        tt = DatetimeIndex(ts)
        ut = tt.tz_convert("US/Eastern")
        expected = Index([9, 9, 9])
        tm.assert_index_equal(ut.hour, expected)

    @pytest.mark.parametrize("freq, n", [("H", 1), ("T", 60), ("S", 3600)])
    def test_dti_tz_convert_trans_pos_plus_1__bug(self, freq, n):
        # Regression test for tslib.tz_convert(vals, tz1, tz2).
        # See https://github.com/pandas-dev/pandas/issues/4496 for details.
        idx = date_range(datetime(2011, 3, 26, 23), datetime(2011, 3, 27, 1), freq=freq)
        idx = idx.tz_localize("UTC")
        idx = idx.tz_convert("Europe/Moscow")

        expected = np.repeat(np.array([3, 4, 5]), np.array([n, n, 1]))
        tm.assert_index_equal(idx.hour, Index(expected))

    def test_dti_tz_convert_dst(self):
        for freq, n in [("H", 1), ("T", 60), ("S", 3600)]:
            # Start DST
            idx = date_range(
                "2014-03-08 23:00", "2014-03-09 09:00", freq=freq, tz="UTC"
            )
            idx = idx.tz_convert("US/Eastern")
            expected = np.repeat(
                np.array([18, 19, 20, 21, 22, 23, 0, 1, 3, 4, 5]),
                np.array([n, n, n, n, n, n, n, n, n, n, 1]),
            )
            tm.assert_index_equal(idx.hour, Index(expected))

            idx = date_range(
                "2014-03-08 18:00", "2014-03-09 05:00", freq=freq, tz="US/Eastern"
            )
            idx = idx.tz_convert("UTC")
            expected = np.repeat(
                np.array([23, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
                np.array([n, n, n, n, n, n, n, n, n, n, 1]),
            )
            tm.assert_index_equal(idx.hour, Index(expected))

            # End DST
            idx = date_range(
                "2014-11-01 23:00", "2014-11-02 09:00", freq=freq, tz="UTC"
            )
            idx = idx.tz_convert("US/Eastern")
            expected = np.repeat(
                np.array([19, 20, 21, 22, 23, 0, 1, 1, 2, 3, 4]),
                np.array([n, n, n, n, n, n, n, n, n, n, 1]),
            )
            tm.assert_index_equal(idx.hour, Index(expected))

            idx = date_range(
                "2014-11-01 18:00", "2014-11-02 05:00", freq=freq, tz="US/Eastern"
            )
            idx = idx.tz_convert("UTC")
            expected = np.repeat(
                np.array([22, 23, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
                np.array([n, n, n, n, n, n, n, n, n, n, n, n, 1]),
            )
            tm.assert_index_equal(idx.hour, Index(expected))

        # daily
        # Start DST
        idx = date_range("2014-03-08 00:00", "2014-03-09 00:00", freq="D", tz="UTC")
        idx = idx.tz_convert("US/Eastern")
        tm.assert_index_equal(idx.hour, Index([19, 19]))

        idx = date_range(
            "2014-03-08 00:00", "2014-03-09 00:00", freq="D", tz="US/Eastern"
        )
        idx = idx.tz_convert("UTC")
        tm.assert_index_equal(idx.hour, Index([5, 5]))

        # End DST
        idx = date_range("2014-11-01 00:00", "2014-11-02 00:00", freq="D", tz="UTC")
        idx = idx.tz_convert("US/Eastern")
        tm.assert_index_equal(idx.hour, Index([20, 20]))

        idx = date_range(
            "2014-11-01 00:00", "2014-11-02 000:00", freq="D", tz="US/Eastern"
        )
        idx = idx.tz_convert("UTC")
        tm.assert_index_equal(idx.hour, Index([4, 4]))

    def test_tz_convert_roundtrip(self, tz_aware_fixture):
        tz = tz_aware_fixture
        idx1 = date_range(start="2014-01-01", end="2014-12-31", freq="M", tz="UTC")
        exp1 = date_range(start="2014-01-01", end="2014-12-31", freq="M")

        idx2 = date_range(start="2014-01-01", end="2014-12-31", freq="D", tz="UTC")
        exp2 = date_range(start="2014-01-01", end="2014-12-31", freq="D")

        idx3 = date_range(start="2014-01-01", end="2014-03-01", freq="H", tz="UTC")
        exp3 = date_range(start="2014-01-01", end="2014-03-01", freq="H")

        idx4 = date_range(start="2014-08-01", end="2014-10-31", freq="T", tz="UTC")
        exp4 = date_range(start="2014-08-01", end="2014-10-31", freq="T")

        for idx, expected in [(idx1, exp1), (idx2, exp2), (idx3, exp3), (idx4, exp4)]:
            converted = idx.tz_convert(tz)
            reset = converted.tz_convert(None)
            tm.assert_index_equal(reset, expected)
            assert reset.tzinfo is None
            expected = converted.tz_convert("UTC").tz_localize(None)
            tm.assert_index_equal(reset, expected)

    def test_dti_tz_convert_tzlocal(self):
        # GH#13583
        # tz_convert doesn't affect to internal
        dti = date_range(start="2001-01-01", end="2001-03-01", tz="UTC")
        dti2 = dti.tz_convert(dateutil.tz.tzlocal())
        tm.assert_numpy_array_equal(dti2.asi8, dti.asi8)

        dti = date_range(start="2001-01-01", end="2001-03-01", tz=dateutil.tz.tzlocal())
        dti2 = dti.tz_convert(None)
        tm.assert_numpy_array_equal(dti2.asi8, dti.asi8)

    @pytest.mark.parametrize(
        "tz",
        [
            "US/Eastern",
            "dateutil/US/Eastern",
            pytz.timezone("US/Eastern"),
            gettz("US/Eastern"),
        ],
    )
    def test_dti_tz_convert_utc_to_local_no_modify(self, tz):
        rng = date_range("3/11/2012", "3/12/2012", freq="H", tz="utc")
        rng_eastern = rng.tz_convert(tz)

        # Values are unmodified
        tm.assert_numpy_array_equal(rng.asi8, rng_eastern.asi8)

        assert timezones.tz_compare(rng_eastern.tz, timezones.maybe_get_tz(tz))

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_tz_convert_unsorted(self, tzstr):
        dr = date_range("2012-03-09", freq="H", periods=100, tz="utc")
        dr = dr.tz_convert(tzstr)

        result = dr[::-1].hour
        exp = dr.hour[::-1]
        tm.assert_almost_equal(result, exp)

    # -------------------------------------------------------------
    # DatetimeIndex.tz_localize

    def test_dti_tz_localize_nonexistent_raise_coerce(self):
        # GH#13057
        times = ["2015-03-08 01:00", "2015-03-08 02:00", "2015-03-08 03:00"]
        index = DatetimeIndex(times)
        tz = "US/Eastern"
        with pytest.raises(pytz.NonExistentTimeError):
            index.tz_localize(tz=tz)

        with pytest.raises(pytz.NonExistentTimeError):
            with tm.assert_produces_warning(FutureWarning):
                index.tz_localize(tz=tz, errors="raise")

        with tm.assert_produces_warning(
            FutureWarning, clear=FutureWarning, check_stacklevel=False
        ):
            result = index.tz_localize(tz=tz, errors="coerce")
        test_times = ["2015-03-08 01:00-05:00", "NaT", "2015-03-08 03:00-04:00"]
        dti = to_datetime(test_times, utc=True)
        expected = dti.tz_convert("US/Eastern")
        tm.assert_index_equal(result, expected)

    @pytest.mark.parametrize("tz", [pytz.timezone("US/Eastern"), gettz("US/Eastern")])
    def test_dti_tz_localize_ambiguous_infer(self, tz):
        # November 6, 2011, fall back, repeat 2 AM hour
        # With no repeated hours, we cannot infer the transition
        dr = date_range(datetime(2011, 11, 6, 0), periods=5, freq=pd.offsets.Hour())
        with pytest.raises(pytz.AmbiguousTimeError):
            dr.tz_localize(tz)

        # With repeated hours, we can infer the transition
        dr = date_range(
            datetime(2011, 11, 6, 0), periods=5, freq=pd.offsets.Hour(), tz=tz
        )
        times = [
            "11/06/2011 00:00",
            "11/06/2011 01:00",
            "11/06/2011 01:00",
            "11/06/2011 02:00",
            "11/06/2011 03:00",
        ]
        di = DatetimeIndex(times)
        localized = di.tz_localize(tz, ambiguous="infer")
        tm.assert_index_equal(dr, localized)
        tm.assert_index_equal(dr, DatetimeIndex(times, tz=tz, ambiguous="infer"))

        # When there is no dst transition, nothing special happens
        dr = date_range(datetime(2011, 6, 1, 0), periods=10, freq=pd.offsets.Hour())
        localized = dr.tz_localize(tz)
        localized_infer = dr.tz_localize(tz, ambiguous="infer")
        tm.assert_index_equal(localized, localized_infer)

    @pytest.mark.parametrize("tz", [pytz.timezone("US/Eastern"), gettz("US/Eastern")])
    def test_dti_tz_localize_ambiguous_times(self, tz):
        # March 13, 2011, spring forward, skip from 2 AM to 3 AM
        dr = date_range(datetime(2011, 3, 13, 1, 30), periods=3, freq=pd.offsets.Hour())
        with pytest.raises(pytz.NonExistentTimeError):
            dr.tz_localize(tz)

        # after dst transition, it works
        dr = date_range(
            datetime(2011, 3, 13, 3, 30), periods=3, freq=pd.offsets.Hour(), tz=tz
        )

        # November 6, 2011, fall back, repeat 2 AM hour
        dr = date_range(datetime(2011, 11, 6, 1, 30), periods=3, freq=pd.offsets.Hour())
        with pytest.raises(pytz.AmbiguousTimeError):
            dr.tz_localize(tz)

        # UTC is OK
        dr = date_range(
            datetime(2011, 3, 13), periods=48, freq=pd.offsets.Minute(30), tz=pytz.utc
        )

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_dti_tz_localize_pass_dates_to_utc(self, tzstr):
        strdates = ["1/1/2012", "3/1/2012", "4/1/2012"]

        idx = DatetimeIndex(strdates)
        conv = idx.tz_localize(tzstr)

        fromdates = DatetimeIndex(strdates, tz=tzstr)

        assert conv.tz == fromdates.tz
        tm.assert_numpy_array_equal(conv.values, fromdates.values)

    @pytest.mark.parametrize("prefix", ["", "dateutil/"])
    def test_dti_tz_localize(self, prefix):
        tzstr = prefix + "US/Eastern"
        dti = pd.date_range(start="1/1/2005", end="1/1/2005 0:00:30.256", freq="L")
        dti2 = dti.tz_localize(tzstr)

        dti_utc = pd.date_range(
            start="1/1/2005 05:00", end="1/1/2005 5:00:30.256", freq="L", tz="utc"
        )

        tm.assert_numpy_array_equal(dti2.values, dti_utc.values)

        dti3 = dti2.tz_convert(prefix + "US/Pacific")
        tm.assert_numpy_array_equal(dti3.values, dti_utc.values)

        dti = pd.date_range(start="11/6/2011 1:59", end="11/6/2011 2:00", freq="L")
        with pytest.raises(pytz.AmbiguousTimeError):
            dti.tz_localize(tzstr)

        dti = pd.date_range(start="3/13/2011 1:59", end="3/13/2011 2:00", freq="L")
        with pytest.raises(pytz.NonExistentTimeError):
            dti.tz_localize(tzstr)

    @pytest.mark.parametrize(
        "tz",
        [
            "US/Eastern",
            "dateutil/US/Eastern",
            pytz.timezone("US/Eastern"),
            gettz("US/Eastern"),
        ],
    )
    def test_dti_tz_localize_utc_conversion(self, tz):
        # Localizing to time zone should:
        #  1) check for DST ambiguities
        #  2) convert to UTC

        rng = date_range("3/10/2012", "3/11/2012", freq="30T")

        converted = rng.tz_localize(tz)
        expected_naive = rng + pd.offsets.Hour(5)
        tm.assert_numpy_array_equal(converted.asi8, expected_naive.asi8)

        # DST ambiguity, this should fail
        rng = date_range("3/11/2012", "3/12/2012", freq="30T")
        # Is this really how it should fail??
        with pytest.raises(pytz.NonExistentTimeError):
            rng.tz_localize(tz)

    def test_dti_tz_localize_roundtrip(self, tz_aware_fixture):
        # note: this tz tests that a tz-naive index can be localized
        # and de-localized successfully, when there are no DST transitions
        # in the range.
        idx = date_range(start="2014-06-01", end="2014-08-30", freq="15T")
        tz = tz_aware_fixture
        localized = idx.tz_localize(tz)
        # cant localize a tz-aware object
        with pytest.raises(TypeError):
            localized.tz_localize(tz)
        reset = localized.tz_localize(None)
        assert reset.tzinfo is None
        tm.assert_index_equal(reset, idx)

    def test_dti_tz_localize_naive(self):
        rng = date_range("1/1/2011", periods=100, freq="H")

        conv = rng.tz_localize("US/Pacific")
        exp = date_range("1/1/2011", periods=100, freq="H", tz="US/Pacific")

        tm.assert_index_equal(conv, exp)

    def test_dti_tz_localize_tzlocal(self):
        # GH#13583
        offset = dateutil.tz.tzlocal().utcoffset(datetime(2011, 1, 1))
        offset = int(offset.total_seconds() * 1000000000)

        dti = date_range(start="2001-01-01", end="2001-03-01")
        dti2 = dti.tz_localize(dateutil.tz.tzlocal())
        tm.assert_numpy_array_equal(dti2.asi8 + offset, dti.asi8)

        dti = date_range(start="2001-01-01", end="2001-03-01", tz=dateutil.tz.tzlocal())
        dti2 = dti.tz_localize(None)
        tm.assert_numpy_array_equal(dti2.asi8 - offset, dti.asi8)

    @pytest.mark.parametrize("tz", [pytz.timezone("US/Eastern"), gettz("US/Eastern")])
    def test_dti_tz_localize_ambiguous_nat(self, tz):
        times = [
            "11/06/2011 00:00",
            "11/06/2011 01:00",
            "11/06/2011 01:00",
            "11/06/2011 02:00",
            "11/06/2011 03:00",
        ]
        di = DatetimeIndex(times)
        localized = di.tz_localize(tz, ambiguous="NaT")

        times = [
            "11/06/2011 00:00",
            np.NaN,
            np.NaN,
            "11/06/2011 02:00",
            "11/06/2011 03:00",
        ]
        di_test = DatetimeIndex(times, tz="US/Eastern")

        # left dtype is datetime64[ns, US/Eastern]
        # right is datetime64[ns, tzfile('/usr/share/zoneinfo/US/Eastern')]
        tm.assert_numpy_array_equal(di_test.values, localized.values)

    @pytest.mark.parametrize("tz", [pytz.timezone("US/Eastern"), gettz("US/Eastern")])
    def test_dti_tz_localize_ambiguous_flags(self, tz):
        # November 6, 2011, fall back, repeat 2 AM hour

        # Pass in flags to determine right dst transition
        dr = date_range(
            datetime(2011, 11, 6, 0), periods=5, freq=pd.offsets.Hour(), tz=tz
        )
        times = [
            "11/06/2011 00:00",
            "11/06/2011 01:00",
            "11/06/2011 01:00",
            "11/06/2011 02:00",
            "11/06/2011 03:00",
        ]

        # Test tz_localize
        di = DatetimeIndex(times)
        is_dst = [1, 1, 0, 0, 0]
        localized = di.tz_localize(tz, ambiguous=is_dst)
        tm.assert_index_equal(dr, localized)
        tm.assert_index_equal(dr, DatetimeIndex(times, tz=tz, ambiguous=is_dst))

        localized = di.tz_localize(tz, ambiguous=np.array(is_dst))
        tm.assert_index_equal(dr, localized)

        localized = di.tz_localize(tz, ambiguous=np.array(is_dst).astype("bool"))
        tm.assert_index_equal(dr, localized)

        # Test constructor
        localized = DatetimeIndex(times, tz=tz, ambiguous=is_dst)
        tm.assert_index_equal(dr, localized)

        # Test duplicate times where inferring the dst fails
        times += times
        di = DatetimeIndex(times)

        # When the sizes are incompatible, make sure error is raised
        with pytest.raises(Exception):
            di.tz_localize(tz, ambiguous=is_dst)

        # When sizes are compatible and there are repeats ('infer' won't work)
        is_dst = np.hstack((is_dst, is_dst))
        localized = di.tz_localize(tz, ambiguous=is_dst)
        dr = dr.append(dr)
        tm.assert_index_equal(dr, localized)

        # When there is no dst transition, nothing special happens
        dr = date_range(datetime(2011, 6, 1, 0), periods=10, freq=pd.offsets.Hour())
        is_dst = np.array([1] * 10)
        localized = dr.tz_localize(tz)
        localized_is_dst = dr.tz_localize(tz, ambiguous=is_dst)
        tm.assert_index_equal(localized, localized_is_dst)

    # TODO: belongs outside tz_localize tests?
    @pytest.mark.parametrize("tz", ["Europe/London", "dateutil/Europe/London"])
    def test_dti_construction_ambiguous_endpoint(self, tz):
        # construction with an ambiguous end-point
        # GH#11626

        with pytest.raises(pytz.AmbiguousTimeError):
            date_range(
                "2013-10-26 23:00", "2013-10-27 01:00", tz="Europe/London", freq="H"
            )

        times = date_range(
            "2013-10-26 23:00", "2013-10-27 01:00", freq="H", tz=tz, ambiguous="infer"
        )
        assert times[0] == Timestamp("2013-10-26 23:00", tz=tz, freq="H")

        if str(tz).startswith("dateutil"):
            # fixed ambiguous behavior
            # see GH#14621
            assert times[-1] == Timestamp("2013-10-27 01:00:00+0100", tz=tz, freq="H")
        else:
            assert times[-1] == Timestamp("2013-10-27 01:00:00+0000", tz=tz, freq="H")

    @pytest.mark.parametrize(
        "tz, option, expected",
        [
            ["US/Pacific", "shift_forward", "2019-03-10 03:00"],
            ["dateutil/US/Pacific", "shift_forward", "2019-03-10 03:00"],
            ["US/Pacific", "shift_backward", "2019-03-10 01:00"],
            pytest.param(
                "dateutil/US/Pacific",
                "shift_backward",
                "2019-03-10 01:00",
                marks=pytest.mark.xfail(reason="GH 24329"),
            ),
            ["US/Pacific", timedelta(hours=1), "2019-03-10 03:00"],
        ],
    )
    def test_dti_construction_nonexistent_endpoint(self, tz, option, expected):
        # construction with an nonexistent end-point

        with pytest.raises(pytz.NonExistentTimeError):
            date_range(
                "2019-03-10 00:00", "2019-03-10 02:00", tz="US/Pacific", freq="H"
            )

        times = date_range(
            "2019-03-10 00:00", "2019-03-10 02:00", freq="H", tz=tz, nonexistent=option
        )
        assert times[-1] == Timestamp(expected, tz=tz, freq="H")

    def test_dti_tz_localize_bdate_range(self):
        dr = pd.bdate_range("1/1/2009", "1/1/2010")
        dr_utc = pd.bdate_range("1/1/2009", "1/1/2010", tz=pytz.utc)
        localized = dr.tz_localize(pytz.utc)
        tm.assert_index_equal(dr_utc, localized)

    @pytest.mark.parametrize("tz", ["Europe/Warsaw", "dateutil/Europe/Warsaw"])
    @pytest.mark.parametrize(
        "method, exp", [["NaT", pd.NaT], ["raise", None], ["foo", "invalid"]]
    )
    def test_dti_tz_localize_nonexistent(self, tz, method, exp):
        # GH 8917
        n = 60
        dti = date_range(start="2015-03-29 02:00:00", periods=n, freq="min")
        if method == "raise":
            with pytest.raises(pytz.NonExistentTimeError):
                dti.tz_localize(tz, nonexistent=method)
        elif exp == "invalid":
            with pytest.raises(ValueError):
                dti.tz_localize(tz, nonexistent=method)
        else:
            result = dti.tz_localize(tz, nonexistent=method)
            expected = DatetimeIndex([exp] * n, tz=tz)
            tm.assert_index_equal(result, expected)

    @pytest.mark.parametrize(
        "start_ts, tz, end_ts, shift",
        [
            ["2015-03-29 02:20:00", "Europe/Warsaw", "2015-03-29 03:00:00", "forward"],
            [
                "2015-03-29 02:20:00",
                "Europe/Warsaw",
                "2015-03-29 01:59:59.999999999",
                "backward",
            ],
            [
                "2015-03-29 02:20:00",
                "Europe/Warsaw",
                "2015-03-29 03:20:00",
                timedelta(hours=1),
            ],
            [
                "2015-03-29 02:20:00",
                "Europe/Warsaw",
                "2015-03-29 01:20:00",
                timedelta(hours=-1),
            ],
            ["2018-03-11 02:33:00", "US/Pacific", "2018-03-11 03:00:00", "forward"],
            [
                "2018-03-11 02:33:00",
                "US/Pacific",
                "2018-03-11 01:59:59.999999999",
                "backward",
            ],
            [
                "2018-03-11 02:33:00",
                "US/Pacific",
                "2018-03-11 03:33:00",
                timedelta(hours=1),
            ],
            [
                "2018-03-11 02:33:00",
                "US/Pacific",
                "2018-03-11 01:33:00",
                timedelta(hours=-1),
            ],
        ],
    )
    @pytest.mark.parametrize("tz_type", ["", "dateutil/"])
    def test_dti_tz_localize_nonexistent_shift(
        self, start_ts, tz, end_ts, shift, tz_type
    ):
        # GH 8917
        tz = tz_type + tz
        if isinstance(shift, str):
            shift = "shift_" + shift
        dti = DatetimeIndex([Timestamp(start_ts)])
        result = dti.tz_localize(tz, nonexistent=shift)
        expected = DatetimeIndex([Timestamp(end_ts)]).tz_localize(tz)
        tm.assert_index_equal(result, expected)

    @pytest.mark.parametrize("offset", [-1, 1])
    @pytest.mark.parametrize("tz_type", ["", "dateutil/"])
    def test_dti_tz_localize_nonexistent_shift_invalid(self, offset, tz_type):
        # GH 8917
        tz = tz_type + "Europe/Warsaw"
        dti = DatetimeIndex([Timestamp("2015-03-29 02:20:00")])
        msg = "The provided timedelta will relocalize on a nonexistent time"
        with pytest.raises(ValueError, match=msg):
            dti.tz_localize(tz, nonexistent=timedelta(seconds=offset))

    @pytest.mark.filterwarnings("ignore::FutureWarning")
    def test_dti_tz_localize_errors_deprecation(self):
        # GH 22644
        tz = "Europe/Warsaw"
        n = 60
        dti = date_range(start="2015-03-29 02:00:00", periods=n, freq="min")
        with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
            with pytest.raises(ValueError):
                dti.tz_localize(tz, errors="foo")
            # make sure errors='coerce' gets mapped correctly to nonexistent
            result = dti.tz_localize(tz, errors="coerce")
            expected = dti.tz_localize(tz, nonexistent="NaT")
            tm.assert_index_equal(result, expected)

    # -------------------------------------------------------------
    # DatetimeIndex.normalize

    def test_normalize_tz(self):
        rng = date_range("1/1/2000 9:30", periods=10, freq="D", tz="US/Eastern")

        result = rng.normalize()
        expected = date_range("1/1/2000", periods=10, freq="D", tz="US/Eastern")
        tm.assert_index_equal(result, expected)

        assert result.is_normalized
        assert not rng.is_normalized

        rng = date_range("1/1/2000 9:30", periods=10, freq="D", tz="UTC")

        result = rng.normalize()
        expected = date_range("1/1/2000", periods=10, freq="D", tz="UTC")
        tm.assert_index_equal(result, expected)

        assert result.is_normalized
        assert not rng.is_normalized

        rng = date_range("1/1/2000 9:30", periods=10, freq="D", tz=tzlocal())
        result = rng.normalize()
        expected = date_range("1/1/2000", periods=10, freq="D", tz=tzlocal())
        tm.assert_index_equal(result, expected)

        assert result.is_normalized
        assert not rng.is_normalized

    @td.skip_if_windows
    @pytest.mark.parametrize(
        "timezone",
        [
            "US/Pacific",
            "US/Eastern",
            "UTC",
            "Asia/Kolkata",
            "Asia/Shanghai",
            "Australia/Canberra",
        ],
    )
    def test_normalize_tz_local(self, timezone):
        # GH#13459
        with tm.set_timezone(timezone):
            rng = date_range("1/1/2000 9:30", periods=10, freq="D", tz=tzlocal())

            result = rng.normalize()
            expected = date_range("1/1/2000", periods=10, freq="D", tz=tzlocal())
            tm.assert_index_equal(result, expected)

            assert result.is_normalized
            assert not rng.is_normalized

    # ------------------------------------------------------------
    # DatetimeIndex.__new__

    @pytest.mark.parametrize("prefix", ["", "dateutil/"])
    def test_dti_constructor_static_tzinfo(self, prefix):
        # it works!
        index = DatetimeIndex([datetime(2012, 1, 1)], tz=prefix + "EST")
        index.hour
        index[0]

    def test_dti_constructor_with_fixed_tz(self):
        off = FixedOffset(420, "+07:00")
        start = datetime(2012, 3, 11, 5, 0, 0, tzinfo=off)
        end = datetime(2012, 6, 11, 5, 0, 0, tzinfo=off)
        rng = date_range(start=start, end=end)
        assert off == rng.tz

        rng2 = date_range(start, periods=len(rng), tz=off)
        tm.assert_index_equal(rng, rng2)

        rng3 = date_range("3/11/2012 05:00:00+07:00", "6/11/2012 05:00:00+07:00")
        assert (rng.values == rng3.values).all()

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_dti_convert_datetime_list(self, tzstr):
        dr = date_range("2012-06-02", periods=10, tz=tzstr, name="foo")
        dr2 = DatetimeIndex(list(dr), name="foo")
        tm.assert_index_equal(dr, dr2)
        assert dr.tz == dr2.tz
        assert dr2.name == "foo"

    def test_dti_construction_univalent(self):
        rng = date_range("03/12/2012 00:00", periods=10, freq="W-FRI", tz="US/Eastern")
        rng2 = DatetimeIndex(data=rng, tz="US/Eastern")
        tm.assert_index_equal(rng, rng2)

    @pytest.mark.parametrize("tz", [pytz.timezone("US/Eastern"), gettz("US/Eastern")])
    def test_dti_from_tzaware_datetime(self, tz):
        d = [datetime(2012, 8, 19, tzinfo=tz)]

        index = DatetimeIndex(d)
        assert timezones.tz_compare(index.tz, tz)

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_dti_tz_constructors(self, tzstr):
        """ Test different DatetimeIndex constructions with timezone
        Follow-up of GH#4229
        """

        arr = ["11/10/2005 08:00:00", "11/10/2005 09:00:00"]

        idx1 = to_datetime(arr).tz_localize(tzstr)
        idx2 = pd.date_range(start="2005-11-10 08:00:00", freq="H", periods=2, tz=tzstr)
        idx3 = DatetimeIndex(arr, tz=tzstr)
        idx4 = DatetimeIndex(np.array(arr), tz=tzstr)

        for other in [idx2, idx3, idx4]:
            tm.assert_index_equal(idx1, other)

    # -------------------------------------------------------------
    # Unsorted

    def test_join_utc_convert(self, join_type):
        rng = date_range("1/1/2011", periods=100, freq="H", tz="utc")

        left = rng.tz_convert("US/Eastern")
        right = rng.tz_convert("Europe/Berlin")

        result = left.join(left[:-5], how=join_type)
        assert isinstance(result, DatetimeIndex)
        assert result.tz == left.tz

        result = left.join(right[:-5], how=join_type)
        assert isinstance(result, DatetimeIndex)
        assert result.tz.zone == "UTC"

    @pytest.mark.parametrize(
        "dtype",
        [None, "datetime64[ns, CET]", "datetime64[ns, EST]", "datetime64[ns, UTC]"],
    )
    def test_date_accessor(self, dtype):
        # Regression test for GH#21230
        expected = np.array([date(2018, 6, 4), pd.NaT])

        index = DatetimeIndex(["2018-06-04 10:00:00", pd.NaT], dtype=dtype)
        result = index.date

        tm.assert_numpy_array_equal(result, expected)

    @pytest.mark.parametrize(
        "dtype",
        [None, "datetime64[ns, CET]", "datetime64[ns, EST]", "datetime64[ns, UTC]"],
    )
    def test_time_accessor(self, dtype):
        # Regression test for GH#21267
        expected = np.array([time(10, 20, 30), pd.NaT])

        index = DatetimeIndex(["2018-06-04 10:20:30", pd.NaT], dtype=dtype)
        result = index.time

        tm.assert_numpy_array_equal(result, expected)

    def test_timetz_accessor(self, tz_naive_fixture):
        # GH21358
        tz = timezones.maybe_get_tz(tz_naive_fixture)

        expected = np.array([time(10, 20, 30, tzinfo=tz), pd.NaT])

        index = DatetimeIndex(["2018-06-04 10:20:30", pd.NaT], tz=tz)
        result = index.timetz

        tm.assert_numpy_array_equal(result, expected)

    def test_dti_drop_dont_lose_tz(self):
        # GH#2621
        ind = date_range("2012-12-01", periods=10, tz="utc")
        ind = ind.drop(ind[-1])

        assert ind.tz is not None

    def test_dti_tz_conversion_freq(self, tz_naive_fixture):
        # GH25241
        t3 = DatetimeIndex(["2019-01-01 10:00"], freq="H")
        assert t3.tz_localize(tz=tz_naive_fixture).freq == t3.freq
        t4 = DatetimeIndex(["2019-01-02 12:00"], tz="UTC", freq="T")
        assert t4.tz_convert(tz="UTC").freq == t4.freq

    def test_drop_dst_boundary(self):
        # see gh-18031
        tz = "Europe/Brussels"
        freq = "15min"

        start = pd.Timestamp("201710290100", tz=tz)
        end = pd.Timestamp("201710290300", tz=tz)
        index = pd.date_range(start=start, end=end, freq=freq)

        expected = DatetimeIndex(
            [
                "201710290115",
                "201710290130",
                "201710290145",
                "201710290200",
                "201710290215",
                "201710290230",
                "201710290245",
                "201710290200",
                "201710290215",
                "201710290230",
                "201710290245",
                "201710290300",
            ],
            tz=tz,
            freq=freq,
            ambiguous=[
                True,
                True,
                True,
                True,
                True,
                True,
                True,
                False,
                False,
                False,
                False,
                False,
            ],
        )
        result = index.drop(index[0])
        tm.assert_index_equal(result, expected)

    def test_date_range_localize(self):
        rng = date_range("3/11/2012 03:00", periods=15, freq="H", tz="US/Eastern")
        rng2 = DatetimeIndex(["3/11/2012 03:00", "3/11/2012 04:00"], tz="US/Eastern")
        rng3 = date_range("3/11/2012 03:00", periods=15, freq="H")
        rng3 = rng3.tz_localize("US/Eastern")

        tm.assert_index_equal(rng, rng3)

        # DST transition time
        val = rng[0]
        exp = Timestamp("3/11/2012 03:00", tz="US/Eastern")

        assert val.hour == 3
        assert exp.hour == 3
        assert val == exp  # same UTC value
        tm.assert_index_equal(rng[:2], rng2)

        # Right before the DST transition
        rng = date_range("3/11/2012 00:00", periods=2, freq="H", tz="US/Eastern")
        rng2 = DatetimeIndex(["3/11/2012 00:00", "3/11/2012 01:00"], tz="US/Eastern")
        tm.assert_index_equal(rng, rng2)
        exp = Timestamp("3/11/2012 00:00", tz="US/Eastern")
        assert exp.hour == 0
        assert rng[0] == exp
        exp = Timestamp("3/11/2012 01:00", tz="US/Eastern")
        assert exp.hour == 1
        assert rng[1] == exp

        rng = date_range("3/11/2012 00:00", periods=10, freq="H", tz="US/Eastern")
        assert rng[2].hour == 3

    def test_timestamp_equality_different_timezones(self):
        utc_range = date_range("1/1/2000", periods=20, tz="UTC")
        eastern_range = utc_range.tz_convert("US/Eastern")
        berlin_range = utc_range.tz_convert("Europe/Berlin")

        for a, b, c in zip(utc_range, eastern_range, berlin_range):
            assert a == b
            assert b == c
            assert a == c

        assert (utc_range == eastern_range).all()
        assert (utc_range == berlin_range).all()
        assert (berlin_range == eastern_range).all()

    def test_dti_intersection(self):
        rng = date_range("1/1/2011", periods=100, freq="H", tz="utc")

        left = rng[10:90][::-1]
        right = rng[20:80][::-1]

        assert left.tz == rng.tz
        result = left.intersection(right)
        assert result.tz == left.tz

    def test_dti_equals_with_tz(self):
        left = date_range("1/1/2011", periods=100, freq="H", tz="utc")
        right = date_range("1/1/2011", periods=100, freq="H", tz="US/Eastern")

        assert not left.equals(right)

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_dti_tz_nat(self, tzstr):
        idx = DatetimeIndex([Timestamp("2013-1-1", tz=tzstr), pd.NaT])

        assert isna(idx[1])
        assert idx[0].tzinfo is not None

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_dti_astype_asobject_tzinfos(self, tzstr):
        # GH#1345

        # dates around a dst transition
        rng = date_range("2/13/2010", "5/6/2010", tz=tzstr)

        objs = rng.astype(object)
        for i, x in enumerate(objs):
            exval = rng[i]
            assert x == exval
            assert x.tzinfo == exval.tzinfo

        objs = rng.astype(object)
        for i, x in enumerate(objs):
            exval = rng[i]
            assert x == exval
            assert x.tzinfo == exval.tzinfo

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_dti_with_timezone_repr(self, tzstr):
        rng = date_range("4/13/2010", "5/6/2010")

        rng_eastern = rng.tz_localize(tzstr)

        rng_repr = repr(rng_eastern)
        assert "2010-04-13 00:00:00" in rng_repr

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_dti_take_dont_lose_meta(self, tzstr):
        rng = date_range("1/1/2000", periods=20, tz=tzstr)

        result = rng.take(range(5))
        assert result.tz == rng.tz
        assert result.freq == rng.freq

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_utc_box_timestamp_and_localize(self, tzstr):
        tz = timezones.maybe_get_tz(tzstr)

        rng = date_range("3/11/2012", "3/12/2012", freq="H", tz="utc")
        rng_eastern = rng.tz_convert(tzstr)

        expected = rng[-1].astimezone(tz)

        stamp = rng_eastern[-1]
        assert stamp == expected
        assert stamp.tzinfo == expected.tzinfo

        # right tzinfo
        rng = date_range("3/13/2012", "3/14/2012", freq="H", tz="utc")
        rng_eastern = rng.tz_convert(tzstr)
        # test not valid for dateutil timezones.
        # assert 'EDT' in repr(rng_eastern[0].tzinfo)
        assert "EDT" in repr(rng_eastern[0].tzinfo) or "tzfile" in repr(
            rng_eastern[0].tzinfo
        )

    def test_dti_to_pydatetime(self):
        dt = dateutil.parser.parse("2012-06-13T01:39:00Z")
        dt = dt.replace(tzinfo=tzlocal())

        arr = np.array([dt], dtype=object)

        result = to_datetime(arr, utc=True)
        assert result.tz is pytz.utc

        rng = date_range("2012-11-03 03:00", "2012-11-05 03:00", tz=tzlocal())
        arr = rng.to_pydatetime()
        result = to_datetime(arr, utc=True)
        assert result.tz is pytz.utc

    def test_dti_to_pydatetime_fizedtz(self):
        dates = np.array(
            [
                datetime(2000, 1, 1, tzinfo=fixed_off),
                datetime(2000, 1, 2, tzinfo=fixed_off),
                datetime(2000, 1, 3, tzinfo=fixed_off),
            ]
        )
        dti = DatetimeIndex(dates)

        result = dti.to_pydatetime()
        tm.assert_numpy_array_equal(dates, result)

        result = dti._mpl_repr()
        tm.assert_numpy_array_equal(dates, result)

    @pytest.mark.parametrize("tz", [pytz.timezone("US/Central"), gettz("US/Central")])
    def test_with_tz(self, tz):
        # just want it to work
        start = datetime(2011, 3, 12, tzinfo=pytz.utc)
        dr = bdate_range(start, periods=50, freq=pd.offsets.Hour())
        assert dr.tz is pytz.utc

        # DateRange with naive datetimes
        dr = bdate_range("1/1/2005", "1/1/2009", tz=pytz.utc)
        dr = bdate_range("1/1/2005", "1/1/2009", tz=tz)

        # normalized
        central = dr.tz_convert(tz)
        assert central.tz is tz
        naive = central[0].to_pydatetime().replace(tzinfo=None)
        comp = conversion.localize_pydatetime(naive, tz).tzinfo
        assert central[0].tz is comp

        # compare vs a localized tz
        naive = dr[0].to_pydatetime().replace(tzinfo=None)
        comp = conversion.localize_pydatetime(naive, tz).tzinfo
        assert central[0].tz is comp

        # datetimes with tzinfo set
        dr = bdate_range(
            datetime(2005, 1, 1, tzinfo=pytz.utc), datetime(2009, 1, 1, tzinfo=pytz.utc)
        )
        with pytest.raises(Exception):
            bdate_range(datetime(2005, 1, 1, tzinfo=pytz.utc), "1/1/2009", tz=tz)

    @pytest.mark.parametrize("prefix", ["", "dateutil/"])
    def test_field_access_localize(self, prefix):
        strdates = ["1/1/2012", "3/1/2012", "4/1/2012"]
        rng = DatetimeIndex(strdates, tz=prefix + "US/Eastern")
        assert (rng.hour == 0).all()

        # a more unusual time zone, #1946
        dr = date_range(
            "2011-10-02 00:00", freq="h", periods=10, tz=prefix + "America/Atikokan"
        )

        expected = Index(np.arange(10, dtype=np.int64))
        tm.assert_index_equal(dr.hour, expected)

    @pytest.mark.parametrize("tz", [pytz.timezone("US/Eastern"), gettz("US/Eastern")])
    def test_dti_convert_tz_aware_datetime_datetime(self, tz):
        # GH#1581
        dates = [datetime(2000, 1, 1), datetime(2000, 1, 2), datetime(2000, 1, 3)]

        dates_aware = [conversion.localize_pydatetime(x, tz) for x in dates]
        result = DatetimeIndex(dates_aware)
        assert timezones.tz_compare(result.tz, tz)

        converted = to_datetime(dates_aware, utc=True)
        ex_vals = np.array([Timestamp(x).value for x in dates_aware])
        tm.assert_numpy_array_equal(converted.asi8, ex_vals)
        assert converted.tz is pytz.utc

    def test_dti_union_aware(self):
        # non-overlapping
        rng = date_range("2012-11-15 00:00:00", periods=6, freq="H", tz="US/Central")

        rng2 = date_range("2012-11-15 12:00:00", periods=6, freq="H", tz="US/Eastern")

        result = rng.union(rng2)
        expected = rng.astype("O").union(rng2.astype("O"))
        tm.assert_index_equal(result, expected)
        assert result[0].tz.zone == "US/Central"
        assert result[-1].tz.zone == "US/Eastern"

    def test_dti_union_mixed(self):
        # GH 21671
        rng = DatetimeIndex([pd.Timestamp("2011-01-01"), pd.NaT])
        rng2 = pd.DatetimeIndex(["2012-01-01", "2012-01-02"], tz="Asia/Tokyo")
        result = rng.union(rng2)
        expected = Index(
            [
                pd.Timestamp("2011-01-01"),
                pd.NaT,
                pd.Timestamp("2012-01-01", tz="Asia/Tokyo"),
                pd.Timestamp("2012-01-02", tz="Asia/Tokyo"),
            ],
            dtype=object,
        )
        tm.assert_index_equal(result, expected)

    @pytest.mark.parametrize(
        "tz", [None, "UTC", "US/Central", dateutil.tz.tzoffset(None, -28800)]
    )
    @pytest.mark.usefixtures("datetime_tz_utc")
    def test_iteration_preserves_nanoseconds(self, tz):
        # GH 19603
        index = DatetimeIndex(
            ["2018-02-08 15:00:00.168456358", "2018-02-08 15:00:00.168456359"], tz=tz
        )
        for i, ts in enumerate(index):
            assert ts == index[i]

class TestDateRange:
    """Tests for date_range with timezones"""

    def test_hongkong_tz_convert(self):
        # GH#1673 smoke test
        dr = date_range("2012-01-01", "2012-01-10", freq="D", tz="Hongkong")

        # it works!
        dr.hour

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_date_range_span_dst_transition(self, tzstr):
        # GH#1778

        # Standard -> Daylight Savings Time
        dr = date_range("03/06/2012 00:00", periods=200, freq="W-FRI", tz="US/Eastern")

        assert (dr.hour == 0).all()

        dr = date_range("2012-11-02", periods=10, tz=tzstr)
        result = dr.hour
        expected = Index([0] * 10)
        tm.assert_index_equal(result, expected)

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_date_range_timezone_str_argument(self, tzstr):
        tz = timezones.maybe_get_tz(tzstr)
        result = date_range("1/1/2000", periods=10, tz=tzstr)
        expected = date_range("1/1/2000", periods=10, tz=tz)

        tm.assert_index_equal(result, expected)

    def test_date_range_with_fixedoffset_noname(self):
        off = fixed_off_no_name
        start = datetime(2012, 3, 11, 5, 0, 0, tzinfo=off)
        end = datetime(2012, 6, 11, 5, 0, 0, tzinfo=off)
        rng = date_range(start=start, end=end)
        assert off == rng.tz

        idx = Index([start, end])
        assert off == idx.tz

    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
    def test_date_range_with_tz(self, tzstr):
        stamp = Timestamp("3/11/2012 05:00", tz=tzstr)
        assert stamp.hour == 5

        rng = date_range("3/11/2012 04:00", periods=10, freq="H", tz=tzstr)

        assert stamp == rng[1]

class TestToDatetime:
    """Tests for the to_datetime constructor with timezones"""

    def test_to_datetime_utc(self):
        arr = np.array([dateutil.parser.parse("2012-06-13T01:39:00Z")], dtype=object)

        result = to_datetime(arr, utc=True)
        assert result.tz is pytz.utc

    def test_to_datetime_fixed_offset(self):
        dates = [
            datetime(2000, 1, 1, tzinfo=fixed_off),
            datetime(2000, 1, 2, tzinfo=fixed_off),
            datetime(2000, 1, 3, tzinfo=fixed_off),
        ]
        result = to_datetime(dates)
        assert result.tz == fixed_off


import unittest, time, sys, copy
sys.path.extend(['.','..','../..','py'])
import h2o, h2o_cmd, h2o_glm, h2o_util, h2o_import as h2i
## Dataset created from this:
#
# from sklearn.datasets import make_hastie_10_2
# import numpy as np
# i = 1000000
# f = 10
# (X,y) = make_hastie_10_2(n_samples=i,random_state=None)
# y.shape = (i,1)
# Y = np.hstack((X,y))
# np.savetxt('./1mx' + str(f) + '_hastie_10_2.data', Y, delimiter=',', fmt='%.2f');

def glm_doit(self, csvFilename, bucket, csvPathname, timeoutSecs=30):
    print "\nStarting parse of", csvFilename
    parseResult = h2i.import_parse(bucket=bucket, path=csvPathname, schema='put', hex_key=csvFilename + ".hex", timeoutSecs=20)
    y = "10"
    # NOTE: hastie has two values, -1 and 1. To make H2O work if two valued and not 0,1 have
    kwargs = {
        'response':  y,
        'max_iter': 10,
        'n_folds': 2,
        'lambda': '1e-8,1e-4,1e-3',
        'alpha': '0,0.25,0.8',
        }

    start = time.time()
    print "\nStarting GLMGrid of", csvFilename
    glmGridResult = h2o_cmd.runGLM(parseResult=parseResult, timeoutSecs=timeoutSecs, **kwargs)
    print "GLMGrid in",  (time.time() - start), "secs (python)"

    # still get zero coeffs..best model is AUC = 0.5 with intercept only.
    h2o_glm.simpleCheckGLMGrid(self,glmGridResult, allowZeroCoeff=True,**kwargs)

class Basic(unittest.TestCase):
    def tearDown(self):
        h2o.check_sandbox_for_errors()

    @classmethod
    def setUpClass(cls):
        h2o.init(2,java_heap_GB=5)
        global SYNDATASETS_DIR
        SYNDATASETS_DIR = h2o.make_syn_dir()

    @classmethod
    def tearDownClass(cls):
        h2o.tear_down_cloud()

    def test_GLM2grid_hastie(self):
        # gunzip it and cat it to create 2x and 4x replications in SYNDATASETS_DIR
        # FIX! eventually we'll compare the 1x, 2x and 4x results like we do
        # in other tests. (catdata?)
        bucket = 'home-0xdiag-datasets'
        csvFilename = "1mx10_hastie_10_2.data.gz"
        csvPathname = 'standard' + '/' + csvFilename
        glm_doit(self, csvFilename, bucket, csvPathname, timeoutSecs=300)

        fullPathname = h2i.find_folder_and_filename('home-0xdiag-datasets', csvPathname, returnFullPath=True)
        filename1x = "hastie_1x.data"
        pathname1x = SYNDATASETS_DIR + '/' + filename1x
        h2o_util.file_gunzip(fullPathname, pathname1x)

        filename2x = "hastie_2x.data"
        pathname2x = SYNDATASETS_DIR + '/' + filename2x
        h2o_util.file_cat(pathname1x,pathname1x,pathname2x)
        glm_doit(self, filename2x, None, pathname2x, timeoutSecs=300)

if __name__ == '__main__':
    h2o.unit_main()