"""
===================
Label image regions
===================

This example shows how to segment an image with image labelling. The following
steps are applied:

1. Thresholding with automatic Otsu method
2. Close small holes with binary closing
3. Remove artifacts touching image border
4. Measure image regions to filter small objects

"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

from skimage import data
from skimage.filters import threshold_otsu
from skimage.segmentation import clear_border
from skimage.measure import label
from skimage.morphology import closing, square
from skimage.measure import regionprops
from skimage.color import label2rgb

image = data.coins()[50:-50, 50:-50]

# apply threshold
thresh = threshold_otsu(image)
bw = closing(image > thresh, square(3))

# remove artifacts connected to image border
cleared = bw.copy()
clear_border(cleared)

# label image regions
label_image = label(cleared)
borders = np.logical_xor(bw, cleared)
label_image[borders] = -1
image_label_overlay = label2rgb(label_image, image=image)

fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(6, 6))
ax.imshow(image_label_overlay)

for region in regionprops(label_image):

    # skip small images
    if region.area < 100:
        continue

    # draw rectangle around segmented coins
    minr, minc, maxr, maxc = region.bbox
    rect = mpatches.Rectangle((minc, minr), maxc - minc, maxr - minr,
                              fill=False, edgecolor='red', linewidth=2)
    ax.add_patch(rect)

plt.show()


""" Example inspired by an example from the scikit-learn project:
http://scikit-learn.org/stable/auto_examples/cluster/plot_cluster_comparison.html
"""

import numpy as np

try:
    from sklearn import cluster, datasets
    from sklearn.preprocessing import StandardScaler
except ImportError:
    raise ImportError('This example requires scikit-learn (conda install sklearn)')

from bokeh.layouts import row, column
from bokeh.plotting import figure, show, output_file

N = 50000
PLOT_SIZE = 400

# generate datasets.
np.random.seed(0)
noisy_circles = datasets.make_circles(n_samples=N, factor=.5, noise=.04)
noisy_moons = datasets.make_moons(n_samples=N, noise=.05)
centers = [(-2, 3), (2, 3), (-2, -3), (2, -3)]
blobs1 = datasets.make_blobs(centers=centers, n_samples=N, cluster_std=0.4, random_state=8)
blobs2 = datasets.make_blobs(centers=centers, n_samples=N, cluster_std=0.7, random_state=8)

colors = np.array([x for x in ('#00f', '#0f0', '#f00', '#0ff', '#f0f', '#ff0')])
colors = np.hstack([colors] * 20)

# create clustering algorithms
dbscan   = cluster.DBSCAN(eps=.2)
birch    = cluster.Birch(n_clusters=2)
means    = cluster.MiniBatchKMeans(n_clusters=2)
spectral = cluster.SpectralClustering(n_clusters=2, eigen_solver='arpack', affinity="nearest_neighbors")
affinity = cluster.AffinityPropagation(damping=.9, preference=-200)

# change here, to select clustering algorithm (note: spectral is slow)
algorithm = dbscan  # <- SELECT ALG

plots =[]
for dataset in (noisy_circles, noisy_moons, blobs1, blobs2):
    X, y = dataset
    X = StandardScaler().fit_transform(X)

    # predict cluster memberships
    algorithm.fit(X)
    if hasattr(algorithm, 'labels_'):
        y_pred = algorithm.labels_.astype(np.int)
    else:
        y_pred = algorithm.predict(X)

    p = figure(output_backend="webgl", title=algorithm.__class__.__name__,
               plot_width=PLOT_SIZE, plot_height=PLOT_SIZE)

    p.scatter(X[:, 0], X[:, 1], color=colors[y_pred].tolist(), alpha=0.1,)

    plots.append(p)

# generate layout for the plots
layout = column(row(plots[:2]), row(plots[2:]))

output_file("clustering.html", title="clustering with sklearn")

show(layout)


#!/usr/bin/env python3

"""
Example: inductive conformal classification using DecisionTreeClassifier
"""

import numpy as np

from sklearn.tree import DecisionTreeClassifier
import Orange

from nonconformist.acp import CrossConformalClassifier, AggregatedCp, CrossSampler
from nonconformist.icp import IcpClassifier
from nonconformist.nc import ProbEstClassifierNc, inverse_probability

# -----------------------------------------------------------------------------
# Setup training, calibration and test indices
# -----------------------------------------------------------------------------
data = Orange.data.Table('iris')
X, y = data.X, data.Y

idx = np.random.permutation(y.size)
train = idx[:idx.size // 3]
calibrate = idx[idx.size // 3:2 * idx.size // 3]
test = idx[2 * idx.size // 3:]

# -----------------------------------------------------------------------------
# Train and calibrate
# -----------------------------------------------------------------------------
icp = IcpClassifier(ProbEstClassifierNc(DecisionTreeClassifier(), inverse_probability))
icp.fit(X[train, :], y[train])
icp.calibrate(X[calibrate, :], y[calibrate])

ccp = CrossConformalClassifier(IcpClassifier(ProbEstClassifierNc(DecisionTreeClassifier(), inverse_probability)))
ccp.fit(X[train, :], y[train])

acp = AggregatedCp(IcpClassifier(ProbEstClassifierNc(DecisionTreeClassifier(), inverse_probability)), CrossSampler())
acp.fit(X[train, :], y[train])

# -----------------------------------------------------------------------------
# Predict
# -----------------------------------------------------------------------------
print('# Inductive')
prediction = icp.predict(X[test, :], significance=0.1)
for pred, actual in zip(prediction[:5], y[test]):
    print(pred, actual)

print('\n# Cross')
prediction = ccp.predict(X[test, :], significance=0.1)
for pred, actual in zip(prediction[:5], y[test]):
    print(pred, actual)

print('\n# Aggre')
prediction = acp.predict(X[test, :], significance=0.1)
for pred, actual in zip(prediction[:5], y[test]):
    print(pred, actual)


#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
=========================================================
Pipelining: chaining a PCA and a logistic regression
=========================================================

The PCA does an unsupervised dimensionality reduction, while the logistic
regression does the prediction.

We use a GridSearchCV to set the dimensionality of the PCA

"""
print(__doc__)

# Code source: Gaël Varoquaux
# Modified for documentation by Jaques Grobler
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt

from sklearn import linear_model, decomposition, datasets
from sklearn.pipeline import Pipeline
from sklearn.grid_search import GridSearchCV

logistic = linear_model.LogisticRegression()

pca = decomposition.PCA()
pipe = Pipeline(steps=[('pca', pca), ('logistic', logistic)])

digits = datasets.load_digits()
X_digits = digits.data
y_digits = digits.target

###############################################################################
# Plot the PCA spectrum
pca.fit(X_digits)

plt.figure(1, figsize=(4, 3))
plt.clf()
plt.axes([.2, .2, .7, .7])
plt.plot(pca.explained_variance_, linewidth=2)
plt.axis('tight')
plt.xlabel('n_components')
plt.ylabel('explained_variance_')

###############################################################################
# Prediction

n_components = [20, 40, 64]
Cs = np.logspace(-4, 4, 3)

#Parameters of pipelines can be set using ‘__’ separated parameter names:

estimator = GridSearchCV(pipe,
                         dict(pca__n_components=n_components,
                              logistic__C=Cs))
estimator.fit(X_digits, y_digits)

plt.axvline(estimator.best_estimator_.named_steps['pca'].n_components,
            linestyle=':', label='n_components chosen')
plt.legend(prop=dict(size=12))
plt.show()


from pylearn2.models.s3c import S3C
from pylearn2.models.s3c import E_Step_Scan
from pylearn2.models.s3c import Grad_M_Step
from pylearn2.models.s3c import E_Step
from theano import function
import numpy as np
import theano.tensor as T
from theano import config
#from pylearn2.utils import serial
import warnings

def broadcast(mat, shape_0):
    rval = mat
    if mat.shape[0] != shape_0:
        assert mat.shape[0] == 1

        rval = np.zeros((shape_0, mat.shape[1]),dtype=mat.dtype)

        for i in xrange(shape_0):
            rval[i,:] = mat[0,:]

    return rval

