# -*- coding: utf-8 -*-
"""
This is a German NER implementation for the dataset GermEval 2014. Also check the folder 'Session 1 - SENNA' for some further information.

This model uses a bi-directional LSTM and runs over an input sentence. The output of the BiLSTM is inputted to a softmax-layer to derive
the probabilities for the different tags.

Single BiLSTM (100 hidden units), after 15 epochs:
Dev-Data: Prec: 0.789, Rec: 0.739, F1: 0.763
Test-Data: Prec: 0.781, Rec: 0.717, F1: 0.747

Stacked BiLSTM (2 layers, 64 hidden states), after 25 epochs:
Dev-Data: Prec: 0.804, Rec: 0.763, F1: 0.783
Test-Data: Prec: 0.795, Rec: 0.738, F1: 0.766

Code was written & tested with:
- Python 2.7
- Theano 0.8.2
- Keras 1.1.1

@author: Nils Reimers
"""
import numpy as np

import random

import time
import gzip
import cPickle as pkl

import BIOF1Validation

import keras
from keras.models import Sequential
from keras.layers.core import *
from keras.layers.wrappers import *
from keras.optimizers import *
from keras.utils import np_utils
from keras.layers.embeddings import Embedding
from keras.layers import LSTM
from docutils.languages.af import labels

f = gzip.open('pkl/embeddings.pkl.gz', 'rb')
embeddings = pkl.load(f)
f.close()

label2Idx = embeddings['label2Idx']
wordEmbeddings = embeddings['wordEmbeddings']
caseEmbeddings = embeddings['caseEmbeddings']

#Inverse label mapping
idx2Label = {v: k for k, v in label2Idx.items()}

f = gzip.open('pkl/data.pkl.gz', 'rb')
train_data = pkl.load(f)
dev_data = pkl.load(f)
test_data = pkl.load(f)
f.close()

#####################################
#
# Create the  Network
#
#####################################

n_out = len(label2Idx)
tokens = Sequential()
tokens.add(Embedding(input_dim=wordEmbeddings.shape[0], output_dim=wordEmbeddings.shape[1],  weights=[wordEmbeddings], trainable=False))

casing = Sequential()
casing.add(Embedding(output_dim=caseEmbeddings.shape[1], input_dim=caseEmbeddings.shape[0], weights=[caseEmbeddings], trainable=False))

model = Sequential();
model.add(Merge([tokens, casing], mode='concat'))
model.add(Bidirectional(LSTM(10, return_sequences=True, dropout_W=0.2)))
model.add(TimeDistributed(Dense(n_out, activation='softmax')))

sgd = SGD(lr=0.1, decay=1e-7, momentum=0.0, nesterov=False, clipvalue=3)
rmsprop = RMSprop(clipvalue=3)
model.compile(loss='sparse_categorical_crossentropy', optimizer=sgd)

model.summary()

##################################
#
# Training of the Network
#
##################################

def iterate_minibatches(dataset, startIdx, endIdx):
    endIdx = min(len(dataset), endIdx)

    for idx in xrange(startIdx, endIdx):
        tokens, casing, labels = dataset[idx]

        labels = np.expand_dims([labels], -1)
        yield labels, np.asarray([tokens]), np.asarray([casing])

def tag_dataset(dataset):
    correctLabels = []
    predLabels = []
    for tokens, casing, labels in dataset:
        tokens = np.asarray([tokens])
        casing = np.asarray([casing])
        pred = model.predict_classes([tokens, casing], verbose=False)[0]
        correctLabels.append(labels)
        predLabels.append(pred)

    return predLabels, correctLabels

number_of_epochs = 20
stepsize = 24000
print "%d epochs" % number_of_epochs

print "%d train sentences" % len(train_data)
print "%d dev sentences" % len(dev_data)
print "%d test sentences" % len(test_data)

for epoch in xrange(number_of_epochs):
    print "--------- Epoch %d -----------" % epoch
    random.shuffle(train_data)
    for startIdx in xrange(0, len(train_data), stepsize):
        start_time = time.time()
        for batch in iterate_minibatches(train_data, startIdx, startIdx+stepsize):
            labels, tokens, casing = batch
            model.train_on_batch([tokens, casing], labels)
        print "%.2f sec for training" % (time.time() - start_time)

        #Train Dataset
        start_time = time.time()
        predLabels, correctLabels = tag_dataset(train_data)
        pre_train, rec_train, f1_train = BIOF1Validation.compute_f1(predLabels, correctLabels, idx2Label)
        print "Train-Data: Prec: %.3f, Rec: %.3f, F1: %.3f" % (pre_train, rec_train, f1_train)

        #Dev Dataset
        predLabels, correctLabels = tag_dataset(dev_data)
        pre_dev, rec_dev, f1_dev = BIOF1Validation.compute_f1(predLabels, correctLabels, idx2Label)
        print "Dev-Data: Prec: %.3f, Rec: %.3f, F1: %.3f" % (pre_dev, rec_dev, f1_dev)

        #Test Dataset
        predLabels, correctLabels = tag_dataset(test_data)
        pre_test, rec_test, f1_test= BIOF1Validation.compute_f1(predLabels, correctLabels, idx2Label)
        print "Test-Data: Prec: %.3f, Rec: %.3f, F1: %.3f" % (pre_test, rec_test, f1_test)

        print "%.2f sec for evaluation" % (time.time() - start_time)
        print ""

2016-11_Seminar/Session 4 - LSTM Sequence Classification/code/NER_BiLSTM.py


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Joe Filippazzo, jfilippazzo@stsci.edu
# !python3
"""
Make nice spectrum objects to pass around SED class
"""
import copy
from functools import wraps, partial
from itertools import groupby
from operator import itemgetter
from multiprocessing import Pool
import os
from pkg_resources import resource_filename

import astropy.constants as ac
import astropy.units as q
import astropy.io.votable as vo
from astropy.io import fits
# from astropy.modeling import fitting, models
from bokeh.plotting import figure, show
from bokeh.models import ColumnDataSource, HoverTool
import numpy as np
from pandas import DataFrame
from svo_filters import Filter

from . import mcmc as mc
from . import utilities as u

def copy_raw(func):
    """A wrapper to copy the raw data to the new Spectrum object"""
    @wraps(func)
    def _copy_raw(*args, **kwargs):
        """Run the function then update the raw data attribute"""
        # Grab the original data
        raw = args[0].raw or args[0].spectrum

        # Run the function and update the raw attribute
        new_spec = func(*args, **kwargs)
        new_spec.raw = raw

        return new_spec

    return _copy_raw

