import numpy as np
import pdb
from scipy.interpolate import interp1d
from scipy.stats import pearsonr
from scipy.stats import norm
import scipy.signal as signal
from matplotlib import gridspec
import matplotlib
matplotlib.use('Qt4Agg')
import matplotlib.pyplot as plt
from matplotlib.widgets import RadioButtons, Button, CheckButtons
import sys
from types import *

class Estimateline:
    '''Class to manually estimate where lines are located'''
    def __init__(self,pspec,ax5,uline):
        print 'If redshift calibration appears correct, hit "Accept and Close". '\
        'Otherwise, "right click" approx. where the '+uline+' line is in the plotted spectrum. '\
        'The program will re-correlate based on this guess.'
        self.ax5 = ax5
        self.cid3 = pspec.figure.canvas.mpl_connect('button_press_event',self.onclick)

    def on_key_press(self,event):
        if event.key == 'shift':
            self.shift_is_held = True

    def on_key_release(self, event):
        if event.key == 'shift':
            self.shift_is_held = False

    def onclick(self,event):
        if event.inaxes == self.ax5:
            if event.button == 3:
                print 'xdata=%f, ydata%f'%(event.xdata, event.ydata)
                self.lam = event.xdata
                plt.close()
            '''
            if event.button == 1:
                #if self.shift_is_held:
                #    print 'xdata=%f, ydata%f'%(event.xdata, event.ydata)
                #    self.lam = event.xdata
                #    plt.close()
                #else:
                plt.close()
            '''
        else: return

class z_est:
    def __init__(self,lower_w=3900.0,upper_w=5500.0,lower_z=0.01,upper_z=0.35,z_res=3.0e-5,skip_initial_priors=False):
        '''
        Initialize redshift estimate parameters
        '''

        #preconditions
        assert lower_w and upper_w, "wavelength bounds must have values"
        assert lower_z and upper_z, "redshift bounds must have values"
        assert (type(lower_w) == IntType or type(lower_w) == FloatType) and (type(upper_w) == IntType or \
                    type(upper_w) == FloatType), "wavelength bounds must be integers or floats"
        assert (type(lower_z) == IntType or type(lower_z) == FloatType) and (type(upper_z) == IntType or \
                    type(upper_z) == FloatType), "wavelength bounds must be integers or floats"
        assert lower_w < upper_w, "lower_w must be < upper_w"
        assert lower_z < upper_z, "lower_z must be < upper_z"

        #set class attributes
        self.lower_w = lower_w
        self.upper_w = upper_w
        self.lower_z = lower_z
        self.upper_z = upper_z
        self.z_res = z_res

        #create redshift array and initialize correlation value array
        self.ztest = np.arange(self.lower_z,self.upper_z,self.z_res)
        self.corr_val_i = np.zeros(self.ztest.size)

        #set redshift prior flag
        if skip_initial_priors:
            self.est_pre_z = '3'
            self.uline_n = 'HK'
            self.uline = 3950.0
            self.z_prior_width = 0.06
        else:
            self.est_pre_z = raw_input('(1) Use a known prior [Examples: median of known redshifts. Galaxy photoz measurements] \n'\
                                    '(2) View spectrum and specify a redshift prior \n'\
                                    '(3) No prior\n')

        #catch and correct false entry
        _est_enter = False
        self.uline_n = raw_input('What is the name of a spectral line you wish to use to identify redshift priors? '\
                                            '[Default: HK]: ')
        if not self.uline_n:
            self.uline_n = 'HK'
        self.uline = raw_input('Please list the approx. rest wavelength (in angstroms) of that line you seek to identify in your spectra '\
                                            '[Default: HK lines are at about 3950]: ')
        if self.uline:
            self.uline = np.float(self.uline)
        else:
            self.uline = 3950.0
        while not _est_enter:
            if self.est_pre_z == '1':
                self.z_prior_width = 0.06
                print 'redshift prior width has been set to',self.z_prior_width
                _est_enter = True
            elif self.est_pre_z == '2':
                self.z_prior_width = 0.06
                print 'redshift prior width has been set to',self.z_prior_width
                _est_enter = True
            elif self.est_pre_z == '3':
                self.z_prior_width = 0.06
                _est_enter = True
            else:
                self.est_pre_z = raw_input('Incorrect entry: Please enter either (1), (2), or (3).')

        #remind user to set the correct values in next step
        if self.est_pre_z == '1':
            print 'Make sure to set the gal_prior argument to the value of the known redshift prior: '\
                '[Example: z_est.redshift_estimate(gal_prior=0.1)]'

        #postconditions
        assert self.est_pre_z, "Must define redshift prior flag"
        assert self.est_pre_z == '1' or self.est_pre_z == '2' or self.est_pre_z == '3', \
                         "Incorrect string value for prior"

    def redshift_estimate(self,early_type_wave,early_type_flux,wave,Flux_science,plotlines=None,template_id='Galaxy',gal_prior=None):
        '''
        estimate redshift for object
        '''
        #manage redshift prior
        self.gal_prior = gal_prior
        self.template_number = template_id
        self.plotlines = plotlines

        #continuum subtract
        Flux_sc = Flux_science - signal.medfilt(Flux_science,171)
        early_type_flux_sc = early_type_flux - signal.medfilt(early_type_flux,171)

        #handle single redshift prior flag
        if self.est_pre_z == '1':
            if self.gal_prior:
                self.pre_z_est = self.gal_prior
            else:
                nospec = raw_input('You said you are either using a spectroscopic or photometric redshift prior. '\
                                        'You need to specify a prior value! Either enter a number in now or type (q) to exit')
                if nospec == 'q':
                    sys.exit()
                elif not nospec:
                    sys.exit()
                else:
                    self.gal_prior = np.float(nospec)
                    self.pre_z_est = self.gal_prior

        #handle user prior flag
        if self.est_pre_z == '2':
            print 'Take a look at the plotted galaxy spectrum and note, approximately, at what wavelength do you see the '+self.uline_n+' line. '\
                    'Then close the plot and enter that wavelength in angstroms.'
            plt.plot(wave,Flux_science)
            plt.xlim(self.lower_w,self.upper_w)
            plt.show()
            line_init = raw_input(self.uline_n+' approx. wavelength (A): ')
            self.pre_z_est = np.float(line_init)/self.uline - 1

        #handle no prior flag
        if self.est_pre_z == '3':
            self.pre_z_est = None

        redshift_est,cor,ztest,corr_val = self._cross_cor(self.pre_z_est,self.z_prior_width,early_type_wave,early_type_flux_sc,wave,Flux_sc)

        self.qualityval = 1
        self.first_pass = True
        self.skip_spec_flag = False
        self._GUI_display(redshift_est,ztest,corr_val,wave,Flux_science)
        #self.line_est = Estimateline(self.pspec,ax)
        #plt.show()
        try:
            self.pre_lam_est = self.line_est.lam
            self.pre_z_est = self.pre_lam_est/3950.0 - 1.0
            self.first_pass = False
            redshift_est,cor,ztest,corr_val = self._cross_cor(self.pre_z_est,self.z_prior_width,early_type_wave,early_type_flux_sc,wave,Flux_sc)
            print 'redshift est:',redshift_est
            self._GUI_display(redshift_est,ztest,corr_val,wave,Flux_science)
            redshift_est = self.spectra2.finalz
        except AttributeError:
            pass
        print 'zpy redshift estimate',redshift_est
        return redshift_est,cor,ztest,corr_val,self.qualityval

    def _cross_cor(self,z_est,unc,early_type_wave,early_type_flux,wave,Flux_sc):
        '''
        This function cross-correlates a continuum subtracted template spectrum with a continuum subtracted observed spectrum.
        It then returns an estimate of the redshift, the correlation value at that redshift, the array of redshifts tested,
        and the unnormalized correlation value.
        '''

        #loop over each possible redshift to compute correlation values
        for i in range(self.ztest.size):
            z = self.ztest[i]
            #redshift the template wavelengths
            wshift = early_type_wave*(1+z)
            #identify the wavelength diff between the lower wave limit and the redshifted template spectrum
            wavediff = np.min(wshift - self.lower_w)

            #if the limit is above the minimum wavelength of the redshifted template spectrum...
            if wavediff < 0:
                wave_range = wave[np.where((wave<self.upper_w)&(wave>self.lower_w))]
                Flux_range = Flux_sc[np.where((wave<self.upper_w)&(wave>self.lower_w))]
            #if the limit is below the minimum wavelength of the redshifted template spectrum...
            else:
                wave_range = wave[np.where((wave<self.upper_w+wavediff)&(wave>self.lower_w+wavediff))]
                Flux_range = Flux_sc[np.where((wave<self.upper_w+wavediff)&(wave>self.lower_w+wavediff))]

            #interpolate the redshifted template spectrum and estimate the flux at the observed spectrum wavelengths
            inter = interp1d(wshift,early_type_flux)
            et_flux_range = inter(wave_range)

            #calculate the pearson r correlation value between the observed and template flux
            self.corr_val_i[i] = pearsonr(et_flux_range,Flux_range)[0]

        #normalize the correlation values as a function of redshift
        corr_val = (self.corr_val_i[np.isfinite(self.corr_val_i)]+1)/np.trapz((self.corr_val_i[np.isfinite(self.corr_val_i)]+1),self.ztest[np.isfinite(self.corr_val_i)])
        self.ztest = self.ztest[np.isfinite(self.corr_val_i)]

        #multiply in prior to likelihood if specified
        self.corr_prior = np.zeros(self.ztest.size)
        if z_est:
            rv = norm(z_est,unc)
            corr_val = corr_val * rv.pdf(self.ztest)
            self.corr_prior = rv.pdf(self.ztest)

        #make redshift estimate
        redshift_est = (self.ztest[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))])[np.where(corr_val[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))] == np.max(corr_val[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))]))]

        #save correlation value at maximum redshift likelihood
        cor = (self.corr_val_i[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))])[np.where(corr_val[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))] == np.max(corr_val[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))]))]

        return redshift_est[0], cor, self.ztest,corr_val

    def _GUI_display(self,redshift_est,ztest,corr_val,wave,flux_sc):
        '''Display the spectrum and reference lines.'''
        self.fig = plt.figure(figsize=(10, 8))
        gs = gridspec.GridSpec(2, 1, height_ratios=[2, 1])
        ax2 = plt.subplot(gs[0])
        ax = plt.subplot(gs[1])

        #maximize window
        figManager = plt.get_current_fig_manager()
        figManager.window.showMaximized()

        plt.subplots_adjust(top=0.96,bottom=0.1,left=0.04,right=0.80)

        ax.plot(ztest,corr_val,'b')
        pspec_corr = ax.axvline(redshift_est,color='k',ls='--')
        ax.fill_between(self.ztest,np.zeros(self.ztest.size),self.corr_prior,facecolor='grey',alpha=0.6)
        ax.set_xlabel('Redshift')
        ax.set_ylabel('Correlation')

        self.pspec, = ax2.plot(wave,flux_sc)
        ax2.set_ylim(np.min(flux_sc),np.max(flux_sc))
        ax2.set_xlim(wave[0],wave[-1])

        #plot lines
        vlin_pspecs = []
        if self.plotlines == None:
            self.plotlines = {'HKlines':[3968.5,3933.7],'emission_lines':[3725.0,4959.0,5007.0],'absorption_lines':[4102.9,4304.0,4862.0,5175.0],'sky_lines':[]}
        else: pass
        try:
            for vlin in self.plotlines['HKlines']:
                vlin_pspecs.append(ax2.axvline(vlin*(1+redshift_est),ls='--',alpha=0.7,c='red'))
        except KeyError: pass
        try:
            for vlin in self.plotlines['emission_lines']:
                vlin_pspecs.append(ax2.axvline(vlin*(1+redshift_est),ls='--',alpha=0.7,c='blue'))
        except KeyError: pass
        try:
            for vlin in self.plotlines['absorption_lines']:
                vlin_pspecs.append(ax2.axvline(vlin*(1+redshift_est),ls='--',alpha=0.7,c='orange'))
        except KeyError: pass
        try:
            for vlin in self.plotlines['sky_lines']:
                vlin_pspecs.append(ax2.axvline(vlin*(1+redshift_est),ls='--',alpha=0.7,c='grey'))
        except KeyError: pass

        if self.first_pass:
            self.line_est = Estimateline(self.pspec,ax2,self.uline_n)

        self.fig.text(0.83, 0.9, 'Template {0}'.format(self.template_number), bbox=dict(facecolor='white', alpha=1.),fontsize=18)

        rax = plt.axes([0.83, 0.43, 0.15, 0.22])
        if self.qualityval == 1:
            radio = RadioButtons(rax, ('1 - No Clue      ','2 - Slight\n    Chance', '3 - Maybe', '4 - Probably', '5 - Clear'))
        else:
            radio = RadioButtons(rax, ('1 - No Clue      ','2 - Slight\n    Chance', '3 - Maybe', '4 - Probably', '5 - Clear'),active=1)
        def qualfunc(label):
            if label == '5 - Clear':
                self.qualityval = 5
            elif label == '4 - Probably':
                self.qualityval = 4
            elif label == '3 - Maybe':
                self.qualityval = 3
            elif label == '2 - Slight\n    Chance':
                self.qualityval = 2
            else:
                self.qualityval = 1

        radio.on_clicked(qualfunc)
        closeax = plt.axes([0.83, 0.3, 0.15, 0.1])
        button = Button(closeax, 'Accept & Close', hovercolor='0.975')
        def closeplot(event):
            plt.close()
        button.on_clicked(closeplot)
        skip_spec_ax = plt.axes([0.83, 0.94, 0.15, 0.04])
        skip_button = Button(skip_spec_ax, 'skip spectra', hovercolor='0.975')
        def skip_spec(event):
            plt.close()
            self.qualityval = 0
            self.skip_spec_flag = True
        skip_button.on_clicked(skip_spec)
        ax2.set_xlim(self.lower_w,self.upper_w)
        ax2.set_xlabel('Wavelength (A)')
        ax2.set_ylabel('Counts')

        plt.draw()

        if not self.first_pass:
            self.spectra2 = DragSpectra(vlin_pspecs,pspec_corr,redshift_est,ax2)
            self.fig.canvas.mpl_connect('motion_notify_event',self.spectra2.on_motion)
            self.fig.canvas.mpl_connect('button_press_event',self.spectra2.on_press)
            self.fig.canvas.mpl_connect('button_release_event',self.spectra2.on_release)
        plt.show()