class Test_S3C_Inference:
    def setUp(self):
        # Temporarily change config.floatX to float64, as s3c inference
        # tests currently fail due to numerical issues for float32.
        self.prev_floatX = config.floatX
        config.floatX = 'float64'

    def tearDown(self):
        # Restore previous value of floatX
        config.floatX = self.prev_floatX

    def __init__(self):
        """ gets a small batch of data
            sets up an S3C model
        """
        # We also have to change the value of config.floatX in __init__.
        self.prev_floatX = config.floatX
        config.floatX = 'float64'

        try:
            self.tol = 1e-5

            #dataset = serial.load('${PYLEARN2_DATA_PATH}/stl10/stl10_patches/data.pkl')

            #X = dataset.get_batch_design(1000)
            #X = X[:,0:5]

            X = np.random.RandomState([1,2,3]).randn(1000,5)

            X -= X.mean()
            X /= X.std()
            m, D = X.shape
            N = 5

            #don't give the model an e_step or learning rate so it won't spend years compiling a learn_func
            self.model = S3C(nvis = D,
                             nhid = N,
                             irange = .1,
                             init_bias_hid = 0.,
                             init_B = 3.,
                             min_B = 1e-8,
                             max_B = 1000.,
                             init_alpha = 1., min_alpha = 1e-8, max_alpha = 1000.,
                             init_mu = 1., e_step = None,
                             m_step = Grad_M_Step(),
                             min_bias_hid = -1e30, max_bias_hid = 1e30,
                            )

            self.model.make_pseudoparams()

            self.h_new_coeff_schedule = [.1, .2, .3, .4, .5, .6, .7, .8, .9, 1. ]

            self.e_step = E_Step_Scan(h_new_coeff_schedule = self.h_new_coeff_schedule)
            self.e_step.register_model(self.model)

            self.X = X
            self.N = N
            self.m = m

        finally:
            config.floatX = self.prev_floatX

    def test_match_unrolled(self):
        """ tests that inference with scan matches result using unrolled loops """

        unrolled_e_step = E_Step(h_new_coeff_schedule = self.h_new_coeff_schedule)
        unrolled_e_step.register_model(self.model)

        V = T.matrix()

        scan_result = self.e_step.infer(V)
        unrolled_result = unrolled_e_step.infer(V)

        outputs = []

        for key in scan_result:
            outputs.append(scan_result[key])
            outputs.append(unrolled_result[key])

        f = function([V], outputs)

        outputs = f(self.X)

        assert len(outputs) % 2 == 0

        for i in xrange(0,len(outputs),2):
            assert np.allclose(outputs[i],outputs[i+1])

    def test_grad_s(self):

        "tests that the gradients with respect to s_i are 0 after doing a mean field update of s_i "

        model = self.model
        e_step = self.e_step
        X = self.X

        assert X.shape[0] == self.m

        model.test_batch_size = X.shape[0]

        init_H = e_step.init_H_hat(V = X)
        init_Mu1 = e_step.init_S_hat(V = X)

        prev_setting = config.compute_test_value
        config.compute_test_value= 'off'
        H, Mu1 = function([], outputs=[init_H, init_Mu1])()
        config.compute_test_value = prev_setting

        H = broadcast(H, self.m)
        Mu1 = broadcast(Mu1, self.m)

        H = np.cast[config.floatX](self.model.rng.uniform(0.,1.,H.shape))
        Mu1 = np.cast[config.floatX](self.model.rng.uniform(-5.,5.,Mu1.shape))

        H_var = T.matrix(name='H_var')
        H_var.tag.test_value = H
        Mu1_var = T.matrix(name='Mu1_var')
        Mu1_var.tag.test_value = Mu1
        idx = T.iscalar()
        idx.tag.test_value = 0

        S = e_step.infer_S_hat(V = X, H_hat = H_var, S_hat = Mu1_var)

        s_idx = S[:,idx]

        s_i_func = function([H_var,Mu1_var,idx],s_idx)

        sigma0 = 1. / model.alpha
        Sigma1 = e_step.infer_var_s1_hat()
        mu0 = T.zeros_like(model.mu)

        #by truncated KL, I mean that I am dropping terms that don't depend on H and Mu1
        # (they don't affect the outcome of this test and some of them are intractable )
        trunc_kl = - model.entropy_hs(H_hat = H_var, var_s0_hat = sigma0, var_s1_hat = Sigma1) + \
                     model.expected_energy_vhs(V = X, H_hat = H_var, S_hat = Mu1_var, var_s0_hat = sigma0, var_s1_hat = Sigma1)

        grad_Mu1 = T.grad(trunc_kl.sum(), Mu1_var)

        grad_Mu1_idx = grad_Mu1[:,idx]

        grad_func = function([H_var, Mu1_var, idx], grad_Mu1_idx)

        for i in xrange(self.N):
            Mu1[:,i] = s_i_func(H, Mu1, i)

            g = grad_func(H,Mu1,i)

            assert not np.any(np.isnan(g))

            g_abs_max = np.abs(g).max()

            if g_abs_max > self.tol:
                raise Exception('after mean field step, gradient of kl divergence wrt mean field parameter should be 0, but here the max magnitude of a gradient element is '+str(g_abs_max)+' after updating s_'+str(i))

    def test_value_s(self):

        "tests that the value of the kl divergence decreases with each update to s_i "

        model = self.model
        e_step = self.e_step
        X = self.X

        assert X.shape[0] == self.m

        init_H = e_step.init_H_hat(V = X)
        init_Mu1 = e_step.init_S_hat(V = X)

        prev_setting = config.compute_test_value
        config.compute_test_value= 'off'
        H, Mu1 = function([], outputs=[init_H, init_Mu1])()
        config.compute_test_value = prev_setting

        H = broadcast(H, self.m)
        Mu1 = broadcast(Mu1, self.m)

        H = np.cast[config.floatX](self.model.rng.uniform(0.,1.,H.shape))
        Mu1 = np.cast[config.floatX](self.model.rng.uniform(-5.,5.,Mu1.shape))

        H_var = T.matrix(name='H_var')
        H_var.tag.test_value = H
        Mu1_var = T.matrix(name='Mu1_var')
        Mu1_var.tag.test_value = Mu1
        idx = T.iscalar()
        idx.tag.test_value = 0

        S = e_step.infer_S_hat( V = X, H_hat = H_var, S_hat = Mu1_var)

        s_idx = S[:,idx]

        s_i_func = function([H_var,Mu1_var,idx],s_idx)

        sigma0 = 1. / model.alpha
        Sigma1 = e_step.infer_var_s1_hat()
        mu0 = T.zeros_like(model.mu)

        #by truncated KL, I mean that I am dropping terms that don't depend on H and Mu1
        # (they don't affect the outcome of this test and some of them are intractable )
        trunc_kl = - model.entropy_hs(H_hat = H_var, var_s0_hat = sigma0, var_s1_hat = Sigma1) + \
                     model.expected_energy_vhs(V = X, H_hat = H_var, S_hat = Mu1_var, var_s0_hat = sigma0, var_s1_hat = Sigma1)

        trunc_kl_func = function([H_var, Mu1_var], trunc_kl)

        for i in xrange(self.N):
            prev_kl = trunc_kl_func(H,Mu1)

            Mu1[:,i] = s_i_func(H, Mu1, i)

            new_kl = trunc_kl_func(H,Mu1)

            increase = new_kl - prev_kl

            mx = increase.max()

            if mx > 1e-3:
                raise Exception('after mean field step in s, kl divergence should decrease, but some elements increased by as much as '+str(mx)+' after updating s_'+str(i))

    def test_grad_h(self):

        "tests that the gradients with respect to h_i are 0 after doing a mean field update of h_i "

        model = self.model
        e_step = self.e_step
        X = self.X

        assert X.shape[0] == self.m

        init_H = e_step.init_H_hat(V = X)
        init_Mu1 = e_step.init_S_hat(V = X)

        prev_setting = config.compute_test_value
        config.compute_test_value= 'off'
        H, Mu1 = function([], outputs=[init_H, init_Mu1])()
        config.compute_test_value = prev_setting

        H = broadcast(H, self.m)
        Mu1 = broadcast(Mu1, self.m)

        H = np.cast[config.floatX](self.model.rng.uniform(0.,1.,H.shape))
        Mu1 = np.cast[config.floatX](self.model.rng.uniform(-5.,5.,Mu1.shape))

        H_var = T.matrix(name='H_var')
        H_var.tag.test_value = H
        Mu1_var = T.matrix(name='Mu1_var')
        Mu1_var.tag.test_value = Mu1
        idx = T.iscalar()
        idx.tag.test_value = 0

        new_H = e_step.infer_H_hat(V = X, H_hat = H_var, S_hat = Mu1_var)
        h_idx = new_H[:,idx]

        updates_func = function([H_var,Mu1_var,idx], h_idx)

        sigma0 = 1. / model.alpha
        Sigma1 = e_step.infer_var_s1_hat()
        mu0 = T.zeros_like(model.mu)

        #by truncated KL, I mean that I am dropping terms that don't depend on H and Mu1
        # (they don't affect the outcome of this test and some of them are intractable )
        trunc_kl = - model.entropy_hs(H_hat = H_var, var_s0_hat = sigma0, var_s1_hat = Sigma1) + \
                     model.expected_energy_vhs(V = X, H_hat = H_var, S_hat = Mu1_var,  var_s0_hat = sigma0,
                             var_s1_hat = Sigma1)

        grad_H = T.grad(trunc_kl.sum(), H_var)

        assert len(grad_H.type.broadcastable) == 2

        #from theano.printing import min_informative_str
        #print min_informative_str(grad_H)

        #grad_H = Print('grad_H')(grad_H)

        #grad_H_idx = grad_H[:,idx]

        grad_func = function([H_var, Mu1_var], grad_H)

        failed = False

        for i in xrange(self.N):
            rval = updates_func(H, Mu1, i)
            H[:,i] = rval

            g = grad_func(H,Mu1)[:,i]

            assert not np.any(np.isnan(g))

            g_abs_max = np.abs(g).max()

            if g_abs_max > self.tol:
                #print "new values of H"
                #print H[:,i]
                #print "gradient on new values of H"
                #print g

                failed = True

                print 'iteration ',i
                #print 'max value of new H: ',H[:,i].max()
                #print 'H for failing g: '
                failing_h = H[np.abs(g) > self.tol, i]
                #print failing_h

                #from matplotlib import pyplot as plt
                #plt.scatter(H[:,i],g)
                #plt.show()

                #ignore failures extremely close to h=1

                high_mask = failing_h > .001
                low_mask = failing_h < .999

                mask = high_mask * low_mask

                print 'masked failures: ',mask.shape[0],' err ',g_abs_max

                if mask.sum() > 0:
                    print 'failing h passing the range mask'
                    print failing_h[ mask.astype(bool) ]
                    raise Exception('after mean field step, gradient of kl divergence'
                            ' wrt freshly updated variational parameter should be 0, '
                            'but here the max magnitude of a gradient element is '
                            +str(g_abs_max)+' after updating h_'+str(i))

        #assert not failed

    def test_value_h(self):

        "tests that the value of the kl divergence decreases with each update to h_i "

        model = self.model
        e_step = self.e_step
        X = self.X

        assert X.shape[0] == self.m

        init_H = e_step.init_H_hat(V = X)
        init_Mu1 = e_step.init_S_hat(V = X)

        prev_setting = config.compute_test_value
        config.compute_test_value= 'off'
        H, Mu1 = function([], outputs=[init_H, init_Mu1])()
        config.compute_test_value = prev_setting

        H = broadcast(H, self.m)
        Mu1 = broadcast(Mu1, self.m)

        H = np.cast[config.floatX](self.model.rng.uniform(0.,1.,H.shape))
        Mu1 = np.cast[config.floatX](self.model.rng.uniform(-5.,5.,Mu1.shape))

        H_var = T.matrix(name='H_var')
        H_var.tag.test_value = H
        Mu1_var = T.matrix(name='Mu1_var')
        Mu1_var.tag.test_value = Mu1
        idx = T.iscalar()
        idx.tag.test_value = 0

        newH = e_step.infer_H_hat(V = X, H_hat = H_var, S_hat = Mu1_var)

        h_idx = newH[:,idx]

        h_i_func = function([H_var,Mu1_var,idx],h_idx)

        sigma0 = 1. / model.alpha
        Sigma1 = e_step.infer_var_s1_hat()
        mu0 = T.zeros_like(model.mu)

        #by truncated KL, I mean that I am dropping terms that don't depend on H and Mu1
        # (they don't affect the outcome of this test and some of them are intractable )
        trunc_kl = - model.entropy_hs(H_hat = H_var, var_s0_hat = sigma0, var_s1_hat = Sigma1) + \
                     model.expected_energy_vhs(V = X, H_hat = H_var, S_hat = Mu1_var, var_s0_hat = sigma0, var_s1_hat = Sigma1)

        trunc_kl_func = function([H_var, Mu1_var], trunc_kl)

        for i in xrange(self.N):
            prev_kl = trunc_kl_func(H,Mu1)

            H[:,i] = h_i_func(H, Mu1, i)
            #we don't update mu, the whole point of the split e step is we don't have to

            new_kl = trunc_kl_func(H,Mu1)

            increase = new_kl - prev_kl

            print 'failures after iteration ',i,': ',(increase > self.tol).sum()

            mx = increase.max()

            if mx > 1e-4:
                print 'increase amounts of failing examples:'
                print increase[increase > self.tol]
                print 'failing H:'
                print H[increase > self.tol,:]
                print 'failing Mu1:'
                print Mu1[increase > self.tol,:]
                print 'failing V:'
                print X[increase > self.tol,:]

                raise Exception('after mean field step in h, kl divergence should decrease, but some elements increased by as much as '+str(mx)+' after updating h_'+str(i))

if __name__ == '__main__':
    obj = Test_S3C_Inference()

    #obj.test_grad_h()
    #obj.test_grad_s()
    #obj.test_value_s()
    obj.test_value_h()


from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six

from nose.tools import assert_raises, assert_equal

import numpy as np
from numpy.testing.utils import assert_array_equal, assert_array_almost_equal

import matplotlib.colors as mcolors
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from matplotlib.testing.decorators import image_comparison, cleanup

def test_colormap_endian():
    """
    Github issue #1005: a bug in putmask caused erroneous
    mapping of 1.0 when input from a non-native-byteorder
    array.
    """
    cmap = cm.get_cmap("jet")
    # Test under, over, and invalid along with values 0 and 1.
    a = [-0.5, 0, 0.5, 1, 1.5, np.nan]
    for dt in ["f2", "f4", "f8"]:
        anative = np.ma.masked_invalid(np.array(a, dtype=dt))
        aforeign = anative.byteswap().newbyteorder()
        #print(anative.dtype.isnative, aforeign.dtype.isnative)
        assert_array_equal(cmap(anative), cmap(aforeign))

def test_BoundaryNorm():
    """
    Github issue #1258: interpolation was failing with numpy
    1.7 pre-release.
    """
    # TODO: expand this into a more general test of BoundaryNorm.
    boundaries = [0, 1.1, 2.2]
    vals = [-1, 0, 2, 2.2, 4]
    expected = [-1, 0, 2, 3, 3]
    # ncolors != len(boundaries) - 1 triggers interpolation
    ncolors = len(boundaries)
    bn = mcolors.BoundaryNorm(boundaries, ncolors)
    assert_array_equal(bn(vals), expected)

def test_LogNorm():
    """
    LogNorm ignored clip, now it has the same
    behavior as Normalize, e.g., values > vmax are bigger than 1
    without clip, with clip they are 1.
    """
    ln = mcolors.LogNorm(clip=True, vmax=5)
    assert_array_equal(ln([1, 6]), [0, 1.0])