class Spectrum:
    """
    A class to store, calibrate, fit, and plot a single spectrum
    """
    def __init__(self, wave, flux, unc=None, snr=None, trim=None, const=1, phot=False,
                 name=None, ref=None, header=None, verbose=False, **kwargs):
        """Initialize the Spectrum object

        Parameters
        ----------
        wave: astropy.units.quantity.Quantity
            The wavelength array
        flux: astropy.units.quantity.Quantity
            The flux density array
        unc: np.ndarray
            The flux density uncertainty array
        snr: float (optional)
            A value to override spectrum SNR
        trim: dict
            A dictionary of the `trim` method arguments to run
        const: int, float
            A multiplicative factor for the flux
        phot: bool
            Photometric spectrum
        name: str
            A name for the spectrum
        ref: str
            A reference for the data
        header: str
            The header for the spectrum file
        verbose: bool
            Print helpful stuff
        """
        # Meta
        self.verbose = verbose
        self.name = name or 'New Spectrum'
        self.ref = ref
        self.header = header
        self.const = const
        self.phot = phot

        # Make sure the arrays are the same shape
        if not wave.shape == flux.shape and ((unc is None) or not (unc.shape == flux.shape)):
            raise TypeError("Wavelength, flux and uncertainty arrays must be the same shape.")

        # Check wave units are length
        if not u.equivalent(wave, q.um):
            raise TypeError("Wavelength array must be in astropy.units.quantity.Quantity length units, e.g. 'um'")

        # Check flux units are flux density
        if not u.equivalent(flux, u.FLAM):
            raise TypeError("Flux array must be in astropy.units.quantity.Quantity flux density units, e.g. 'erg/s/cm2/A'")

        # Generate uncertainty array
        if unc is None and isinstance(snr, (int, float)):
            unc = flux / snr

        # Make sure the uncertainty array is in the correct units
        if unc is not None:
            if not u.equivalent(unc, u.FLAM):
                raise TypeError("Uncertainty array must be in astropy.units.quantity.Quantity flux density units, e.g. 'erg/s/cm2/A'")

        # Replace negatives, zeros, and infs with nans
        spectrum = [wave, flux]
        spectrum += [unc] if unc is not None else []
        spectrum = u.scrub(spectrum, fill_value=np.nan)

        # Store history info and raw data
        self.history = {}
        self.raw = None

        # Strip and store units
        self._wave_units = wave.unit
        self._flux_units = flux.unit
        self.units = [self._wave_units, self._flux_units]
        self.units += [self._flux_units] if unc is not None else []
        spectrum = [i.value for i in spectrum]

        # Add the data
        self.wave = spectrum[0]
        self._flux = spectrum[1]
        self._unc = None if unc is None else spectrum[2]

        # Store components if added
        self.components = None
        self.best_fit = {}

        # Store kwargs
        for key, val in kwargs.items():
            setattr(self, key, val)

        # Trim
        if trim is not None:
            self.trim(**trim)

    @copy_raw
    def __add__(self, spec):
        """Add the spectra of this and another Spectrum object

        Parameters
        ----------
        spec: sedkit.spectrum.Spectrum
            The spectrum object to add

        Returns
        -------
        sedkit.spectrum.Spectrum
            A new spectrum object with the input spectra stitched together
        """
        # If None is added, just return a copy
        if spec is None:
            return Spectrum(*self.spectrum, name=self.name)

        if not isinstance(type(spec), type(Spectrum)):
            raise TypeError('spec must be sedkit.spectrum.Spectrum')

        # Make spec the same units
        spec.wave_units = self.wave_units
        spec.flux_units = self.flux_units

        # Get the two spectra to stitch
        s1 = self.data
        s2 = spec.data

        # Determine if overlapping
        overlap = True
        try:
            if s1[0][-1] > s1[0][0] > s2[0][-1] or s2[0][-1] > s2[0][0] > s1[0][-1]:
                overlap = False
        except IndexError:
            overlap = False

        # Concatenate and order two segments if no overlap
        if not overlap:

            # Drop uncertainties on both spectra if one is missing
            if self.unc is None or spec.unc is None:
                s1 = [self.wave, self.flux]
                s2 = [spec.wave, spec.flux]

            # Concatenate arrays and sort by wavelength
            new_spec = np.concatenate([s1, s2], axis=1).T
            new_spec = new_spec[np.argsort(new_spec[:, 0])].T

        # Otherwise there are three segments, (left, overlap, right)
        else:

            # Get the left segemnt
            left = s1[:, s1[0] <= s2[0][0]]
            if not np.any(left):
                left = s2[:, s2[0] <= s1[0][0]]

            # Get the right segment
            right = s1[:, s1[0] >= s2[0][-1]]
            if not np.any(right):
                right = s2[:, s2[0] >= s1[0][-1]]

            # Get the overlapping segements
            o1 = s1[:, np.where((s1[0] < right[0][0]) & (s1[0] > left[0][-1]))].squeeze()
            o2 = s2[:, np.where((s2[0] < right[0][0]) & (s2[0] > left[0][-1]))].squeeze()

            # Get the resolutions
            r1 = s1.shape[1] / (max(s1[0]) - min(s1[0]))
            r2 = s2.shape[1] / (max(s2[0]) - min(s2[0]))

            # Make higher resolution s1
            if r1 < r2:
                o1, o2 = o2, o1

            # Interpolate s2 to s1
            o2_flux = np.interp(o1[0], s2[0], s2[1])

            # Get the average
            o_flux = np.nanmean([o1[1], o2_flux], axis=0)

            # Calculate uncertainties if possible
            if len(s2) == len(o1) == 3:
                o2_unc = np.interp(o1[0], s2[0], s2[2])
                o_unc = np.sqrt(o1[2]**2 + o2_unc**2)
                overlap = np.array([o1[0], o_flux, o_unc])
            else:
                overlap = np.array([o1[0], o_flux])
                left = left[:2]
                right = right[:2]

            # Make sure it is 2D
            if overlap.shape == (3,):
                overlap.shape = 3, 1
            if overlap.shape == (2,):
                overlap.shape = 2, 1

            # Concatenate the segments
            new_spec = np.concatenate([left, overlap, right], axis=1)

        # Add units
        new_spec = [i * Q for i, Q in zip(new_spec, self.units)]

        # Make the new spectrum object
        new_spec = Spectrum(*new_spec)

        # Store the components
        new_spec.components = Spectrum(*self.spectrum), spec

        return new_spec

    def best_fit_model(self, modelgrid, report=None, name=None, **kwargs):
        """Perform simple fitting of the spectrum to all models in the given
        modelgrid and store the best fit

        Parameters
        ----------
        modelgrid: sedkit.modelgrid.ModelGrid
            The model grid to fit
        report: str
            The name of the parameter to plot versus the
            Goodness-of-fit statistic
        name: str
            A name for the fit
        """
        # Prepare modelgrid
        modelgrid.wave_units = self.wave_units

        # Prepare data
        name = name or '{} fit'.format(modelgrid.name)
        spectrum = Spectrum(*self.spectrum)
        rows = [row for n, row in modelgrid.index.iterrows()]

        # Iterate over entire model grid
        pool = Pool(8)
        func = partial(fit_model, fitspec=spectrum, resample=not modelgrid.phot)
        fit_rows = pool.map(func, rows)
        pool.close()
        pool.join()

        # Turn the results into a DataFrame and sort
        models = DataFrame(fit_rows)
        models = models.sort_values('gstat')

        # Get the best fit
        bf = copy.copy(models.iloc[0])

        # Make into a dictionary
        bdict = dict(bf)

        # Add full model
        bdict['wave_units'] = self.wave_units
        bdict['flux_units'] = self.flux_units
        bdict['fit_to'] = 'phot' if modelgrid.phot else 'spec'
        full_model = modelgrid.get_spectrum(const=bdict['const'], snr=5, **{par: val for par, val in bdict.items() if par in modelgrid.parameters})
        full_model.phot = True
        bdict['full_model'] = full_model

        self.message(bf[modelgrid.parameters])

        if report is not None:

            # Configure plot
            tools = "pan, wheel_zoom, box_zoom, reset"
            rep = figure(tools=tools, x_axis_label=report, y_axis_label='Goodness-of-fit', plot_width=600, plot_height=400)

            # Single out best fit
            best = ColumnDataSource(data=models.iloc[:1])
            others = ColumnDataSource(data=models.iloc[1:])

            # Add hover tool
            hover = HoverTool(tooltips=[('label', '@label'), ('gstat', '@gstat')])
            rep.add_tools(hover)

            # Plot the fits
            rep.circle(report, 'gstat', source=best, color='red', legend_label=bf['label'])
            rep.circle(report, 'gstat', source=others)

            # Show the plot
            show(rep)

        self.best_fit[name] = bdict

    def convolve_filter(self, filter, **kwargs):
        """
        Convolve the spectrum with a filter

        Parameters
        ----------
        filter: svo_filters.svo.Filter, str
            The filter object or name

        Returns
        -------
        sedkit.spectrum.Spectrum
            The convolved spectrum object
        """
        # Ensure filter object
        if isinstance(filter, str):
            filter = Filter(filter)

        # Get the wavelengths and throughput
        flx, unc = filter.apply(self.spectrum)

        # Make the new spectrum object
        new_spec = Spectrum(filter.wave[0], flx, unc=unc)

        return new_spec

    @property
    def data(self):
        """
        Store the spectrum without units
        """
        if self.unc is None:
            data = np.stack([self.wave, self.flux])
        else:
            data = np.stack([self.wave, self.flux, self.unc])

        return data

    def export(self, filepath, header=None):
        """
        Export the spectrum to file

        Parameters
        ----------
        filepath: str
            The path for the exported file
        """
        # Get target directory
        dirname = os.path.dirname(filepath) or '.'
        name = self.name.replace(' ', '_')

        # Check the parent directory
        if not os.path.exists(dirname):
            raise IOError('{}: No such directory'.format(dirname))

        # Write the file with a header
        head = '{}\nWavelength [{}], Flux Density [{}]'.format(name, self.wave_units, self.flux_units)
        if isinstance(header, str):
            head += '\n{}'.format(header)
        t_data = np.asarray(self.spectrum).T
        np.savetxt(filepath, t_data, header=head)

    def fit(self, spec, weights=None, wave_units=None, scale=True, resample=True, plot=False):
        """Determine the goodness of fit between this and another spectrum

        Parameters
        ----------
        spec: sedkit.spectrum.Spectrum, np.ndarray
            The spectrum object or [W, F] array to fit
        wave_units: astropy.units.quantity.Quantity
            The wavelength units of the input spectrum if
            it is a numpy array
        scale: bool
            Scale spec when measuring the goodness of fit

        Returns
        -------
        tuple
            The fit statistic, and normalization for the fit
        """
        # In case the wavelength units are different
        xnorm = 1
        wav = self.wave

        if hasattr(spec, 'spectrum'):

            # Resample spec onto self wavelength
            spec2 = spec.resamp(self.spectrum[0])
            flx2 = spec2.flux
            err2 = np.ones_like(spec2.flux) if spec2.unc is None else spec2.unc

        elif isinstance(spec, (list, tuple, np.ndarray)):

            spec2 = copy.copy(spec)

            # Convert wave units
            wave_units = wave_units or q.AA
            xnorm = q.Unit(wave_units).to(self.wave_units)
            spec2[0] = spec2[0] * xnorm

            # Resample spec onto self wavelength
            if resample:
                spec2 = u.spectres(self.wave, *spec2)
            wav = spec2[0]
            flx2 = spec2[1]
            err2 = np.ones_like(flx2) if len(spec2) == 2 else spec2[2]

        else:
            raise TypeError("Only an sedkit.spectrum.Spectrum or numpy.ndarray can be fit.")

        # Get the self data
        flx1 = self.flux
        err1 = np.ones_like(flx1) if self.unc is None else self.unc

        # Make default weights the bin widths, excluding gaps in spectra
        if weights is None:
            weights = np.ones_like(wav)
            # weights = np.gradient(wav)
            # weights[weights > np.std(weights)] = 1

        # Run the fitting and get the normalization
        gstat, ynorm = u.goodness(flx1, flx2, err1, err2, weights)

        # Run it again with the scaling removed
        if scale:
            gstat, _ = u.goodness(flx1, flx2 * ynorm, err1, err2 * ynorm, weights)

        if plot:
            fig = self.plot(best_fit=False)
            fig.line(spec.wave, spec.flux * ynorm, legend_label='Fit')
            show(fig)

        return gstat, ynorm, xnorm

    # def fit_blackbody(self, init=8000, epsilon=0.0001, acc=1, maxiter=500, **kwargs):
    #     """
    #     Fit a blackbody spectrum to the spectrum
    #
    #     Returns
    #     -------
    #     int
    #         The best fit blackbody temperature
    #     """
    #     # Determine optimal parameters for data
    #     wav, flx, err = self.spectrum
    #
    #     @models.custom_model
    #     def blackbody(wavelength, temperature=2000):
    #         wavelength *= q.um
    #         temperature *= q.K
    #
    #         bb = models.BlackBody(temperature=temperature)
    #         flux = (bb(wavelength) * q.sr / bb.bolometric_flux.value).to(u.FLAM, q.spectral_density(wavelength)) * 1E-8
    #
    #         max_val = blackbody_lambda((ac.b_wien / temperature).to(q.um), temperature).value
    #         return blackbody_lambda(wavelength, temperature).value / max_val
    #
    #     bb = blackbody(temperature=init)
    #     fit = fitting.LevMarLSQFitter()
    #     bb_fit = fit(bb, wav.to(q.AA).value/10000, flx.to(u.FLAM).value)
    #     teff = int(bb_fit.temperature.value) * q.K
    #
    #     self.message('{} blackbody fit to {}'.format(teff, self.name))
    #
    #     fig = figure()
    #     fig.line(wav, flx)
    #     fig.line(wav, blackbody_lambda(wav.to(q.AA).value, teff), color='red')
    #     show(fig)
    #
    #     return teff

    @property
    def flux(self):
        """Getter for the flux"""
        return self._flux * self.const

    @copy_raw
    def flux_calibrate(self, distance, target_distance=10 * q.pc, flux_units=None):
        """Flux calibrate the spectrum from the given distance to the target distance

        Parameters
        ----------
        distance: astropy.unit.quantity.Quantity, sequence
            The current distance or (distance, uncertainty) of the spectrum
        target_distance: astropy.unit.quantity.Quantity
            The distance to flux calibrate the spectrum to
        flux_units: astropy.unit.quantity.Quantity
            The desired flux units of the output

        Returns
        -------
        sedkit.spectrum.Spectrum
            The flux calibrated spectrum object
        """
        # Set target flux units
        if flux_units is None:
            flux_units = self.flux_units

        # Calculate the scaled flux
        flux = (self.spectrum[1] * (distance[0] / target_distance)**2).to(flux_units)

        # Calculate the scaled uncertainty
        if self.unc is None:
            unc = None
        else:
            term1 = (self.spectrum[2] * distance[0] / target_distance).to(flux_units)
            term2 = (2 * self.spectrum[1] * (distance[1] * distance[0] / target_distance**2)).to(flux_units)
            unc = np.sqrt(term1**2 + term2**2)

        return Spectrum(self.spectrum[0], flux, unc, name=self.name)

    @property
    def flux_units(self):
        """A property for flux_units"""
        return self._flux_units

    @flux_units.setter
    def flux_units(self, flux_units):
        """A setter for flux_units

        Parameters
        ----------
        flux_units: astropy.units.quantity.Quantity
            The astropy units of the SED wavelength
        """
        # Check the units
        if not u.equivalent(flux_units, u.FLAM):
            raise TypeError("flux_units must be in flux density units, e.g. 'erg/s/cm2/A'")

        # Update the flux and unc arrays
        self._flux = self._flux * self.flux_units.to(flux_units)
        if self.unc is not None:
            self._unc = self._unc * self.flux_units.to(flux_units)

        # Set the flux_units
        self._flux_units = flux_units
        self._set_units()

    def integrate(self, units=q.erg / q.s / q.cm**2):
        """Calculate the area under the spectrum

        Parameters
        ----------
        units: astropy.units.quantity.Quantity
            The target units for the integral

        Returns
        -------
        sequence
            The integrated flux and uncertainty
        """
        # Make sure the target units are flux units
        if not u.equivalent(units, q.erg / q.s / q.cm**2):
            raise TypeError("units must be in flux units, e.g. 'erg/s/cm2'")

        # Calculate the factor for the given units
        m = self.flux_units * self.wave_units

        # Scrub the spectrum
        spec = u.scrub(self.data)
        val = (np.trapz(spec[1], x=spec[0]) * m).to(units)

        if self.unc is None:
            unc = None
        else:
            unc = np.sqrt(np.nansum((spec[2] * np.gradient(spec[0]) * m)**2)).to(units)

        return val, unc

    @copy_raw
    def interpolate(self, wave):
        """Interpolate the spectrum to another wavelength array

        Parameters
        ----------
        wave: astropy.units.quantity.Quantity, sedkit.spectrum.Spectrum
            The wavelength array to interpolate to

        Returns
        -------
        sedkit.spectrum.Spectrum
            The interpolated spectrum object
        """
        # Pull out wave if its a Spectrum object
        if hasattr(wave, 'spectrum'):
            wave = wave.spectrum[0]

        # Test units
        if not u.equivalent(wave, q.um):
            raise ValueError("New wavelength array must be in units of length.")

        # Get the data and make into same wavelength units
        w0 = self.wave * self.wave_units.to(wave.unit)
        f0 = self.spectrum[1]
        if len(self.spectrum) > 2:
            e0 = self.spectrum[2]
        else:
            e0 = np.zeros_like(f0)

        # Interpolate self to new wavelengths
        f1 = np.interp(wave.value, w0, f0.value, left=np.nan, right=np.nan) * self.flux_units
        e1 = np.interp(wave.value, w0, e0.value, left=np.nan, right=np.nan) * self.flux_units

        return Spectrum(wave, f1, e1, name=self.name)

    def message(self, msg, pre='[sedkit]'):
        """
        Only print message if verbose=True

        Parameters
        ----------
        msg: str
            The message to print
        pre: str
            The stuff to print before
        """
        if self.verbose:
            if pre is None:
                print(msg)
            else:
                print("{} {}".format(pre, msg))

    def mcmc_fit(self, model_grid, params=['teff'], walkers=5, steps=20, name=None, report=None):
        """
        Produces a marginalized distribution plot of best fit parameters from the specified model_grid

        Parameters
        ----------
        model_grid: sedkit.modelgrid.ModelGrid
            The model grid to use
        params: list
            The list of model grid parameters to fit
        walkers: int
            The number of walkers to deploy
        steps: int
            The number of steps for each walker to take
        name: str
            Name for the fit
        plot: bool
            Make plots
        """
        # Specify the parameter space to be walked
        for param in params:
            if param not in model_grid.parameters:
                raise ValueError("'{}' not a parameter in this model grid, {}".format(param, model_grid.parameters))

        # A name for the fit
        name = name or model_grid.name

        # Ensure modelgrid and spectruym are the same wave_units
        model_grid.wave_units = self.wave_units

        # Set up the sampler object
        self.sampler = mc.SpecSampler(self, model_grid, params)

        # Run the mcmc method
        self.sampler.mcmc_go(nwalk_mult=walkers, nstep_mult=steps)

        # Save the chi-sq best fit
        self.best_fit[name + ' (chi2)'] = self.sampler.spectrum.best_fit['best']

        # Make plots
        if report is not None:
            self.sampler.plot_chains()

        # Generate best fit spectrum the 50th quantile value
        best_fit_params = {k: v for k, v in zip(self.sampler.all_params, self.sampler.all_quantiles.T[1])}
        params_with_unc = self.sampler.get_error_and_unc()
        for param, quant in zip(self.sampler.all_params, params_with_unc):
            best_fit_params['{}_unc'.format(param)] = np.mean([quant[0], quant[2]])

        # Add missing parameters
        for param in model_grid.parameters:
            if param not in best_fit_params:
                best_fit_params[param] = getattr(model_grid, '{}_vals'.format(param))[0]

        # Get best fit model and scale to spectrum
        model = model_grid.get_spectrum(**{param: best_fit_params[param] for param in model_grid.parameters})
        model = model.norm_to_spec(self)
        model.phot = model_grid.phot

        # Make dict for best fit model
        best_fit_params['label'] = model.name
        best_fit_params['filepath'] = None
        best_fit_params['spectrum'] = np.array(model.spectrum)
        best_fit_params['full_model'] = model
        best_fit_params['const'] = 1.
        best_fit_params['fit_to'] = 'phot' if model_grid.phot else 'spec'

        self.best_fit[name] = best_fit_params

    @copy_raw
    def norm_to_mags(self, photometry, force=False, exclude=[], include=[]):
        """
        Normalize the spectrum to the given bandpasses

        Parameters
        ----------
        photometry: astropy.table.QTable
            A table of the photometry
        force: bool
            Force the normalization even if bandpass is not completely covered by spectrum
        exclude: sequence (optional)
            A list of bands to exclude from the normalization
        include: sequence (optional)
            A list of bands to include in the normalization

        Returns
        -------
        sedkit.spectrum.Spectrum
            The normalized spectrum object
        """
        # Default norm
        norm = 1

        # Compile list of photometry to include
        keep = []
        for band in photometry['band']:

            # Keep only explicitly included bands...
            if include:
                if band in include:
                    keep.append(band)

            # ...or just keep non excluded bands
            else:
                if band not in exclude:
                    keep.append(band)

        # Trim the table
        photometry = photometry[[idx for idx, bnd in enumerate(photometry['band']) if bnd in keep]]

        if len(photometry) == 0:
            self.message('No photometry to normalize this spectrum.')

        else:
            # Calculate all the synthetic magnitudes
            data = []
            for row in photometry:

                try:
                    bp = row['bandpass']
                    syn_flx, syn_unc = self.synthetic_flux(bp, force=force)
                    if syn_flx is not None:
                        flx = row['app_flux']
                        unc = row['app_flux_unc']
                        weight = bp.fwhm.value
                        unc = unc.value if hasattr(unc, 'unit') else None
                        syn_unc = syn_unc.value if hasattr(syn_unc, 'unit') else None
                        data.append([flx.value, unc, syn_flx.value, syn_unc, weight])
                except IOError:
                    pass

            # Check if there is overlap
            if len(data) == 0:
                self.message('No photometry in the range {} to normalize this spectrum.'.format([self.wave_min, self.wave_max]))

            # Calculate the weighted normalization
            else:
                f1, e1, f2, e2, weights = np.array(data).T
                if not any(e1):
                    e1 = None
                if not any(e2):
                    e2 = None
                gstat, norm = u.goodness(f1, f2, e1, e2, weights)

        # Make new spectrum
        spectrum = self.spectrum
        spectrum[1] *= norm
        if self.unc is not None:
            spectrum[2] *= norm

        return Spectrum(*spectrum, name=self.name)

    @copy_raw
    def norm_to_spec(self, spec, add=False, plot=False, **kwargs):
        """Normalize the spectrum to another spectrum

        Parameters
        ----------
        spec: sedkit.spectrum.Spectrum
            The spectrum to normalize to
        add: bool
            Add spec to self after normalization
        plot: bool
            Plot the spectra

        Returns
        -------
        sedkit.spectrum.Spectrum
          The normalized spectrum
        """
        # Resample self onto spec wavelengths
        w0 = self.wave * self.wave_units.to(spec.wave_units)
        slf = self.resamp(spec.spectrum[0])

        # Trim both to just overlapping wavelengths
        idx = u.idx_overlap(w0, spec.wave, inclusive=True)
        spec0 = slf.data[:, idx]
        spec1 = spec.data[:, idx]

        # Find the normalization factor
        norm = u.minimize_norm(spec1[1], spec0[1], **kwargs)

        # Make new spectrum
        spectrum = self.spectrum
        spectrum[1] *= norm
        if self.unc is not None:
            spectrum[2] *= norm

        # Make the new spectrum
        new_spec = Spectrum(*spectrum, name=self.name)

        # Add them together if necessary
        if add:
            new_spec = new_spec + spec

        if plot:
            # Rename and plot each
            new_spec.name = 'Normalized'
            self.name = 'Input'
            spec.name = 'Target'
            fig = new_spec.plot()
            fig = self.plot(fig=fig)
            fig = spec.plot(fig=fig)
            show(fig)

        return new_spec

    def plot(self, fig=None, components=False, best_fit=True, scale='log', draw=False, const=1., **kwargs):
        """Plot the spectrum

        Parameters
        ----------
        fig: bokeh.figure (optional)
            The figure to plot on
        components: bool
            Plot all components of the spectrum
        best_fit: bool
            Plot the best fit model if available
        scale: str
            The scale of the x and y axes, ['linear', 'log']
        draw: bool
            Draw the plot rather than just return it

        Returns
        -------
        bokeh.figure
            The figure
        """
        # Make the figure
        if fig is None:
            # fig = figure(x_axis_type=scale, y_axis_type=scale)
            fig = figure()
            fig.xaxis.axis_label = "Wavelength [{}]".format(self.wave_units)
            fig.yaxis.axis_label = "Flux Density [{}]".format(self.flux_units)

        # Plot the spectrum
        c = kwargs.get('color', next(u.COLORS))

        # Scatter plot
        if self.phot:

            idx = np.arange(len(self.wave))

            # Uncertainties
            if self.unc is not None:

                # Plot points with errors
                idx, = np.where(self.unc > 0)
                if len(idx) > 0:
                    fig.circle(self.wave[idx], self.flux[idx] * const, color=c, fill_alpha=0.7, size=8, legend_label=self.name)
                    y_err_x = [(px, px) for px in self.wave[idx]]
                    y_err_y = [(py - err, py + err) for py, err in zip(self.flux[idx], self.unc[idx])]
                    fig.multi_line(y_err_x, y_err_y, color=c)

                # Plot points without errors
                idn, = np.where(self.unc <= 0)
                if len(idn) > 0:
                    fig.circle(self.wave[idn], self.flux[idn] * const, color=c, fill_alpha=0, size=8)

            else:
                fig.circle(self.wave, self.flux * const, color=c, fill_alpha=0.7, size=8, legend_label=self.name)

            # Plot the best fit
            if best_fit:
                for name, bf in self.best_fit.items():
                    spec = bf['spectrum']
                    fig.square(spec[0], spec[1] * const, alpha=0.3, color=next(u.COLORS), legend_label=bf['label'])

        # Line plot
        else:
            fig.line(self.wave, self.flux * const, color=c, alpha=0.8, legend_label=self.name)

            # Plot the uncertainties
            if self.unc is not None:

                # Upper and lower bounds
                lower = (self.flux - self.unc) * const
                upper = (self.flux + self.unc) * const

                # Change x values to NaN if y value is NaN
                xpatch = copy.copy(self.wave)
                xpatch[np.argwhere(np.isnan(upper)) | np.argwhere(np.isnan(lower))] = np.nan

                # Make top and bottom of patch
                xpatch = np.concatenate([xpatch, xpatch[::-1]])
                ypatch = np.concatenate([upper, lower[::-1]])

                # Add the patch
                fig.patch(xpatch, ypatch, color=c, fill_alpha=0.1, line_alpha=0)

            # Plot the components
            if components and self.components is not None:
                for spec in self.components:
                    fig.line(spec.wave, spec.flux * const, color=next(u.COLORS), legend_label=spec.name)

            # Plot the best fit
            if best_fit:
                for name, bf in self.best_fit.items():
                    spec = bf['spectrum']
                    fig.line(spec[0], spec[1] * const, alpha=0.3, color=next(u.COLORS), legend_label=bf['label'])

        if draw:
            show(fig)
        else:
            return fig

    def renormalize(self, mag, bandpass, system='vegamag', force=False, no_spec=False, name=None):
        """Renormalize the spectrum to the given magnitude

        Parameters
        ----------
        mag: float
            The target magnitude
        bandpass: svo_filters.svo.Filter, str
            The bandpass to use
        system: str
            The magnitude system to use
        force: bool
            Force the synthetic photometry even if incomplete coverage
        no_spec: bool
            Return the normalization constant only

        Returns
        -------
        float, sedkit.spectrum.Spectrum
            The normalization constant or normalized spectrum object
        """
        # My solution
        norm = u.mag2flux(bandpass, mag)[0] / self.synthetic_flux(bandpass, force=force)[0]

        # Just return the normalization factor
        if no_spec:
            return float(norm.value)

        # Scale the spectrum
        spectrum = self.spectrum
        spectrum[1] *= norm
        if self.unc is not None:
            spectrum[2] *= norm

        return Spectrum(*spectrum, name=name or self.name)

    @copy_raw
    def resamp(self, wave=None, resolution=None, name=None):
        """Resample the spectrum onto a new wavelength array or to a new
        resolution

        Parameters
        ----------
        wave: astropy.units.quantity.Quantity (optional)
            The wavelength array to resample onto
        resolution: int (optional)
            The new resolution to resample to, keeping the same
            wavelength range

        Returns
        -------
        sedkit.spectrum.Spectrum
            The resampled spectrum
        """
        # Generate wavelength if resolution is set
        if resolution is not None:

            # Make the wavelength array
            mn = np.nanmin(self.wave)
            mx = np.nanmax(self.wave)
            d_lam = (mx - mn) / resolution
            wave = np.arange(mn, mx, d_lam) * self.wave_units

        if not u.equivalent(wave, q.um):
            raise TypeError("wave must be in length units")

        # Convert wave to target units
        self.wave_units = wave.unit
        wave = wave.value

        # Bin the spectrum
        binned = u.spectres(wave, self.wave, self.flux, self.unc)

        # Update the spectrum
        spectrum = [i * Q for i, Q in zip(binned, self.units)]

        return Spectrum(*spectrum, name=name or self.name)

    def restore(self):
        """
        Restore the spectrum to the original raw data
        """
        return Spectrum(*(self.raw or self.spectrum), name=self.name)

    def _set_units(self):
        """Set the units for convenience"""
        self.units = [self._wave_units, self._flux_units]
        self.units += [self._flux_units] if self.unc is not None else []

    @property
    def size(self):
        """The length of the data"""
        return len(self.wave)

    @copy_raw
    def smooth(self, beta, window=11):
        """
        Smooths the spectrum using a Kaiser-Bessel smoothing window of
        narrowness *beta* (~1 => very smooth, ~100 => not smooth)

        Parameters
        ----------
        beta: float, int
            The narrowness of the window
        window: int
            The length of the window

        Returns
        -------
        sedkit.spectrum.Spectrum
            The smoothed spectrum
        """
        s = np.r_[self.flux[window - 1:0:-1], self.flux, self.flux[-1:-window:-1]]
        w = np.kaiser(window, beta)
        y = np.convolve(w / w.sum(), s, mode='valid')
        smoothed = y[5:len(y) - 5]

        # Replace with smoothed spectrum
        spectrum = self.spectrum
        spectrum[1] = smoothed * self.flux_units

        return Spectrum(*spectrum, name=self.name)

    @property
    def spectrum(self):
        """Store the spectrum with units
        """
        return [i * Q for i, Q in zip(self.data, self.units)]

    def synthetic_flux(self, bandpass, force=False, plot=False):
        """
        Calculate the magnitude in a bandpass

        Parameters
        ----------
        bandpass: svo_filters.svo.Filter
            The bandpass to use
        force: bool
            Force the magnitude calculation even if
            overlap is only partial
        plot: bool
            Plot the spectrum and the flux point

        Returns
        -------
        float
            The magnitude
        """
        # Test overlap
        overlap = bandpass.overlap(self.spectrum)

        # Initialize
        flx = unc = None

        if overlap == 'full' or (overlap == 'partial' and force):

            # Convert self to bandpass units
            self.wave_units = bandpass.wave_units

            # Calculate the bits
            wav = bandpass.wave[0]
            rsr = bandpass.throughput
            grad = np.gradient(wav).value

            # Interpolate the spectrum to the filter wavelengths
            f = np.interp(wav.value, self.wave, self.flux, left=0, right=0) * self.flux_units

            # Filter out NaNs
            idx = np.where([not np.isnan(i) for i in f])[0]

            # Calculate the flux
            flx = (np.trapz(f[idx] * rsr[0][idx], x=wav[idx]) / np.trapz(rsr[0][idx], x=wav[idx])).to(self.flux_units)

            # Calculate uncertainty
            if self.unc is not None:
                sig_f = np.interp(wav.value, self.wave, self.unc, left=0, right=0) * self.flux_units
                unc = np.sqrt(np.sum(((sig_f * rsr * grad)**2).to(self.flux_units**2)))
            else:
                unc = None

            # Plot it
            if plot:
                fig = figure()
                fig.line(self.wave, self.flux, color='navy')
                fig.circle([bandpass.wave_eff], [flx], color='red')
                show(fig)

        return flx, unc

    def synthetic_magnitude(self, bandpass, force=False):
        """
        Calculate the synthetic magnitude in the given bandpass

        Parameters
        ----------
        bandpass: svo_filters.svo.Filter
            The bandpass to use
        force: bool
            Force the magnitude calculation even if
            overlap is only partial

        Returns
        -------
        tuple
            The flux and uncertainty
        """
        flx, flx_unc = self.synthetic_flux(bandpass, force=force)

        # Calculate the magnitude
        mag, mag_unc = None, None
        if flx is not None:
            mag, mag_unc = u.flux2mag((flx, flx_unc), bandpass)

        return mag, mag_unc

    def trim(self, include=[], exclude=[], concat=False):
        """Trim the spectrum in the given wavelength ranges

        Parameters
        ----------
        include: sequence
            The (min_wave, max_wave) ranges to include from the spectrum
        exclude: sequence
            The (min_wave, max_wave) ranges to exclude from the spectrum
        """
        if not isinstance(include, (list, tuple)):
            print("""Please provide a list of (lower,upper) bounds with units to include in trim, e.g. [(0*q.um,0.8*q.um)]""")

        if not isinstance(exclude, (list, tuple)):
            print("""Please provide a list of (lower,upper) bounds with units to exclude in trim, e.g. [(0*q.um,0.8*q.um)]""")

        # Default to include everything
        if len(include) == 0:
            include.append([self.wave_min, self.wave_max])

        # Get element indexes of included ranges
        idx_include = []
        for mn, mx in include:
            inc, = np.where(np.logical_and(self.spectrum[0] >= mn, self.spectrum[0] <= mx))
            idx_include.append(inc)

        # Get element indexes of excluded ranges
        idx_exclude = []
        for mn, mx in exclude:
            exc, = np.where(np.logical_and(self.spectrum[0] >= mn, self.spectrum[0] <= mx))
            idx_exclude.append(exc)

        # Get difference of each included set with each excluded set
        segments = []
        for inc in idx_include:
            set_inc = set(inc)

            for exc in idx_exclude:
                set_inc = set_inc.difference(set(exc))

            # Sort and split the indexes by missing elements
            set_inc = sorted(set_inc)
            for k, g in groupby(enumerate(list(set_inc)), lambda i_x: i_x[0] - i_x[1]):

                # Make the spectra
                group = list(map(itemgetter(1), g))
                data = [i[group] for i in copy.copy(self.spectrum)]
                spectrum = Spectrum(*data, name=self.name)
                spectrum.wave_units = self.wave_units
                segments.append(spectrum)

        return np.sum(segments) if concat else segments

    @property
    def unc(self):
        """Getter for the unc"""
        if self._unc is None:
            return None
        else:
            return self.flux / (self._flux / self._unc)

    @property
    def wave_max(self):
        """The minimum wavelength"""
        return max(self.wave) * self.wave_units

    @property
    def wave_min(self):
        """The minimum wavelength"""
        return min(self.wave) * self.wave_units

    @property
    def wave_units(self):
        """A property for wave_units"""
        return self._wave_units

    @wave_units.setter
    def wave_units(self, wave_units):
        """A setter for wave_units

        Parameters
        ----------
        wave_units: astropy.units.quantity.Quantity
            The astropy units of the SED wavelength
        """
        # Make sure the values are in length units
        if not u.equivalent(wave_units, q.um):
            raise TypeError("wave_units must be a unit of length, e.g. 'um'")

        # Update the wavelength array
        self.wave = self.wave * self.wave_units.to(wave_units)

        # Set the wave_units
        self._wave_units = wave_units
        self._set_units()