class DragSpectra:
    '''Class to drage the spectra back and forth to match lines of interest'''
    def __init__(self,vlin_spectra,corr_spec,redshift_estimate,ax5):
        self.ax5 = ax5
        self.corr_spec = corr_spec
        self.yzs = self.corr_spec.get_data()[1]
        print 'begin shift'
        self.vlin_spectra = vlin_spectra
        self.vline_ys = vlin_spectra[0].get_data()[1]
        self.pressed = False
        self.finalz = redshift_estimate
        #figure.canvas.mpl_connect('motion_notify_event',self.on_motion)
        #figure.canvas.mpl_connect('button_press_event',self.on_press)
        #figure.canvas.mpl_connect('button_release_event',self.on_release)

    def on_motion(self,evt):
        if self.pressed:
            #dx = evt.xdata - self.mouse_x
            #print "%d %d" % (evt.xdata,self.mouse_x)
            newz = ((evt.xdata/self.mouse_x)*(1.+self.z_on_press))-1.  #((1. + (dx/self.mouse_x))*(1.+self.z0))-1.
            newxs = self.vline_lams*(evt.xdata/self.mouse_x) # equivalent to spec_x*((1+newz)/(1+z0))
            for i in np.arange(len(self.vlin_spectra)):
                self.vlin_spectra[i].set_data([newxs[i], newxs[i]], self.vline_ys)

            self.corr_spec.set_data([newz, newz], self.yzs)
            plt.draw()

    def on_press(self,evt):
        if evt.inaxes == self.ax5:
            self.mouse_x = evt.xdata
            self.z_on_press = self.corr_spec.get_data()[0][0]
            self.vline_lams = np.array([self.vlin_spectra[x].get_data()[0][0] for x in np.arange(len(self.vlin_spectra))])
            self.pressed = True
        else: return

    def on_release(self,evt):
        if evt.inaxes == self.ax5:
            self.pressed = False
            try:
                self.finalz = self.corr_spec.get_data()[0][0]
            except AttributeError:
                self.finalz = self.finalz
        else: return

if __name__ == '__main__':
    R = z_est()


import torch
import math
import numpy
import os
import PIL
import PIL.Image
import sys
try:
	from correlation import correlation # the custom cost volume layer
except:
	sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), 'correlation')); import correlation # you should consider upgrading python