def test_PowerNorm():
    a = np.array([0, 0.5, 1, 1.5], dtype=np.float)
    pnorm = mcolors.PowerNorm(1)
    norm = mcolors.Normalize()
    assert_array_almost_equal(norm(a), pnorm(a))

    a = np.array([-0.5, 0, 2, 4, 8], dtype=np.float)
    expected = [0, 0, 1/16, 1/4, 1]
    pnorm = mcolors.PowerNorm(2, vmin=0, vmax=8)
    assert_array_almost_equal(pnorm(a), expected)
    assert_equal(pnorm(a[0]), expected[0])
    assert_equal(pnorm(a[2]), expected[2])
    assert_array_almost_equal(a[1:], pnorm.inverse(pnorm(a))[1:])

    # Clip = True
    a = np.array([-0.5, 0, 1, 8, 16], dtype=np.float)
    expected = [0, 0, 0, 1, 1]
    pnorm = mcolors.PowerNorm(2, vmin=2, vmax=8, clip=True)
    assert_array_almost_equal(pnorm(a), expected)
    assert_equal(pnorm(a[0]), expected[0])
    assert_equal(pnorm(a[-1]), expected[-1])

    # Clip = True at call time
    a = np.array([-0.5, 0, 1, 8, 16], dtype=np.float)
    expected = [0, 0, 0, 1, 1]
    pnorm = mcolors.PowerNorm(2, vmin=2, vmax=8, clip=False)
    assert_array_almost_equal(pnorm(a, clip=True), expected)
    assert_equal(pnorm(a[0], clip=True), expected[0])
    assert_equal(pnorm(a[-1], clip=True), expected[-1])

def test_Normalize():
    norm = mcolors.Normalize()
    vals = np.arange(-10, 10, 1, dtype=np.float)
    _inverse_tester(norm, vals)
    _scalar_tester(norm, vals)
    _mask_tester(norm, vals)

def test_SymLogNorm():
    """
    Test SymLogNorm behavior
    """
    norm = mcolors.SymLogNorm(3, vmax=5, linscale=1.2)
    vals = np.array([-30, -1, 2, 6], dtype=np.float)
    normed_vals = norm(vals)
    expected = [0., 0.53980074, 0.826991, 1.02758204]
    assert_array_almost_equal(normed_vals, expected)
    _inverse_tester(norm, vals)
    _scalar_tester(norm, vals)
    _mask_tester(norm, vals)

    # Ensure that specifying vmin returns the same result as above
    norm = mcolors.SymLogNorm(3, vmin=-30, vmax=5, linscale=1.2)
    normed_vals = norm(vals)
    assert_array_almost_equal(normed_vals, expected)

def _inverse_tester(norm_instance, vals):
    """
    Checks if the inverse of the given normalization is working.
    """
    assert_array_almost_equal(norm_instance.inverse(norm_instance(vals)), vals)

def _scalar_tester(norm_instance, vals):
    """
    Checks if scalars and arrays are handled the same way.
    Tests only for float.
    """
    scalar_result = [norm_instance(float(v)) for v in vals]
    assert_array_almost_equal(scalar_result, norm_instance(vals))

def _mask_tester(norm_instance, vals):
    """
    Checks mask handling
    """
    masked_array = np.ma.array(vals)
    masked_array[0] = np.ma.masked
    assert_array_equal(masked_array.mask, norm_instance(masked_array).mask)

@image_comparison(baseline_images=['levels_and_colors'],
                  extensions=['png'])
def test_cmap_and_norm_from_levels_and_colors():
    data = np.linspace(-2, 4, 49).reshape(7, 7)
    levels = [-1, 2, 2.5, 3]
    colors = ['red', 'green', 'blue', 'yellow', 'black']
    extend = 'both'
    cmap, norm = mcolors.from_levels_and_colors(levels, colors, extend=extend)

    ax = plt.axes()
    m = plt.pcolormesh(data, cmap=cmap, norm=norm)
    plt.colorbar(m)

    # Hide the axes labels (but not the colorbar ones, as they are useful)
    for lab in ax.get_xticklabels() + ax.get_yticklabels():
        lab.set_visible(False)

def test_cmap_and_norm_from_levels_and_colors2():
    levels = [-1, 2, 2.5, 3]
    colors = ['red', (0, 1, 0), 'blue', (0.5, 0.5, 0.5), (0.0, 0.0, 0.0, 1.0)]
    clr = mcolors.colorConverter.to_rgba_array(colors)
    bad = (0.1, 0.1, 0.1, 0.1)
    no_color = (0.0, 0.0, 0.0, 0.0)
    masked_value = 'masked_value'

    # Define the test values which are of interest.
    # Note: levels are lev[i] <= v < lev[i+1]
    tests = [('both', None, {-2: clr[0],
                             -1: clr[1],
                             2: clr[2],
                             2.25: clr[2],
                             3: clr[4],
                             3.5: clr[4],
                             masked_value: bad}),

             ('min', -1, {-2: clr[0],
                          -1: clr[1],
                          2: clr[2],
                          2.25: clr[2],
                          3: no_color,
                          3.5: no_color,
                          masked_value: bad}),

             ('max', -1, {-2: no_color,
                          -1: clr[0],
                          2: clr[1],
                          2.25: clr[1],
                          3: clr[3],
                          3.5: clr[3],
                          masked_value: bad}),

             ('neither', -2, {-2: no_color,
                              -1: clr[0],
                              2: clr[1],
                              2.25: clr[1],
                              3: no_color,
                              3.5: no_color,
                              masked_value: bad}),
             ]

    for extend, i1, cases in tests:
        cmap, norm = mcolors.from_levels_and_colors(levels, colors[0:i1],
                                                    extend=extend)
        cmap.set_bad(bad)
        for d_val, expected_color in cases.items():
            if d_val == masked_value:
                d_val = np.ma.array([1], mask=True)
            else:
                d_val = [d_val]
            assert_array_equal(expected_color, cmap(norm(d_val))[0],
                               'Wih extend={0!r} and data '
                               'value={1!r}'.format(extend, d_val))

    assert_raises(ValueError, mcolors.from_levels_and_colors, levels, colors)

def test_rgb_hsv_round_trip():
    for a_shape in [(500, 500, 3), (500, 3), (1, 3), (3,)]:
        np.random.seed(0)
        tt = np.random.random(a_shape)
        assert_array_almost_equal(tt,
            mcolors.hsv_to_rgb(mcolors.rgb_to_hsv(tt)))
        assert_array_almost_equal(tt,
            mcolors.rgb_to_hsv(mcolors.hsv_to_rgb(tt)))

@cleanup
def test_autoscale_masked():
    # Test for #2336. Previously fully masked data would trigger a ValueError.
    data = np.ma.masked_all((12, 20))
    plt.pcolor(data)
    plt.draw()

def test_colors_no_float():
    # Gray must be a string to distinguish 3-4 grays from RGB or RGBA.

    def gray_from_float_rgb():
        return mcolors.colorConverter.to_rgb(0.4)

    def gray_from_float_rgba():
        return mcolors.colorConverter.to_rgba(0.4)

    assert_raises(ValueError, gray_from_float_rgb)
    assert_raises(ValueError, gray_from_float_rgba)

def test_light_source_shading_color_range():
    # see also
    #http://matplotlib.org/examples/pylab_examples/shading_example.html

    from matplotlib.colors import LightSource
    from matplotlib.colors import Normalize

    refinput = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
    norm = Normalize(vmin=0, vmax=50)
    ls = LightSource(azdeg=0, altdeg=65)
    testoutput = ls.shade(refinput, plt.cm.jet, norm=norm)
    refoutput = np.array([
        [[0., 0., 0.58912656, 1.],
        [0., 0., 0.67825312, 1.],
        [0., 0., 0.76737968, 1.],
        [0., 0., 0.85650624, 1.]],
        [[0., 0., 0.9456328, 1.],
        [0., 0., 1., 1.],
        [0., 0.04901961, 1., 1.],
        [0., 0.12745098, 1., 1.]],
        [[0., 0.22156863, 1., 1.],
        [0., 0.3, 1., 1.],
        [0., 0.37843137, 1., 1.],
        [0., 0.45686275, 1., 1.]]
        ])
    assert_array_almost_equal(refoutput, testoutput)

if __name__ == '__main__':
    import nose
    nose.runmodule(argv=['-s', '--with-doctest'], exit=False)


"""
showspec is my homegrown spectrum plotter, meant to somewhat follow STARLINK's
SPLAT and have functionality similar to GAIA, but with an emphasis on producing
publication-quality plots (which, while splat may do, it does unreproducibly)

.. todo::
    -add spectrum arithmetic tools
        (as is, you can use numpy.interp with sp.vind and sp.spectrum pretty
        easily)
    -implement other fitters
        -e.g., NH3 hyperfine, Voigt
    -move to object-oriented pylab/pyplot implementation (for bulk / speedup work)
    -allow for non-plotting fitting work (curious... I've never needed that yet)
    -Equivalent Width measurement without gaussian fitting
        -probably should be part of the baseline code
    -write documentation other people can read

12/21/2011 - ALL of the above to do is IS DONE!  It's now hosted at <http://pyspeckit.bitbucket.org>

"""

import math

import pylab
from pylab import *
for k,v in pylab.__dict__.iteritems():
    if hasattr(v,'__module__'):
        if v.__module__ is None:
            locals()[k].__module__ = 'pylab'
import matplotlib

from agpy.mpfit import mpfit

from collapse_gaussfit import *
from ratosexagesimal import *
import pyfits
import gaussfitter

import numpy
from numpy import isnan
from mad import MAD,nanmedian

def steppify(arr,isX=False,interval=0,sign=+1.0):
    """
    *support function*
    Converts an array to double-length for step plotting
    """
    if isX and interval==0:
        interval = numpy.abs(arr[1]-arr[0]) / 2.0
    newarr = pylab.array(zip(arr-sign*interval,arr+sign*interval)).ravel()
    return newarr