class Blackbody(Spectrum):
    """A spectrum object specifically for blackbodies"""
    def __init__(self, wavelength, teff, radius=None, distance=None, **kwargs):
        """
        Given a wavelength array and temperature, returns an array of Planck
        function values in [erg s-1 cm-2 A-1]

        Parameters
        ----------
        wavelength: array-like
            The array of wavelength values to evaluate the Planck function
        teff: astropy.unit.quantity.Quantity, sequence
            The effective temperature and (optional) uncertainty
        radius: astropy.unit.quantity.Quantity, sequence
            The radius and (optional) uncertainty
        distance: astropy.unit.quantity.Quantity, sequence
            The distance and (optional) uncertainty
        """
        # Check wavelength units
        if not u.equivalent(wavelength, q.um):
            raise TypeError("Wavelength must be in astropy units of length, e.g. 'um'")

        # Check teff
        if not u.issequence(teff, length=2):
            teff = teff, None
        if not u.equivalent(teff[0], q.K):
            raise TypeError("teff must be in astropy units of temperature, eg. 'K'")
        if not u.equivalent(teff[1], q.K) and teff[1] is not None:
            raise TypeError("teff_unc must be in astropy units of temperature, eg. 'K'")
        self.teff, self.teff_unc = teff[:2]

        # Check radius
        if u.issequence(radius, length=3):
            radius = radius
        if not u.issequence(radius):
            radius = radius, None, None
        if not u.equivalent(radius[0], q.R_jup) and radius[0] is not None:
            raise TypeError("radius must be in astropy units of length, eg. 'R_jup'")
        if not u.equivalent(radius[1], q.R_jup) and radius[1] is not None:
            raise TypeError("radius_unc must be in astropy units of length, eg. 'R_jup'")
        self.radius, self.radius_unc = radius[:2]

        # Check distance
        if u.issequence(distance, length=3):
            distance = distance
        if not u.issequence(distance):
            distance = distance, None, None
        if not u.equivalent(distance[0], q.pc) and distance[0] is not None:
            raise TypeError("distance must be in astropy units of length, eg. 'pc'")
        if not u.equivalent(distance[1], q.pc) and distance[1] is not None:
            raise TypeError("distance_unc must be in astropy units of length, eg. 'pc'")
        self.distance, self.distance_unc = distance[:2]

        # Evaluate
        I, I_unc = self.eval(wavelength)

        # Inherit from Spectrum
        super().__init__(wavelength, I, I_unc, **kwargs)

        self.name = '{} Blackbody'.format(teff)

    def eval(self, wavelength, Flam=False):
        """Evaluate the blackbody at the given wavelengths

        Parameters
        ----------
        wavelength: sequence
            The wavelength values

        Returns
        -------
        list
            The blackbody flux and uncertainty arrays
        """
        if not u.equivalent(wavelength, q.um):
            raise TypeError("Wavelength must be in astropy units, e.g. 'um'")

        units = q.erg / q.s / q.cm**2 / (1 if Flam else q.AA)

        # Check for radius and distance
        if self.radius is not None and self.distance is not None:
            scale = (self.radius**2 / self.distance**2).decompose()
        else:
            scale = 1.

        # Get numerator and denominator
        const = ac.h * ac.c / (wavelength * ac.k_B)
        numer = 2 * np.pi * ac.h * ac.c**2 * scale / (wavelength**(4 if Flam else 5))
        denom = np.exp((const / self.teff)).decompose()

        # Calculate intensity
        I = (numer / (denom - 1.)).to(units)

        # Calculate dI/dr
        if self.radius is not None and self.radius_unc is not None:
            dIdr = (self.radius_unc * 2 * I / self.radius).to(units)
        else:
            dIdr = 0. * units

        # Calculate dI/dd
        if self.distance is not None and self.distance_unc is not None:
            dIdd = (self.distance_unc * 2 * I / self.distance).to(units)
        else:
            dIdd = 0. * units

        # Calculate dI/dT
        if self.teff is not None and self.teff_unc is not None:
            dIdT = (self.teff_unc * I * ac.h * ac.c / wavelength / ac.k_B / self.teff**2).to(units)
        else:
            dIdT = 0. * units

        # Calculate sigma_I from derivative terms
        I_unc = np.sqrt(dIdT**2 + dIdr**2 + dIdd**2)
        if isinstance(I_unc.value, float):
            I_unc = None

        return I, I_unc