Backward_tensorGrid={}
def Backward(tensorInput, tensorFlow):
	if str(tensorFlow.size()) not in Backward_tensorGrid:
		tensorHorizontal = torch.linspace(-1.0, 1.0, tensorFlow.size(3)).view(1, 1, 1, tensorFlow.size(3)).expand(tensorFlow.size(0), -1, tensorFlow.size(2), -1)
		tensorVertical = torch.linspace(-1.0, 1.0, tensorFlow.size(2)).view(1, 1, tensorFlow.size(2), 1).expand(tensorFlow.size(0), -1, -1, tensorFlow.size(3))

		Backward_tensorGrid[str(tensorFlow.size())] = torch.cat([ tensorHorizontal, tensorVertical ], 1).cuda()
	# end

	tensorFlow = torch.cat([ tensorFlow[:, 0:1, :, :] / ((tensorInput.size(3) - 1.0) / 2.0), tensorFlow[:, 1:2, :, :] / ((tensorInput.size(2) - 1.0) / 2.0) ], 1)

	return torch.nn.functional.grid_sample(input=tensorInput, grid=(Backward_tensorGrid[str(tensorFlow.size())] + tensorFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='zeros', align_corners=True)
# end

##########################################################

class LiteFlowNet(torch.nn.Module):
	def __init__(self, model_path):
		super(LiteFlowNet, self).__init__()

		class Features(torch.nn.Module):
			def __init__(self):
				super(Features, self).__init__()

				self.moduleOne = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=3, out_channels=32, kernel_size=7, stride=1, padding=3),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
				)

				self.moduleTwo = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=2, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
				)

				self.moduleThr = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
				)

				self.moduleFou = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=64, out_channels=96, kernel_size=3, stride=2, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
				)

				self.moduleFiv = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=96, out_channels=128, kernel_size=3, stride=2, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
				)

				self.moduleSix = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=128, out_channels=192, kernel_size=3, stride=2, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
				)
			# end

			def forward(self, tensorInput):
				tensorOne = self.moduleOne(tensorInput)
				tensorTwo = self.moduleTwo(tensorOne)
				tensorThr = self.moduleThr(tensorTwo)
				tensorFou = self.moduleFou(tensorThr)
				tensorFiv = self.moduleFiv(tensorFou)
				tensorSix = self.moduleSix(tensorFiv)

				return [ tensorOne, tensorTwo, tensorThr, tensorFou, tensorFiv, tensorSix ]
			# end
		# end

		class Matching(torch.nn.Module):
			def __init__(self, intLevel):
				super(Matching, self).__init__()

				self.dblBackward = [ 0.0, 0.0, 10.0, 5.0, 2.5, 1.25, 0.625 ][intLevel]

				if intLevel != 2:
					self.moduleFeat = torch.nn.Sequential()

				elif intLevel == 2:
					self.moduleFeat = torch.nn.Sequential(
						torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=1, stride=1, padding=0),
						torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
					)

				# end

				if intLevel == 6:
					self.moduleUpflow = None

				elif intLevel != 6:
					self.moduleUpflow = torch.nn.ConvTranspose2d(in_channels=2, out_channels=2, kernel_size=4, stride=2, padding=1, bias=False, groups=2)

				# end

				if intLevel >= 4:
					self.moduleUpcorr = None

				elif intLevel < 4:
					self.moduleUpcorr = torch.nn.ConvTranspose2d(in_channels=49, out_channels=49, kernel_size=4, stride=2, padding=1, bias=False, groups=49)

				# end

				self.moduleMain = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=49, out_channels=128, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=32, out_channels=2, kernel_size=[ 0, 0, 7, 5, 5, 3, 3 ][intLevel], stride=1, padding=[ 0, 0, 3, 2, 2, 1, 1 ][intLevel])
				)
			# end

			def forward(self, tensorFirst, tensorSecond, tensorFeaturesFirst, tensorFeaturesSecond, tensorFlow):
				tensorFeaturesFirst = self.moduleFeat(tensorFeaturesFirst)
				tensorFeaturesSecond = self.moduleFeat(tensorFeaturesSecond)

				if tensorFlow is not None:
					tensorFlow = self.moduleUpflow(tensorFlow)
				# end

				if tensorFlow is not None:
					tensorFeaturesSecond = Backward(tensorInput=tensorFeaturesSecond, tensorFlow=tensorFlow * self.dblBackward)
				# end

				if self.moduleUpcorr is None:
					tensorCorrelation = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tensorFirst=tensorFeaturesFirst, tensorSecond=tensorFeaturesSecond, intStride=1), negative_slope=0.1, inplace=False)

				elif self.moduleUpcorr is not None:
					tensorCorrelation = self.moduleUpcorr(torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tensorFirst=tensorFeaturesFirst, tensorSecond=tensorFeaturesSecond, intStride=2), negative_slope=0.1, inplace=False))

				# end

				return (tensorFlow if tensorFlow is not None else 0.0) + self.moduleMain(tensorCorrelation)
			# end
		# end

		class Subpixel(torch.nn.Module):
			def __init__(self, intLevel):
				super(Subpixel, self).__init__()

				self.dblBackward = [ 0.0, 0.0, 10.0, 5.0, 2.5, 1.25, 0.625 ][intLevel]

				if intLevel != 2:
					self.moduleFeat = torch.nn.Sequential()

				elif intLevel == 2:
					self.moduleFeat = torch.nn.Sequential(
						torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=1, stride=1, padding=0),
						torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
					)

				# end

				self.moduleMain = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=[ 0, 0, 130, 130, 194, 258, 386 ][intLevel], out_channels=128, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=32, out_channels=2, kernel_size=[ 0, 0, 7, 5, 5, 3, 3 ][intLevel], stride=1, padding=[ 0, 0, 3, 2, 2, 1, 1 ][intLevel])
				)
			# end

			def forward(self, tensorFirst, tensorSecond, tensorFeaturesFirst, tensorFeaturesSecond, tensorFlow):
				tensorFeaturesFirst = self.moduleFeat(tensorFeaturesFirst)
				tensorFeaturesSecond = self.moduleFeat(tensorFeaturesSecond)

				if tensorFlow is not None:
					tensorFeaturesSecond = Backward(tensorInput=tensorFeaturesSecond, tensorFlow=tensorFlow * self.dblBackward)
				# end

				return (tensorFlow if tensorFlow is not None else 0.0) + self.moduleMain(torch.cat([ tensorFeaturesFirst, tensorFeaturesSecond, tensorFlow ], 1))
			# end
		# end

		class Regularization(torch.nn.Module):
			def __init__(self, intLevel):
				super(Regularization, self).__init__()

				self.dblBackward = [ 0.0, 0.0, 10.0, 5.0, 2.5, 1.25, 0.625 ][intLevel]

				self.intUnfold = [ 0, 0, 7, 5, 5, 3, 3 ][intLevel]

				if intLevel >= 5:
					self.moduleFeat = torch.nn.Sequential()

				elif intLevel < 5:
					self.moduleFeat = torch.nn.Sequential(
						torch.nn.Conv2d(in_channels=[ 0, 0, 32, 64, 96, 128, 192 ][intLevel], out_channels=128, kernel_size=1, stride=1, padding=0),
						torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
					)

				# end

				self.moduleMain = torch.nn.Sequential(
					torch.nn.Conv2d(in_channels=[ 0, 0, 131, 131, 131, 131, 195 ][intLevel], out_channels=128, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
				)

				if intLevel >= 5:
					self.moduleDist = torch.nn.Sequential(
						torch.nn.Conv2d(in_channels=32, out_channels=[ 0, 0, 49, 25, 25, 9, 9 ][intLevel], kernel_size=[ 0, 0, 7, 5, 5, 3, 3 ][intLevel], stride=1, padding=[ 0, 0, 3, 2, 2, 1, 1 ][intLevel])
					)

				elif intLevel < 5:
					self.moduleDist = torch.nn.Sequential(
						torch.nn.Conv2d(in_channels=32, out_channels=[ 0, 0, 49, 25, 25, 9, 9 ][intLevel], kernel_size=([ 0, 0, 7, 5, 5, 3, 3 ][intLevel], 1), stride=1, padding=([ 0, 0, 3, 2, 2, 1, 1 ][intLevel], 0)),
						torch.nn.Conv2d(in_channels=[ 0, 0, 49, 25, 25, 9, 9 ][intLevel], out_channels=[ 0, 0, 49, 25, 25, 9, 9 ][intLevel], kernel_size=(1, [ 0, 0, 7, 5, 5, 3, 3 ][intLevel]), stride=1, padding=(0, [ 0, 0, 3, 2, 2, 1, 1 ][intLevel]))
					)

				# end

				self.moduleScaleX = torch.nn.Conv2d(in_channels=[ 0, 0, 49, 25, 25, 9, 9 ][intLevel], out_channels=1, kernel_size=1, stride=1, padding=0)
				self.moduleScaleY = torch.nn.Conv2d(in_channels=[ 0, 0, 49, 25, 25, 9, 9 ][intLevel], out_channels=1, kernel_size=1, stride=1, padding=0)
			# eny

			def forward(self, tensorFirst, tensorSecond, tensorFeaturesFirst, tensorFeaturesSecond, tensorFlow):
				tensorDifference = (tensorFirst - Backward(tensorInput=tensorSecond, tensorFlow=tensorFlow * self.dblBackward)).pow(2.0).sum(1, True).sqrt().detach()

				tensorDist = self.moduleDist(self.moduleMain(torch.cat([ tensorDifference, tensorFlow - tensorFlow.view(tensorFlow.size(0), 2, -1).mean(2, True).view(tensorFlow.size(0), 2, 1, 1), self.moduleFeat(tensorFeaturesFirst) ], 1)))
				tensorDist = tensorDist.pow(2.0).neg()
				tensorDist = (tensorDist - tensorDist.max(1, True)[0]).exp()

				tensorDivisor = tensorDist.sum(1, True).reciprocal()

				tensorScaleX = self.moduleScaleX(tensorDist * torch.nn.functional.unfold(input=tensorFlow[:, 0:1, :, :], kernel_size=self.intUnfold, stride=1, padding=int((self.intUnfold - 1) / 2)).view_as(tensorDist)) * tensorDivisor
				tensorScaleY = self.moduleScaleY(tensorDist * torch.nn.functional.unfold(input=tensorFlow[:, 1:2, :, :], kernel_size=self.intUnfold, stride=1, padding=int((self.intUnfold - 1) / 2)).view_as(tensorDist)) * tensorDivisor

				return torch.cat([ tensorScaleX, tensorScaleY ], 1)
			# end
		# end

		self.moduleFeatures = Features()
		self.moduleMatching = torch.nn.ModuleList([ Matching(intLevel) for intLevel in [ 2, 3, 4, 5, 6 ] ])
		self.moduleSubpixel = torch.nn.ModuleList([ Subpixel(intLevel) for intLevel in [ 2, 3, 4, 5, 6 ] ])
		self.moduleRegularization = torch.nn.ModuleList([ Regularization(intLevel) for intLevel in [ 2, 3, 4, 5, 6 ] ])

		self.load_state_dict(torch.load(model_path))
	# end

	def forward(self, tensorFirst, tensorSecond):
		tensorFirst[:, 0, :, :] = tensorFirst[:, 0, :, :] - 0.411618
		tensorFirst[:, 1, :, :] = tensorFirst[:, 1, :, :] - 0.434631
		tensorFirst[:, 2, :, :] = tensorFirst[:, 2, :, :] - 0.454253

		tensorSecond[:, 0, :, :] = tensorSecond[:, 0, :, :] - 0.410782
		tensorSecond[:, 1, :, :] = tensorSecond[:, 1, :, :] - 0.433645
		tensorSecond[:, 2, :, :] = tensorSecond[:, 2, :, :] - 0.452793

		tensorFeaturesFirst = self.moduleFeatures(tensorFirst)
		tensorFeaturesSecond = self.moduleFeatures(tensorSecond)

		tensorFirst = [ tensorFirst ]
		tensorSecond = [ tensorSecond ]

		for intLevel in [ 1, 2, 3, 4, 5 ]:
			tensorFirst.append(torch.nn.functional.interpolate(input=tensorFirst[-1], size=(tensorFeaturesFirst[intLevel].size(2), tensorFeaturesFirst[intLevel].size(3)), mode='bilinear', align_corners=False))
			tensorSecond.append(torch.nn.functional.interpolate(input=tensorSecond[-1], size=(tensorFeaturesSecond[intLevel].size(2), tensorFeaturesSecond[intLevel].size(3)), mode='bilinear', align_corners=False))
		# end

		tensorFlow = None

		for intLevel in [ -1, -2, -3, -4, -5 ]:
			tensorFlow = self.moduleMatching[intLevel](tensorFirst[intLevel], tensorSecond[intLevel], tensorFeaturesFirst[intLevel], tensorFeaturesSecond[intLevel], tensorFlow)
			tensorFlow = self.moduleSubpixel[intLevel](tensorFirst[intLevel], tensorSecond[intLevel], tensorFeaturesFirst[intLevel], tensorFeaturesSecond[intLevel], tensorFlow)
			tensorFlow = self.moduleRegularization[intLevel](tensorFirst[intLevel], tensorSecond[intLevel], tensorFeaturesFirst[intLevel], tensorFeaturesSecond[intLevel], tensorFlow)
		# end

		return tensorFlow * 20.0
	# end
# end


import io
import os
import librosa
import torch
import scipy
import numpy as np
import soundfile as sf
from TTS.utils.text import text_to_sequence
from TTS.utils.generic_utils import load_config
from TTS.utils.audio import AudioProcessor
from TTS.models.tacotron import Tacotron
from matplotlib import pylab as plt

class Synthesizer(object):
    def load_model(self, model_path, model_name, model_config, use_cuda):
        model_config = os.path.join(model_path, model_config)
        self.model_file = os.path.join(model_path, model_name)
        print(" > Loading model ...")
        print(" | > model config: ", model_config)
        print(" | > model file: ", self.model_file)
        config = load_config(model_config)
        self.config = config
        self.use_cuda = use_cuda
        self.model = Tacotron(config.embedding_size, config.num_freq,
                              config.num_mels, config.r)
        self.ap = AudioProcessor(
            config.sample_rate,
            config.num_mels,
            config.min_level_db,
            config.frame_shift_ms,
            config.frame_length_ms,
            config.preemphasis,
            config.ref_level_db,
            config.num_freq,
            config.power,
            griffin_lim_iters=60)
        # load model state
        if use_cuda:
            cp = torch.load(self.model_file)
        else:
            cp = torch.load(
                self.model_file, map_location=lambda storage, loc: storage)
        # load the model
        self.model.load_state_dict(cp['model'])
        if use_cuda:
            self.model.cuda()
        self.model.eval()

    def save_wav(self, wav, path):
        wav *= 32767 / max(1e-8, np.max(np.abs(wav)))
        librosa.output.write_wav(path, wav.astype(np.int16),
                                 self.config.sample_rate)

    def tts(self, text):
        text_cleaner = [self.config.text_cleaner]
        wavs = []
        for sen in text.split('.'):
            if len(sen) < 3:
                continue
            sen = sen.strip()
            sen += '.'
            print(sen)
            sen = sen.strip()
            seq = np.array(text_to_sequence(text, text_cleaner))
            chars_var = torch.from_numpy(seq).unsqueeze(0).long()
            if self.use_cuda:
                chars_var = chars_var.cuda()
            mel_out, linear_out, alignments, stop_tokens = self.model.forward(
                chars_var)
            linear_out = linear_out[0].data.cpu().numpy()
            wav = self.ap.inv_spectrogram(linear_out.T)
            # wav = wav[:self.ap.find_endpoint(wav)]
            out = io.BytesIO()
            wavs.append(wav)
            wavs.append(np.zeros(10000))
        self.save_wav(wav, out)
        return out


from pickle import load
from numpy import argmax
from keras.preprocessing.sequence import pad_sequences
from keras.applications.vgg16 import VGG16
from keras.preprocessing.image import load_img
from keras.preprocessing.image import img_to_array
from keras.applications.vgg16 import preprocess_input
from keras.models import Model
from keras.models import load_model

def ext_feats(filename):
    # load in the model
    model = VGG16()
    # re-structure the model
    model = Model(inputs = model.inputs, outputs = model.layers[-2].output)
    # laod the photo
    img = load_img(filename, target_size = (224, 224))
    # convert the image pixels to numpy array
    img = img_to_array(img)
    # reshape data
    img = img.reshape((1, img.shape[0], img.shape[1], img.shape[2]))
    # prep the img for  vgg model
    img = preprocess_input(img)
    # get feats
    feats = model.predict(img, verbose = 0)
    return feats

def word_for_id(integer, token):
    for word, index, in token.word_index.items():
        if index == integer:
            return word
        return None

def gene_desc(model, token, photo, max_len):
    # seed the genertation process
    in_text = 'startseq'
    # iterate over the whole len of the seq
    for i in range(max_len):
        # int encode input seq
        seq = token.texts_to_sequences([in_text])[0]
        # pad input
        seq = pad_sequences([seq], maxlen = max_len)
        # predict the next word
        yhat = model.predict([photo, seq], verbose = 0)
        # convert probability to a int
        yhat = argmax(yhat)
        # map the int to word
        word = word_for_id(yhat, token)
        # stop if the word cannot be mapped
        if word is None:
            break
        #append as each word is being generated
        in_text += ' ' + word
        # stop if predicts end of seq
        if word == 'endseq':
            break
    return in_text

#load the tokenizer
token = load(open('token.pkl', 'rb'))
max_len = 34
# load the model
model = load_model('model-ep020-loss3.431-val_loss3.118.h5')
# load & prepare the photo to be identified
photo = ext_feats('beach.jpg')
# gene desc
desc = gene_desc(model, token, photo, max_len)
print(desc)


import numpy as np
from localization import CheckPoint, Position

class Param:
    def __init__(self):
        self.base_speed = 18
        self.min_speed = 13

        self.speed_decay = 2
        self.max_steer_angle = 60.0

        self.steer_angle_scale = 1
        self.middle_pos_scale = 1

        self.sign_size_1 = 33
        self.sign_size_2 = 37

        self.speed_slow_down = 1

        self.delay_time = 1

        # set up check-points
        self.checkpoint1 = CheckPoint("1", Position(2, -36, 116))
        self.checkpoint2 = CheckPoint("2", Position(-36, -126, 838))
        self.checkpoint3 = CheckPoint("3", Position(418, -163, 1100))

        # add neighbors check-points
        self.checkpoint1.add_neighbors(self.checkpoint2)
        self.checkpoint2.add_neighbors(self.checkpoint1, self.checkpoint3)
        self.checkpoint3.add_neighbors(self.checkpoint2)


"""Guitar Resynthesiser using chord annotation and synthesize strums
"""

import librosa
import numpy as np
import pretty_midi
import scipy

from massage.resynth import util
from massage.core import Resynthesizer
from massage import VOICING_FILE

class GuitarStrummer(Resynthesizer):
    """ Resynthesize a stereo audio file that simulates a strumming guitar.
        Takes a guitar stem and associated JAMS object containing a chord
        annotation and generates strummed chords according to the JAMS chord
        annotation and onsets detected from the audio signal.
    """

    def run(self, y, fs, jam=None, instrument_label=None):
        """ run the resynthesizer

        Parameters
        ----------
        y : np.array
            an audio signal.
            shape[0] = number of channels
            shape[1] = number of samples
        fs : int
            the sample rate of the audio signal y
        jam : JAMS
            a jams object that contains 'chord' annotation for y
        instrument_label : str
            Label specifying the instrument of the track. Must be one of:
                -'acoustic guitar'
        """
        # input guarding
        if len(y) == 0:
            raise ValueError('len(y) is 0!')

        # Get Chord annotation
        chord_sequence = []
        try:
            time_intervals, chord_labs = jam.search(
                namespace='chord')[0].to_interval_values()
            for interval, lab in zip(time_intervals, chord_labs):
                chord_sequence.append([interval[0], interval[1], lab])
        except IndexError:
            raise ValueError('no "chord" annotation in the JAMS object.')

        envelope = util.get_energy_envelope(y)
        y_mono = librosa.core.to_mono(y)

        onsets, offsets, energies = util.onset_offset(y=y_mono, sr=fs)
        velocities = util.amplitude_to_velocity(energies)
        sound_font, program = util.pick_sf(y_mono, fs, instrument_label)

        env_times = librosa.samples_to_time(np.arange(len(envelope[0])), sr=fs)

        # list of interp1d objects, len(.) = number of channels
        energy_interp = []

        num_ch = 1 if envelope.ndim == 1 else envelope.shape[0]
        for ch in range(num_ch):
            energy_interp_ch = scipy.interpolate.interp1d(
                env_times, envelope[0], bounds_error=False,
                fill_value='extrapolate')
            energy_interp.append(energy_interp_ch)

        # Generating midi for synthesis and also jams file
        midi_data = self._generate_chord_midi(
            chord_sequence, onsets, offsets, velocities, program,
            energy_interp, VOICING_FILE)
        # Convert midi_data -> jams file
        jams_out = util.midi_to_jams(midi_data) # audio duration

        # Resynth
        y = midi_data.fluidsynth(sf2_path=sound_font, fs=fs)

        # Normalize
        if np.max(np.abs(y)) > 0:
            y = y / np.max(np.abs(y))

        y_len = len(y)
        env_len = len(envelope[0])
        if y_len > env_len:
            y_stereo = np.array(
                [envelope[0] * y[:env_len], envelope[1] * y[:env_len]])
        else:
            y_stereo = np.array(
                [envelope[0][:y_len] * y, envelope[1][:y_len] * y])

        return y_stereo, jams_out

    @classmethod
    def get_id(cls):
        """Method to get the id of the pitch tracker
        """
        return 'guitar_strummer'

    def _get_strum(self, start_t, end_t, chord, voicings, prev_voicing,
                   backwards=False, velocity=None, max_strum_duration=2.5):
        """
        Parameters
        ----------
        start_t : float
            starting time of the strum
        end_t : float
            ending time of the strum
        chord : str
            Chord-Harte Format chord labels, 'Ab:maj6', 'D#:min', etc
        voicings : dict
            keys are Chorde-Harte chord labels, values are lists of midi
            notes corresponding to different chord voicings.
        prev_voicing : list of int
            a list of midi numbers to represent the previous voicing
        backwards : bool, default=False
            - False : strumming down
            - True  : strumming up
        velocity : int
            midi velocity value of the strum to be created
        max_strum_duration : float default=2.5
            the maximum time in seconds that a strum will be held

        Returns
        -------
        strum : list of pretty_midi.Note
            a list of pretty_midi.Note object with small delays to simulate
            strumming
        current_voicing : list of int
            a list of the pitches in midi number to specify voicing
        """
        strum = []

        notes = util.choose_voicing(chord, voicings, prev_voicing)
        strum_duration = end_t - start_t

        if strum_duration > max_strum_duration:
            end_t = start_t + max_strum_duration
            strum_duration = max_strum_duration

        current_voicing = notes
        shifted_start = start_t

        # USING FIXED VELOCITY RANGE HERE BECAUSE WE'RE ADAPTING
        # THE OVERALL ENVELOPE LATER!
        if velocity is None:
            velocity = np.random.choice(range(70, 100))
        elif velocity < 70:
            velocity = 70
        elif velocity > 100:
            velocity = 100

        max_delay = strum_duration / float(len(notes))
        if max_delay > 0.05:
            all_choices = [0.02, 0.03, 0.04, 0.05]
        elif max_delay > 0.04:
            all_choices = [0.02, 0.03, 0.04]
        elif max_delay > 0.03:
            all_choices = [0.01, 0.02, 0.03]
        elif max_delay > 0.02:
            all_choices = [0.01, 0.02]
        elif max_delay >= 0.01:
            all_choices = [0.01]
        else:
            raise ValueError(
                "Something terrible happened! max_delay is {}".format(max_delay))

        if backwards:
            notes = notes[::-1]
            choices = all_choices[:np.random.choice(range(1, 3))]
        else:
            choices = all_choices[:np.random.choice(range(1, 5))]

        for note in notes:
            note = pretty_midi.Note(
                velocity=velocity + np.random.choice(range(-5, 5)),
                pitch=note, start=shifted_start, end=end_t
            )
            strum.append(note)
            shifted_start = shifted_start + np.random.choice(choices)

        return strum, current_voicing

    def _generate_chord_midi(self, chord_sequence, onsets, offsets, velocities,
                             program, energy_interp, voicing_file,
                             onset_lag=0.01, energy_thresh=0.03):
        """ Given list of triples of the form
            (start_time, end_time, chord_name), generate midi file.

        Parameters
        ----------
        chord_sequence : list
            list of triples of the form (start_time, end_time, chord_name),
            with start_time, end_time in seconds, and chord names of the
            form 'A:min6', 'Bb:maj', etc
        onsets : ndarray of float
            a list of onset times
        offsets : ndarray of float
            a list of offset times
        velocities : ndarray of int
            array of midi velocity values for each chord
        program : int
            midi program number
        energy_interp : list of scipy.interpolate.interp1d
            a list of interpolator objects (1 per channel) to facilitate energy
            interpolation
        voicing_file : str
            Path to a json file containing chord voicings
        onset_lag : float default=2.5
            number of secs to move the strum forward in order to compensate
            for strumming time.
        energy_thresh : float default=0.03
            threshold the energy needs to exceed in order to be voiced.

        Returns
        midi_chords : pretty_midi.PrettyMIDI
            pretty_midi object ready to be synthesized

        """
        voicings = util.get_all_voicings(voicing_file)
        midi_chords = pretty_midi.PrettyMIDI()

        chords = pretty_midi.Instrument(program=program)

        prev_voicing = None
        for triple in chord_sequence:
            start_t = triple[0]
            end_t = triple[1]

            if triple[2] == 'N' or triple[2] == 'X':
                continue

            onsets_idx_in_range = np.where(
                np.logical_and(onsets >= start_t, onsets < end_t))[0]

            # shift onset back by onset_lag (sec) to adjust for strum start
            start_times = [onsets[i] - onset_lag for i in onsets_idx_in_range]

            energies = [np.max([energy_interp[0](s), energy_interp[1](s)])
                        for s in start_times]

            end_times = [offsets[i] for i in onsets_idx_in_range]
            if len(start_times) == 0:
                continue

            vel_values = [velocities[i] for i in onsets_idx_in_range]

            next_start_times = start_times[1:] + [end_t]
            zipper = zip(start_times, end_times, next_start_times, energies)

            for i, (start_t, end_t, next_start_t, energy) in enumerate(zipper):
                if energy < energy_thresh:
                    continue

                if i % 2 == 0:
                    backwards = False
                else:
                    backwards = True

                min_duration = 0.06
                if next_start_t - start_t < min_duration:
                    continue
                elif end_t - start_t < min_duration:
                    end_t = np.random.uniform(
                        start_t + min_duration, next_start_t)
                else:
                    end_t = np.random.uniform(end_t, next_start_t)

                strum, prev_voicing = self._get_strum(
                    start_t, end_t, triple[2], voicings, prev_voicing,
                    backwards=backwards, velocity=vel_values[i]
                )
                chords.notes.extend(strum)

        midi_chords.instruments.append(chords)

        return midi_chords


#!/usr/bin/env python
import json
import cv2
import os
import tensorflow as tf
import time
import numpy as np
import keras
import datetime
import argparse
from keras import backend as K
from keras.models import load_model
from keras.models import model_from_json
from copy import copy
from threading import Thread
from tqdm import tqdm

"""
Demonstrates live webcam background segmentation on average CPU workstations.

"""

class FixedDropout(keras.layers.Dropout):
    """
      Custom dropout layer defined in EfficientNet; the definition is needed while loading the graph
   """

    def _get_noise_shape(self, inputs):
        """Noise shape for input

         Args:
         ----
            inputs(np.array)

         Returns:
         -------
            tuple(noise): noise based on input dimesnions
      """
        if self.noise_shape is None:
            return self.noise_shape

        symbolic_shape = keras.backend.shape(inputs)
        noise_shape = [
            symbolic_shape[axis] if shape is None else shape
            for axis, shape in enumerate(self.noise_shape)
        ]
        return tuple(noise_shape)

class WebStream:
    """Class to acquire webcam frames in a threaded fashion
   """

    def __init__(self, src=0):
        """Initialize the webcam stream

         Args:
         -----
            src(int or video_file_path): source to capture frames

         Returns:
         -------
            None
      """
        self.stream = cv2.VideoCapture(src)
        (self.grabbed, self.frame) = self.stream.read()
        self.stopped = False

    def start(self):
        """Start separate thread to acquire frames

         Args:
         -----
            None

         Returns:
         --------
            None
      """
        Thread(target=self.update, args=()).start()
        return self

    def update(self):
        """Update the frame which is captured

         Args:
         -----
            None

         Returns:
         --------
            None
      """
        while True:
            if self.stopped:
                self.stream.release()
                return

            (self.grabbed, self.frame) = self.stream.read()

    def read(self):
        """Update the frame which is currently read

         Args:
         -----
            None

         Returns:
         --------
            frame(np.array): input image frame
      """
        return self.frame

    def stop(self):
        """Stop capturing frames from the input stream

         Args:
         -----
            None

         Returns:
         --------
            None
      """
        self.stopped = True

class Segmentation(object):
    def __init__(self, args):
        self.split_h, self.split_w = (
            int(args["webcam_height"]) - 150,
            int(args["webcam_width"] // 2),
        )
        self.h, self.w, self.c = 128, 128, 3
        self.threshold = 0.5
        self.base_model_path = args["base_model_path"]

    def load_model(self):
        """Load keras model based on path
               Args:
               -----
                  base_model_path(str): path to the frozen model

               Returns:
               --------
                  model(keras.model): model loaded with weights
      """
        full_dir_path = os.path.abspath(os.path.dirname(__file__))

        # load the model from graph & setup the weights
        print("===============Loading Model==============")
        base_model_full_path = full_dir_path + self.base_model_path
        with open(base_model_full_path + "Unet_EB0_128_graph.json", "r") as f:
            model_json = json.load(f)

        model = model_from_json(
            model_json,
            custom_objects={"swish": tf.nn.swish, "FixedDropout": FixedDropout},
        )
        model.load_weights(base_model_full_path + "Unet_EB0_128_weights.h5")
        print("===============Model loaded==============")
        return model

    def preprocess(self, frame):
        """Preprocess input image
               Args:
               -----
                  frame(np.array): input image

               Returns:
               --------
                  images(tuple): tuple of preprocessed images
      """
        orig = cv2.resize(frame, (self.split_h, self.split_w))
        orig_normal = np.divide(orig, 255, dtype=np.float32)
        orig_blur = cv2.blur(
            orig_normal, (int(self.split_h / 16), int(self.split_w / 16)), 0
        )
        image = cv2.resize(orig, (self.h, self.w), interpolation=cv2.INTER_AREA)
        image = image[..., ::-1]  # switch BGR to RGB
        image = np.divide(image, 255, dtype=np.float32)
        image = image[np.newaxis, ...]
        return image, orig_normal, orig_blur

    def postprocess(self, mask, orig_normal, orig_blur):
        """Preprocess input image
               Args:
               -----
                  mask(np.array): input masked image
                  orig_normal(np.array): input normalized image
                  orig_blur(np.array): input blurred image

               Returns:
               --------
                  new_image(np.array): final background segmented masked image
      """
        # remove background and apply background blur for better visually appealing stream
        mask_dst = cv2.resize(
            mask, (self.split_h, self.split_w), interpolation=cv2.INTER_CUBIC
        )
        mask_dst = cv2.blur(mask_dst, (15, 15), 0)
        new_image = np.multiply(orig_normal, mask_dst[:, :, None], dtype=np.float32)
        mask = np.dstack((mask_dst, mask_dst, mask_dst))
        new_image = np.where(mask > self.threshold, new_image, orig_blur)
        return new_image

if __name__ == "__main__":
    # assumes wecam resolution is 1920x1080

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-I",
        "--input-device",
        type=int,
        default=0,
        dest="input_device",
        help="Device number input",
    )

    parser.add_argument(
        "-F",
        "--full-screen",
        type=int,
        default=0,
        dest="full_screen",
        help="Full screen",
    )

    parser.add_argument(
        "-H",
        "--webcam-height",
        type=int,
        default=1080,
        dest="webcam_height",
        help="Webcam height resolution",
    )

    parser.add_argument(
        "-W",
        "--webcam-width",
        type=int,
        default=1920,
        dest="webcam_width",
        help="Webcam width resolution",
    )

    parser.add_argument(
        "-P",
        "--model_path",
        type=str,
        default="/./data/assets/saved_models/",
        dest="base_model_path",
        help="Path to frozen model",
    )

    args = vars(parser.parse_args())
    h, w, c = 128, 128, 3

    seg = Segmentation(args)
    model = seg.load_model()

    print("===============Warming up the graph==============")
    r = np.random.rand(1, h, w, c)
    for i in tqdm(range(500)):
        pr_mask = model.predict(r)
    print("===============Graph warmed up===============")

    ws = WebStream(src=args["input_device"])
    cam = ws.start()

    total_time = []
    frames = 0
    myframe = None
    try:
        while True:
            frame = cam.read()
            myframe = copy(frame)
            t1 = time.time()

            # preprocess
            image, orig_normal, orig_blur = seg.preprocess(myframe)

            # model prediction
            pr_mask = model.predict(image)
            mask = pr_mask[..., 0].squeeze()

            # postprocess
            new_image = seg.postprocess(mask, orig_normal, orig_blur)

            # display the frame
            color_and_mask = np.concatenate((orig_normal, new_image), axis=1)
            total_time.append(time.time() - t1)

            if args["full_screen"]:
                cv2.namedWindow("Segmentation", cv2.WND_PROP_FULLSCREEN)
                cv2.setWindowProperty(
                    "Segmentation", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN
                )

            cv2.imshow("Segmentation", color_and_mask)
            if cv2.waitKey(1) & 0xFF == ord("q"):
                cam.stop()
                cv2.destroyAllWindows()
                break

            frames += 1
    except Exception as e:
        print(e)
        cam.stop()
        cv2.destroyAllWindows()

    print("FPS {t}".format(t=(frames / sum(total_time))))
    print("Total Frames {f}".format(f=frames))


########################################################################################################################
__doc__ = \
    """
This is inherited by both place and combine via _MonsterMerge
    """

########################################################################################################################

from rdkit import Chem
from rdkit.Geometry.rdGeometry import Point3D
from typing import Tuple, List, Dict, Optional, Union
import numpy as np
from warnings import warn
from .bond_provenance import BondProvenance
from ._communal import _MonsterCommunal

class _MonsterJoinNeigh(_MonsterCommunal):
    def join_neighboring_mols(self, mol_A: Chem.Mol, mol_B: Chem.Mol):
        """
        Joins two molecules by first calling _find_closest to find closest.
        That method does all the thinking.
        then by calling _join_atoms.

        :param mol_A:
        :param mol_B:
        :return:
        """
        # get closets atoms
        combo, candidates = self._find_all_closest(mol_A, mol_B)  # _find_all_closest is in communal
        anchor_A, anchor_B, distance = candidates[0]
        mol = self._join_atoms(combo, anchor_A, anchor_B, distance, linking=True)
        for anchor_A, anchor_B, distance in candidates[1:]:
            mol = self._join_atoms(combo, anchor_A, anchor_B, distance, linking=False)
        mol.SetProp('_Name', mol_A.GetProp('_Name') + '~' + mol_B.GetProp('_Name'))
        return mol

    def _join_atoms(self,
                    combo: Chem.RWMol,
                    anchor_A: int,
                    anchor_B: int,
                    distance: float,
                    linking: bool=True):
        """
        extrapolate positions between. by adding linkers if needed.
        """
        conf = combo.GetConformer()
        pos_A = conf.GetAtomPosition(anchor_A)
        pos_B = conf.GetAtomPosition(anchor_B)
        n_new = int(round(distance / 1.22) - 1)
        xs = np.linspace(pos_A.x, pos_B.x, n_new + 2)[1:-1]
        ys = np.linspace(pos_A.y, pos_B.y, n_new + 2)[1:-1]
        zs = np.linspace(pos_A.z, pos_B.z, n_new + 2)[1:-1]

        # correcting for ring marker atoms
        def is_ring_atom(anchor: int) -> bool:
            atom = combo.GetAtomWithIdx(anchor)
            if atom.HasProp('_ori_i') and atom.GetIntProp('_ori_i') == -1:
                return True
            else:
                return False

        if is_ring_atom(anchor_A):
            distance -= 1.35 + 0.2 # Arbitrary + 0.2 to compensate for the ring not reaching (out of plane).
            n_new -= 1
            xs = xs[1:]
            ys = ys[1:]
            zs = zs[1:]

        if is_ring_atom(anchor_B):
            distance -= 1.35 + 0.2 # Arbitrary + 0.2 to compensate for the ring not reaching  (out of plane).
            n_new -= 1
            xs = xs[:-1]
            ys = ys[:-1]
            zs = zs[:-1]

        # notify that things could be leary.
        if distance < 0:
            self.journal.debug(f'Two ring atoms detected to be close. Joining for now.'+
                               ' They will be bonded/fused/spiro afterwards')
        # check if valid.
        if distance > self.joining_cutoff:
            msg = f'Atoms {anchor_A}+{anchor_B} are {distance} Å away. Cutoff is {self.joining_cutoff}.'
            self.journal.warning(msg)
            raise ConnectionError(msg)
        # place new atoms
        self.journal.debug(f'Molecules will be joined via atoms {anchor_A}+{anchor_B} ({distance} Å) via the addition of {n_new} atoms.')
        previous = anchor_A
        if linking is False and n_new > 0:
            self.journal.warning(f'Was going to bond {anchor_A} and {anchor_B} but reconsidered.')
        elif linking is True and n_new <= 0:
            combo.AddBond(previous, anchor_B, Chem.BondType.SINGLE)
            new_bond = combo.GetBondBetweenAtoms(previous, anchor_B)
            BondProvenance.set_bond(new_bond, 'main_novel')
        elif linking is False and n_new <= 0:
            combo.AddBond(previous, anchor_B, Chem.BondType.SINGLE)
            new_bond = combo.GetBondBetweenAtoms(previous, anchor_B)
            BondProvenance.set_bond(new_bond, 'other_novel')
        elif linking is True and n_new > 0:
            for i in range(n_new):
                # make oxygen the first and last bridging atom.
                if i == 0 and combo.GetAtomWithIdx(anchor_A).GetSymbol() == 'C':
                    new_atomic = 8
                elif i > 2 and i == n_new -1 and combo.GetAtomWithIdx(anchor_B).GetSymbol() == 'C':
                    new_atomic = 8
                else:
                    new_atomic = 6
                idx = combo.AddAtom(Chem.Atom(new_atomic))
                new = combo.GetAtomWithIdx(idx)
                new.SetBoolProp('_Novel', True)
                new.SetIntProp('_ori_i', 999)
                conf.SetAtomPosition(idx, Point3D(float(xs[i]), float(ys[i]), float(zs[i])))
                combo.AddBond(idx, previous, Chem.BondType.SINGLE)
                new_bond = combo.GetBondBetweenAtoms(idx, previous)
                BondProvenance.set_bond(new_bond, 'linker')
                previous = idx
            combo.AddBond(previous, anchor_B, Chem.BondType.SINGLE)
            new_bond = combo.GetBondBetweenAtoms(previous, anchor_B)
            BondProvenance.set_bond(new_bond, 'linker')
        else:
            raise ValueError('Impossible')
        return combo.GetMol()

    def get_largest_fragment(self, mol):
        frags = Chem.GetMolFrags(mol, asMols=True, sanitizeFrags=False)
        frags = sorted(frags, key=lambda mol: mol.GetNumAtoms(), reverse=True)
        discarded_origins = {a.GetProp('_ori_name') for m in frags[1:] for a in m.GetAtoms() if a.HasProp('_ori_name')}
        self.unmatched.extend(discarded_origins)
        return frags[0]


import cv2
import numpy as np

class Detector:
    """
        Detectar personas en la imagen
    """
    SCALE = 1 / 255.0
    CONF_THRESHOLD = 0.9
    NMS_THRESHOLD = 0.4

    def __init__(self, yolo_weights_path, yolo_config_path):

        self._net = cv2.dnn.readNet(yolo_weights_path, yolo_config_path)
        self._net.setPreferableTarget(cv2.dnn.DNN_TARGET_OPENCL_FP16)
        self._net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        self._frame_width = 0
        self._frame_height = 0

    def detect_image(self, frame):
        """
            Detectar personas en una imagen

        :param frame -> Imagen
        :return bboxes -> Lista con los bboxes detectados
        """
        bboxes = []

        self._frame_width = frame.shape[1]
        self._frame_height = frame.shape[0]

        blob = cv2.dnn.blobFromImage(frame, self.SCALE, (416, 416), (0, 0, 0), True, crop=False)
        self._net.setInput(blob)
        outs = self._net.forward(self.get_output_layers())

        class_ids = []
        confidences = []

        for out in outs:
            for detection in out:
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                if class_id == 0 and confidence > self.CONF_THRESHOLD:
                    center_x = int(detection[0] * self._frame_width)
                    center_y = int(detection[1] * self._frame_height)
                    w = int(detection[2] * self._frame_width)
                    h = int(detection[3] * self._frame_height)
                    x = center_x - w / 2
                    y = center_y - h / 2
                    class_ids.append(class_id)
                    confidences.append(float(confidence))
                    bboxes.append((int(x), int(y), int(w), int(h)))

        cv2.dnn.NMSBoxes(bboxes, confidences, self.CONF_THRESHOLD, self.NMS_THRESHOLD)

        return bboxes

    def get_output_layers(self):
        """
            Obtener capas de la red

        :return:  output_layers -> Lista con las capas de la red
        """
        layer_names = self._net.getLayerNames()
        output_layers = [layer_names[i[0] - 1] for i in self._net.getUnconnectedOutLayers()]
        return output_layers


import cv2
import numpy as np

mat=np.genfromtxt("train.txt",dtype=str,delimiter=';')
mat_x,mat_y=mat[:,0],mat[:,1]
train_x,train_y=[],[]
name=["YueBH","ZhaoRF","BaiQL"]
for i in mat_x:
    #print type(i)
    img=cv2.imread(i,0)
    #img_1=np.float32(img)
    train_x.append(img)
for j in mat_y:
    train_y.append(int(j))
#train_data=np.asarray(train_x)
#train_label=np.asarray(train_y)
#print train_x,train_y.shape
np.save("train_data",train_x)
np.save("train_labels",train_y)
'''

a=np.load("train_data.npy")
b=np.load("train_labels.npy")
c=b.T
d=a.T
print a.shape,b.shape,c.shape,d.shape
#print a.size(),b.size()

model=cv2.ml.SVM_create()
model.train(np.asarray(train_x),np.asarray(train_y))
a=np.load('tain.txt')
a[0]'''


# Standard library
from collections import namedtuple
import warnings
import re

# Third-party
import astropy.coordinates as coord
from astropy.coordinates import representation as r
import astropy.units as u
from astropy.utils.compat.misc import override__dir__
import numpy as np

# Project
from . import representation_nd as rep_nd
from .plot import plot_projections
from ..io import quantity_to_hdf5, quantity_from_hdf5
from ..units import UnitSystem, DimensionlessUnitSystem, _greek_letters
from ..util import atleast_2d

__all__ = ['PhaseSpacePosition']

_RepresentationMappingBase = \
    namedtuple('RepresentationMapping',
               ('repr_name', 'new_name', 'default_unit'))

class RepresentationMapping(_RepresentationMappingBase):
    """
    This `~collections.namedtuple` is used to override the representation and
    differential class component names in the `PhaseSpacePosition` and `Orbit`
    classes.
    """

    def __new__(cls, repr_name, new_name, default_unit='recommended'):
        # this trick just provides some defaults
        return super().__new__(cls, repr_name, new_name, default_unit)

class RegexRepresentationMapping(RepresentationMapping):
    """
    A representation mapping that uses a regex to map the original attribute
    name to the new attribute name.
    """
    pass

class PhaseSpacePosition(object):

    representation_mappings = {
        r.CartesianRepresentation: [
            RepresentationMapping('xyz', 'xyz')
        ],
        r.SphericalCosLatDifferential: [
            RepresentationMapping('d_lon_coslat', 'pm_lon_coslat', u.mas/u.yr),
            RepresentationMapping('d_lat', 'pm_lat', u.mas/u.yr),
            RepresentationMapping('d_distance', 'radial_velocity')
        ],
        r.SphericalDifferential: [
            RepresentationMapping('d_lon', 'pm_lon', u.mas/u.yr),
            RepresentationMapping('d_lat', 'pm_lat', u.mas/u.yr),
            RepresentationMapping('d_distance', 'radial_velocity')
        ],
        r.PhysicsSphericalDifferential: [
            RepresentationMapping('d_phi', 'pm_phi', u.mas/u.yr),
            RepresentationMapping('d_theta', 'pm_theta', u.mas/u.yr),
            RepresentationMapping('d_r', 'radial_velocity')
        ],
        r.CartesianDifferential: [
            RepresentationMapping('d_x', 'v_x'),
            RepresentationMapping('d_y', 'v_y'),
            RepresentationMapping('d_z', 'v_z'),
            RepresentationMapping('d_xyz', 'v_xyz')
        ],
        r.CylindricalDifferential: [
            RepresentationMapping('d_rho', 'v_rho'),
            RepresentationMapping('d_phi', 'pm_phi'),
            RepresentationMapping('d_z', 'v_z')
        ],
        rep_nd.NDCartesianRepresentation: [
            RepresentationMapping('xyz', 'xyz')
        ],
        rep_nd.NDCartesianDifferential: [
            RepresentationMapping('d_xyz', 'v_xyz'),
            RegexRepresentationMapping('d_x([0-9])', 'v_x{0}')
        ],
    }
    representation_mappings[r.UnitSphericalCosLatDifferential] = \
        representation_mappings[r.SphericalCosLatDifferential]
    representation_mappings[r.UnitSphericalDifferential] = \
        representation_mappings[r.SphericalDifferential]

    def __init__(self, pos, vel=None, frame=None):
        """
        Represents phase-space positions, i.e. positions and conjugate momenta
        (velocities).

        The class can be instantiated with Astropy representation objects (e.g.,
        :class:`~astropy.coordinates.CartesianRepresentation`), Astropy
        :class:`~astropy.units.Quantity` objects, or plain Numpy arrays.

        If passing in representation objects, the default representation is
        taken to be the class that is passed in.

        If passing in Quantity or Numpy array instances for both position and
        velocity, they are assumed to be Cartesian. Array inputs are interpreted
        as dimensionless quantities. The input position and velocity objects can
        have an arbitrary number of (broadcastable) dimensions. For Quantity or
        array inputs, the first axis (0) has special meaning:

            - `axis=0` is the coordinate dimension (e.g., x, y, z for Cartesian)

        So if the input position array, `pos`, has shape `pos.shape = (3, 100)`,
        this would represent 100 3D positions (`pos[0]` is `x`, `pos[1]` is `y`,
        etc.). The same is true for velocity.

        Parameters
        ----------
        pos : representation, quantity_like, or array_like
            Positions. If a numpy array (e.g., has no units), this will be
            stored as a dimensionless :class:`~astropy.units.Quantity`. See
            the note above about the assumed meaning of the axes of this object.
        vel : differential, quantity_like, or array_like
            Velocities. If a numpy array (e.g., has no units), this will be
            stored as a dimensionless :class:`~astropy.units.Quantity`. See
            the note above about the assumed meaning of the axes of this object.
        frame : :class:`~gala.potential.FrameBase` (optional)
            The reference frame of the input phase-space positions.

        """

        if isinstance(pos, coord.Galactocentric):
            pos = pos.data

        if not isinstance(pos, coord.BaseRepresentation):
            # assume Cartesian if not specified
            if not hasattr(pos, 'unit'):
                pos = pos * u.one

            # 3D coordinates get special treatment
            ndim = pos.shape[0]
            if ndim == 3:
                # TODO: HACK: until this stuff is in astropy core
                if isinstance(pos, coord.BaseRepresentation):
                    kw = [(k, getattr(pos, k)) for k in pos.components]
                    pos = getattr(coord, pos.__class__.__name__)(**kw)

                else:
                    pos = coord.CartesianRepresentation(pos)

            else:
                pos = rep_nd.NDCartesianRepresentation(pos)

        else:
            ndim = 3

        if vel is None:
            if 's' not in pos.differentials:
                raise TypeError("You must specify velocity data when creating "
                                "a {0} object.".format(self.__class__.__name__))
            else:
                vel = pos.differentials.get('s', None)

        if not isinstance(vel, coord.BaseDifferential):
            # assume representation is same as pos if not specified
            if not hasattr(vel, 'unit'):
                vel = vel * u.one

            if ndim == 3:
                name = pos.__class__.get_name()
                Diff = coord.representation.DIFFERENTIAL_CLASSES[name]
                vel = Diff(*vel)
            else:
                Diff = rep_nd.NDCartesianDifferential
                vel = Diff(vel)

        # make sure shape is the same
        if pos.shape != vel.shape:
            raise ValueError("Position and velocity must have the same shape "
                             f"{pos.shape} vs. {vel.shape}")

        from ..potential.frame import FrameBase
        if frame is not None and not isinstance(frame, FrameBase):
            raise TypeError("Input reference frame must be a FrameBase "
                            "subclass instance.")

        self.pos = pos
        self.vel = vel
        self.frame = frame
        self.ndim = ndim

    def __getitem__(self, slyce):
        return self.__class__(pos=self.pos[slyce],
                              vel=self.vel[slyce],
                              frame=self.frame)

    def get_components(self, which):
        """
        Get the component name dictionary for the desired object.

        The returned dictionary maps component names on this class to component
        names on the desired object.

        Parameters
        ----------
        which : str
            Can either be ``'pos'`` or ``'vel'`` to get the components for the
            position or velocity object.
        """
        mappings = self.representation_mappings.get(
            getattr(self, which).__class__, [])

        old_to_new = dict()
        for name in getattr(self, which).components:
            for m in mappings:
                if isinstance(m, RegexRepresentationMapping):
                    pattr = re.match(m.repr_name, name)
                    old_to_new[name] = m.new_name.format(*pattr.groups())

                elif m.repr_name == name:
                    old_to_new[name] = m.new_name

        mapping = dict()
        for name in getattr(self, which).components:
            mapping[old_to_new.get(name, name)] = name

        return mapping

    @property
    def pos_components(self):
        return self.get_components('pos')

    @property
    def vel_components(self):
        return self.get_components('vel')

    def _get_extra_mappings(self, which):
        mappings = self.representation_mappings.get(
            getattr(self, which).__class__, [])

        extra = dict()
        for m in mappings:
            if (m.new_name not in self.get_components(which) and
                    not isinstance(m, RegexRepresentationMapping)):
                extra[m.new_name] = m.repr_name
        return extra

    @override__dir__
    def __dir__(self):
        """
        Override the builtin `dir` behavior to include representation and
        differential names.
        """
        dir_values = set(self.pos_components.keys())
        dir_values |= set(self.vel_components.keys())
        dir_values |= set(self._get_extra_mappings('pos').keys())
        dir_values |= set(self._get_extra_mappings('vel').keys())
        dir_values |= set(r.REPRESENTATION_CLASSES.keys())
        return dir_values

    def __getattr__(self, attr):
        """
        Allow access to attributes on the ``pos`` and ``vel`` representation and
        differential objects.
        """

        # Prevent infinite recursion here.
        if attr.startswith('_'):
            return self.__getattribute__(attr)  # Raise AttributeError.

        # TODO: with >3.5 support, can do:
        # pos_comps = {**self.pos_components,
        #              **self._get_extra_mappings('pos')}
        pos_comps = self.pos_components.copy()
        pos_comps.update(self._get_extra_mappings('pos'))
        if attr in pos_comps:
            val = getattr(self.pos, pos_comps[attr])
            return val

        # TODO: with >3.5 support, can do:
        # pos_comps = {**self.vel_components,
        #              **self._get_extra_mappings('vel')}
        vel_comps = self.vel_components.copy()
        vel_comps.update(self._get_extra_mappings('vel'))
        if attr in vel_comps:
            val = getattr(self.vel, vel_comps[attr])
            return val

        if attr in r.REPRESENTATION_CLASSES:
            return self.represent_as(attr)

        return self.__getattribute__(attr)  # Raise AttributeError.

    @property
    def data(self):
        return self.pos.with_differentials(self.vel)

    # ------------------------------------------------------------------------
    # Convert from Cartesian to other representations
    #
    def represent_as(self, new_pos, new_vel=None):
        """
        Represent the position and velocity of the orbit in an alternate
        coordinate system. Supports any of the Astropy coordinates
        representation classes.

        Parameters
        ----------
        new_pos : :class:`~astropy.coordinates.BaseRepresentation`
            The type of representation to generate. Must be a class (not an
            instance), or the string name of the representation class.
        new_vel : :class:`~astropy.coordinates.BaseDifferential` (optional)
            Class in which any velocities should be represented. Must be a class
            (not an instance), or the string name of the differential class. If
            None, uses the default differential for the new position class.

        Returns
        -------
        new_psp : `gala.dynamics.PhaseSpacePosition`
        """

        if self.ndim != 3:
            raise ValueError("Can only change representation for "
                             "ndim=3 instances.")

        # get the name of the desired representation
        if isinstance(new_pos, str):
            pos_name = new_pos
        else:
            pos_name = new_pos.get_name()

        if isinstance(new_vel, str):
            vel_name = new_vel
        elif new_vel is None:
            vel_name = pos_name
        else:
            vel_name = new_vel.get_name()

        Representation = coord.representation.REPRESENTATION_CLASSES[pos_name]
        Differential = coord.representation.DIFFERENTIAL_CLASSES[vel_name]

        new_pos = self.pos.represent_as(Representation)
        new_vel = self.vel.represent_as(Differential, self.pos)

        return self.__class__(pos=new_pos,
                              vel=new_vel,
                              frame=self.frame)

    def to_frame(self, frame, current_frame=None, **kwargs):
        """
        Transform to a new reference frame.

        Parameters
        ----------
        frame : `~gala.potential.FrameBase`
            The frame to transform to.
        current_frame : `gala.potential.CFrameBase`
            The current frame the phase-space position is in.
        **kwargs
            Any additional arguments are passed through to the individual frame
            transformation functions (see:
            `~gala.potential.frame.builtin.transformations`).

        Returns
        -------
        psp : `gala.dynamics.PhaseSpacePosition`
            The phase-space position in the new reference frame.

        """

        from ..potential.frame.builtin import transformations as frame_trans

        if self.frame is None and current_frame is None:
            raise ValueError(f"If no frame was specified when this {self} was "
                             "initialized, you must pass the current frame in "
                             "via the current_frame argument to transform to a "
                             "new frame.")

        elif self.frame is not None and current_frame is None:
            current_frame = self.frame

        name1 = current_frame.__class__.__name__.rstrip('Frame').lower()
        name2 = frame.__class__.__name__.rstrip('Frame').lower()
        func_name = f"{name1}_to_{name2}"

        if not hasattr(frame_trans, func_name):
            raise ValueError("Unsupported frame transformation: {} to {}"
                             .format(current_frame, frame))
        else:
            trans_func = getattr(frame_trans, func_name)

        pos, vel = trans_func(current_frame, frame, self, **kwargs)
        return PhaseSpacePosition(pos=pos, vel=vel, frame=frame)

    def to_coord_frame(self, frame, galactocentric_frame=None, **kwargs):
        """
        Transform the orbit from Galactocentric, cartesian coordinates to
        Heliocentric coordinates in the specified Astropy coordinate frame.

        Parameters
        ----------
        frame : :class:`~astropy.coordinates.BaseCoordinateFrame`
            The frame instance specifying the desired output frame.
            For example, :class:`~astropy.coordinates.ICRS`.
        galactocentric_frame : :class:`~astropy.coordinates.Galactocentric`
            This is the assumed frame that the position and velocity of this
            object are in. The ``Galactocentric`` instand should have parameters
            specifying the position and motion of the sun in the Galactocentric
            frame, but no data.

        Returns
        -------
        c : :class:`~astropy.coordinates.BaseCoordinateFrame`
            An instantiated coordinate frame containing the positions and
            velocities from this object transformed to the specified coordinate
            frame.

        """

        if self.ndim != 3:
            raise ValueError("Can only change representation for "
                             "ndim=3 instances.")

        if galactocentric_frame is None:
            galactocentric_frame = coord.Galactocentric()

        pos_keys = list(self.pos_components.keys())
        vel_keys = list(self.vel_components.keys())
        if (getattr(self, pos_keys[0]).unit == u.one or
                getattr(self, vel_keys[0]).unit == u.one):
            raise u.UnitConversionError("Position and velocity must have "
                                        "dimensioned units to convert to a "
                                        "coordinate frame.")

        # first we need to turn the position into a Galactocentric instance
        gc_c = galactocentric_frame.realize_frame(
            self.pos.with_differentials(self.vel))
        c = gc_c.transform_to(frame)
        return c

    # Pseudo-backwards compatibility
    def w(self, units=None):
        """
        This returns a single array containing the phase-space positions.

        Parameters
        ----------
        units : `~gala.units.UnitSystem` (optional)
            The unit system to represent the position and velocity in
            before combining into the full array.

        Returns
        -------
        w : `~numpy.ndarray`
            A numpy array of all positions and velocities, without units.
            Will have shape ``(2*ndim, ...)``.

        """
        if self.ndim == 3:
            cart = self.cartesian
        else:
            cart = self

        xyz = cart.xyz
        d_xyz = cart.v_xyz

        x_unit = xyz.unit
        v_unit = d_xyz.unit
        if ((units is None or isinstance(units, DimensionlessUnitSystem)) and
                (x_unit == u.one and v_unit == u.one)):
            units = DimensionlessUnitSystem()

        elif units is None:
            raise ValueError("A UnitSystem must be provided.")

        x = xyz.decompose(units).value
        if x.ndim < 2:
            x = atleast_2d(x, insert_axis=1)

        v = d_xyz.decompose(units).value
        if v.ndim < 2:
            v = atleast_2d(v, insert_axis=1)

        return np.vstack((x, v))

    @classmethod
    def from_w(cls, w, units=None, **kwargs):
        """
        Create a {name} object from a single array specifying positions
        and velocities. This is mainly for backwards-compatibility and
        it is not recommended for new users.

        Parameters
        ----------
        w : array_like
            The array of phase-space positions.
        units : `~gala.units.UnitSystem` (optional)
            The unit system that the input position+velocity array, ``w``,
            is represented in.
        **kwargs
            Any aditional keyword arguments passed to the class initializer.

        Returns
        -------
        obj : `~gala.dynamics.{name}`

        """.format(name=cls.__name__)

        w = np.array(w)

        ndim = w.shape[0]//2
        pos = w[:ndim]
        vel = w[ndim:]

        # TODO: this is bad form - UnitSystem should know what to do with a
        # Dimensionless
        if units is not None and not isinstance(units, DimensionlessUnitSystem):
            units = UnitSystem(units)
            pos = pos*units['length']
            vel = vel*units['length']/units['time']  # from _core_units

        return cls(pos=pos, vel=vel, **kwargs)

    # ------------------------------------------------------------------------
    # Input / output
    #
    def to_hdf5(self, f):
        """
        Serialize this object to an HDF5 file.

        Requires ``h5py``.

        Parameters
        ----------
        f : str, :class:`h5py.File`
            Either the filename or an open HDF5 file.
        """

        if isinstance(f, str):
            import h5py
            f = h5py.File(f, mode='r')

        if self.frame is not None:
            frame_group = f.create_group('frame')
            frame_group.attrs['module'] = self.frame.__module__
            frame_group.attrs['class'] = self.frame.__class__.__name__

            units = [str(x).encode('utf8')
                     for x in self.frame.units.to_dict().values()]
            frame_group.create_dataset('units', data=units)

            d = frame_group.create_group('parameters')
            for k, par in self.frame.parameters.items():
                quantity_to_hdf5(d, k, par)

        cart = self.represent_as('cartesian')
        quantity_to_hdf5(f, 'pos', cart.xyz)
        quantity_to_hdf5(f, 'vel', cart.v_xyz)

        return f

    @classmethod
    def from_hdf5(cls, f):
        """
        Load an object from an HDF5 file.

        Requires ``h5py``.

        Parameters
        ----------
        f : str, :class:`h5py.File`
            Either the filename or an open HDF5 file.
        """
        if isinstance(f, str):
            import h5py
            f = h5py.File(f, mode='r')

        pos = quantity_from_hdf5(f['pos'])
        vel = quantity_from_hdf5(f['vel'])

        frame = None
        if 'frame' in f:
            g = f['frame']

            frame_mod = g.attrs['module']
            frame_cls = g.attrs['class']
            frame_units = [u.Unit(x.decode('utf-8')) for x in g['units']]

            if u.dimensionless_unscaled in frame_units:
                units = DimensionlessUnitSystem()
            else:
                units = UnitSystem(*frame_units)

            pars = dict()
            for k in g['parameters']:
                pars[k] = quantity_from_hdf5(g['parameters/'+k])

            exec("from {0} import {1}".format(frame_mod, frame_cls))
            frame_cls = eval(frame_cls)

            frame = frame_cls(units=units, **pars)

        return cls(pos=pos, vel=vel, frame=frame)

    # ------------------------------------------------------------------------
    # Computed dynamical quantities
    #
    def kinetic_energy(self):
        r"""
        The kinetic energy *per unit mass*:

        .. math::

            E_K = \frac{1}{2} \, |\boldsymbol{v}|^2

        Returns
        -------
        E : :class:`~astropy.units.Quantity`
            The kinetic energy.
        """
        return 0.5 * self.vel.norm()**2

    def potential_energy(self, potential):
        r"""
        The potential energy *per unit mass*:

        .. math::

            E_\Phi = \Phi(\boldsymbol{q})

        Parameters
        ----------
        potential : `gala.potential.PotentialBase`
            The potential object to compute the energy from.

        Returns
        -------
        E : :class:`~astropy.units.Quantity`
            The potential energy.
        """
        # TODO: check that potential ndim is consistent with here
        return potential.energy(self)

    def energy(self, hamiltonian):
        r"""
        The total energy *per unit mass* (e.g., kinetic + potential):

        Parameters
        ----------
        hamiltonian : `gala.potential.Hamiltonian`
            The Hamiltonian object to evaluate the energy.

        Returns
        -------
        E : :class:`~astropy.units.Quantity`
            The total energy.
        """
        from ..potential import PotentialBase
        if isinstance(hamiltonian, PotentialBase):
            from ..potential import Hamiltonian

            warnings.warn("This function now expects a `Hamiltonian` instance "
                          "instead of  a `PotentialBase` subclass instance. If "
                          "you are using a static reference frame, you just "
                          "need to pass your potential object in to the "
                          "Hamiltonian constructor to use, e.g., "
                          "Hamiltonian(potential).", DeprecationWarning)

            hamiltonian = Hamiltonian(hamiltonian)

        return hamiltonian(self)

    def angular_momentum(self):
        r"""
        Compute the angular momentum for the phase-space positions contained
        in this object::

        .. math::

            \boldsymbol{{L}} = \boldsymbol{{q}} \times \boldsymbol{{p}}

        See :ref:`shape-conventions` for more information about the shapes of
        input and output objects.

        Returns
        -------
        L : :class:`~astropy.units.Quantity`
            Array of angular momentum vectors.

        Examples
        --------

            >>> import numpy as np
            >>> import astropy.units as u
            >>> pos = np.array([1., 0, 0]) * u.au
            >>> vel = np.array([0, 2*np.pi, 0]) * u.au/u.yr
            >>> w = PhaseSpacePosition(pos, vel)
            >>> w.angular_momentum() # doctest: +FLOAT_CMP
            <Quantity [0.        ,0.        ,6.28318531] AU2 / yr>
        """
        cart = self.represent_as(coord.CartesianRepresentation)
        return cart.pos.cross(cart.vel).xyz

    # ------------------------------------------------------------------------
    # Misc. useful methods
    #
    def _plot_prepare(self, components, units):
        """
        Prepare the ``PhaseSpacePosition`` or subclass for passing to a plotting
        routine to plot all projections of the object.
        """

        # components to plot
        if components is None:
            components = self.pos.components
        n_comps = len(components)

        # if units not specified, get units from the components
        if units is not None:
            if isinstance(units, u.UnitBase):
                units = [units] * n_comps  # global unit

            elif len(units) != n_comps:
                raise ValueError('You must specify a unit for each axis, or a '
                                 'single unit for all axes.')

        labels = []
        x = []
        for i, name in enumerate(components):
            val = getattr(self, name)

            if units is not None:
                val = val.to(units[i])
                unit = units[i]
            else:
                unit = val.unit

            if val.unit != u.one:
                uu = unit.to_string(format='latex_inline')
                unit_str = ' [{}]'.format(uu)
            else:
                unit_str = ''

            # Figure out how to fancy display the component name
            if name.startswith('d_'):
                dot = True
                name = name[2:]
            else:
                dot = False

            if name in _greek_letters:
                name = r"\{}".format(name)

            if dot:
                name = r"\dot{{{}}}".format(name)

            labels.append('${}$'.format(name) + unit_str)
            x.append(val.value)

        return x, labels

    def plot(self, components=None, units=None, auto_aspect=True, **kwargs):
        """
        Plot the positions in all projections. This is a wrapper around
        `~gala.dynamics.plot_projections` for fast access and quick
        visualization. All extra keyword arguments are passed to that function
        (the docstring for this function is included here for convenience).

        Parameters
        ----------
        components : iterable (optional)
            A list of component names (strings) to plot. By default, this is the
            Cartesian positions ``['x', 'y', 'z']``. To plot Cartesian
            velocities, pass in the velocity component names
            ``['d_x', 'd_y', 'd_z']``.
        units : `~astropy.units.UnitBase`, iterable (optional)
            A single unit or list of units to display the components in.
        auto_aspect : bool (optional)
            Automatically enforce an equal aspect ratio.
        relative_to : bool (optional)
            Plot the values relative to this value or values.
        autolim : bool (optional)
            Automatically set the plot limits to be something sensible.
        axes : array_like (optional)
            Array of matplotlib Axes objects.
        subplots_kwargs : dict (optional)
            Dictionary of kwargs passed to :func:`~matplotlib.pyplot.subplots`.
        labels : iterable (optional)
            List or iterable of axis labels as strings. They should correspond to
            the dimensions of the input orbit.
        plot_function : callable (optional)
            The ``matplotlib`` plot function to use. By default, this is
            :func:`~matplotlib.pyplot.scatter`, but can also be, e.g.,
            :func:`~matplotlib.pyplot.plot`.
        **kwargs
            All other keyword arguments are passed to the ``plot_function``.
            You can pass in any of the usual style kwargs like ``color=...``,
            ``marker=...``, etc.

        Returns
        -------
        fig : `~matplotlib.Figure`

        """

        try:
            import matplotlib.pyplot as plt
        except ImportError:
            msg = 'matplotlib is required for visualization.'
            raise ImportError(msg)

        if components is None:
            components = self.pos.components

        x, labels = self._plot_prepare(components=components,
                                       units=units)

        default_kwargs = {
            'marker': '.',
            'labels': labels,
            'plot_function': plt.scatter,
            'autolim': False
        }

        for k, v in default_kwargs.items():
            kwargs[k] = kwargs.get(k, v)

        fig = plot_projections(x, **kwargs)

        if self.pos.get_name() == 'cartesian' and \
                all([not c.startswith('d_') for c in components]) and \
                auto_aspect:
            for ax in fig.axes:
                ax.set(aspect='equal', adjustable='datalim')

        return fig

    # ------------------------------------------------------------------------
    # Display
    #
    def __repr__(self):
        return "<{} {}, dim={}, shape={}>".format(self.__class__.__name__,
                                                  self.pos.get_name(),
                                                  self.ndim,
                                                  self.pos.shape)

    def __str__(self):
        return "pos={}\nvel={}".format(self.pos, self.vel)

    # ------------------------------------------------------------------------
    # Shape and size
    #

    @property
    def shape(self):
        """
        This is *not* the shape of the position or velocity arrays. That is
        accessed by doing, e.g., ``obj.x.shape``.
        """
        return self.pos.shape

    #
    # Compatibility with other packages
    #

    def to_galpy_orbit(self, ro=None, vo=None):
        """Convert this object to a ``galpy.Orbit`` instance.

        Parameters
        ----------
        ro : `astropy.units.Quantity` or `astropy.units.UnitBase`
            "Natural" length unit.
        vo : `astropy.units.Quantity` or `astropy.units.UnitBase`
            "Natural" velocity unit.

        Returns
        -------
        galpy_orbit : `galpy.orbit.Orbit`

        """
        from galpy.orbit import Orbit
        from galpy.util.config import __config__ as galpy_config

        if self.frame is not None:
            from ..potential import StaticFrame
            w = self.to_frame(StaticFrame(self.frame.units))
        else:
            w = self

        if ro is None:
            ro = galpy_config.getfloat('normalization', 'ro')
            ro = ro * u.kpc

        if vo is None:
            vo = galpy_config.getfloat('normalization', 'vo')
            vo = vo * u.km/u.s

        # PhaseSpacePosition or Orbit:
        cyl = w.cylindrical

        R = cyl.rho.to_value(ro).T
        phi = cyl.phi.to_value(u.rad).T
        z = cyl.z.to_value(ro).T

        vR = cyl.v_rho.to_value(vo).T
        vT = (cyl.rho * cyl.pm_phi).to_value(vo, u.dimensionless_angles()).T
        vz = cyl.v_z.to_value(vo).T

        o = Orbit(np.array([R, vR, vT, z, vz, phi]).T, ro=ro, vo=vo)

        if hasattr(w, 't') and w.t is not None:
            o.t = w.t.to_value(ro / vo)

        return o

import numpy,import astropy,from astropy