class SpecPlotter:
    """
    SpecPlotter class.  Takes in a spectrum or data cube, plotting properties,
    and a velocity axis determination function.  Look at splat_1d for a wrapper
    that might actually be useful.

    Whew, needs more documentation
    """

    def __init__(self,  cube,  axis=None,  xtol=None,  ytol=None, vconv=lambda
            x: x, xtora=lambda x: x, ytodec=lambda x: x, specname=None,
            dv=None, color='k', hdr=None, errspec=None,  maskspec=None,
            fig=None, fignum=1, clear=True, title=None, xunits='km/s',
            erralpha=0.2, ivconv=None, autorefresh=True, reffreq=None,
            gaiafignum=0, gaiafig=None, clickid=None, **kwargs ):
      self.vconv = vconv
      self.xtora = xtora
      self.ytodec = ytodec
      self.cube = cube # where(numpy.isnan(cube),0,cube)
      if len(self.cube.shape) > 1:
          self.spectrum = cube[:,0,0] # spectrum is what's plotted; cube is the "raw data"
      else:
          self.spectrum = cube # spectrum is what's plotted; cube is the "raw data"
      self.specname=specname
      self.dv=dv
      self.reffreq=reffreq
      self.scale=1.0
      self.units='K'
      self.xunits=xunits
      self.voff=0.0
      self.offset=0.0
      self.continuum=0.0
      self.errspec = errspec
      self.erralpha=erralpha
      self.plotcolor=color
      self.specfit = Specfit(self)
      self.fitspec = self.specfit
      self.baseline = Baseline(self)
      #self.fft = FFT(self)
      #self.psd = self.fft.psd
      self.vmin=None
      self.vmax=None
      self.title=title
      self.ivconv=ivconv
      self.autorefresh=autorefresh
      self.spectrumplot=None
      self.errorplot=None
      self.gaiafignum = gaiafignum
      self.gaiafig = gaiafig
      self.clickid = clickid
      self.plotkwargs = kwargs
      if maskspec is not None:
          self.maskspec = maskspec
      else:
          self.maskspec = zeros(self.cube.shape)
      self.linecollections =[]
      self.texts =[]
      if hdr: self.header = hdr

      # figure out where to put the plot
      if fig is None and axis is None:
          fig=figure(fignum)
          if clear: fig.clf()
          self.axis = fig.gca()
      elif fig is None and axis is None:
          self.axis = pylab.gca()
      elif fig is not None and axis is None:
          if clear: fig.clf()
          self.axis = fig.gca()
      elif fig is None and axis is not None:
          self.axis = axis
      else: # if figure and axis are both set, just use axis
          self.axis = axis
      if clear: self.axis.clear()

    def __call__(self, event):
      """
      Connects map cube to specplotter...
      """
      if event.inaxes:
        clickX = event.xdata
        clickY = event.ydata
        tb = get_current_fig_manager().toolbar
        #if ((self.axis is None) or (self.axis==event.inaxes)) and tb.mode=='':
        if event.button==1 and tb.mode=='':
            print "OverPlotting spectrum from point %i,%i" % (clickX,clickY)
            self.plotspec(clickY,clickX,button=event.button,cube=True)
        elif event.button==2:
            print "Plotting spectrum from point %i,%i" % (clickX,clickY)
            self.plotspec(clickY,clickX,button=event.button,cube=True,clear=True)
        elif event.button==3:
            print "Disconnecting GAIA-like tool"
            self.gaiafig.canvas.mpl_disconnect(self.clickid)
        else:
            print "Call failed for some reason: "
            print "event: ",event

    def plotspec(self, i=0, j=0, cube=False, title=None,
            clear=False, color=None, continuum=None,
            axis=None, offset=None, scale=None, voff=None, vmin=None,
            vmax=None, units=None, xunits=None, erralpha=None, plotpix=False,
            errstyle='fill', autorefresh=None, button=None, **kwargs):
      """
      Plot a spectrum
      Originally written to plot spectra from data cubes, hence the i,j parameter
      to specify the location in the cube

      Now, cube defaults to False, but you can still pass in a data cube.

      Inputs:
        title,color, kwargs - semi-obvious plot-related comands
        axis - You can pass in a Matplotlib axis instance and it will plot on that
        clear - Clear the axis before plotting?
        continuum - if you've already subtracted out a continuum, you can add
            it back in (only if it is a constant offset).  It will be included in
            the spectrum
        offset - Like continuum, but ONLY for plotting purposes.  Will move the
            plot vertically but will NOT include values in the .spectrum
        scale - multiplicative factor to scale the data by (NOT for plotting
            purposes; modifies spectrum)
        voff - Shift the spectrum on the velocity axis by this amount
        vmin,vmax - only plot within this range
            (note that these keywords passed to splat_1d MAY crop the spectrum)
        units - units of the data.  At the moment, no conversions are done
        xunits - units of the Y axis.  Can affect other procedures, like show_lines,
            and some unit conversion (Hz to GHz) is done
        erralpha - Transparency of the errorbars if plotted
        errstyle - style of errorbars if plotted
        plotpix - if set, will plot against a pixel (channel) axis instead of a
            physical axis
        autorefresh - automatically update the plot when fitting gaussians, labeling,
            etc?
      """

      if kwargs.has_key('fignum'): kwargs.pop('fignum')  # HACK because I want __init__ to accept different kwargs
      if kwargs.has_key('fig'): kwargs.pop('fig')        # is there a better workaround?
      if scale  is not None: self.scale = scale
      if units  is not None: self.units = units
      if xunits is not None: self.xunits= xunits
      if voff   is not None: self.voff  = voff
      if offset is not None: self.offset= offset
      if continuum is not None: self.continuum= continuum
      if color  is not None: self.plotcolor=color
      if erralpha is not None: self.erralpha= erralpha
      if vmax  is not None: self.vmax = vmax
      if vmin  is not None: self.vmin = vmin
      if title is not None: self.title = title
      if autorefresh is not None: self.autorefresh = autorefresh
      if axis is None: axis=self.axis # allow spectrum to be plotted on other axis
      if clear: axis.clear()

      if plotpix:
          self.vind = arange(self.cube.shape[0])
      else:
          self.vind = self.vconv(arange(self.cube.shape[0])) + self.voff

      if kwargs.has_key('fignum'): kwargs.pop('fignum')
      if kwargs.has_key('linewidth'):
          linewidth = kwargs.pop('linewidth')
      else:
          linewidth=0.5

      if cube or len(self.cube.shape) == 3:
          self.spectrum = self.cube[:,i,j]*self.scale+self.continuum-self.baseline.basespec
          self.spectrumplot = axis.plot(self.vind,self.spectrum+self.offset,color=self.plotcolor,
                  linestyle='steps-mid',linewidth=linewidth,
                  **kwargs)
      else:
          if self.maskspec.sum() > 0:
              nanmask = where(self.maskspec,numpy.nan,1)
              self.spectrum = self.cube*self.scale*nanmask+self.continuum-self.baseline.basespec
              self.spectrumplot = axis.plot(self.vind,self.spectrum+self.offset,color=self.plotcolor,
                      linestyle='steps-mid',linewidth=linewidth,
                      **kwargs)
          else:
              self.spectrum = self.cube*self.scale+self.continuum-self.baseline.basespec
              self.spectrumplot = axis.plot(self.vind,self.spectrum+self.offset,color=self.plotcolor,
                      linestyle='steps-mid',linewidth=linewidth,
                      **kwargs)
          if self.errspec is not None:
              if errstyle == 'fill':
                  self.errorplot = [axis.fill_between(steppify(self.vind,isX=True,sign=sign(self.dv)),
                          steppify(self.spectrum+self.offset-self.errspec*self.scale),
                          steppify(self.spectrum+self.offset+self.errspec*self.scale),
                          facecolor=self.plotcolor, alpha=self.erralpha, **kwargs)]
              elif errstyle == 'bars':
                  self.errorplot = axis.errorbar(self.vind, self.spectrum+self.offset,
                          yerr=self.errspec*self.scale, ecolor=self.plotcolor, fmt=None,
                          **kwargs)

      if vmin is not None: xlo = self.vmin
      else: xlo=self.vind.min()
      if vmax is not None: xhi = self.vmax
      else: xhi=self.vind.max()
      axis.set_xlim(xlo,xhi)

      if self.title is not None:
          axis.set_title(self.title)
      elif self.xtora and self.ytodec:
          axis.set_title("Spectrum at %s %s" %
                  (ratos(self.xtora(i)),dectos(self.ytodec(j))))
      elif self.specname:
          axis.set_title("Spectrum of %s" % self.specname)
      if isinstance(self.xunits,str):
          axis.set_xlabel(self.xunits)
      else:
          axis.set_xlabel("V$_{LSR}$ (km s$^{-1}$)")
          self.xunits = 'km/s'
      if units in ['Ta*','Tastar','K']:
        axis.set_ylabel("$T_A^*$ (K)")
      elif units == 'mJy':
        axis.set_ylabel("$S_\\nu$ (mJy)")
      elif units == 'Jy':
        axis.set_ylabel("$S_\\nu$ (Jy)")
      else:
        axis.set_ylabel(self.units)
      if self.autorefresh: self.refresh()

    def save(self,fname,**kwargs):
      """
      Save the current spectrum (useful for saving baselined data)
      """
      newfile = pyfits.PrimaryHDU(data=self.cube,header=self.header)
      newfile.writeto(fname,**kwargs)

    def savefig(self,fname,bbox_inches='tight',**kwargs):
        """
        simple wrapper of maplotlib's savefig.
        """
        self.axis.figure.savefig(fname,bbox_inches=bbox_inches,**kwargs)

    def showlines(self,linefreqs,linenames,ctype='freq',cunit='hz',yscale=0.8,vofflines=0.0,
            voffunit='km/s',**kwargs):
        """
        Overplot vertical lines and labels at the frequencies (or velocities) of each line

        yscale - fraction of maximum at which to label
        """

        self.clearlines()

        if ctype != 'freq':
            print "Sorry, non-frequency units not implemented yet."
            return

        speedoflight=2.99792458e5
        if self.reffreq and self.xunits in ('km/s','m/s'):
            linefreqs = -(array(linefreqs)-self.reffreq)/self.reffreq * speedoflight

        if 'hz' in cunit or 'Hz' in cunit:
            linefreqs *= (1.0 + vofflines / speedoflight)
        else:
            linefreqs += vofflines

        ymax = (self.spectrum[self.spectrum==self.spectrum]).max()
        for lf,ln in zip(linefreqs,linenames):
            if lf < self.vind.max() and lf > self.vind.min():
                self.linecollections.append(vlines(lf,0,ymax,**kwargs))
                self.texts.append(text(lf,ymax*yscale,ln,rotation='vertical',**kwargs))

        if self.autorefresh: self.refresh()

    def clearlines(self):
        if len(self.texts) > 0:
            for T in self.texts:
                if T in self.axis.texts:
                    self.axis.texts.remove(T)
        if len(self.linecollections) > 0:
            for LC in self.linecollections:
                if LC in self.axis.collections:
                    self.axis.collections.remove(LC)

    def refresh(self):
        self.axis.figure.canvas.draw()

class FFT:
    def __init__(self,specplotter,fignum=3,axis=None, color='k'):
        self.specplotter=specplotter
        if axis is None:
            self.fignum=fignum
            self.figure=figure(self.fignum)
            self.axis=gca()
        else:
            self.axis=axis
            self.figure=self.axis.figure
            self.fignum=None
        #self.axis.clear()
        self.color=color
        self.fftplot=None
        self.setspec()
        self.setshift()
        self.clear()

    def __call__(self,psd=False,shift=True):
        self.setspec()
        if psd:
            self.psd(shift=shift)
        else:
            self.fft(shift=shift)

    def fft(self,shift=True,logplot=False,**kwargs):
        self.clear()
        self.setshift(shift)
        if logplot: self.axis.set_yscale('log')
        else: self.axis.set_yscale('linear')
        self.fftspec = fft(self.spectofft)
        self.realfft = self.fftspec.real
        self.imagfft = self.fftspec.imag
        self.fftplot = self.axis.plot(self.shiftfunc(self.realfft),
                drawstyle='steps-mid',color=self.color,**kwargs)
        self.refresh()

    def psd(self,logplot=True,shift=True,**kwargs):
        self.clear()
        if logplot: self.axis.set_yscale('log')
        else: self.axis.set_yscale('linear')
        self.setshift(shift)
        self.psdspec = fft(self.spectofft) * fft(self.spectofft[::-1])
        self.psdreal = abs(self.psdspec)
        self.fftplot = self.axis.plot(self.shiftfunc(self.psdreal),
                drawstyle='steps-mid',color=self.color,**kwargs)
        if logplot: self.axis.set_yscale('log')
        else: self.axis.set_yscale('linear')
        self.refresh()

    def setshift(self,shift=True):
        if shift: self.shiftfunc = fftshift
        else: self.shiftfunc = lambda x: x

    def setspec(self):
        self.spectofft = copy(self.specplotter.spectrum)
        OKmask = (self.spectofft==self.spectofft)
        self.spectofft[(True-OKmask)] = 0

    def clear(self):
        if self.fftplot is not None:
            for p in self.fftplot:
                p.set_visible(False)
                if p in self.axis.lines: self.axis.lines.remove(p)
        self.axis.clear()
        self.refresh()

    def refresh(self):
        self.axis.figure.canvas.draw()

class PSD(FFT):
    def __call__(self,shift=True):
        self.setspec()
        self.setshift(shift)
        self.clear()
        self.psd()
        self.refresh()