class FileSpectrum(Spectrum):
    def __init__(self, file, wave_units=None, flux_units=None, ext=0,
                 survey=None, name=None, **kwargs):
        """Create a spectrum from an ASCII or FITS file

        Parameters
        ----------
        file: str
            The path to the ascii or FITS file
        wave_units: astropy.units.quantity.Quantity
            The wavelength units
        flux_units: astropy.units.quantity.Quantity
            The flux units
        ext: int, str
            The FITS extension name or index
        survey: str (optional)
            The name of the survey
        """
        # Read the fits data...
        if file.endswith('.fits'):

            if file.endswith('.fits'):
                data = u.spectrum_from_fits(file, ext=ext)

            elif survey == 'SDSS':
                head = fits.getheader(file)
                raw = fits.getdata(file)
                flux_units = 1E-17 * u.FLAM
                wave_units = q.AA
                log_w = head['COEFF0'] + head['COEFF1'] * np.arange(len(raw.flux))
                data = [10**log_w, raw.flux, raw.ivar]

            # Check if it is a recarray
            elif isinstance(raw, fits.fitsrec.FITS_rec):

                # Check if it's an SDSS spectrum
                raw = fits.getdata(file, ext=ext)
                data = raw['WAVELENGTH'], raw['FLUX'], raw['ERROR']

            # Otherwise just an array
            else:
                print("Sorry, I cannot read the file at", file)

        # ...or the ascii data...
        elif file.endswith('.txt'):
            data = np.genfromtxt(file, unpack=True)

        # ...or the VO Table
        elif file.endswith('.xml'):
            vot = vo.parse_single_table(file)
            data = np.array([list(i) for i in vot.array]).T

        else:
            raise IOError('The file needs to be ASCII, XML, or FITS.')

        # Make sure units are astropy quantities
        if isinstance(wave_units, str):
            wave_units = q.Unit(wave_units)
        if isinstance(flux_units, str):
            flux_units = q.Unit(flux_units)

        # Sanity check for wave_units
        if data[0].min() > 100 and wave_units == q.um:
            print("WARNING: Your wavelength range ({} - {}) looks like Angstroms. Are you sure it's {}?".format(data[0].min(), data[0].max(), wave_units))

        # Apply units
        wave = data[0] * wave_units
        flux = data[1] * flux_units
        if len(data) > 2:
            unc = data[2] * flux_units
        else:
            unc = None

        if name is None:
            name = file

        super().__init__(wave, flux, unc, name=name, **kwargs)