class Baseline:
    def __init__(self,specplotter):
        self.baselinepars  = None
        self.order = None
        self.basespec = zeros(specplotter.spectrum.shape[0])
        self.excludemask = zeros(specplotter.spectrum.shape[0],dtype='bool')
        self.OKmask = ones(specplotter.spectrum.shape[0],dtype='bool')
        self.specplotter = specplotter
        self.blleg = None
        self.click = 0
        self.nclicks_b1 = 0
        self.nclicks_b2 = 0
        self.fitregion=[]
        self.excludevelo = []
        self.excludepix  = []

    def __call__(self, order=1, annotate=False, excludefit=False, save=True,
            exclude=None, exclusionlevel=0.01,
            interactive=False, **kwargs):
        """
        Fit and remove a polynomial from the spectrum.
        It will be saved in the variable "self.basespec"
        and the fit parameters will be saved in "self.order"

        function baseline(spectrum,xarr=None,xmin=None,xmax=None,order=1,quiet=True,exclude=None):
        Subtract a baseline from a spectrum
        If xmin,xmax are not specified, defaults to ignoring first and last 10% of spectrum

        exclude is a set of start/end indices to ignore when baseline fitting
        (ignored by setting error to infinite in fitting procedure)

        excludefit creates a mask based on the fitted gaussian model (assuming
        that it has a zero-height) using an exclusion level of (exclusionlevel)
        * the smallest gaussian peak that was fit

        "basespec" is added back to the spectrum before fitting so you can run this
        procedure multiple times without losing information
        """
        specfit = self.specplotter.specfit
        self.order = order
        fitp = zeros(self.order+1)
        self.spectofit = self.specplotter.spectrum+self.basespec
        self.OKmask = (self.spectofit==self.spectofit)
        if exclude == 'interactive' or interactive:
            self.excludemask[:] = True
            self.excludevelo = []
            self.excludepix  = []
            self.click = self.specplotter.axis.figure.canvas.mpl_connect('button_press_event',self.selectregion)
        else:
            if excludefit and specfit.modelpars is not None:
                #vlo = self.specplotter.specfit.modelpars[1] - 2*self.specplotter.specfit.modelpars[2]
                #vhi = self.specplotter.specfit.modelpars[1] + 2*self.specplotter.specfit.modelpars[2]
                #exclude = [argmin(abs(self.specplotter.vind-vlo)),argmin(abs(self.specplotter.vind-vhi))]
                specfit.fullsizemodel() # make sure the spectrum is the right size
                self.excludemask = abs(specfit.model) > exclusionlevel*abs(min(specfit.modelpars[0::3]))
            else:
                self.excludemask[:] = False
            self.dofit(exclude=exclude,annotate=annotate,**kwargs)
        if save: self.savefit()

    def dofit(self, exclude=None, excludeunits='velo', annotate=False,
            **kwargs):
        """
        Do the baseline fitting and save and plot the results.

        Can specify a region to exclude using velocity units or pixel units
        """
        if exclude is not None and excludeunits in ['velo','km/s']:
            if len(exclude) % 2 == 0:
                self.excludevelo = exclude
                self.excludepix = []
                for vl,vu in zip(exclude[::2],exclude[1::2]):
                    xl = argmin(abs(self.specplotter.vind-vl))
                    xu = argmin(abs(self.specplotter.vind-vu))
                    if xl > xu: xl,xu=xu,xl
                    self.excludemask[xl:xu] = True
                    self.excludepix += [xl,xu]
        elif excludeunits in ['pix','pixel','chan','channel']:
            if len(exclude) % 2 == 0:
                self.excludepix = []
                for xl,xu in zip(exclude[::2],exclude[1::2]):
                    if xl > xu: xl,xu=xu,xl
                    self.excludemask[xl:xu] = True
                    self.excludepix += [xl,xu]
        self.specplotter.spectrum, self.baselinepars = baseline(
                self.spectofit,
                xarr=self.specplotter.vind,
                order=self.order, exclude=None,
                mask=(True-self.OKmask)+self.excludemask,
                **kwargs)
        self.basespec = poly1d(self.baselinepars)(self.specplotter.vind)
        if self.specplotter.spectrumplot is not None:
            [self.specplotter.axis.lines.remove(p) for p in self.specplotter.spectrumplot]
        if self.specplotter.errorplot is not None:
            [self.specplotter.axis.collections.remove(p) for p in self.specplotter.errorplot if isinstance(p,matplotlib.collections.PolyCollection)]
            [self.specplotter.axis.lines.remove(p) for p in self.specplotter.errorplot if isinstance(p,matplotlib.lines.Line2D)]
        self.specplotter.plotspec(**self.specplotter.plotkwargs)
        self.specplotter.axis.set_ylim(
                abs(self.specplotter.spectrum[self.OKmask].min())*1.1*sign(self.specplotter.spectrum[self.OKmask].min()),
                abs(self.specplotter.spectrum[self.OKmask].max())*1.1*sign(self.specplotter.spectrum[self.OKmask].max()))
        if annotate: self.annotate() # refreshes automatically
        elif self.specplotter.autorefresh: self.specplotter.refresh()

    def selectregion(self,event,annotate=False):
        """
        select regions for baseline fitting
        """
        if event.button == 1:
            if self.nclicks_b1 == 0:
                self.bx1 = argmin(abs(event.xdata-self.specplotter.vind))
                self.excludevelo += [self.specplotter.vind]
                self.excludepix  += [self.bx1]
                self.nclicks_b1 += 1
            elif self.nclicks_b1 == 1:
                self.bx2 = argmin(abs(event.xdata-self.specplotter.vind))
                self.nclicks_b1 -= 1
                if self.bx1 > self.bx2: self.bx1,self.bx2 = self.bx2,self.bx1
                self.fitregion += self.specplotter.axis.plot(
                        self.specplotter.vind[self.bx1:self.bx2],
                        self.specplotter.spectrum[self.bx1:self.bx2]+self.specplotter.offset,
                        drawstyle='steps-mid',
                        color='g',alpha=0.5)
                self.specplotter.refresh()
                self.excludemask[self.bx1:self.bx2] = False
                self.excludevelo += [self.specplotter.vind]
                self.excludepix  += [self.bx2]
        if event.button in [2,3]:
            disconnect(self.click)
            self.dofit(exclude=None,annotate=annotate)
            for p in self.fitregion:
                p.set_visible(False)
                self.specplotter.axis.lines.remove(p)
            self.fitregion=[] # I should be able to just remove from the list... but it breaks the loop...
            self.specplotter.refresh()

    def annotate(self,loc='upper left'):
        bltext = "bl: $y=$"+"".join(["$%+6.3gx^{%i}$" % (f,self.order-i)
            for i,f in enumerate(self.baselinepars)])
        #self.blleg = text(xloc,yloc     ,bltext,transform = self.specplotter.axis.transAxes)
        self.clearlegend()
        pl = matplotlib.collections.CircleCollection([0],edgecolors=['k'])
        self.blleg = self.specplotter.axis.legend(
                (pl,),
                (bltext,),loc=loc,markerscale=0.001,
                borderpad=0.1, handlelength=0.1, handletextpad=0.1
                )
        self.specplotter.axis.add_artist(self.blleg)
        if self.specplotter.autorefresh: self.specplotter.refresh()

    def clearlegend(self):
        if self.blleg is not None:
            self.blleg.set_visible(False)
            if self.blleg in self.specplotter.axis.artists:
                self.specplotter.axis.artists.remove(self.blleg)
        if self.specplotter.autorefresh: self.specplotter.refresh()

    def savefit(self):
        if self.baselinepars is not None:
            for ii,p in enumerate(self.baselinepars):
                self.specplotter.header.update('BLCOEF%0.2i' % (ii),p,comment="Baseline power-law best-fit coefficient x^%i" % (self.order-ii-1))

class Specfit:

    def __init__(self,specplotter):
        self.model = None
        self.modelpars = None
        self.modelerrs = None
        self.modelplot = None
        self.guessplot = []
        self.fitregion = []
        self.ngauss = 0
        self.nclicks_b1 = 0
        self.nclicks_b2 = 0
        self.gx1 = 0
        self.gx2 = specplotter.spectrum.shape[0]
        self.guesses = []
        self.click = 0
        self.fitkwargs = {}
        self.auto = False
        self.autoannotate = True
        self.specplotter = specplotter
        self.gaussleg=None
        self.residuals=None
        self.setfitspec()
        #self.seterrspec()

    def __call__(self, interactive=False, usemoments=True, fitcolor='r',
            multifit=False, guesses=None, annotate=True, save=True,
            **kwargs):
        """
        Fit gaussians to a spectrum

        guesses = [height,amplitude,center,width]
        """

        self.fitcolor = fitcolor
        self.clear()

        self.ngauss = 0
        self.fitkwargs = kwargs
        if interactive:
            print "Left-click twice to select a fitting range, then middle-click twice to select a peak and width"
            self.nclicks_b1 = 0
            self.nclicks_b2 = 0
            self.guesses = []
            self.click = self.specplotter.axis.figure.canvas.mpl_connect('button_press_event',self.makeguess)
            self.autoannotate = annotate
        elif multifit:
            if guesses is None:
                print "You must input guesses when using multifit.  Also, baseline first!"
            else:
                self.guesses = guesses
                self.multifit()
                self.autoannotate = annotate
        else:
            #print "Non-interactive, 1D fit with automatic guessing"
            if self.specplotter.baseline.order is None:
                self.specplotter.baseline.order=0
                self.onedfit(usemoments=usemoments,annotate=annotate,**kwargs)
            else:
                self.onedfit(usemoments=usemoments,annotate=annotate,
                        vheight=False,height=0.0,**kwargs)
            if self.specplotter.autorefresh: self.specplotter.refresh()
        if save: self.savefit()

    def seterrspec(self,usestd=None,useresiduals=True):
        if self.specplotter.errspec is not None and not usestd:
            self.errspec = self.specplotter.errspec
        elif self.residuals is not None and useresiduals:
            self.errspec = ones(self.spectofit.shape[0]) * self.residuals.std()
        else: self.errspec = ones(self.spectofit.shape[0]) * self.spectofit.std()

    def setfitspec(self):
        self.spectofit = copy(self.specplotter.spectrum)
        OKmask = (self.spectofit==self.spectofit)
        self.spectofit[(True-OKmask)] = 0
        self.seterrspec()
        self.errspec[(True-OKmask)] = 1e10

    def multifit(self):
        self.ngauss = len(self.guesses)/3
        self.setfitspec()
        if self.fitkwargs.has_key('negamp'): self.fitkwargs.pop('negamp')
        mpp,model,mpperr,chi2 = gaussfitter.multigaussfit(
                self.specplotter.vind[self.gx1:self.gx2],
                self.spectofit[self.gx1:self.gx2],
                err=self.errspec[self.gx1:self.gx2],
                ngauss=self.ngauss,
                params=self.guesses,
                **self.fitkwargs)
        self.chi2 = chi2
        self.dof  = self.gx2-self.gx1-self.ngauss*3
        self.model = model
        self.modelpars = mpp.tolist()
        self.modelerrs = mpperr.tolist()
        self.modelplot = self.specplotter.axis.plot(
                self.specplotter.vind[self.gx1:self.gx2],
                self.model+self.specplotter.offset, color=self.fitcolor, linewidth=0.5)
        self.residuals = self.spectofit[self.gx1:self.gx2] - self.model
        if self.autoannotate:
            self.annotate()

    def onedfit(self, usemoments=True, annotate=True, vheight=True, height=0, negamp=None,**kwargs):
        self.ngauss = 1
        self.auto = True
        self.setfitspec()
        if usemoments: # this can be done within gaussfit but I want to save them
            self.guesses = gaussfitter.onedmoments(
                    self.specplotter.vind[self.gx1:self.gx2],
                    self.spectofit[self.gx1:self.gx2],
                    vheight=vheight,negamp=negamp,**kwargs)
            if vheight is False: self.guesses = [height]+self.guesses
        else:
            if negamp: self.guesses = [height,-1,0,1]
            else:  self.guesses = [height,1,0,1]
        mpp,model,mpperr,chi2 = gaussfitter.onedgaussfit(
                self.specplotter.vind[self.gx1:self.gx2],
                self.spectofit[self.gx1:self.gx2],
                err=self.errspec[self.gx1:self.gx2],
                vheight=vheight,
                params=self.guesses,
                **self.fitkwargs)
        self.chi2 = chi2
        self.dof  = self.gx2-self.gx1-self.ngauss*3-vheight
        if vheight:
            self.specplotter.baseline.baselinepars = mpp[:1] # first item in list form
            self.model = model - mpp[0]
        else: self.model = model
        self.residuals = self.spectofit[self.gx1:self.gx2] - self.model
        self.modelpars = mpp[1:].tolist()
        self.modelerrs = mpperr[1:].tolist()
        self.modelplot = self.specplotter.axis.plot(
                self.specplotter.vind[self.gx1:self.gx2],
                self.model+self.specplotter.offset, color=self.fitcolor, linewidth=0.5)
        if annotate:
            self.annotate()
            if vheight: self.specplotter.baseline.annotate()

    def fullsizemodel(self):
        """
        If the gaussian was fit to a sub-region of the spectrum,
        expand it (with zeros) to fill the spectrum.  You can
        always recover the original by:
        origmodel = model[gx1:gx2]
        """

        if self.model.shape != self.specplotter.spectrum.shape:
            temp = zeros(self.specplotter.spectrum.shape)
            temp[self.gx1:self.gx2] = self.model
            self.model = temp
            self.residuals = self.spectofit - self.model

    def plotresiduals(self,fig=None,axis=None,clear=True,**kwargs):
        """
        Plot residuals of the fit.  Specify a figure or
        axis; defaults to figure(2).

        kwargs are passed to matplotlib plot
        """
        if axis is None:
            fig=figure(2)
            self.residualaxis = gca()
            if clear: self.residualaxis.clear()
        else:
            self.residualaxis = axis
            if clear: self.residualaxis.clear()
        self.residualplot = self.residualaxis.plot(self.specplotter.vind[self.gx1:self.gx2],
                self.residuals,drawstyle='steps-mid',
                linewidth=0.5, color='k', **kwargs)
        if self.specplotter.vmin is not None and self.specplotter.vmax is not None:
            self.residualaxis.set_xlim(self.specplotter.vmin,self.specplotter.vmax)
        self.residualaxis.figure.canvas.draw()

    def annotate(self,loc='upper right'):
        #text(xloc,yloc     ,"c=%g" % self.modelpars[1],transform = self.specplotter.axis.transAxes)
        #text(xloc,yloc-0.05,"w=%g" % self.modelpars[2],transform = self.specplotter.axis.transAxes)
        #text(xloc,yloc-0.10,"a=%g" % self.modelpars[0],transform = self.specplotter.axis.transAxes)
        self.clearlegend()
        pl = matplotlib.collections.CircleCollection([0],edgecolors=['k'])
        self.gaussleg = self.specplotter.axis.legend(
                tuple([pl]*3*self.ngauss),
                tuple(flatten(
                    [("c%i=%6.4g $\\pm$ %6.4g" % (jj,self.modelpars[1+jj*3],self.modelerrs[1+jj*3]),
                      "w%i=%6.4g $\\pm$ %6.4g" % (jj,self.modelpars[2+jj*3],self.modelerrs[2+jj*3]),
                      "a%i=%6.4g $\\pm$ %6.4g" % (jj,self.modelpars[0+jj*3],self.modelerrs[0+jj*3]))
                      for jj in range(self.ngauss)])),
                loc=loc,markerscale=0.01,
                borderpad=0.1, handlelength=0.1, handletextpad=0.1
                )
        self.gaussleg.draggable(True)
        self.specplotter.axis.add_artist(self.gaussleg)
        if self.specplotter.autorefresh: self.specplotter.refresh()

    def selectregion(self,event):
        if self.nclicks_b1 == 0:
            self.gx1 = argmin(abs(event.xdata-self.specplotter.vind))
            self.nclicks_b1 += 1
        elif self.nclicks_b1 == 1:
            self.gx2 = argmin(abs(event.xdata-self.specplotter.vind))
            self.nclicks_b1 -= 1
            if self.gx1 > self.gx2: self.gx1,self.gx2 = self.gx2,self.gx1
            if abs(self.gx1-self.gx2) > 3: # can't fit w/ fewer data than pars
                self.fitregion = self.specplotter.axis.plot(
                        self.specplotter.vind[self.gx1:self.gx2],
                        self.specplotter.spectrum[self.gx1:self.gx2]+self.specplotter.offset,
                        drawstyle='steps-mid',
                        color='c')
                if self.guesses == []:
                    self.guesses = gaussfitter.onedmoments(
                            self.specplotter.vind[self.gx1:self.gx2],
                            self.spectofit[self.gx1:self.gx2],
                            vheight=0)
                    self.ngauss = 1
                    self.auto = True
            else:
                print "Fitting region is too small (channels %i:%i).  Try again." % (self.gx1,self.gx2)

    def guesspeakwidth(self,event):
        if self.nclicks_b2 % 2 == 0:
            if self.auto:
                self.guesses[:2] = [event.ydata,event.xdata]
            else:
                self.guesses += [event.ydata,event.xdata,1]
                self.ngauss += 1
            self.nclicks_b2 += 1
            self.guessplot += [self.specplotter.axis.scatter(event.xdata,event.ydata,marker='x',c='r')]
        elif self.nclicks_b2 % 2 == 1:
            self.guesses[-1] = abs(event.xdata-self.guesses[-2])
            self.nclicks_b2 += 1
            self.guessplot += self.specplotter.axis.plot([event.xdata,
                2*self.guesses[-2]-event.xdata],[event.ydata]*2,
                color='r')
            if self.auto:
                self.auto = False
            if self.nclicks_b2 / 2 > self.ngauss:
                print "There have been %i middle-clicks but there are only %i gaussians" % (self.nclicks_b2,self.ngauss)
                self.ngauss += 1

    def clear(self,legend=True):
        if self.modelplot is not None:
            for p in self.modelplot:
                p.set_visible(False)
        if legend: self.clearlegend()

    def makeguess(self,event):
        if event.button == 1:
            self.selectregion(event)
        elif event.button == 2:
            self.guesspeakwidth(event)
        elif event.button == 3:
            disconnect(self.click)
            if self.ngauss > 0:
                print len(self.guesses)/3," Guesses: ",self.guesses," X channel range: ",self.gx1,self.gx2
                if len(self.guesses) % 3 == 0:
                    self.multifit()
                    for p in self.guessplot + self.fitregion:
                        p.set_visible(False)
                else:
                    print "error, wrong # of pars"
        if self.specplotter.autorefresh: self.specplotter.refresh()

    def clearlegend(self):
        if self.gaussleg is not None:
            self.gaussleg.set_visible(False)
            if self.gaussleg in self.specplotter.axis.artists:
                self.specplotter.axis.artists.remove(self.gaussleg)
        if self.specplotter.autorefresh: self.specplotter.refresh()

    def savefit(self):
        if self.modelpars is not None:
            for ii,p in enumerate(self.modelpars):
                if ii % 3 == 0: self.specplotter.header.update('AMP%1i' % (ii/3),p,comment="Gaussian best fit amplitude #%i" % (ii/3))
                if ii % 3 == 1: self.specplotter.header.update('CEN%1i' % (ii/3),p,comment="Gaussian best fit center #%i" % (ii/3))
                if ii % 3 == 2: self.specplotter.header.update('WID%1i' % (ii/3),p,comment="Gaussian best fit width #%i" % (ii/3))

def mapplot(plane,cube,vconv=lambda x: x,xtora=lambda x: x,ytodec=lambda x: x, gaiafignum=0, specfignum=1):

    gaiafig = figure(gaiafignum)
    gaiafig.clf()
    gaiaax = gaiafig.add_subplot(111)
    gaiaax.imshow(plane)

    sp = SpecPlotter(cube, vconv=vconv, xtora=xtora, ytodec=ytodec,
            gaiafignum=gaiafignum, fignum=specfignum, gaiafig=gaiafig)
    sp.clickid = gaiafig.canvas.mpl_connect('button_press_event',sp)
    #connect('button_press_event',sp)

def splat_3d(filename,xi=0,yi=0,vmin=None,vmax=None,button=1,dobaseline=False,exclude=None,
        smooth=None,smoothto=None,smoothtype='gaussian',order=1,savepre=None,**kwargs):
    """
    Inputs:
        vmin,vmax - range over which to baseline and plottransform = ax.transAxes
        exclude - (internal) range to exclude from baseline fit
    """
    dv,v0,p3,hdr,cube,xtora,ytodec,vconv,xunits,conversion_factor,units = open_3d(filename)

    if units is None: units="UNITS"
    if xunits is None: xunits="km/s"
    if conversion_factor == 0 or conversion_factor is None: conversion_factor=1.0
    sp = splat_1d(vpars=[dv, v0, p3], hdr=hdr, spec=cube[:, yi, xi],
            xtora=xtora, ytodec=ytodec, vconv=vconv, units=units,
            conversion_factor=conversion_factor, xunits=xunits, **kwargs)

    sp.cube = cube

    return sp

def gaia(filename,estimator='max',axis=0):
    f = pyfits.open(filename)
    hdr = f[0].header
    cube = f[0].data
    dv,v0,p3 = hdr.get('CD3_3'),hdr.get('CRVAL3'),hdr.get('CRPIX3')
    dr,r0,p1 = hdr.get('CD1_1'),hdr.get('CRVAL1'),hdr.get('CRPIX1')
    dd,d0,p2 = hdr.get('CD2_2'),hdr.get('CRVAL2'),hdr.get('CRPIX2')
    if dv is None: dv = hdr.get('CDELT3')
    if dr is None: dr = hdr.get('CDELT1')
    if dd is None: dd = hdr.get('CDELT2')
    xtora = lambda x: (x-p1+1)*dr+r0    # convert pixel coordinates to RA/Dec/Velocity
    ytodec = lambda y: (y-p2+1)*dd+d0
    vconv = lambda v: (v-p3+1)*dv+v0

    if axis > 0:
        cube = cube.swapaxes(0,axis)

    if estimator == 'max':
        p = where(isnan(cube),0,cube).max(axis=0)
    elif estimator == 'int':
        p = where(isnan(cube),0,cube).sum(axis=0) * dv
    elif estimator == 'intdivmax':
        cut = MAD(cube.ravel()) + nanmedian(cube.ravel())
        if cut < 0:
            cut = 0
        m = where(isnan(cube),0,cube).max(axis=0)
        i = where(isnan(cube),0,cube).sum(axis=0) * dv
        p = where(i<0,0,i)/where(m<=cut,numpy.inf,m)
    elif estimator[-5:] == ".fits":
        p = pyfits.open(estimator)[0].data

    mapplot(p,cube,vconv,xtora,ytodec)

def baseline_file(filename,outfilename,vmin=None,vmax=None,order=1,crop=False):
    f = pyfits.open(filename)
    hdr = f[0].header
    cube = f[0].data.squeeze()
    dv,v0,p3 = hdr.get('CD3_3'),hdr.get('CRVAL3'),hdr.get('CRPIX3')
    dr,r0,p1 = hdr.get('CD1_1'),hdr.get('CRVAL1'),hdr.get('CRPIX1')
    dd,d0,p2 = hdr.get('CD2_2'),hdr.get('CRVAL2'),hdr.get('CRPIX2')
    if dv is None: dv = hdr.get('CDELT3')
    if dr is None: dr = hdr.get('CDELT1')
    if dd is None: dd = hdr.get('CDELT2')
    vconv = lambda v: (v-p3+1)*dv+v0
    varr = vconv(arange(cube.shape[-1]))
    if vmin is None: argvmin = None
    else: argvmin = argmin(abs(varr-vmin))
    if vmax is None: argvmax = None
    else: argvmax = argmin(abs(varr-vmax))

    bspec,bfit = baseline(cube,vmin=argvmin,vmax=argvmax,order=order)

def baseline(spectrum,xarr=None,xmin='default',xmax='default',order=1,quiet=True,exclude=None,
        mask=None):
    """
    Subtract a baseline from a spectrum
    If xmin,xmax are not specified, defaults to ignoring first and last 10% of spectrum
    *unless* order > 1, in which case ignoring the ends tends to cause strange effects

    exclude is a set of start/end indices to ignore when baseline fitting
    (ignored by setting error to infinite in fitting procedure)
    """
    if xmin == 'default':
        if order <= 1: xmin = floor( spectrum.shape[-1]*0.1 )
        else:          xmin = 0
    elif xmin is None:
        xmin = 0
    if xmax == 'default':
        if order <= 1: xmax = ceil( spectrum.shape[-1]*0.9 )
        else:          xmax = spectrum.shape[-1]
    elif xmax is None:
        xmax = spectrum.shape[-1]

    pguess = [1]*(order+1)

    if xarr is None:
        xarr = indices(spectrum.shape).squeeze()

    subxarr = xarr[xmin:xmax]
    def mpfitfun(data,err):
        def f(p,fjac=None): return [0,numpy.ravel((poly1d(p)(subxarr)-data)/err)]
        return f

    err = ones(spectrum.shape)
    if exclude is not None:
        err[exclude[0]:exclude[1]] = 1e10
    if mask is not None:
        if mask.dtype.name != 'bool': mask = mask.astype('bool')
        err[mask] = 1e10
        spectrum[mask] = 0
    if (spectrum!=spectrum).sum() > 0:
        print "There is an error in baseline: some values are NaN"
        import pdb; pdb.set_trace()

    mp = mpfit(mpfitfun(spectrum[xmin:xmax],err[xmin:xmax]),xall=pguess,quiet=quiet)
    fitp = mp.params
    bestfit = poly1d(fitp)(xarr).squeeze()

    return (spectrum-bestfit),fitp