class Vega(Spectrum):
    """A Spectrum object of Vega"""
    def __init__(self, wave_units=q.AA, flux_units=u.FLAM, **kwargs):
        """Initialize the Spectrum object

        Parameters
        ----------
        wave_units: astropy.units.quantity.Quantity
            The desired wavelength units
        flux_units: astropy.units.quantity.Quantity
            The desired flux units
        """
        # Get the data and apply units
        vega_file = resource_filename('sedkit', 'data/STScI_Vega.txt')
        ref = '2007ASPC..364..315B, 2004AJ....127.3508B, 2005MSAIS...8..189K'
        wave, flux = np.genfromtxt(vega_file, unpack=True)
        wave *= q.AA
        flux *= u.FLAM

        # Convert to target units
        wave = wave.to(wave_units)
        flux = flux.to(flux_units)

        # Make the Spectrum object
        super().__init__(wave, flux, ref=ref, **kwargs)

        self.name = 'Vega'

def fit_model(row, fitspec, resample=True, wave_units=q.AA):
    """Fit the model grid row to the spectrum with the given parameters

    Parameters
    ----------
    row: pandas.Row
        The dataframe row to fit
    fitspec: sedkit.spectrum.Spectrum
        The spectrum to fit

    Returns
    -------
    pandas.Row
        The input row with the normalized spectrum and additional gstat
    """
    gstat, yn, xn = list(fitspec.fit(row['spectrum'], resample=resample, wave_units=wave_units, weights=row.get('weights')))
    spectrum = np.array([row['spectrum'][0] * xn, row['spectrum'][1] * yn])
    row['spectrum'] = spectrum
    row['gstat'] = gstat
    row['const'] = yn

    return row

import numpy,import astropy,from astropy


"""API for NatGeo-Microsoft AI4 Earth NatCap."""
import argparse
import logging
import os
import sys
import uuid

from keras_retinanet import models
from keras_retinanet.utils.gpu import setup_gpu
from keras_retinanet.utils.keras_version import check_keras_version
from keras_retinanet.utils.tf_version import check_tf_version
import cv2
import flask
import keras
import PIL
import numpy
import shapely.geometry

APP = flask.Flask(__name__)

logging.basicConfig(
    level=logging.DEBUG,
    format=(
        '%(asctime)s (%(relativeCreated)d) %(levelname)s %(name)s'
        ' [%(funcName)s:%(lineno)d] %(message)s'),
    stream=sys.stdout)
LOGGER = logging.getLogger(__name__)

def compute_resize_scale(image_shape, min_side=800, max_side=1333):
    """ Compute an image scale size to constrained to min_side/max_side.
    Args
        min_side: The image's min side will be equal to min_side after
            resizing.
        max_side: If after resizing the image's max side is above max_side,
            resize until the max side is equal to max_side.
    Returns
        A resizing scale.
    """
    (rows, cols, _) = image_shape

    smallest_side = min(rows, cols)

    # rescale the image so the smallest side is min_side
    scale = min_side / smallest_side

    # check if the largest side is now greater than max_side, which can happen
    # when images have a large aspect ratio
    largest_side = max(rows, cols)
    if largest_side * scale > max_side:
        scale = max_side / largest_side

    return scale

def read_image_bgr(path):
    """ Read an image in BGR format.
    Args
        path: Path to the image.
    """
    image = numpy.asarray(PIL.Image.open(path).convert('RGB'))
    return image[:, :, ::-1].copy()

def preprocess_image(x, mode='caffe'):
    """ Preprocess an image by subtracting the ImageNet mean.
    Args
        x: numpy.array of shape (None, None, 3) or (3, None, None).
        mode: One of "caffe" or "tf".
            - caffe: will zero-center each color channel with
                respect to the ImageNet dataset, without scaling.
            - tf: will scale pixels between -1 and 1, sample-wise.
    Returns
        The input with the ImageNet mean subtracted.
    """
    # covert always to float32 to keep compatibility with opencv
    x = x.astype(numpy.float32)
    if mode == 'tf':
        x /= 127.5
        x -= 1.
    elif mode == 'caffe':
        x -= [103.939, 116.779, 123.68]
    return x

@APP.route('/api/v1/process_image', methods=['POST'])
def process_image():
    """Invoked to do inference on a binary image.

    Body of the post includes a url to the stored .zip file of the archive.

    Returns
        None.

    """
    try:
        image_id = uuid.uuid4().hex
        payload = flask.request.data
        image_path = os.path.join('/', 'usr', 'local', 'images', '%s.png' % image_id)
        try:
            os.makedirs(os.path.dirname(image_path))
        except OSError:
            pass
        with open(image_path, 'wb') as png_file:
            png_file.write(payload)
        LOGGER.debug('payload written to: %s', image_path)
        LOGGER.debug('preprocess image')
        raw_image = read_image_bgr(image_path)
        image = preprocess_image(raw_image)
        scale = compute_resize_scale(
            image.shape, min_side=800, max_side=1333)
        image = cv2.resize(image, None, fx=scale, fy=scale)
        if keras.backend.image_data_format() == 'channels_first':
            image = image.transpose((2, 0, 1))

        LOGGER.debug('run inference on image %s', str(image.shape))
        result = model.predict_on_batch(
            numpy.expand_dims(image, axis=0))
        # correct boxes for image scale
        LOGGER.debug('inference complete')
        boxes, scores, labels = result
        boxes /= scale

        # convert box to a list from a numpy array and score to a value
        # from a single element array
        non_max_supression_box_list = []
        box_score_tuple_list = [
            (list(box), score) for box, score in zip(boxes[0], scores[0])
            if score > 0.3]
        while box_score_tuple_list:
            box, score = box_score_tuple_list.pop()
            shapely_box = shapely.geometry.box(*box)
            keep = True
            # this list makes a copy
            for test_box, test_score in list(box_score_tuple_list):
                shapely_test_box = shapely.geometry.box(*test_box)
                if shapely_test_box.intersects(shapely_box):
                    if test_score > score:
                        # keep the new one
                        keep = False
                        break
            if keep:
                non_max_supression_box_list.append((
                    [float(x) for x in box], float(score)))

        LOGGER.debug(non_max_supression_box_list)
        return flask.jsonify(non_max_supression_box_list)
    except Exception as e:
        LOGGER.exception('error on processing image')
        return str(e), 500

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='NCI NDR Analysis.')
    parser.add_argument(
        '--app_port', type=int, default=8080,
        help='port to listen on for requests')
    args = parser.parse_args()

    LOGGER.debug('set up model')

    check_keras_version()
    check_tf_version()

    model = models.load_model(
        'natgeo_dams_model_resnet50_csv_64.h5', backbone_name='resnet50')

    LOGGER.debug('start the APP')
    APP.run(
        host='0.0.0.0',
        port=args.app_port)


#!/usr/bin/env python3
# Copyright 2019 Mycroft AI Inc.
#
# 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.
import numpy as np
from os.path import join
from precise_runner import PreciseRunner
from precise_runner.runner import ListenerEngine
from prettyparse import Usage
from random import randint
from shutil import get_terminal_size
from threading import Event

from precise.network_runner import Listener
from precise.scripts.base_script import BaseScript
from precise.util import save_audio, buffer_to_audio, activate_notify

from .sequence_nlp_service import nlp_task

class ListenScript(BaseScript):
    usage = Usage('''
        Run a model on microphone audio input

        :model str
            Either Keras (.net) or TensorFlow (.pb) model to run

        :-c --chunk-size int 2048
            Samples between inferences

        :-l --trigger-level int 3
            Number of activated chunks to cause an activation

        :-s --sensitivity float 0.5
            Network output required to be considered activated

        :-b --basic-mode
            Report using . or ! rather than a visual representation

        :-d --save-dir str -
            Folder to save false positives

        :-p --save-prefix str -
            Prefix for saved filenames
    ''')

    def __init__(self, args):
        super().__init__(args)
        self.listener = Listener(args.model, args.chunk_size)
        self.audio_buffer = np.zeros(self.listener.pr.buffer_samples, dtype=float)
        self.engine = ListenerEngine(self.listener, args.chunk_size)
        self.engine.get_prediction = self.get_prediction
        self.runner = PreciseRunner(self.engine, args.trigger_level, sensitivity=args.sensitivity,
                                    on_activation=self.on_activation, on_prediction=self.on_prediction)
        self.session_id, self.chunk_num = '%09d' % randint(0, 999999999), 0
        self.event = Event()

    def on_activation(self):
        activate_notify()
        # self.event.set()
        nlp_task(self.runner.stream)
        # self.event.set()
        # self.runner.stream.start_stream()

        if self.args.save_dir:
            nm = join(self.args.save_dir, self.args.save_prefix + self.session_id + '.' + str(self.chunk_num) + '.wav')
            save_audio(nm, self.audio_buffer)
            print()
            print('Saved to ' + nm + '.')
            self.chunk_num += 1

    def on_prediction(self, conf):
        if self.args.basic_mode:
            print('!' if conf > 0.7 else '.', end='', flush=True)
        else:
            max_width = 80
            width = min(get_terminal_size()[0], max_width)
            units = int(round(conf * width))
            bar = 'X' * units + '-' * (width - units)
            cutoff = round((1.0 - self.args.sensitivity) * width)
            print(bar[:cutoff] + bar[cutoff:].replace('X', 'x'))
            # print(conf)

    def get_prediction(self, chunk):
        audio = buffer_to_audio(chunk)
        self.audio_buffer = np.concatenate((self.audio_buffer[len(audio):], audio))
        return self.listener.update(chunk)

    def run(self):
        self.runner.start()
        self.event.wait()  # Wait forever

main = ListenScript.run_main

if __name__ == '__main__':
    main()


from collections import OrderedDict
from copy import deepcopy

from ml_collections import ConfigDict

import numpy as np
import torch
import torch.optim as optim
from torch import nn as nn
import torch.nn.functional as F

from .model import Scalar

def soft_target_update(network, target_network, soft_target_update_rate):
    target_network_params = {k: v for k, v in target_network.named_parameters()}
    for k, v in network.named_parameters():
        target_network_params[k].data = (
            (1 - soft_target_update_rate) * target_network_params[k].data
            + soft_target_update_rate * v.data
        )

class SAC(object):

    @staticmethod
    def get_default_config(updates=None):
        config = ConfigDict()
        config.discount = 0.99
        config.reward_scale = 1.0
        config.alpha_multiplier = 1.0
        config.use_automatic_entropy_tuning = True
        config.backup_entropy = True
        config.target_entropy = 0.0
        config.policy_lr = 3e-4
        config.qf_lr = 3e-4
        config.optimizer_type = 'adam'
        config.soft_target_update_rate = 5e-3
        config.target_update_period = 1

        if updates is not None:
            config.update(updates)
        return config

    def __init__(self, config, policy, qf1, qf2, target_qf1, target_qf2):
        self.config = SAC.get_default_config(config)
        self.policy = policy
        self.qf1 = qf1
        self.qf2 = qf2
        self.target_qf1 = target_qf1
        self.target_qf2 = target_qf2

        optimizer_class = {
            'adam': torch.optim.Adam,
            'sgd': torch.optim.SGD,
        }[self.config.optimizer_type]

        self.policy_optimizer = optimizer_class(
            self.policy.parameters(), self.config.policy_lr,
        )
        self.qf_optimizer = optimizer_class(
            list(self.qf1.parameters()) + list(self.qf2.parameters()), self.config.qf_lr
        )

        if self.config.use_automatic_entropy_tuning:
            self.log_alpha = Scalar(0.0)
            self.alpha_optimizer = optimizer_class(
                self.log_alpha.parameters(),
                lr=self.config.policy_lr,
            )
        else:
            self.log_alpha = None

        self.update_target_network(1.0)
        self._total_steps = 0

    def update_target_network(self, soft_target_update_rate):
        soft_target_update(self.qf1, self.target_qf1, soft_target_update_rate)
        soft_target_update(self.qf2, self.target_qf2, soft_target_update_rate)

    def train(self, batch):
        self._total_steps += 1

        observations = batch['observations']
        actions = batch['actions']
        rewards = batch['rewards']
        next_observations = batch['next_observations']
        dones = batch['dones']

        new_actions, log_pi = self.policy(observations)

        if self.config.use_automatic_entropy_tuning:
            alpha_loss = -(self.log_alpha() * (log_pi + self.config.target_entropy).detach()).mean()
            alpha = self.log_alpha().exp() * self.config.alpha_multiplier
        else:
            alpha_loss = observations.new_tensor(0.0)
            alpha = observations.new_tensor(self.config.alpha_multiplier)

        """ Policy loss """
        q_new_actions = torch.min(
            self.qf1(observations, new_actions),
            self.qf2(observations, new_actions),
        )
        policy_loss = (alpha*log_pi - q_new_actions).mean()

        """ Q function loss """
        q1_pred = self.qf1(observations, actions)
        q2_pred = self.qf2(observations, actions)

        new_next_actions, next_log_pi = self.policy(next_observations)
        target_q_values = torch.min(
            self.target_qf1(next_observations, new_next_actions),
            self.target_qf2(next_observations, new_next_actions),
        )

        if self.config.backup_entropy:
            target_q_values = target_q_values - alpha * next_log_pi

        q_target = self.config.reward_scale * rewards + (1. - dones) * self.config.discount * target_q_values
        qf1_loss = F.mse_loss(q1_pred, q_target.detach())
        qf2_loss = F.mse_loss(q2_pred, q_target.detach())
        qf_loss = qf1_loss + qf2_loss

        if self.config.use_automatic_entropy_tuning:
            self.alpha_optimizer.zero_grad()
            alpha_loss.backward()
            self.alpha_optimizer.step()

        self.policy_optimizer.zero_grad()
        policy_loss.backward()
        self.policy_optimizer.step()

        self.qf_optimizer.zero_grad()
        qf_loss.backward()
        self.qf_optimizer.step()

        if self.total_steps % self.config.target_update_period == 0:
            self.update_target_network(
                self.config.soft_target_update_rate
            )

        return dict(
            log_pi=log_pi.mean().item(),
            policy_loss=policy_loss.item(),
            qf1_loss=qf1_loss.item(),
            qf2_loss=qf2_loss.item(),
            alpha_loss=alpha_loss.item(),
            alpha=alpha.item(),
            average_qf1=q1_pred.mean().item(),
            average_qf2=q2_pred.mean().item(),
            average_target_q=target_q_values.mean().item(),
            total_steps=self.total_steps,
        )

    def torch_to_device(self, device):
        for module in self.modules:
            module.to(device)

    @property
    def modules(self):
        modules = [self.policy, self.qf1, self.qf2, self.target_qf1, self.target_qf2]
        if self.config.use_automatic_entropy_tuning:
            modules.append(self.log_alpha)
        return modules

    @property
    def total_steps(self):
        return self._total_steps


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