def open_3d(filename):
    f = pyfits.open(filename)
    hdr = f[0].header
    cube = f[0].data
    if len(cube.shape) == 4: cube=cube[0,:,:,:]
    #cube = reshape(cube.mean(axis=2).mean(axis=1),[cube.shape[0],1,1])
    dv,v0,p3 = hdr.get('CD3_3'),hdr.get('CRVAL3'),hdr.get('CRPIX3')
    dr,r0,p1 = hdr.get('CD1_1'),hdr.get('CRVAL1'),hdr.get('CRPIX1')
    dd,d0,p2 = hdr.get('CD2_2'),hdr.get('CRVAL2'),hdr.get('CRPIX2')
    if dv is None: dv = hdr.get('CDELT3')
    if dr is None: dr = hdr.get('CDELT1')
    if dd is None: dd = hdr.get('CDELT2')
    xtora = lambda x: (x-p1+1)*dr+r0    # convert pixel coordinates to RA/Dec/Velocity
    ytodec = lambda y: (y-p2+1)*dd+d0
    vconv = lambda v: (v-p3+1)*dv+v0
    if hdr.get('CUNIT3') in ['m/s','M/S']:
        conversion_factor = 1000.0
        xunits = 'km/s' # change to km/s because you're converting units
    else:
        xunits = hdr.get('CUNIT3')
        if xunits in ("hz","Hz"):
            print "Converting from Hz to GHz"
            xunits = "GHz"
            conversion_factor = 1.0e9
        else:
            conversion_factor = 1.0
    units = hdr.get('BUNIT')

    return dv,v0,p3,hdr,cube,xtora,ytodec,vconv,xunits,conversion_factor,units

def open_1d(filename,specnum=0,wcstype='',errspecnum=None,maskspecnum=None):
    """
    Grabs all the relevant pieces of a 1d spectrum for plotting
    wcstype is the suffix on the WCS type to get to velocity/frequency/whatever
    """
    f = pyfits.open(filename)
    hdr = f[0].header
    spec = f[0].data
    errspec  = None
    maskspec = None
    if hdr.get('NAXIS') == 2:
        if errspecnum is not None:
            errspec = spec[errspecnum,:]
        if maskspecnum is not None:
            maskspec = spec[maskspecnum,:]
        if isinstance(specnum,list):
            spec = spec[specnum,:].mean(axis=0)
        elif isinstance(specnum,int):
            spec = spec[specnum,:]
        else:
            raise TypeError("Specnum is of wrong type (not a list of integers or an integer).  Type: %s" %
                    str(type(specnum)))
    elif hdr.get('NAXIS') > 2:
        raise ValueError("Too many axes for open_1d (splat_1d) - use cube instead")
    if hdr.get('CD1_1'+wcstype):
        dv,v0,p3 = hdr['CD1_1'+wcstype],hdr['CRVAL1'+wcstype],hdr['CRPIX1'+wcstype]
    else:
        dv,v0,p3 = hdr['CDELT1'+wcstype],hdr['CRVAL1'+wcstype],hdr['CRPIX1'+wcstype]
    if hdr.get('OBJECT'):
        specname = hdr.get('OBJECT')
    elif hdr.get('GLON') and hdr.get('GLAT'):
        specname = "%s %s" % (hdr.get('GLON'),hdr.get('GLAT'))
    else:
        specname = filename.rstrip(".fits")
    if hdr.get('CUNIT1'+wcstype) in ['m/s','M/S']:
        conversion_factor = 1000.0
        xunits = 'km/s' # change to km/s because you're converting units
    else:
        xunits = hdr.get('CUNIT1'+wcstype)
        if xunits in ("hz","Hz"):
            print "Converting from Hz to GHz"
            xunits = "GHz"
            conversion_factor = 1.0e9
        else:
            conversion_factor = 1.0
    vconv = lambda v: ((v-p3+1)*dv+v0)/conversion_factor
    xtora=None
    ytodec=None
    units = hdr.get('BUNIT').strip()
    if hdr.get('CTYPE1'+wcstype):
        xtype = hdr.get('CTYPE1'+wcstype)
    else:
        xtype = 'VLSR'
    if hdr.get('REFFREQ'+wcstype):
        reffreq = hdr.get('REFFREQ'+wcstype)
    else:
        reffreq = None

    return dv,v0,p3,conversion_factor,hdr,spec,vconv,xtora,ytodec,specname,units,xunits,errspec,maskspec,reffreq

def splat_1d(filename=None,vmin=None,vmax=None,button=1,dobaseline=False,
        exclude=None,smooth=None,order=1,savepre=None,vcrop=True,
        vconv=None,vpars=None,hdr=None,spec=None,xtora=None,ytodec=None,
        specname=None,quiet=True,specnum=0,errspecnum=None,wcstype='',
        offset=0.0, continuum=0.0, annotatebaseline=False, plotspectrum=True,
        smoothto=None, xunits=None, units=None, conversion_factor=None,
        smoothtype='gaussian',convmode='valid',maskspecnum=None,**kwargs):
    """
    Wrapper for specplotter creation.  Works nicely with 1D spectra with well-defined
    FITS headers (i.e., CRVAL1, CRPIX1, CDELT1, and optionally CUNIT1 and CTYPE1)

    This documentation needs to be updated a lot... I implemented a lot of features
    without documenting them, which was a mistake

    Inputs:
        vmin,vmax - range over which to baseline and plot
        exclude - (internal) range to exclude from baseline fit
        vcrop - will vmin/vmax crop out data, or just set the plot limits?
    """
    if (vpars and vconv and hdr and spec is not None and xtora and ytodec
            and units and xunits and conversion_factor):
        dv,v0,p3 = vpars
        errspec = None
        maskspec = None
        reffreq = None
        if units is None and kwargs.has_key('units'): units = kwargs.pop('units')
    else:
        dv,v0,p3,conversion_factor,hdr,spec,vconv,xtora,ytodec,specname_file,units,xunits,errspec,maskspec,reffreq = \
                open_1d(filename,specnum=specnum,wcstype=wcstype,errspecnum=errspecnum,maskspecnum=maskspecnum)
        if specname is None: specname=specname_file
        if units is None and kwargs.has_key('units'): units = kwargs.pop('units')

    if type(continuum)==type('str'):
        if hdr.get(continuum) is not None:
            continuum = hdr.get(continuum)
        else:
            raise ValueError("Continuum specified but none present.")

    varr = vconv(arange(spec.shape[0]))
    if vmin is None or vcrop==False: argvmin = 0
    else:
      argvmin = argmin(abs(varr-vmin))
      if dv > 0:
        hdr.update('CRPIX1'+wcstype,p3-argvmin)
    if vmax is None or vcrop==False: argvmax = spec.shape[0]
    else:
      argvmax = argmin(abs(varr-vmax))
      if dv < 0:
        hdr.update('CRPIX1'+wcstype,p3-argvmax)

    if argvmin > argvmax:
        argvmin,argvmax = argvmax,argvmin
        #if exclude is not None: exclude = exclude[::-1]
    elif argvmin == argvmax:
        raise Exception("Error: no data in velocity range %g:%g for source %s."
                % (vmin,vmax,filename))

    # these lines were meant to automatically put "exclude" into velocity
    # units; this is now done in the baseline code
    #if exclude is not None:
    #    exclude[0] = argmin(abs(varr-exclude[0]))
    #    exclude[1] = argmin(abs(varr-exclude[1]))
    #    exclude = array(exclude) - argvmin

    vconv = lambda v: ((v-p3+argvmin+1)*dv+v0) / conversion_factor
    ivconv = lambda V: p3-1-argvmin+(V*conversion_factor-v0)/dv

    specplot = spec[argvmin:argvmax]
    if errspec is not None: errspec=errspec[argvmin:argvmax]
    if maskspec is not None: maskspec=maskspec[argvmin:argvmax]

    if smoothto:
        smooth = abs(smoothto/dv)

    if smooth:
        roundsmooth = round(smooth) # can only downsample by integers
        # change fitter first
        if smoothtype == 'hanning':
            specplot = convolve(specplot,hanning(2+roundsmooth)/hanning(2+roundsmooth).sum(),convmode)[::roundsmooth]
            kernsize = smooth
            ones_sameshape = zeros(smooth+2)
            ones_sameshape[1:-1] = 1
        elif smoothtype == 'boxcar':
            specplot = convolve(specplot,ones(roundsmooth)/float(roundsmooth),convmode)[::roundsmooth]
            kernsize = roundsmooth
            ones_sameshape = ones(roundsmooth)
        elif smoothtype == 'gaussian':
            speclen = specplot.shape[0]
            xkern  = linspace(-1*smooth,smooth,smooth*3)
            kernel = exp(-xkern**2/(2*(smooth/sqrt(8*log(2)))**2))
            kernel /= kernel.sum()
            kernsize = len(kernel)
            specplot = convolve(specplot,kernel,convmode)[::roundsmooth]
            ones_sameshape = zeros(roundsmooth*3)
            ones_sameshape[roundsmooth:-roundsmooth] = 1
        if errspec is not None:
            errspec = sqrt(convolve(errspec**2,ones_sameshape,convmode)[::roundsmooth]) / float(roundsmooth)
        if maskspec is not None:
            maskspec = array(convolve(maskspec,ones_sameshape,convmode)[::roundsmooth],dtype='bool')
            if maskspec.shape != specplot.shape: import pdb; pdb.set_trace()
        # this bit of code may also make sense, but I'm shifting the center pixel instead
        # b/c it's easier (?) to deal with velocity range
        #v0 += (abs(dv)*smooth - abs(dv))/2.0 # pixel center moves by half the original pixel size
        dv *= roundsmooth
        if convmode == 'same':
          newrefpix = (p3-argvmin)/roundsmooth
        elif convmode == 'full':
          newrefpix = (p3-0.5-argvmin+kernsize/2.0)/roundsmooth
        elif convmode == 'valid':
          newrefpix = (p3-0.5-argvmin-kernsize/2.0)/roundsmooth
        # this was resolved by advanced guess-and check
        # but also, sort of makes sense: FITS refers to the *center* of a pixel.  You want to
        # shift 1/2 pixel to the right so that the first pixel goes from 0 to 1
        vconv = lambda v: ((v-newrefpix)*dv+v0)/conversion_factor
        ivconv = lambda V: newrefpix+(V*conversion_factor-v0)/dv
        hdr.update('CRPIX1'+wcstype,newrefpix+1)
        hdr.update('CDELT1'+wcstype,dv)

    sp = SpecPlotter(specplot, vconv=vconv, xtora=xtora, ytodec=ytodec,
            specname=specname, dv=dv/conversion_factor, hdr=hdr, reffreq=reffreq,
            errspec=errspec, maskspec=maskspec, xunits=xunits, **kwargs)

    if plotspectrum:
        sp.plotspec(button=button, cube=False, vmin=vmin, vmax=vmax,
                units=units, offset=offset, continuum=continuum,
                **kwargs)

    if dobaseline:
        sp.baseline(exclude=exclude,order=order,quiet=quiet,annotate=annotatebaseline)
    if plotspectrum: sp.refresh()

    if hdr.get('GLON') and hdr.get('GLAT'):
        sp.glon = hdr.get('GLON')
        sp.glat = hdr.get('GLAT')

    if savepre is not None:
        glon,glat = sp.glon,sp.glat
        if glat < 0: pm=""
        else: pm = "+"
        savename = savepre + "G%07.3f%0s%07.3f_" % (glon,pm,glat) + hdr.get('MOLECULE').replace(' ','') + hdr.get('TRANSITI').replace(' ','')
        savefig(savename+'.png')

    return sp

def splat_tspec(filename,specnum=0,**kwargs):
    """
    Same as splat_1d for tspec data
    """

    tdata = pyfits.getdata(filename)
    theader = pyfits.getheader(filename)
    if len(tdata.shape) == 3:
        tdata = tdata[specnum,:,:]
    wavelength = tdata[0,:]
    spectrum   = tdata[1,:]
    error      = tdata[2,:]

    vconv = lambda x: wavelength[x]
    ivconv = lambda x: argmin(abs(wavelength-x))

    specname='TSPEC'
    dv = median(wavelength[1:] - wavelength[:-1])

    sp = SpecPlotter(spectrum,vconv=vconv,specname=specname,dv=dv,hdr=theader)

    sp.plotspec(cube=False,units=theader.get('YUNITS'),xunits=theader.get('XUNITS'),**kwargs)

    return sp


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for learn.dataframe.tensorflow_dataframe."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import csv
import math
import tempfile

import numpy as np