#
#
# # Earth Analytics Education - EA  Python Course Spring 2021

# ## Important  - Assignment Guidelines
#
# 1. Before you submit your assignment to GitHub, make sure to run the entire notebook with a fresh kernel. To do this first, **restart the kernel** (in the menubar, select Kernel$\rightarrow$Restart & Run All)
# 2. Always replace the `raise NotImplementedError()` code with your code that addresses the activity challenge. If you don't replace that code, your notebook will not run.
#
# ```
# # YOUR CODE HERE
# raise NotImplementedError()
# ```
#
# 3. Any open ended questions will have a "YOUR ANSWER HERE" within a markdown cell. Replace that text with your answer also formatted using Markdown.
# 4. **DO NOT RENAME THIS NOTEBOOK File!** If the file name changes, the autograder will not grade your assignment properly.
# 6. When you create a figure, comment out `plt.show()` to ensure the autograder can grade your plots. For figure cells, DO NOT DELETE the code that says `DO NOT REMOVE LINE BELOW`.
#
# ```
# ### DO NOT REMOVE LINE BELOW ###
# student_plot1_ax = nb.convert_axes(plt)
# ```
#
# * Only include the package imports, code, and outputs that are required to run your homework assignment.
# * Be sure that your code can be run on any operating system. This means that:
#    1. the data should be downloaded in the notebook to ensure it's reproducible
#    2. all paths should be created dynamically using the `os.path.join`
#
# ## Follow to PEP 8 Syntax Guidelines & Documentation
#
# * Run the `autopep8` tool on all cells prior to submitting (HINT: hit shift + the tool to run it on all cells at once!
# * Use clear and expressive names for variables.
# * Organize your code to support readability.
# * Check for code line length
# * Use comments and white space sparingly where it is needed
# * Make sure all python imports are at the top of your notebook and follow PEP 8 order conventions
# * Spell check your Notebook before submitting it.
#
# For all of the plots below, be sure to do the following:
#
# * Make sure each plot has a clear TITLE and, where appropriate, label the x and y axes. Be sure to include UNITS in your labels.
#

# ### Add Your Name Below
# **Your Name: Thomas Schoenrock-Rossiter**

#

# ---

# # Week 04 and 05 Homework - Automate NDVI Workflow
#
# For this assignment, you will write code to generate a plot of the mean normalized difference vegetation index (NDVI) for two different sites in the United States across one year of data:
#
# * San Joaquin Experimental Range (SJER) in Southern California, United States
# * Harvard Forest (HARV) in the Northeastern United States
#
# The data that you will use for this week is available from **earthpy** using the following download:
#
# `et.data.get_data('ndvi-automation')`
#
# ## Assignment Goals
#
# Your goal in this assignment is to create the most efficient and concise workflow that you can that allows for:
#
# 1. The code to scale if you added new sites or more time periods to the analysis.
# 2. Someone else to understand your workflow.
# 3. The LEAST and most efficient (i.e. runs fast, minimize repetition) amount of code that completes the task.
#
# ### HINTS
#
# * Remove values outside of the landsat valid range of values as specified in the metadata, as needed.
# * Keep any output files SEPARATE FROM input files. Outputs should be created in an outputs directory that is created in the code (if needed) and/or tested for.
# * Use the functions that we demonstrated during class to make your workflow more efficient.
# * BONUS - if you  chose - you can export your data as a csv file. You will get bonus points for doing this.
#
#
# ## Assignment Requirements
#
# Your submission to the GitHub repository should include:
# * This Jupyter Notebook file (.ipynb) with:
#     * The code to create a plot of mean NDVI across a year for  2 NEON Field Sites:
#         * NDVI on the x axis and formatted dates on the y for both NEON sites on one figure/axis object
#     * The **data should be cleaned to remove the influence of clouds**. See the [earthdatascience website for an example of what your plot might look like with and without removal of clouds](https://www.earthdatascience.org/courses/earth-analytics-python/create-efficient-data-workflows/).
# * BONUS: Create one output `.csv` file that has 3 columns - NDVI, Date and Site Name - with values for SJER and HARV.
#
# Your notebook should:
# * Have *at least* 2 well documented and well named functions with docstrings.
# * Include a Markdown cell at the top of the notebook that outlines the overall workflow using pseudocode (i.e. plain language, not code)
# * Include additional Markdown cells throughout the notebook to describe:
#     * the data that you used - and where it is from
#     * how data are being processing
#     * how the code is optimized to run fast and be more concise

# # Replace this cell with your pseudocode  for this workflow
#
# If you happen to be a diagram person a diagram is ok too
#
#

# Workflow pseudocode
#
# Run first for one month at one site - next run for all months in one site - last run for all months in both sites
#
# Preparation -
# - Import packages
# - Set working directory to 'data'
# - Download data from earthpy
#
# Access Landsat band files for single scene -
# - Create path to Landsat data
# - Create list of Landsat bands in scene (into df?)
# - Subset bands into just RGB, NIR
# - Stack bands into single raster
# - clip stacked raster by boundary layer
# - Output raster for individual month
#
# Access Landsat band files for all scenes in a site -
# - Get list of all directories
# - For bands in each scene directory perform the above steps
# - Output, into a list or df, the average ndvi and data for each scene
#
# Access Landsat band files for all scenes in both sites -
#
# Calculate ndvi -
# - Run ndvi expression on red and NIR bands
# - Calculate mean ndvi for site
# - Output ndvi df/raster
# - Output ndvi values to spreadsheet
#
# Plot ndvi -
# - Create ndvi graph
#

# In[1]:

# Autograding imports - do not modify this cell
import matplotcheck.autograde as ag
import matplotcheck.notebook as nb
import matplotcheck.timeseries as ts
from datetime import datetime

# In[2]:

# Import needed packages in PEP 8 order
# and no unused imports listed (10 points total)

import os
from glob import glob

import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from matplotlib import patches as mpatches
from matplotlib import colors
from matplotlib.dates import DateFormatter

import seaborn as sns
import numpy as np
import pandas as pd
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()
from datetime import datetime
import pyproj
pyproj.set_use_global_context(True)

import rioxarray as rxr
import rasterio
from rasterio.plot import plotting_extent
import xarray as xr
import geopandas as gpd

import earthpy as et
import earthpy.spatial as es
import earthpy.plot as ep
import earthpy.mask as em

from shapely.geometry import mapping, box

sns.set(font_scale = 1.2, style = "darkgrid")

sns.set_style("ticks")

landsat_data = et.data.get_data('ndvi-automation')
os.chdir(os.path.join(et.io.HOME, 'earth-analytics', 'data'))

# In[3]:

# DO NOT MODIFY THIS CELL
# Tests that the working directory is set to earth-analytics/data

path = os.path.normpath(os.getcwd())
student_wd_parts = path.split(os.sep)

if student_wd_parts[-2:] == ['earth-analytics', 'data']:
    print("\u2705 Great - it looks like your working directory is set correctly to ~/earth-analytics/data")
else:
    print("\u274C Oops, the autograder will not run unless your working directory is set to earth-analytics/data")

# # Figure 1: Plot 1 - Mean NDVI For Each Site Across the Year (50 points)
#
# Create a plot of the mean normalized difference vegetation index (NDVI) for the two different sites in the United States across the year:
#
# * NDVI on the x axis and formatted dates on the y for both NEON sites on one figure/axis object.
# * Each site should be identified with a different color in the plot and legend.
# * The final plot **data should be cleaned to remove the influence of clouds**.
# * Be sure to include appropriate title and axes labels.
#
# Add additional cells as needed for processing data (e.g. defining functions, etc), but be sure to:
# * follow the instructions in the code cells that have been provided to ensure that you are able to use the sanity check tests that are provided.
# * include only the plot code in the cell identified for the final plot code below

# ## Task 1:
#
# In the cell below, create a single dataframe containing MEAN NDVI, the site name,
# and the date of the data for the HARV site
# scene `HARV/landsat-crop/LC080130302017031701T1-SC20181023151837`.  The column names for the  final
# DataFrame should be`mean_ndvi`, and `site`, and the data should be **indexed on the date**.
#
# Use the functions that we reviewed in class (or create your own versions of them) to implement your code
#
# ### In the Cell below Place  All Functions Needed to Run this Notebook (20 points)

# In[4]:

### DO NOT REMOVE THIS LINE OR EDIT / MOVE THIS CELL ###
start_time = datetime.now()

# The cell below contains the three functions that were used to streamline the amount of code and repitition used. The third function takes the previous two functions, one of which opens a sinlge band and crops it and the other takes the cropped band and masks it using a qa layer, and combines them so that they can be called in the subsequent loop. The loop applies this code to each landsat scene file for both of the sites before outputting a list.

# In[5]:

# In this cell place all of the functions needed to run your notebook
# You will be graded here on function application, docstrings, efficiency so ensure
# All functions are placed here!

#Create paths to both sites

site_path = os.path.join("ndvi-automation",
                         "sites")

# path to sjer and harv sites
study_sites = glob(os.path.join(site_path + "/*/"))
print(study_sites)

#function to open single landsat band, mask using qa and crop
def open_clean_bands(band_path,
                     crop_extent,
                     valid_range=None,):
    """Open and mask a single landsat band using a pixel_qa layer.

    Parameters
    -----------
    band_path : string
        A path to the array to be opened
    crop_extent : vector
        A path to the vector file to use as crop boundary
    valid_range : tuple (optional)
        A tuple of min and max range of values for the data. Default = None

    Returns
    -----------
    arr : xarray DataArray
        An xarray DataArray with values that should be masked set to 1 for
        True (Boolean)
    """
    band = rxr.open_rasterio(band_path, masked = True).rio.clip    (crop_extent.geometry, from_disk = True).squeeze()

    # Only run this step if a valid range tuple is provided
    if valid_range:
        mask = ((band < valid_range[0]) | (band > valid_range[1]))
        band = band.where(~xr.where(mask, True, False))

    return band

# function to apply cloud mask and caluclate ndvi (mean)
def mask_crop_ndvi(all_bands,
                   crop_bound,
                   pixel_qa,
                   masked_values):
    """Open and mask a single landsat band using a pixel_qa layer.

    Parameters
    -----------
    all_bands : list
        a list containing the xarray objects for landsat  bands 4 and  5
    crop_bound: geopandas GeoDataFrame
        A geopandas dataframe to be used to crop the raster data using
        rasterio mask().
    pixel_qa: xarray DataArray
        An xarray DataArray with pixel qa values that have not yet been
        turned into a mask (0s and 1s)
    masked_values: list
        A list of cloud mask values

    Returns
    -----------
    ndvi_crop : Xarray Dataset
        a cropped and masked xarray object containing NDVI values
    """

    crop_json = crop_bound.geometry

    # Clip pixel qa cloud mask layer
    cl_mask_crop = pixel_qa.rio.clip(crop_json)

    # Calculate NDVI
    ndvi_xr = (all_bands[1]-all_bands[0]) / (all_bands[1]+all_bands[0])
    # Clip NDVI layer
    ndvi_crop = ndvi_xr.rio.clip(crop_json)

    # Apply cloud mask to NDVI
    ndvi_crop = ndvi_crop.where(~pixel_qa.isin(masked_values))

    return ndvi_crop.mean().values

#third function containing the above two
def output_ndvi(scene_path, crop_extent):
    """Function that combines the above two functions to open, stack, crop
    and mask Landsat scenes.

    Parameters
    -----------
    scene_path : string
        A path to the array to be opened
    crop_extent: vector
        A path to the vector file to use as crop boundary

    Returns
    -----------
    ndvi_crop : ndarray Dataset
        a cropped and masked ndarray object containing mean NDVI values
    """

    band_paths = glob(os.path.join(scene_path, "*band*[4-5].tif"))
    pixel_qa_path = glob(os.path.join(scene_path, "*qa*"))
    landsat_qa = rxr.open_rasterio(pixel_qa_path[0], masked = True).squeeze()

    high_cloud_confidence = em.pixel_flags["pixel_qa"]["L8"]    ["High Cloud Confidence"]
    cloud = em.pixel_flags["pixel_qa"]["L8"]["Cloud"]
    cloud_shadow = em.pixel_flags["pixel_qa"]["L8"]["Cloud Shadow"]

    all_masked_values = cloud_shadow + cloud + high_cloud_confidence

    #run clean_bands function
    bands = []
    for band_path in band_paths:
        cleaned_band = open_clean_bands(band_path = band_path,
                                    crop_extent = crop_extent,
                                    valid_range = (0, 10000))
        bands.append(cleaned_band)

    #run mask_crop function
    ndvi = mask_crop_ndvi(all_bands = bands,
                        crop_bound = crop_extent,
                        pixel_qa = landsat_qa,
                        masked_values = all_masked_values)

    #print(ndvi)
    return ndvi

ndvi_list = []

for site_file in study_sites:
    all_dirs = glob(os.path.join(site_file, 'landsat-crop', '*'))
    site_name = os.path.basename(os.path.normpath(site_file))
    crop_bound_path = os.path.join(site_file, 'vector', site_name +
                                  "-crop.shp")
    crop_bound = gpd.read_file(crop_bound_path)

    for adir in all_dirs:

        date = adir[-29 : -21]
        ndvi = output_ndvi(adir, crop_bound)
        ndvi_list.append([site_name, date, ndvi])

ndvi_final_df = pd.DataFrame(ndvi_list,
                    columns = ["site", "date", "mean_ndvi"])

ndvi_final_df['date'] = pd.to_datetime(ndvi_final_df['date'])

ndvi_final_df.set_index('date', inplace = True)

#this was used as .dropna did not seem to work
ndvi_df = ndvi_final_df[ndvi_final_df['mean_ndvi'] > 0]

ndvi_df

# In[6]:

# This cell  is testing your data output above

student_ndvi_ts_single_site = _

single_scene_points = 0

# Ensure the data is stored in a dataframe.
if isinstance(student_ndvi_ts_single_site, pd.DataFrame):
    print('\u2705 Your data is stored in a DataFrame!')
    single_scene_points += 1
else:
    print('\u274C It appears your data is not stored in a DataFrame. ',
          'To see what type of object your data is stored in, check its type with type(object)')

# Ensure that the date column is the index
if isinstance(student_ndvi_ts_single_site.index, pd.core.indexes.datetimes.DatetimeIndex):
    print('\u2705 You have the index set to the date column!')
    single_scene_points += 2
else:
    print('\u274C You do not have the index set to the date column.')

# Ensure that the date column is datetime
if isinstance(student_ndvi_ts_single_site.index[0], pd._libs.tslibs.timestamps.Timestamp):
    print('\u2705 The data in your date column is datetime!')
    single_scene_points += 2
else:
    print('\u274C The data in your date column is not datetime.')

# Ensure the site name is correct
if student_ndvi_ts_single_site.site.values[0] == 'HARV':
    print('\u2705 You have the correct site name!')
    single_scene_points += 5
else:
    print('\u274C You do not have the correct site name.')

if np.allclose(0.281131628228094, student_ndvi_ts_single_site.mean_ndvi.values[0]):
    print('\u2705 You have the correct mean NDVI value!')
    single_scene_points += 5
else:
    print('\u274C You do not have the correct mean ndvi value.')

print("\n \u27A1 You received {} out of 15 points for creating a dataframe.".format(
    single_scene_points))
single_scene_points

# In[7]:

# Create dataframe of mean NDVI in this cell using the functions created above
# Important: to use the ungraded tests below as a sanity check,
# name your columns: mean_ndvi and site
# Call the dataframe at the end of the cell so the tests run on it!
# Be sure that the date column is an index of type date
# HINT: the time series lessons may help you remember how to do this!

# ## Task 2:
#
# In the cell below, process all of the landsat scenes. Create a DataFrame that contains the following
# information for each scene
#
#
# |   | index  | site  | mean_ndvi  |
# |---|---|---|---|
# | Date  |   |   |   |
# |  2017-01-07  | 0  | SJER  | .4  |
#
# Be sure to call your dataframe at the end of the cell to ensure autograding works.
# HINT: FOR THIS STEP, leave any rows containing missing values (`NAN`).

# In[8]:

# Create dataframe of NDVI including the cleaning data to deal with clouds

# Important: to use the ungraded tests below as a sanity check,
# name your columns: mean_ndvi and site
# Don't forget to set date as the index and make the values of type datetime

#export csv file
csv_output_path = os.path.join('ndvi-automation', 'outputs', 'mean_ndvi.csv')
ndvi_df.to_csv(csv_output_path)

ndvi_df

# In[9]:

# Last sanity check before creating your plot (10 points)

# Ensure that you call your dataframe at the bottom of the cell above
# and that it has columns called: mean_ndvi and site

# Ensure the data is stored in a dataframe.
student_ndvi_df = _

df_points = 0

if isinstance(student_ndvi_df, pd.DataFrame):
    print('\u2705 Your data is stored in a DataFrame!')
    df_points +=2
else:
    print('\u274C It appears your data is not stored in a DataFrame. ',
          'To see what type of object your data is stored in, check its type with type(object)')

# Check that dataframe contains the appropriate number of NAN values
if student_ndvi_df.mean_ndvi.isna().sum() == 15:
    print('\u2705 Correct number of masked data values!')
    df_points +=2
else:
    print('\u274C The amount of null data in your dataframe is incorrect.')

# Ensure that the date column is the index
if isinstance(student_ndvi_df.index, pd.core.indexes.datetimes.DatetimeIndex):
    print('\u2705 You have the index set to the date column!')
    df_points +=3
else:
    print('\u274C You do not have the index set to the date column.')

# Ensure that the date column is datetime
if isinstance(student_ndvi_df.index[0], pd._libs.tslibs.timestamps.Timestamp):
    print('\u2705 The data in your date column is datetime!')
    df_points +=3
else:
    print('\u274C The data in your date column is not datetime.')

# Output for timer, # DO NOT MODIFY
end_time = datetime.now()
total_time = end_time - start_time
print(
    "Your total run time for processing the data was {0}.".format(total_time))

print("\n \u27A1 You received {} out of 10 points for creating a dataframe.".format(
    df_points))

df_points

# In[10]:

# Add only the plot code to this cell

# This is the final figure of mean NDVI
# for both sites across the year
# with data cleaned to deal with clouds

fig, ax = plt.subplots(figsize = (14, 10))

for site, df in ndvi_df.groupby("site"):
    if site in ["HARV"]:
        site_name = 'HARV'
        color = 'red'
    else:
        site_name = 'SJER'
        color = 'green'

    ax.plot(df.index, df.mean_ndvi, label = site_name, marker = 'o',
            color = color)

ax.set(title = "Mean Normalized Difference Vegetation Index (NDVI) for two sites (HARV & SJER) \n Mar 2017 - Dec 2017 (Cleaned data)",
       xlabel = "Month", ylabel = "Mean NDVI")

ax.legend(title = "Site")

### DO NOT REMOVE LINES BELOW ###
final_masked_solution = nb.convert_axes(plt, which_axes="current")

# In[11]:

# Ignore this cell for the autograding tests

# In[12]:

# Ignore this cell for the autograding tests

# # Question 1 (10 points)
#
# Imagine that you are planning NEON’s upcoming flight season to capture remote sensing data in these locations and want to ensure that you fly the area when the vegetation is the most green.
#
# When would you recommend the flights take place for each site?
#
# Answer the question in 2-3 sentences in the Markdown cell below.

# For HARV I would advise flying in July where NDVI values are maximum. However there appears to be several no data landsat scenes here so it may be more difficult to capture information. For SJER I would advise flying around April.

# # Question 2 (10 points)
#
# How could you modify your workflow to look at vegetation changes over time in each site?
#
# Answer the question in 2-3 sentences in the Markdown cell below.

# First, I would monitor annual datasets, rather than monthly to remove the effects of seasonal variation on ndvi response. Then I would look over a longer time range (years) to assess any variation in ndvi. Perhaps one could also assess seasonal ndvi variation across the years by taking the average ndvi from blocks of three months (to represtent the seasons) and then comparing this accross several years of data.

# # Do not edit this cell! (10 points)
#
# The notebook includes:
# * additional Markdown cells throughout the notebook to describe:
#     * the data that you used - and where it is from
#     * how data are being processing
#     * how the code is optimized to run fast and be more concise

# # Do not edit this cell! (20 points)
#
# The notebook will also be checked for overall clean code requirements as specified at the **top** of this notebook. Some of these requirements include (review the top cells for more specifics):
#
# * Notebook begins at cell [1] and runs on any machine in its entirety.
# * PEP 8 format is applied throughout (including lengths of comment and code lines).
# * No additional code or imports in the notebook that is not needed for the workflow.
# * Notebook is fully reproducible. This means:
#    * reproducible paths using the os module.
#    * data downloaded using code in the notebook.
#    * all imports at top of notebook.

# ## BONUS - Export a  .CSV File to Share (10 points possible)
#
# This is optional - if you export a **.csv** file with the columns specified above: Site, Date and NDVI Value you can get an additional 10 points.
#
# * FULL CREDIT: File exists in csv format and contains the columns specified.
# We will check your github repo for this file!
#