from tensorflow.contrib.learn.python.learn.dataframe import tensorflow_dataframe as df
from tensorflow.contrib.learn.python.learn.dataframe.transforms import densify
from tensorflow.core.example import example_pb2
from tensorflow.python.framework import dtypes
from tensorflow.python.lib.io import tf_record
from tensorflow.python.ops import parsing_ops
from tensorflow.python.platform import test

# pylint: disable=g-import-not-at-top
try:
  import pandas as pd
  HAS_PANDAS = True
except ImportError:
  HAS_PANDAS = False

def _assert_df_equals_dict(expected_df, actual_dict):
  for col in expected_df:
    if expected_df[col].dtype in [np.float32, np.float64]:
      assertion = np.testing.assert_allclose
    else:
      assertion = np.testing.assert_array_equal

    if expected_df[col].dtype.kind in ["O", "S", "U"]:
      # Python 2/3 compatibility
      # TensorFlow always returns bytes, so we just convert the unicode
      # expectations to bytes also before comparing.
      expected_values = [x.encode("utf-8") for x in expected_df[col].values]
    else:
      expected_values = expected_df[col].values

    assertion(
        expected_values,
        actual_dict[col],
        err_msg="Expected {} in column '{}'; got {}.".format(expected_values,
                                                             col,
                                                             actual_dict[col]))

def _make_test_csv():
  f = tempfile.NamedTemporaryFile(
      dir=test.get_temp_dir(), delete=False, mode="w")
  w = csv.writer(f)
  w.writerow(["int", "float", "bool", "string"])
  for _ in range(100):
    intvalue = np.random.randint(-10, 10)
    floatvalue = np.random.rand()
    boolvalue = int(np.random.rand() > 0.3)
    stringvalue = "S: %.4f" % np.random.rand()

    row = [intvalue, floatvalue, boolvalue, stringvalue]
    w.writerow(row)
  f.close()
  return f.name

def _make_test_csv_sparse():
  f = tempfile.NamedTemporaryFile(
      dir=test.get_temp_dir(), delete=False, mode="w")
  w = csv.writer(f)
  w.writerow(["int", "float", "bool", "string"])
  for _ in range(100):
    # leave columns empty; these will be read as default value (e.g. 0 or NaN)
    intvalue = np.random.randint(-10, 10) if np.random.rand() > 0.5 else ""
    floatvalue = np.random.rand() if np.random.rand() > 0.5 else ""
    boolvalue = int(np.random.rand() > 0.3) if np.random.rand() > 0.5 else ""
    stringvalue = (("S: %.4f" % np.random.rand()) if np.random.rand() > 0.5 else
                   "")

    row = [intvalue, floatvalue, boolvalue, stringvalue]
    w.writerow(row)
  f.close()
  return f.name

def _make_test_tfrecord():
  f = tempfile.NamedTemporaryFile(dir=test.get_temp_dir(), delete=False)
  w = tf_record.TFRecordWriter(f.name)
  for i in range(100):
    ex = example_pb2.Example()
    ex.features.feature["var_len_int"].int64_list.value.extend(range((i % 3)))
    ex.features.feature["fixed_len_float"].float_list.value.extend(
        [float(i), 2 * float(i)])
    w.write(ex.SerializeToString())
  return f.name

class TensorFlowDataFrameTestCase(test.TestCase):
  """Tests for `TensorFlowDataFrame`."""

  def _assert_pandas_equals_tensorflow(self, pandas_df, tensorflow_df,
                                       num_batches, batch_size):
    self.assertItemsEqual(
        list(pandas_df.columns) + ["index"], tensorflow_df.columns())

    for batch_num, batch in enumerate(tensorflow_df.run(num_batches)):
      row_numbers = [
          total_row_num % pandas_df.shape[0]
          for total_row_num in range(batch_size * batch_num, batch_size * (
              batch_num + 1))
      ]
      expected_df = pandas_df.iloc[row_numbers]
      _assert_df_equals_dict(expected_df, batch)

  def testInitFromPandas(self):
    """Test construction from Pandas DataFrame."""
    if not HAS_PANDAS:
      return
    pandas_df = pd.DataFrame({"sparrow": range(10), "ostrich": 1})
    tensorflow_df = df.TensorFlowDataFrame.from_pandas(
        pandas_df, batch_size=10, shuffle=False)

    batch = tensorflow_df.run_one_batch()

    np.testing.assert_array_equal(pandas_df.index.values, batch["index"],
                                  "Expected index {}; got {}".format(
                                      pandas_df.index.values, batch["index"]))
    _assert_df_equals_dict(pandas_df, batch)

  def testBatch(self):
    """Tests `batch` method.

    `DataFrame.batch()` should iterate through the rows of the
    `pandas.DataFrame`, and should "wrap around" when it reaches the last row.
    """
    if not HAS_PANDAS:
      return
    pandas_df = pd.DataFrame({
        "albatross": range(10),
        "bluejay": 1,
        "cockatoo": range(0, 20, 2),
        "penguin": list("abcdefghij")
    })
    tensorflow_df = df.TensorFlowDataFrame.from_pandas(pandas_df, shuffle=False)

    # Rebatch `df` into the following sizes successively.
    batch_sizes = [4, 7]
    num_batches = 3

    final_batch_size = batch_sizes[-1]

    for batch_size in batch_sizes:
      tensorflow_df = tensorflow_df.batch(batch_size, shuffle=False)

    self._assert_pandas_equals_tensorflow(
        pandas_df,
        tensorflow_df,
        num_batches=num_batches,
        batch_size=final_batch_size)

  def testFromNumpy(self):
    x = np.eye(20)
    tensorflow_df = df.TensorFlowDataFrame.from_numpy(x, batch_size=10)
    for batch in tensorflow_df.run(30):
      for ind, val in zip(batch["index"], batch["value"]):
        expected_val = np.zeros_like(val)
        expected_val[ind] = 1
        np.testing.assert_array_equal(expected_val, val)

  def testFromCSV(self):
    if not HAS_PANDAS:
      return
    num_batches = 100
    batch_size = 8
    enqueue_size = 7

    data_path = _make_test_csv()
    default_values = [0, 0.0, 0, ""]

    pandas_df = pd.read_csv(data_path)
    tensorflow_df = df.TensorFlowDataFrame.from_csv(
        [data_path],
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        shuffle=False,
        default_values=default_values)
    self._assert_pandas_equals_tensorflow(
        pandas_df,
        tensorflow_df,
        num_batches=num_batches,
        batch_size=batch_size)

  def testFromCSVLimitEpoch(self):
    batch_size = 8
    num_epochs = 17
    expected_num_batches = (num_epochs * 100) // batch_size

    data_path = _make_test_csv()
    default_values = [0, 0.0, 0, ""]

    tensorflow_df = df.TensorFlowDataFrame.from_csv(
        [data_path],
        batch_size=batch_size,
        shuffle=False,
        default_values=default_values)
    result_batches = list(tensorflow_df.run(num_epochs=num_epochs))
    actual_num_batches = len(result_batches)
    self.assertEqual(expected_num_batches, actual_num_batches)

    # TODO(soergel): figure out how to dequeue the final small batch
    expected_rows = 1696  # num_epochs * 100
    actual_rows = sum([len(x["int"]) for x in result_batches])
    self.assertEqual(expected_rows, actual_rows)

  def testFromCSVWithFeatureSpec(self):
    if not HAS_PANDAS:
      return
    num_batches = 100
    batch_size = 8

    data_path = _make_test_csv_sparse()
    feature_spec = {
        "int": parsing_ops.FixedLenFeature(None, dtypes.int16, np.nan),
        "float": parsing_ops.VarLenFeature(dtypes.float16),
        "bool": parsing_ops.VarLenFeature(dtypes.bool),
        "string": parsing_ops.FixedLenFeature(None, dtypes.string, "")
    }

    pandas_df = pd.read_csv(data_path, dtype={"string": object})
    # Pandas insanely uses NaN for empty cells in a string column.
    # And, we can't use Pandas replace() to fix them because nan != nan
    s = pandas_df["string"]
    for i in range(0, len(s)):
      if isinstance(s[i], float) and math.isnan(s[i]):
        pandas_df.set_value(i, "string", "")
    tensorflow_df = df.TensorFlowDataFrame.from_csv_with_feature_spec(
        [data_path],
        batch_size=batch_size,
        shuffle=False,
        feature_spec=feature_spec)

    # These columns were sparse; re-densify them for comparison
    tensorflow_df["float"] = densify.Densify(np.nan)(tensorflow_df["float"])
    tensorflow_df["bool"] = densify.Densify(np.nan)(tensorflow_df["bool"])

    self._assert_pandas_equals_tensorflow(
        pandas_df,
        tensorflow_df,
        num_batches=num_batches,
        batch_size=batch_size)

  def testFromExamples(self):
    num_batches = 77
    enqueue_size = 11
    batch_size = 13

    data_path = _make_test_tfrecord()
    features = {
        "fixed_len_float":
            parsing_ops.FixedLenFeature(
                shape=[2], dtype=dtypes.float32, default_value=[0.0, 0.0]),
        "var_len_int":
            parsing_ops.VarLenFeature(dtype=dtypes.int64)
    }

    tensorflow_df = df.TensorFlowDataFrame.from_examples(
        data_path,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        features=features,
        shuffle=False)

    # `test.tfrecord` contains 100 records with two features: var_len_int and
    # fixed_len_float. Entry n contains `range(n % 3)` and
    # `float(n)` for var_len_int and fixed_len_float,
    # respectively.
    num_records = 100

    def _expected_fixed_len_float(n):
      return np.array([float(n), 2 * float(n)])

    def _expected_var_len_int(n):
      return np.arange(n % 3)

    for batch_num, batch in enumerate(tensorflow_df.run(num_batches)):
      record_numbers = [
          n % num_records
          for n in range(batch_num * batch_size, (batch_num + 1) * batch_size)
      ]
      for i, j in enumerate(record_numbers):
        np.testing.assert_allclose(
            _expected_fixed_len_float(j), batch["fixed_len_float"][i])
        var_len_int = batch["var_len_int"]
      for i, ind in enumerate(var_len_int.indices):
        val = var_len_int.values[i]
        expected_row = _expected_var_len_int(record_numbers[ind[0]])
        expected_value = expected_row[ind[1]]
        np.testing.assert_array_equal(expected_value, val)

  def testSplitString(self):
    batch_size = 8
    num_epochs = 17
    expected_num_batches = (num_epochs * 100) // batch_size

    data_path = _make_test_csv()
    default_values = [0, 0.0, 0, ""]

    tensorflow_df = df.TensorFlowDataFrame.from_csv(
        [data_path],
        batch_size=batch_size,
        shuffle=False,
        default_values=default_values)

    a, b = tensorflow_df.split("string", 0.7)  # no rebatching

    total_result_batches = list(tensorflow_df.run(num_epochs=num_epochs))
    a_result_batches = list(a.run(num_epochs=num_epochs))
    b_result_batches = list(b.run(num_epochs=num_epochs))

    self.assertEqual(expected_num_batches, len(total_result_batches))
    self.assertEqual(expected_num_batches, len(a_result_batches))
    self.assertEqual(expected_num_batches, len(b_result_batches))

    total_rows = sum([len(x["int"]) for x in total_result_batches])
    a_total_rows = sum([len(x["int"]) for x in a_result_batches])
    b_total_rows = sum([len(x["int"]) for x in b_result_batches])

    print("Split rows: %s => %s, %s" % (total_rows, a_total_rows, b_total_rows))

    # TODO(soergel): figure out how to dequeue the final small batch
    expected_total_rows = 1696  # (num_epochs * 100)

    self.assertEqual(expected_total_rows, total_rows)
    self.assertEqual(1087, a_total_rows)  # stochastic but deterministic
    # self.assertEqual(int(total_rows * 0.7), a_total_rows)
    self.assertEqual(609, b_total_rows)  # stochastic but deterministic
    # self.assertEqual(int(total_rows * 0.3), b_total_rows)

    # The strings used for hashing were all unique in the original data, but
    # we ran 17 epochs, so each one should appear 17 times.  Each copy should
    # be hashed into the same partition, so there should be no overlap of the
    # keys.
    a_strings = set([s for x in a_result_batches for s in x["string"]])
    b_strings = set([s for x in b_result_batches for s in x["string"]])
    self.assertEqual(frozenset(), a_strings & b_strings)

if __name__ == "__main__":
  test.main()