#ebtel_plot.py

#Will Barnes
#7 May 2015

#Import necessary modules
try:
    import __builtin__
except ImportError:
    import builtins as __builtin__

import logging
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import seaborn.apionly as sns
from matplotlib.ticker import MaxNLocator
from scipy.optimize import curve_fit

#Resolve Python 2/3 exception problem
exc = getattr(__builtin__,"IOError","FileNotFoundError")

class Plotter(object):

    def __init__(self,lvl0_filename=None,dpi=1000,fontsize=18,alfs=0.65,figsize=(8,8),fformat='eps',two_fluid=True,**kwargs):
        #configure logger
        self.logger = logging.getLogger(type(self).__name__)
        #configure keyword arguments
        self.dpi = dpi
        self.fontsize = fontsize
        self.alfs = alfs
        self.figsize = figsize
        self.fformat = fformat
        self.two_fluid = two_fluid
        #plotting helpers
        self.linestyles = ('-','--','-.',':')
        #load variables
        if lvl0_filename is not None:
            self.load_variables(lvl0_filename)
        else:
            self.logger.warning("No file specified. Variable namespace will not be populated.")

    def load_variables(self,lvl0_filename,**kwargs):
        #load plasma parameters
        try:
            index_offset = 0
            data = np.loadtxt(lvl0_filename+'.txt')
            if self.two_fluid:
                index_offset += 1
                self.temp_i = data[:,2]
                self.temp_apex_i = data[:,8]
            self.time = data[:,0]
            self.temp_e = data[:,1]
            self.dens = data[:,2+index_offset]
            self.temp_apex_e = data[:,5+2*index_offset]
            self.dens_apex = data[:,6+3*index_offset]
            self.heat = data[:,10+5*index_offset]
        except exc:
            self.logger.warning("Unable to load plasma parameters from %s."%(lvl0_filename+'.txt'))
            pass

        #load dem parameters
        try:
            data = np.loadtxt(lvl0_filename+'_dem.txt')
            self.temp_dem = data[:,0]
            self.dem_tr = data[:,1]
            self.dem_cor = data[:,2]
            self.dem_tot = data[:,3]
            self.em_cor = data[:,4]
        except exc:
            self.logger.warning("Unable to load DEM parameters from %s."%(lvl0_filename+'_dem.txt'))
            pass

        #load heat parameters
        try:
            self.events = np.loadtxt(lvl0_filename+'_heat_amp.txt')
        except exc:
            self.logger.warning("Unable to load heating event amplitudes from %s."%(lvl0_filename+'_heat_amp.txt'))
            pass

    def plot_params(self,print_fig_filename=None,**kwargs):
        #set up figure
        fig,ax = plt.subplots(3,1,figsize=(1.5*self.figsize[0],self.figsize[1]),sharex=True)
        ax_n = ax[1].twinx()
        ax_na = ax[2].twinx()

        if self.two_fluid:
            tlab = r'$T_e$'
        else:
            tlab = r'$T$'

        #plot heating
        ax[0].plot(self.time,self.heat,color=sns.color_palette('deep')[0])
        ax[0].set_ylabel(r'$h$ (erg cm$^{-3}$ s$^{-1}$)',fontsize=self.fontsize)
        ax[0].set_xlim([self.time[0],self.time[-1]])
        ax[0].locator_params(nbins=5)
        ax[0].ticklabel_format(axis='y', style='sci', scilimits=(-2,2) )
        ax[0].tick_params(axis='both',labelsize=self.alfs*self.fontsize,pad=8)
        #plot average temperature and density
        line_te = ax[1].plot(self.time,self.temp_e/10**6,label=tlab,color=sns.color_palette('deep')[0])
        if self.two_fluid:
            line_ti = ax[1].plot(self.time,self.temp_i/10**6,color=sns.color_palette('deep')[2],label=r'$T_i$')
        ax[1].set_ylabel(r'$T$ (MK)',fontsize=self.fontsize)
        ax[1].yaxis.set_major_locator(MaxNLocator(prune='lower'))
        ax[1].locator_params(nbins=5)
        ax[1].ticklabel_format(axis='y', style='sci', scilimits=(-2,2) )
        ax[1].tick_params(axis='both',labelsize=self.alfs*self.fontsize,pad=8)
        line_n = ax_n.plot(self.time,self.dens/10**8,color=sns.color_palette('deep')[0],linestyle='--',label=r'$n$')
        ax_n.set_ylabel(r'$n$ (10$^8$ cm$^{-3}$)',fontsize=self.fontsize)
        ax_n.yaxis.set_major_locator(MaxNLocator(prune='lower'))
        ax_n.locator_params(nbins=5)
        ax_n.ticklabel_format(axis='y', style='sci', scilimits=(-2,2) )
        ax_n.tick_params(axis='both',labelsize=self.alfs*self.fontsize,pad=8)
        ax[1].set_xlim([self.time[0],self.time[-1]])
        #plot apex temperature and density
        ax[2].plot(self.time,self.temp_apex_e/10**6,color=sns.color_palette('deep')[0])
        if self.two_fluid:
            ax[2].plot(self.time,self.temp_apex_i/10**6,color=sns.color_palette('deep')[2])
        ax[2].set_ylabel(r'$T_a$ (MK)',fontsize=self.fontsize)
        ax[2].yaxis.set_major_locator(MaxNLocator(prune='lower'))
        ax[2].locator_params(nbins=5)
        ax[2].ticklabel_format(axis='y', style='sci', scilimits=(-2,2) )
        ax[2].tick_params(axis='both',labelsize=self.alfs*self.fontsize,pad=8)
        ax_na.plot(self.time,self.dens_apex/10**8,color=sns.color_palette('deep')[0],linestyle='--')
        ax_na.set_ylabel(r'$n_a$ (10$^8$ cm$^{-3}$)',fontsize=self.fontsize)
        ax_na.yaxis.set_major_locator(MaxNLocator(prune='lower'))
        ax_na.locator_params(nbins=5)
        ax_na.ticklabel_format(axis='y', style='sci', scilimits=(-2,2) )
        ax_na.tick_params(axis='both',labelsize=self.alfs*self.fontsize,pad=8)
        ax[2].set_xlim([self.time[0],self.time[-1]])
        ax[2].set_xlabel(r'$t$ (s)',fontsize=self.fontsize)

        #configure legend
        lines = line_te + line_n
        if self.two_fluid:
            lines = line_te + line_ti + line_n
        labels = [l.get_label() for l in lines]
        ax[1].legend(lines,labels,loc=1)

        #Check if output filename is specified
        if print_fig_filename is not None:
            plt.savefig(print_fig_filename+'.'+self.fformat,format=self.fformat,dpi=self.dpi)
        else:
            plt.show()

    def plot_dem(self,print_fig_filename=None,**kwargs):
        #set up figure
        fig = plt.figure(figsize=self.figsize)
        ax = fig.gca()

        #plot dem curves
        ax.plot(self.temp_dem,self.dem_tr,label=r'TR')
        ax.plot(self.temp_dem,self.dem_cor,label=r'corona')
        ax.plot(self.temp_dem,self.dem_tot,label=r'total')
        ax.plot(self.temp_dem,self.em_cor,label=r'EM$_{corona}$')
        ax.legend()
        ax.set_xlabel(r'$\log{T}$ (K)',fontsize=self.fontsize)
        ax.set_ylabel(r'$\log{\mathrm{DEM}}$ (cm$^{-5}$ K$^{-1}$)',fontsize=self.fontsize)
        ax.set_xlim([5.5,7.5])

        #Check if output filename is specified
        if print_fig_filename is not None:
            plt.savefig(print_fig_filename+'.'+self.fformat,format=self.fformat,dpi=self.dpi)
        else:
            plt.show()

    def plot_event_distribution(self,print_fig_filename=None,noise_thresh=0.01,return_params=True,show_plot=True,xmin=None,**kwargs):
        """Fit event energy distribution with a power-law and plot it."""

        #set up figure
        fig = plt.figure(figsize=self.figsize)
        ax = fig.gca()

        #Create a histogram
        num_bins = self._freedman_diaconis()
        n,bins,patches = ax.hist(self.events,num_bins,histtype='stepfilled',facecolor='blue',alpha=0.25,label=r'Events')
        bin_centers = np.log10(np.diff(bins)/2.0+bins[0:-1])

        #Fit with "graphical method"
        #check for bins with no entries in them; below these entries (if they exist), don't calculate fit
        noise = np.where(n <= int(np.max(n)*noise_thresh))
        if len(noise[0]) > 0:
            n = n[0:noise[0][0]]
            bin_centers = bin_centers[0:noise[0][0]]

        #calculate fit
        pars,covar = curve_fit(self._power_law_curve,bin_centers,np.log10(n),sigma=np.sqrt(np.log10(n)))
        pl_fit = self._power_law_curve(bin_centers,*pars)

        #exception for when uncertainty calculation fails
        try:
            sigma = np.sqrt(np.diag(covar))
        except:
            sigma = [0.0,0.0]
            print("Uncertainty calculation failed. Resulting value is a placeholder.")
            pass

        #estimate power-law fit using maximum likelihood estimation (see D'Huys et al., 2016, Sol. Phys.)
        if xmin is None:
            xmin = np.min(self.events)
        alpha_mle = 1. + len(self.events)*1.0/(np.sum(np.log([e/xmin for e in self.events])))
        sigma_mle = (alpha_mle - 1.)/np.sqrt(len(self.events))

        #plot fit
        ax.plot(10**bin_centers,10**pl_fit,'--r',label=r'Fit',linewidth=2.0)
        ax.set_xlabel(r'$E_H$ (erg cm$^{-3}$ s$^{-1}$)',fontsize=self.fontsize)
        ax.set_ylabel(r'Number of Events',fontsize=self.fontsize)
        ax.set_title(r'Graphical: $\alpha$ = %.2f $\pm$ %.2e, MLE: $\alpha$= %.2f $\pm$ %.2e' % (pars[1], sigma[1], alpha_mle, sigma_mle),fontsize=self.fontsize)
        ax.set_yscale('log',nonposy='clip')
        ax.set_xscale('log')
        ax.set_xlim([np.min(self.events),np.max(self.events)])
        ax.tick_params(axis='both',labelsize=0.75*self.fontsize)
        ax.legend(fontsize=0.75*self.fontsize,loc=1)

        #Check if output filename is specified
        if print_fig_filename is not None:
            plt.savefig(print_fig_filename+'.'+self.fformat,format=self.fformat,dpi=self.dpi)
            plt.close('all')
        elif show_plot:
            plt.show()
        else:
            plt.close('all')

        if return_params:
            return {'graphical':{'alpha':pars[1],'sigma':sigma[1]},'mle':{'alpha':alpha_mle,'sigma':sigma_mle}}

    def _power_law_curve(self,x,a,b):
        return a + b*x

    def _freedman_diaconis(self,**kwargs):
        q75,q25 = np.percentile(self.events,[75,25])
        iqr = q75 - q25
        w = 2.0*iqr*(len(self.events))**(-1.0/3.0)
        return int((np.max(np.array(self.events)) - np.min(np.array(self.events)))/w)

    def plot_surface(self,param_1,param_2,surf_list,**kwargs):
        #set up figure
        fig = plt.figure(figsize=self.figsize)
        ax = fig.gca()

        #set colorbar limits
        if 'vmin' in kwargs:
            vmin = kwargs['vmin']
        else:
            vmin = np.min(np.array(surf_list))
        if 'vmax' in kwargs:
            vmax = kwargs['vmax']
        else:
            vmax = np.max(np.array(surf_list))

        #set up mesh
        p1_mesh,p2_mesh = np.meshgrid(np.array(param_1),np.array(param_2))
        surf = ax.pcolormesh(p1_mesh,p2_mesh,np.array(surf_list),cmap='hot',vmin=vmin,vmax=vmax)
        fig.colorbar(surf,ax=ax)

        #set limits
        if 'xlim' in kwargs:
            ax.set_xlim([kwargs['xlim'][0],kwargs['xlim'][1]])
        else:
            ax.set_xlim([param_1[0],param_1[-1]])
        if 'ylim' in kwargs:
            ax.set_ylim([kwargs['ylim'][0],kwargs['ylim'][1]])
        else:
            ax.set_ylim([param_2[0],param_2[-1]])

        #set labels
        if 'ylab' in kwargs:
            ax.set_ylabel(kwargs['ylab'],fontsize=self.fontsize)
        if 'xlab' in kwargs:
            ax.set_xlabel(kwargs['xlab'],fontsize=self.fontsize)
        if 'plot_title' in kwargs:
            ax.set_title(kwargs['plot_title'],fontsize=self.fontsize)

        #Check if output filename is specified
        if 'print_fig_filename' in kwargs:
            plt.savefig(kwargs['print_fig_filename']+'.'+self.fformat,format=self.fformat,dpi=self.dpi)
        else:
            plt.show()


#!/usr/bin/python

import os
import os.path
import numpy as np
import ROOT
import argparse
import subprocess
import resource
import glob
import sys
sys.path.append("..")
from lib import colourLogger
from lib import generalUtil as gUtil

def calcLC(source, configAnal):

    logStdout.info([['wb', 'Producing a lightcurve for'],
                    ['p', source]])

    anasumFile = configAnal['anasumFiles'][source]
    sourceSpectralParam = configAnal['spectralParameters'][source]
    dirNow = os.path.join(os.getcwd(), source)

    minEnergyFluxCalc = sourceSpectralParam[0]
    normalizationEnergy = sourceSpectralParam[1]
    spectralIndex = sourceSpectralParam[2]
    maxEnergyFluxCalc = sourceSpectralParam[3]

    # Read energy threshold from file
    headersType = {'names': ('Source', 'z', 'tau = 1',
                             'tau = 2', 'tau = 3'),
                   'formats': ('U20', 'f8',
                               'f8', 'f8', 'f8')}

    sourceEnergyThresholds = np.loadtxt('../../{}_sourcesThresholds.txt'.
                                        format(configAnal['EBL']['model']), dtype=headersType)
    sourceIndex = np.where(sourceEnergyThresholds['Source'] == source)
    # first two entries are source and z
    thresholdsForThisSource = list(sourceEnergyThresholds[sourceIndex[0][0]])[2:4]
    thresholdsForThisSource.insert(0, minEnergyFluxCalc)

    binningDict = {'nightly': 1, 'weekly': 7, 'monthly': 28, 'yearly': 365}
    mjd_min = configAnal['dates']['veritas'][0]
    mjd_max = configAnal['dates']['veritas'][1]

    thresholds = list()
    # Build a list of pairs of min/max energies to run over
    for i_thr, thresholdNow in enumerate(thresholdsForThisSource):

        minEnergy = thresholdNow
        if i_thr == len(thresholdsForThisSource) - 1:
            maxEnergy = 30.
        else:
            maxEnergy = thresholdsForThisSource[i_thr + 1]

        prefix = 'LightCurve_'
        interfix = 'tau-{}_'.format(i_thr + 1)
        suffix = 'range_{0:1.0f}_{1:1.0f}_GeV.txt'.format(minEnergy*1000., maxEnergy*1000.)
        fileName = prefix + interfix + suffix
        thresholds.append({'minEnergy': thresholdNow,
                           'maxEnergy': maxEnergy,
                           'fileName': fileName})

    # Make two versions of the lightcurves, one with negative fluxes
    # and one without. The former is only for the luminosity function study.
    suffixNegativeFlux = ['', '_unbound']

    for suffixNegativeFluxNow in suffixNegativeFlux:

        prefix = 'LightCurve_'
        interfix = 'fullEnergyRange_'
        suffix = 'range_{0:1.0f}_{1:1.0f}_GeV{2}.txt'.format(minEnergyFluxCalc*1000.,
                                                             30*1000.,
                                                             suffixNegativeFluxNow)
        fileName = prefix + interfix + suffix
        thresholds.append({'minEnergy': minEnergyFluxCalc,
                           'maxEnergy': maxEnergy,
                           'fileName': fileName})

    for thresholdsNow in thresholds:

        for binning, timeBinLength_days in binningDict.items():

            logStdout.info([['p', source],
                            ['bb', binning],
                            ['wb', 'lightcurve,'],
                            ['g', '{} < E < {} TeV'.format(thresholdsNow['minEnergy'],
                                                           thresholdsNow['maxEnergy'])]])

            iLightCurve = ROOT.VLightCurve()

            iLightCurve.initializeTeVLightCurve(anasumFile, timeBinLength_days,
                                                mjd_min, mjd_max)

            # set spectral parameters (start of flux calculation [TeV],
            # normalization at energy E0 [TeV], spectral index, end of flux calculation [TeV])
            iLightCurve.setSpectralParameters(minEnergyFluxCalc, normalizationEnergy,
                                              spectralIndex, maxEnergyFluxCalc)

            # plot ULs for poins with <2 sigma significance; to get rid of
            # significance limits for upper limits of flux use (-999, 999);
            iLightCurve.setSignificanceParameters(-999, -999)
            # Avoid negative fluxes (if True it avoids them)
            iLightCurve.setFluxCalculationMethod('unbound' not in thresholdsNow['fileName'])
            # calculate fluxes and upper flux limits for energies in the range given
            iLightCurve.fill(thresholdsNow['minEnergy'], thresholdsNow['maxEnergy'])
            iLightCurve.writeASCIIFile(os.path.join(dirNow,
                                                    binning +
                                                    thresholdsNow['fileName']))
            del iLightCurve

def moveFileToLumiFunctionStudyDir(source):

    files = glob.glob('{}/*fullEnergyRange*'.format(source))
    for fileNow in files:
        if 'unbound' in fileNow:
            command = 'mv'
            fileDest = fileNow.replace('_unbound', '')
        else:
            command = 'cp'
            fileDest = fileNow.replace('.txt', '_bounded.txt')

        fileDest = fileDest.replace('fullEnergyRange_range', 'energyRange')
        # The 195 GeV threshold in ED is essentially equivalent to 200 GeV.
        fileDest = fileDest.replace('195', '200')

        subprocess.call('mkdir -p lightcurvesForLumiFunctionStudy/{}'.format(source), shell=True)
        fullCommand = '/bin/{} {} lightcurvesForLumiFunctionStudy/{}'.format(command,
                                                                             fileNow,
                                                                             fileDest)
        subprocess.call(fullCommand, shell=True)

    return

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Produce light-curve.')
    parser.add_argument('source', action='store',
                        help='Source to produce the light-curve for')
    parser.add_argument('configAnalFile', action='store',
                        help='YAML file with the analysis configuration')

    args = parser.parse_args()

    logStdout = colourLogger.initStdOutLogger()

    # A horrible hack to avoid the following ROOT error
    # SysError in <TFile::TFile>: file  can not be opened for reading (Too many open files)
    logStdout.info([['wb', 'getrlimit before:'],
                    ['bb', resource.getrlimit(resource.RLIMIT_NOFILE)]])
    resource.setrlimit(resource.RLIMIT_NOFILE, (4096, 4096))
    logStdout.info([['wb', 'getrlimit after:'],
                    ['bb', resource.getrlimit(resource.RLIMIT_NOFILE)]])

    ROOT.gROOT.SetBatch(True)
    ROOT.TFile.Open._creates = True

    # load shared library
    ROOT.gSystem.Load("$EVNDISPSYS/lib/libVAnaSum.so")

    configAnal = gUtil.readYamlFile(logStdout, args.configAnalFile)
    calcLC(args.source, configAnal)

    moveFileToLumiFunctionStudyDir(args.source)


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

from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from past.utils import old_div
from builtins import range
from __future__ import (division, print_function, absolute_import, unicode_literals)

import sys
import numpy as np
import fsps
from cloudyfsps.ASCIItools import (writeASCII, compileASCII, checkCompiled, compiledExists)

# this code snippet goes through every step needed
# to integrate FSPS into Cloudy.
# This example uses stellar pops with a constant SFH
# as the input ionizing source.
# 1. Write an ascii file in Cloudy format with grid
#    of FSPS spectra in all available ages and
#    metallicities
# 2. Compile asii file into binary format required
#    for Cloudy use. Assumes $CLOUDY_EXE is set to
#    your /path/to/cloudy.exe
# 3. Writes Cloudy input files for a subset of grid
#    parameters.
# 4. Runs Cloudy on the *.in files
# 5. Formats the various output files

zsun = 0.0142 # this is solar metallicity for the MIST isochrones

exec_write_ascii = True

# Function to write the ascii file.
# This is where you set the properties of the
# ionizing spectrum (SSP/CSFH, IMF, FBHB, etc)

def mist_ascii(fileout, **kwargs):
    # change these parameters to modify the ionizing source grid
    # default mode is to produce an ascii grid in age and Z,
    # though different variables and more dimensions are possible.
    sp_dict = dict(zcontinuous=1,
                   imf_type=2,
                   sfh=0,
                   const=0.0,
                   sf_start=0.0)
    sp = fsps.StellarPopulation(**sp_dict)
    # all ages and Zs
    ages = 10.**sp.log_age
    logZs = np.log10(old_div(sp.zlegend,zsun))
    modpars = [(age, logZ) for age in ages for logZ in logZs]
    lam = sp.wavelengths
    all_fluxs = []
    for logZ in logZs:
        sp.params['logzsol'] = logZ
        all_fluxs.append(sp.get_spectrum()[1]) #lsun per hz
    nmod = len(modpars)
    # flatten flux for writing
    flat_flux = np.array([all_fluxs[j][i]
                          for i in range(len(ages))
                          for j in range(len(logZs))])
    # this function is flexible, ndim can be 3/4/n.
    # in this example, however, ndim is 2 (age, logz).
    writeASCII(fileout, lam, flat_flux, modpars,
               nx=len(lam), ndim=2, npar=2, nmod=nmod)
    return
#---------------------------------------------------------------------
# ASCII FILE: WRITE AND COMPILE
#---------------------------------------------------------------------
# assumes you have $CLOUDY_EXE and $CLOUDY_DATA_PATH set as sys vars.

# name of ascii file
ascii_file = 'FSPS_MIST_SSP.ascii'

# the ascii file takes a while to generate, so if an already-compiled
# version exists, the code will not overwrite it.

compiled_ascii = '{}.mod'.format(ascii_file.split('.')[0])
if exec_write_ascii:
    print("Executing write ascii sequence...")
    if not compiledExists(ascii_file):
        print("No compiled model exists...Writing.")
        mist_ascii(ascii_file)
        print("Compiling {} with Cloudy".format(ascii_file))
        compileASCII(ascii_file)
        print("Checking to see if compilation was successful...")
        if checkCompiled(ascii_file):
            print("Your model {} is ready to run.".format(compiled_ascii))
        else:
            sys.exit()
    else:
        print("{} already exists.".format(compiled_ascii))


# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst

import numpy as np

from astropy import units as u
from astropy.utils.state import ScienceState
from astropy.utils.decorators import format_doc
from astropy.coordinates.angles import Angle
from astropy.coordinates.matrix_utilities import rotation_matrix, matrix_product, matrix_transpose
from astropy.coordinates import representation as r
from astropy.coordinates.baseframe import (BaseCoordinateFrame,
                                           frame_transform_graph,
                                           base_doc)
from astropy.coordinates.attributes import (CoordinateAttribute,
                                            QuantityAttribute,
                                            DifferentialAttribute)
from astropy.coordinates.transformations import AffineTransform
from astropy.coordinates.errors import ConvertError

from .icrs import ICRS

__all__ = ['Galactocentric']

# Measured by minimizing the difference between a plane of coordinates along
#   l=0, b=[-90,90] and the Galactocentric x-z plane
# This is not used directly, but accessed via `get_roll0`.  We define it here to
# prevent having to create new Angle objects every time `get_roll0` is called.
_ROLL0 = Angle(58.5986320306*u.degree)

class galactocentric_frame_defaults(ScienceState):
    """This class controls the global setting of default values for the frame
    attributes in the `~astropy.coordinates.Galactocentric` frame, which may be
    updated in future versions of ``astropy``. Note that when using
    `~astropy.coordinates.Galactocentric`, changing values here will not affect
    any attributes that are set explicitly by passing values in to the
    `~astropy.coordinates.Galactocentric` initializer. Modifying these defaults
    will only affect the frame attribute values when using the frame as, e.g.,
    ``Galactocentric`` or ``Galactocentric()`` with no explicit arguments.

    This class controls the parameter settings by specifying a string name,
    which can be one of:

    - 'pre-v4.0': The current default value, which sets the default frame
      attribute values to their original (pre-astropy-v4.0) values.
    - 'v4.0': The attribute values as updated in Astropy version 4.0.
    - 'latest': An alias of the most recent parameter set (currently: 'v4.0')

    See :ref:`astropy-coordinates-galactocentric-defaults` for more information.

    Examples
    --------
    The default `~astropy.coordinates.Galactocentric` frame parameters can be
    modified globally::

        >>> from astropy.coordinates import galactocentric_frame_defaults
        >>> _ = galactocentric_frame_defaults.set('v4.0') # doctest: +SKIP
        >>> Galactocentric() # doctest: +SKIP
        <Galactocentric Frame (galcen_coord=<ICRS Coordinate: (ra, dec) in deg
            (266.4051, -28.936175)>, galcen_distance=8.122 kpc, galcen_v_sun=(12.9, 245.6, 7.78) km / s, z_sun=20.8 pc, roll=0.0 deg)>
        >>> _ = galactocentric_frame_defaults.set('pre-v4.0') # doctest: +SKIP
        >>> Galactocentric() # doctest: +SKIP
        <Galactocentric Frame (galcen_coord=<ICRS Coordinate: (ra, dec) in deg
            (266.4051, -28.936175)>, galcen_distance=8.3 kpc, galcen_v_sun=(11.1, 232.24, 7.25) km / s, z_sun=27.0 pc, roll=0.0 deg)>

    The default parameters can also be updated by using this class as a context
    manager::

        >>> with galactocentric_frame_defaults.set('pre-v4.0'):
        ...     print(Galactocentric()) # doctest: +FLOAT_CMP
        <Galactocentric Frame (galcen_coord=<ICRS Coordinate: (ra, dec) in deg
            (266.4051, -28.936175)>, galcen_distance=8.3 kpc, galcen_v_sun=(11.1, 232.24, 7.25) km / s, z_sun=27.0 pc, roll=0.0 deg)>

    Again, changing the default parameter values will not affect frame
    attributes that are explicitly specified::

        >>> import astropy.units as u
        >>> with galactocentric_frame_defaults.set('pre-v4.0'):
        ...     print(Galactocentric(galcen_distance=8.0*u.kpc)) # doctest: +FLOAT_CMP
        <Galactocentric Frame (galcen_coord=<ICRS Coordinate: (ra, dec) in deg
            (266.4051, -28.936175)>, galcen_distance=8.0 kpc, galcen_v_sun=(11.1, 232.24, 7.25) km / s, z_sun=27.0 pc, roll=0.0 deg)>

    """

    _latest_value = 'v4.0'
    _references = None
    _value = None

    @classmethod
    def get_solar_params_from_string(cls, arg):
        """Return Galactocentric solar parameters given string names for the
        parameter sets.
        """

        # Resolve the meaning of 'latest': The latest parameter set is from v4.0
        # - update this as newer parameter choices are added
        if arg == 'latest':
            arg = cls._latest_value

        params = dict()
        references = dict()

        # Currently, all versions use the same sky position for Sgr A*:
        params['galcen_coord'] = ICRS(ra=266.4051*u.degree,
                                      dec=-28.936175*u.degree)
        references['galcen_coord'] = \
            'http://adsabs.harvard.edu/abs/2004ApJ...616..872R'

        # The roll angle is the same for both frames:
        params['roll'] = 0 * u.deg

        if arg == 'pre-v4.0':
            params['galcen_distance'] = 8.3 * u.kpc
            references['galcen_distance'] = \
                'https://ui.adsabs.harvard.edu/#abs/2009ApJ...692.1075G'

            params['galcen_v_sun'] = r.CartesianDifferential([11.1,
                                                              220+12.24,
                                                              7.25]*u.km/u.s)
            references['galcen_v_sun'] = \
                ['https://ui.adsabs.harvard.edu/#abs/2010MNRAS.403.1829S',
                 'https://ui.adsabs.harvard.edu/#abs/2015ApJS..216...29B']

            params['z_sun'] = 27.0 * u.pc
            references['z_sun'] = \
                'https://ui.adsabs.harvard.edu/#abs/2001ApJ...553..184C'

        elif arg == 'v4.0':
            params['galcen_distance'] = 8.122 * u.kpc
            references['galcen_distance'] = \
                'https://ui.adsabs.harvard.edu/abs/2018A%26A...615L..15G'

            params['galcen_v_sun'] = r.CartesianDifferential([12.9,
                                                              245.6,
                                                              7.78]*u.km/u.s)
            references['galcen_v_sun'] = \
                ['https://ui.adsabs.harvard.edu/abs/2018RNAAS...2..210D',
                 'https://ui.adsabs.harvard.edu/abs/2018A%26A...615L..15G',
                 'https://ui.adsabs.harvard.edu/abs/2004ApJ...616..872R']

            params['z_sun'] = 20.8 * u.pc
            references['z_sun'] = \
                'https://ui.adsabs.harvard.edu/abs/2019MNRAS.482.1417B'

        else:
            raise ValueError(f'Invalid string input to retrieve solar '
                             f'parameters for Galactocentric frame: "{arg}"')

        return params, references

    @classmethod
    def validate(cls, value):
        if value is None:
            value = cls._latest_value

        if isinstance(value, str):
            params, refs = cls.get_solar_params_from_string(value)
            cls._references = refs
            return params

        elif isinstance(value, dict):
            return value

        elif isinstance(value, Galactocentric):
            # turn the frame instance into a dict of frame attributes
            attrs = dict()
            for k in value.frame_attributes:
                attrs[k] = getattr(value, k)
            cls._references = value.frame_attribute_references()
            return attrs

        else:
            raise ValueError("Invalid input to retrieve solar parameters for "
                             "Galactocentric frame: input must be a string, "
                             "dict, or Galactocentric instance")

doc_components = """
    x : `~astropy.units.Quantity`, optional
        Cartesian, Galactocentric :math:`x` position component.
    y : `~astropy.units.Quantity`, optional
        Cartesian, Galactocentric :math:`y` position component.
    z : `~astropy.units.Quantity`, optional
        Cartesian, Galactocentric :math:`z` position component.

    v_x : `~astropy.units.Quantity`, optional
        Cartesian, Galactocentric :math:`v_x` velocity component.
    v_y : `~astropy.units.Quantity`, optional
        Cartesian, Galactocentric :math:`v_y` velocity component.
    v_z : `~astropy.units.Quantity`, optional
        Cartesian, Galactocentric :math:`v_z` velocity component.
"""

doc_footer = """
    Other parameters
    ----------------
    galcen_coord : `ICRS`, optional, must be keyword
        The ICRS coordinates of the Galactic center.
    galcen_distance : `~astropy.units.Quantity`, optional, must be keyword
        The distance from the sun to the Galactic center.
    galcen_v_sun : `~astropy.coordinates.representation.CartesianDifferential`, optional, must be keyword
        The velocity of the sun *in the Galactocentric frame* as Cartesian
        velocity components.
    z_sun : `~astropy.units.Quantity`, optional, must be keyword
        The distance from the sun to the Galactic midplane.
    roll : `~astropy.coordinates.Angle`, optional, must be keyword
        The angle to rotate about the final x-axis, relative to the
        orientation for Galactic. For example, if this roll angle is 0,
        the final x-z plane will align with the Galactic coordinates x-z
        plane. Unless you really know what this means, you probably should
        not change this!

    Examples
    --------

    To transform to the Galactocentric frame with the default
    frame attributes, pass the uninstantiated class name to the
    ``transform_to()`` method of a `~astropy.coordinates.SkyCoord` object::

        >>> import astropy.units as u
        >>> import astropy.coordinates as coord
        >>> c = coord.SkyCoord(ra=[158.3122, 24.5] * u.degree,
        ...                    dec=[-17.3, 81.52] * u.degree,
        ...                    distance=[11.5, 24.12] * u.kpc,
        ...                    frame='icrs')
        >>> c.transform_to(coord.Galactocentric) # doctest: +FLOAT_CMP
        <SkyCoord (Galactocentric: galcen_coord=<ICRS Coordinate: (ra, dec) in deg
            (266.4051, -28.936175)>, galcen_distance=8.122 kpc, galcen_v_sun=(12.9, 245.6, 7.78) km / s, z_sun=20.8 pc, roll=0.0 deg): (x, y, z) in kpc
            [( -9.43489286, -9.40062188, 6.51345359),
             (-21.11044918, 18.76334013, 7.83175149)]>

    To specify a custom set of parameters, you have to include extra keyword
    arguments when initializing the Galactocentric frame object::

        >>> c.transform_to(coord.Galactocentric(galcen_distance=8.1*u.kpc)) # doctest: +FLOAT_CMP
        <SkyCoord (Galactocentric: galcen_coord=<ICRS Coordinate: (ra, dec) in deg
            (266.4051, -28.936175)>, galcen_distance=8.1 kpc, galcen_v_sun=(12.9, 245.6, 7.78) km / s, z_sun=20.8 pc, roll=0.0 deg): (x, y, z) in kpc
            [( -9.41284763, -9.40062188, 6.51346272),
             (-21.08839478, 18.76334013, 7.83184184)]>

    Similarly, transforming from the Galactocentric frame to another coordinate frame::

        >>> c = coord.SkyCoord(x=[-8.3, 4.5] * u.kpc,
        ...                    y=[0., 81.52] * u.kpc,
        ...                    z=[0.027, 24.12] * u.kpc,
        ...                    frame=coord.Galactocentric)
        >>> c.transform_to(coord.ICRS) # doctest: +FLOAT_CMP
        <SkyCoord (ICRS): (ra, dec, distance) in (deg, deg, kpc)
            [( 88.22423301, 29.88672864,  0.17813456),
             (289.72864549, 49.9865043 , 85.93949064)]>

    Or, with custom specification of the Galactic center::

        >>> c = coord.SkyCoord(x=[-8.0, 4.5] * u.kpc,
        ...                    y=[0., 81.52] * u.kpc,
        ...                    z=[21.0, 24120.0] * u.pc,
        ...                    frame=coord.Galactocentric,
        ...                    z_sun=21 * u.pc, galcen_distance=8. * u.kpc)
        >>> c.transform_to(coord.ICRS) # doctest: +FLOAT_CMP
        <SkyCoord (ICRS): (ra, dec, distance) in (deg, deg, kpc)
            [( 86.2585249 , 28.85773187, 2.75625475e-05),
             (289.77285255, 50.06290457, 8.59216010e+01)]>

"""

@format_doc(base_doc, components=doc_components, footer=doc_footer)
class Galactocentric(BaseCoordinateFrame):
    r"""
    A coordinate or frame in the Galactocentric system.

    This frame allows specifying the Sun-Galactic center distance, the height of
    the Sun above the Galactic midplane, and the solar motion relative to the
    Galactic center. However, as there is no modern standard definition of a
    Galactocentric reference frame, it is important to pay attention to the
    default values used in this class if precision is important in your code.
    The default values of the parameters of this frame are taken from the
    original definition of the frame in 2014. As such, the defaults are somewhat
    out of date relative to recent measurements made possible by, e.g., Gaia.
    The defaults can, however, be changed at runtime by setting the parameter
    set name in `~astropy.coordinates.galactocentric_frame_defaults`.

    The current default parameter set is ``"pre-v4.0"``, indicating that the
    parameters were adopted before ``astropy`` version 4.0. A regularly-updated
    parameter set can instead be used by setting
    ``galactocentric_frame_defaults.set ('latest')``, and other parameter set
    names may be added in future versions. To find out the scientific papers
    that the current default parameters are derived from, use
    ``galcen.frame_attribute_references`` (where ``galcen`` is an instance of
    this frame), which will update even if the default parameter set is changed.

    The position of the Sun is assumed to be on the x axis of the final,
    right-handed system. That is, the x axis points from the position of
    the Sun projected to the Galactic midplane to the Galactic center --
    roughly towards :math:`(l,b) = (0^\circ,0^\circ)`. For the default
    transformation (:math:`{\rm roll}=0^\circ`), the y axis points roughly
    towards Galactic longitude :math:`l=90^\circ`, and the z axis points
    roughly towards the North Galactic Pole (:math:`b=90^\circ`).

    For a more detailed look at the math behind this transformation, see
    the document :ref:`coordinates-galactocentric`.

    The frame attributes are listed under **Other Parameters**.
    """

    default_representation = r.CartesianRepresentation
    default_differential = r.CartesianDifferential

    # frame attributes
    galcen_coord = CoordinateAttribute(frame=ICRS)
    galcen_distance = QuantityAttribute(unit=u.kpc)

    galcen_v_sun = DifferentialAttribute(
        allowed_classes=[r.CartesianDifferential])

    z_sun = QuantityAttribute(unit=u.pc)
    roll = QuantityAttribute(unit=u.deg)

    def __init__(self, *args, **kwargs):
        # Set default frame attribute values based on the ScienceState instance
        # for the solar parameters defined above
        default_params = galactocentric_frame_defaults.get()
        self.frame_attribute_references = \
            galactocentric_frame_defaults._references.copy()

        for k in default_params:
            if k in kwargs:
                # If a frame attribute is set by the user, remove its reference
                self.frame_attribute_references.pop(k, None)

            # Keep the frame attribute if it is set by the user, otherwise use
            # the default value
            kwargs[k] = kwargs.get(k, default_params[k])

        super().__init__(*args, **kwargs)

    @classmethod
    def get_roll0(cls):
        """
        The additional roll angle (about the final x axis) necessary to align
        the final z axis to match the Galactic yz-plane.  Setting the ``roll``
        frame attribute to  -this method's return value removes this rotation,
        allowing the use of the `Galactocentric` frame in more general contexts.
        """
        # note that the actual value is defined at the module level.  We make at
        # a property here because this module isn't actually part of the public
        # API, so it's better for it to be accessable from Galactocentric
        return _ROLL0

# ICRS to/from Galactocentric ----------------------->

def get_matrix_vectors(galactocentric_frame, inverse=False):
    """
    Use the ``inverse`` argument to get the inverse transformation, matrix and
    offsets to go from Galactocentric to ICRS.
    """
    # shorthand
    gcf = galactocentric_frame

    # rotation matrix to align x(ICRS) with the vector to the Galactic center
    mat1 = rotation_matrix(-gcf.galcen_coord.dec, 'y')
    mat2 = rotation_matrix(gcf.galcen_coord.ra, 'z')
    # extra roll away from the Galactic x-z plane
    mat0 = rotation_matrix(gcf.get_roll0() - gcf.roll, 'x')

    # construct transformation matrix and use it
    R = matrix_product(mat0, mat1, mat2)

    # Now need to translate by Sun-Galactic center distance around x' and
    # rotate about y' to account for tilt due to Sun's height above the plane
    translation = r.CartesianRepresentation(gcf.galcen_distance * [1., 0., 0.])
    z_d = gcf.z_sun / gcf.galcen_distance
    H = rotation_matrix(-np.arcsin(z_d), 'y')

    # compute total matrices
    A = matrix_product(H, R)

    # Now we re-align the translation vector to account for the Sun's height
    # above the midplane
    offset = -translation.transform(H)

    if inverse:
        # the inverse of a rotation matrix is a transpose, which is much faster
        #   and more stable to compute
        A = matrix_transpose(A)
        offset = (-offset).transform(A)
        offset_v = r.CartesianDifferential.from_cartesian(
            (-gcf.galcen_v_sun).to_cartesian().transform(A))
        offset = offset.with_differentials(offset_v)

    else:
        offset = offset.with_differentials(gcf.galcen_v_sun)

    return A, offset

def _check_coord_repr_diff_types(c):
    if isinstance(c.data, r.UnitSphericalRepresentation):
        raise ConvertError("Transforming to/from a Galactocentric frame "
                           "requires a 3D coordinate, e.g. (angle, angle, "
                           "distance) or (x, y, z).")

    if ('s' in c.data.differentials and
            isinstance(c.data.differentials['s'],
                       (r.UnitSphericalDifferential,
                        r.UnitSphericalCosLatDifferential,
                        r.RadialDifferential))):
        raise ConvertError("Transforming to/from a Galactocentric frame "
                           "requires a 3D velocity, e.g., proper motion "
                           "components and radial velocity.")

@frame_transform_graph.transform(AffineTransform, ICRS, Galactocentric)
def icrs_to_galactocentric(icrs_coord, galactocentric_frame):
    _check_coord_repr_diff_types(icrs_coord)
    return get_matrix_vectors(galactocentric_frame)

@frame_transform_graph.transform(AffineTransform, Galactocentric, ICRS)
def galactocentric_to_icrs(galactocentric_coord, icrs_frame):
    _check_coord_repr_diff_types(galactocentric_coord)
    return get_matrix_vectors(galactocentric_coord, inverse=True)

import numpy,import astropy,from astropy


from __future__ import division, print_function
import sys
from pkg_resources import Requirement, resource_stream
import numpy as np
import itertools
import logging

nt = { 'a': 0, 'c': 1, 'g': 2, 't': 3 }
tops = lambda s: 4*s[:,:-1]+s[:,1:]

class ToeholdSpecificationError(ValueError):
    '''
    Error raised when stickydesign cannot satisfy user's toehold specifications.
    '''
    def __init__(self, message):
        #self.expression = expression
        self.message = message

def exceptionhook(exception_type, exception, traceback, default_hook = sys.excepthook):
    if 'Toehold' in exception_type.__name__ :
        print("{}: {}".format('RuntimeError', exception.message))
    else:
        default_hook(exception_type, exception, traceback)

sys.excepthook = exceptionhook

class energyfuncs:
    """
    Energy functions based on SantaLucia's 2004 paper.

    mismatchtype is one of 'max', 'loop', or 'dangle', specifying how to
    consider mismatches.  'max' is probably the best choice, but is slowest -
    it takes the maximum interaction of the 'loop' and 'dangle' options.
    """
    def __init__(self, targetdG=7.7, length=7, deviation=0.5, max_spurious=0.4):
        import os
        try:
            dsb = resource_stream('stickydesign', 'stickydesign/params/dnastackingbig.csv')
        except:
            try:
                dsb = resource_stream('stickydesign', 'params/dnastackingbig.csv')
            except IOError:
                raise IOError("Error loading dnastackingbig.csv")
        try:
            dgl = resource_stream('piperine', 'data/dnadangle.csv')
        except:
            try:
                this_dir, this_filename = os.path.split(__file__)
                dgl = open( os.path.join(this_dir, "data", "dnadangle.csv") )
            except IOError:
                raise IOError("Error loading dnadangle.csv")
        self.targetdG=targetdG
        self.alphabet='h'
        self.adjs=['c', 'g']
        self.deviation=deviation
        self.max_spurious=max_spurious
        self.length=length
        self.nndG_full = -np.loadtxt(dsb ,delimiter=',')
        self.dgldG_full = -np.loadtxt(dgl ,delimiter=',')
        self.taildG = 1.3
        dsb.close()
        dgl.close()
        self.initdG = 0.0 # 1.96 DISABLED FOR NOW
        self.nndG = self.nndG_full[np.arange(0,16),15-np.arange(0,16)]
        # 30-01-15: The only dangle contexts we are interested in are 3' dangle
        # s. Select those from the Santa Lucia table. We'll have to flip the
        # order of the vector, though, to mach the 5->3 orientation of the gene
        # rated toeholds. To flip, we need to count up to 15 with the opposite-
        # endian order, in terms of quaternary representation
        indcs = 4*np.tile(np.arange(4), 4) + np.repeat(np.arange(4), 4)
        self.dgldG = self.dgldG_full[1, indcs]
        # 30-01-15: As of now, the dangle base is set to C, so make a lookup ta
        # ble ordered by terminating toehold base
        self.dgldG_fixedC = self.dgldG_full[1, np.arange(4) + 4 * nt['c']]
        self.uniform = lambda x,y: np.maximum( self.uniform_loopmismatch(x,y), \
                                               self.uniform_danglemismatch(x,y) \
                                             )

    def th_external_dG(self, seqs):
        # Convert nearest-neighbor stacks to dG-table lookup indices.
        # Sum up the near-neighbor energy contributions
        # Add context-specific dG values, eg tail or dangle contributions
        seqs_len = np.size(seqs, 1)
        # The external context involves a 3' dangle, so exclude the 3' flank
        # base.
        cols_external = np.arange(seqs_len-1)
        tops_external = tops(seqs[:, cols_external])
        nndG_external = np.sum(self.nndG[tops_external], 1)
        # The external-context dangle is fixed at C.
        dgldG_external = self.dgldG_fixedC[seqs[:, seqs_len-2]]
        return nndG_external + dgldG_external - self.taildG - self.initdG

    def th_internal_dG(self, seqs):
        # Convert nearest-neighbor stacks to dG-table lookup indices
        # Sum up and return the near-neighbor energy contributions
        # Add context-specific dG values, eg tail or dangle contributions
        seqs_len = np.size(seqs, 1)
        # The internal context involves a truncated toehold. Remove first 3' to
        # ehold base.
        cols_internal = np.concatenate((np.arange(seqs_len-2), [seqs_len-1]))
        tops_internal = tops(seqs[:, cols_internal])
        nndG_internal = np.sum(self.nndG[tops_internal], 1)
        return nndG_internal - self.taildG - self.initdG

    def matching_uniform(self, seqs):
        # Make a boolean vector representing which toeholds' external context dG
        # is further from the target dG than than their internal context dG
        dG_external = self.th_external_dG(seqs)
        dG_internal = self.th_internal_dG(seqs)
        external_further_bool = np.abs(dG_external - self.targetdG) >\
                                np.abs(dG_internal - self.targetdG)
        return np.choose(external_further_bool, [dG_internal, dG_external])

    def uniform_loopmismatch(self, seqs1, seqs2):
        if seqs1.shape != seqs2.shape:
            if seqs1.ndim == 1:
                seqs1 = endarray( np.repeat(np.array([seqs1]),seqs2.shape[0],0), seqs1.endtype )
            else:
                raise InputError("Lengths of sequence arrays are not acceptable.")
        assert seqs1.endtype == seqs2.endtype
        endtype = seqs1.endtype

        endlen = seqs1.endlen
        plen = endlen-1

        # Run through the
        # TODO: replace this with cleaner code
        if endtype=='DT':
            ps1 = seqs1[:,1:-1]*4+seqs1[:,2:]
            pa1 = seqs1[:,0]*4+seqs1[:,1]
            pac1 = (3-seqs1[:,0])*4+seqs2[:,-1]
            ps2 = seqs2[:,::-1][:,:-2]*4+seqs2[:,::-1][:,1:-1]
            pa2 = seqs2[:,0]*4+seqs2[:,1]
            pac2 = (3-seqs2[:,0])*4+seqs1[:,-1]
        if endtype=='TD':
            ps1 = seqs1[:,:-2]*4+seqs1[:,1:-1]
            pa1 = seqs1[:,-2]*4+seqs1[:,-1]
            pac1 = seqs2[:,0]*4+(3-seqs1[:,-1])
            ps2 = seqs2[:,::-1][:,1:-1]*4+seqs2[:,::-1][:,2:]
            pa2 = seqs2[:,-2]*4+seqs2[:,-1]
            pac2 = (seqs1[:,0])*4+(3-seqs2[:,-1])

        # Shift here is considering the first strand as fixed, and the second one as
        # shifting.  The shift is the offset of the bottom one in terms of pair
        # sequences (thus +2 and -1 instead of +1 and 0).
        en = np.zeros( (ps1.shape[0], 2*plen) )
        for shift in range(-plen+1,plen):
            #import pdb
            #pdb.set_trace()
            en[:,plen+shift-1] = np.sum( \
                    self.nndG_full[ ps1[:,max(shift,0):plen+shift], \
                               ps2[:,max(-shift,0):plen-shift] ], \
                               axis=1)
        en[:,plen-1] = en[:,plen-1] + self.nndG_full[pa1,pac1] + self.nndG_full[pa2,pac2]
        return np.amax(en,1) - self.initdG

    def uniform_danglemismatch(self, seqs1,seqs2,fast=True):
        if seqs1.shape != seqs2.shape:
            if seqs1.ndim == 1:
                seqs1 = endarray( np.repeat(np.array([seqs1]),seqs2.shape[0],0), seqs1.endtype )
            else:
                raise InputError("Lengths of sequence arrays are not acceptable.")
        assert seqs1.endtype == seqs2.endtype
        endtype = seqs1.endtype
        s1 = tops(seqs1)
        s2 = tops(seqs2)
        l = s1.shape[1]
        s2r = np.fliplr(np.invert(s2)%16)
        s2r = s2r//4 + 4*(s2r%4)
        m = np.zeros((s1.shape[0],2*np.sum(np.arange(2,l+1))+l+1))
        r = np.zeros(m.shape[0])
        z = 0;
        if endtype == 'TD':
            s1c = s1[:,0:-1]
            s2rc = s2r[:,1:]
            s1l = np.hstack(( (4*(s2r[:,0]//4) + s1[:,0]//4).reshape(-1,1) , s1 ))
            s2rl = np.hstack(( s2r , (4*(s2r[:,-1]%4) + s1[:,-1]%4).reshape(-1,1) ))
        elif endtype == 'DT':
            s1c = s1[:,1:]
            s2rc = s2r[:,0:-1]
            s2rl = np.hstack(( (4*(s1[:,0]//4) + s2r[:,0]//4).reshape(-1,1) , s2r ))
            s1l = np.hstack(( s1 , (4*(s1[:,-1]%4) + s2r[:,-1]%4).reshape(-1,1) ))
        for o in range(1,l-1):
            zn = l-1-o
            m[:,z:z+zn] = ( s1c[:,:-o]==s2rc[:,o:] ) * self.nndG[s1c[:,:-o]]
            z = z+zn+2
            m[:,z:z+zn] = ( s2rc[:,:-o]==s1c[:,o:] ) * self.nndG[s2rc[:,:-o]]
            z = z+zn+2
        m[:,z:z+l+1] = (s1l == s2rl) * self.nndG[s1l]
        i = 0
        im = len(m)
        # This needs to be changed to something faster
        if not fast:
            for xi in range(0,m.shape[0]):
                gm = 0
                g = 0
                for y in m[xi,:]:
                    if y == 0:
                        g = 0
                    else:
                        g += y
                        if gm > g:
                            gm = g
                r[xi] = gm
                i+=1
                if not i%1000:
                    print("%d/%d" % (i,im))
        else:
            from stickydesign import _stickyext
            x = m
            _stickyext.fastsub(x,r)

        return r-self.initdG

    def score_toeholds(self, toeholds):
        import stickydesign as sd
        toeholds = [ th_set[0] for th_set in toeholds]
        toeholds_flanked = [ 'c' + th.lower() + 'c' for th in toeholds]
        ends = sd.endarray(toeholds_flanked, 'TD')
        e_vec_ext = self.th_external_dG(ends)
        e_vec_int = self.th_internal_dG(ends)
        e_vec_all = np.concatenate( (e_vec_int, e_vec_ext))
        e_err = np.abs(e_vec_all.mean() - self.targetdG)
        e_rng = e_vec_all.max() - e_vec_all.min()
        return (e_err, e_rng)

    def calculate_unrestricted_toehold_characteristics(self):
        import stickydesign as sd
        ends = sd.easyends('TD',
                           self.length,
                           alphabet=self.alphabet,
                           adjs=self.adjs,
                           energetics=self)
        n_ends = len(ends)
        e_array = sd.energy_array_uniform(ends, self)
        e_array = e_array[n_ends:, :n_ends]
        for i in range(n_ends):
            e_array[i,i] = 0
        e_spr = e_array.max()/self.targetdG
        e_vec_ext = self.th_external_dG(ends)
        e_vec_int = self.th_internal_dG(ends)
        e_vec_all = np.concatenate( (e_vec_int, e_vec_ext))
        e_avg = e_vec_all.mean()
        e_dev = np.max(np.abs(e_vec_all - self.targetdG))
        return e_avg, e_dev, e_spr, n_ends

    def get_toeholds(self, n_ths=6, timeout=8):
        from  time import time
        import stickydesign as sd
        """ Generate specified stickyends for the Soloveichik DSD approach

        A given run of stickydesign may not generate toeholds that match the
        Soloveichik approach. This function reports whether the run was successful
        or not. Otherwise, the toeholds are matched to respect the backwards
        strand back-to-back toeholds.

        Args:
            n_ths: Number of toeholds to generate. (6)
            timeout: Time duration allowed for finding toeholds in seconds. (8)
        Returns:
            List of toehold strings
        """
        # Give StickyDesign a set of trivial, single-nucleotide toeholds to avoid poor
        # designs. I'm not sure if this helps now, but it did once.
        avoid_list = [i * int(self.length + 2) for i in ['a', 'c', 't']]

        # Generate toeholds
        fdev = self.deviation / self.targetdG
        notoes = True
        startime = time()
        while notoes:
            try:
                ends = sd.easyends('TD',
                                   self.length,
                                   interaction=self.targetdG,
                                   fdev=fdev,
                                   alphabet=self.alphabet,
                                   adjs=self.adjs,
                                   maxspurious=self.max_spurious,
                                   energetics=self,
                                   oldends=avoid_list)
                notoes = len(ends) < n_ths + len(avoid_list)
                if (time() - startime) > timeout:
                    e_avg, e_spr, e_dev, n_ends = self.calculate_unrestricted_toehold_characteristics()
                    msg = "Cannot make toeholds to user specification! Try target energy:{:.2}, maxspurious:{:.2}, deviation:{:.2}, which makes {:d} toeholds."
                    exception = ToeholdSpecificationError(msg.format(e_avg, e_spr, e_dev, n_ends))
                    raise exception
            except ValueError:
                e_avg, e_spr, e_dev, n_ends = self.calculate_unrestricted_toehold_characteristics()
                msg = "Cannot make toeholds to user specification! Try target energy:{:.2}, maxspurious:{:.2}, deviation:{:.2}, which makes {:d} toeholds."
                exception = ToeholdSpecificationError(msg.format(e_avg, e_spr, e_dev, n_ends))
                raise exception
        th_cands = ends.tolist()
        # remove "avoid" sequences
        th_cands = th_cands[len(avoid_list):]
        # Make as many end in c as possible
        th_cands = th_cands[:n_ths]
        th_all = [ th[1:-1] for th in th_cands]
        ends_full = sd.endarray(th_cands, 'TD')
        ends_all = sd.endarray(th_all, 'TD')
        return ends_all.tolist()


from abc import ABC, abstractmethod

import numpy as np
from astropy.nddata import StdDevUncertainty, VarianceUncertainty, InverseVariance
from astropy.units import Quantity
from scipy.interpolate import CubicSpline

from ..spectra import Spectrum1D, SpectralAxis

__all__ = ['ResamplerBase', 'FluxConservingResampler',
           'LinearInterpolatedResampler', 'SplineInterpolatedResampler']

class ResamplerBase(ABC):
    """
    Base class for resample classes.  The algorithms and needs for difference
    resamples will vary quite a bit, so this class is relatively sparse.

    Parameters
    ----------
    extrapolation_treatment : str
        What to do when resampling off the edge of the spectrum.  Can be
        ``'nan_fill'`` to have points beyond the edges by set to NaN, or
        ``'zero_fill'`` to be set to zero.
    """
    def __init__(self, extrapolation_treatment='nan_fill'):
        if extrapolation_treatment not in ('nan_fill', 'zero_fill'):
            raise ValueError('invalid extrapolation_treatment value: ' + str(extrapolation_treatment))
        self.extrapolation_treatment = extrapolation_treatment

    def __call__(self, orig_spectrum, fin_spec_axis):
        """
        Return the resulting `~specutils.Spectrum1D` of the resampling.
        """
        return self.resample1d(orig_spectrum, fin_spec_axis)

    @abstractmethod
    def resample1d(self, orig_spectrum, fin_spec_axis):
        """
        Workhorse method that will return the resampled Spectrum1D
        object.
        """
        return NotImplemented

class FluxConservingResampler(ResamplerBase):
    """
    This resampling algorithm conserves overall integrated flux (as opposed to
    flux density).
    Algorithm based on the equations documented in the following paper:
    https://ui.adsabs.harvard.edu/abs/2017arXiv170505165C/abstract

    Parameters
    ----------
    extrapolation_treatment : str
        What to do when resampling off the edge of the spectrum.  Can be
        ``'nan_fill'`` to have points beyond the edges by set to NaN, or
        ``'zero_fill'`` to be set to zero.

    Examples
    --------

    To resample an input spectrum to a user specified spectral grid using
    a flux conserving algorithm:

    >>> import numpy as np
    >>> import astropy.units as u
    >>> from specutils import Spectrum1D
    >>> from specutils.manipulation import FluxConservingResampler
    >>> input_spectra = Spectrum1D(
    ...     flux=np.array([1, 3, 7, 6, 20]) * u.mJy,
    ...     spectral_axis=np.array([2, 4, 12, 16, 20]) * u.nm)
    >>> resample_grid = [1, 5, 9, 13, 14, 17, 21, 22, 23]  *u.nm
    >>> fluxc_resample = FluxConservingResampler()
    >>> output_spectrum1D = fluxc_resample(input_spectra, resample_grid) # doctest: +IGNORE_OUTPUT

    """

    def _resample_matrix(self, orig_spec_axis, fin_spec_axis):
        """
        Create a re-sampling matrix to be used in re-sampling spectra in a way
        that conserves flux. This code was heavily influenced by Nick Earl's
        resample rough draft: nmearl@0ff6ef1.

        Parameters
        ----------
        orig_spec_axis : SpectralAxis
            The original spectral axis array.
        fin_spec_axis : SpectralAxis
            The desired spectral axis array.

        Returns
        -------
        resample_mat : ndarray
            An [[N_{fin_spec_axis}, M_{orig_spec_axis}]] matrix.
        """
        # Lower bin and upper bin edges
        orig_edges = orig_spec_axis.bin_edges
        fin_edges = fin_spec_axis.bin_edges

        # I could get rid of these alias variables,
        # but it does add readability
        orig_low = orig_edges[:-1]
        fin_low = fin_edges[:-1]
        orig_upp = orig_edges[1:]
        fin_upp = fin_edges[1:]

        # Here's the real work in figuring out the bin overlaps
        # i.e., contribution of each original bin to the resampled bin
        l_inf = np.where(orig_low > fin_low[:, np.newaxis],
                         orig_low, fin_low[:, np.newaxis])
        l_sup = np.where(orig_upp < fin_upp[:, np.newaxis],
                         orig_upp, fin_upp[:, np.newaxis])

        resamp_mat = (l_sup - l_inf).clip(0)
        resamp_mat = resamp_mat * (orig_upp - orig_low)

        # set bins that don't overlap 100% with original bins
        # to zero by checking edges, and applying generated mask
        left_clip = np.where(fin_edges[:-1] - orig_edges[0] < 0, 0, 1)
        right_clip = np.where(orig_edges[-1] - fin_edges[1:] < 0, 0, 1)
        keep_overlapping_matrix = left_clip * right_clip

        resamp_mat *= keep_overlapping_matrix[:, np.newaxis]

        return resamp_mat.value

    def resample1d(self, orig_spectrum, fin_spec_axis):
        """
        Create a re-sampling matrix to be used in re-sampling spectra in a way
        that conserves flux. If an uncertainty is present in the input spectra
        it will be propagated through to the final resampled output spectra
        as an InverseVariance uncertainty.

        Parameters
        ----------
        orig_spectrum : `~specutils.Spectrum1D`
            The original 1D spectrum.
        fin_spec_axis :  Quantity
            The desired spectral axis array.

        Returns
        -------
        resample_spectrum : `~specutils.Spectrum1D`
            An output spectrum containing the resampled `~specutils.Spectrum1D`
        """

        # Check if units on original spectrum and new wavelength (if defined)
        # match
        if isinstance(fin_spec_axis, Quantity):
            if orig_spectrum.spectral_axis.unit != fin_spec_axis.unit:
                raise ValueError("Original spectrum spectral axis grid and new"
                                 "spectral axis grid must have the same units.")

        if not isinstance(fin_spec_axis, SpectralAxis):
            fin_spec_axis = SpectralAxis(fin_spec_axis)

        # todo: Would be good to return uncertainty in type it was provided?
        # todo: add in weighting options

        # Get provided uncertainty into variance
        if orig_spectrum.uncertainty is not None:
            if isinstance(orig_spectrum.uncertainty, StdDevUncertainty):
                pixel_uncer = np.square(orig_spectrum.uncertainty.array)
            elif isinstance(orig_spectrum.uncertainty, VarianceUncertainty):
                pixel_uncer = orig_spectrum.uncertainty.array
            elif isinstance(orig_spectrum.uncertainty, InverseVariance):
                pixel_uncer = np.reciprocal(orig_spectrum.uncertainty.array)
        else:
            pixel_uncer = None

        orig_axis_in_fin = orig_spectrum.spectral_axis.to(fin_spec_axis.unit)
        resample_grid = self._resample_matrix(orig_axis_in_fin, fin_spec_axis)

        # Now for some broadcasting magic to handle multi dimensional flux inputs
        # Essentially this part is inserting length one dimensions as fillers
        # For example, if we have a (5,6,10) input flux, and an output grid
        # of 3, flux will be broadcast to (5,6,1,10) and resample_grid will
        # Be broadcast to (1,1,3,10).  The sum then reduces down the 10, the
        # original dispersion grid, leaving 3, the new dispersion grid, as
        # the last index.
        new_flux_shape = list(orig_spectrum.flux.shape)
        new_flux_shape.insert(-1, 1)
        in_flux = orig_spectrum.flux.reshape(new_flux_shape)

        ones = [1] * len(orig_spectrum.flux.shape[:-1])
        new_shape_resample_grid = ones + list(resample_grid.shape)
        resample_grid = resample_grid.reshape(new_shape_resample_grid)

        # Calculate final flux
        out_flux = np.sum(in_flux * resample_grid, axis=-1) / np.sum(
            resample_grid, axis=-1)

        # Calculate output uncertainty
        if pixel_uncer is not None:
            pixel_uncer = pixel_uncer.reshape(new_flux_shape)

            out_variance = np.sum(pixel_uncer * resample_grid**2, axis=-1) / np.sum(
                resample_grid**2, axis=-1)
            out_uncertainty = InverseVariance(np.reciprocal(out_variance))
        else:
            out_uncertainty = None

        # nan-filling happens by default - replace with zeros if requested:
        if self.extrapolation_treatment == 'zero_fill':
            origedges = orig_spectrum.spectral_axis.bin_edges
            off_edges = (fin_spec_axis < origedges[0]) | (origedges[-1] < fin_spec_axis)
            out_flux[off_edges] = 0
            if out_uncertainty is not None:
                out_uncertainty.array[off_edges] = 0

        # todo: for now, use the units from the pre-resampled
        # spectra, although if a unit is defined for fin_spec_axis and it doesn't
        # match the input spectrum it won't work right, will have to think
        # more about how to handle that... could convert before and after
        # calculation, which is probably easiest. Matrix math algorithm is
        # geometry based, so won't work to just let quantity math handle it.
        resampled_spectrum = Spectrum1D(flux=out_flux,
                                        spectral_axis=np.array(fin_spec_axis) * orig_spectrum.spectral_axis.unit,
                                        uncertainty=out_uncertainty)

        return resampled_spectrum

class LinearInterpolatedResampler(ResamplerBase):
    """
    Resample a spectrum onto a new ``spectral_axis`` using linear interpolation.

    Parameters
    ----------
    extrapolation_treatment : str
        What to do when resampling off the edge of the spectrum.  Can be
        ``'nan_fill'`` to have points beyond the edges by set to NaN, or
        ``'zero_fill'`` to be set to zero.

    Examples
    --------

    To resample an input spectrum to a user specified dispersion grid using
    linear interpolation:

    >>> import numpy as np
    >>> import astropy.units as u
    >>> from specutils import Spectrum1D
    >>> from specutils.manipulation import LinearInterpolatedResampler
    >>> input_spectra = Spectrum1D(
    ...     flux=np.array([1, 3, 7, 6, 20]) * u.mJy,
    ...     spectral_axis=np.array([2, 4, 12, 16, 20]) * u.nm)
    >>> resample_grid = [1, 5, 9, 13, 14, 17, 21, 22, 23] * u.nm
    >>> fluxc_resample = LinearInterpolatedResampler()
    >>> output_spectrum1D = fluxc_resample(input_spectra, resample_grid) # doctest: +IGNORE_OUTPUT
    """
    def __init__(self, extrapolation_treatment='nan_fill'):
        super().__init__(extrapolation_treatment)

    def resample1d(self, orig_spectrum, fin_spec_axis):
        """
        Call interpolation, repackage new spectra

        Parameters
        ----------
        orig_spectrum : `~specutils.Spectrum1D`
            The original 1D spectrum.
        fin_spec_axis : ndarray
            The desired spectral axis array.

        Returns
        -------
        resample_spectrum : `~specutils.Spectrum1D`
            An output spectrum containing the resampled `~specutils.Spectrum1D`
        """

        fill_val = np.nan  # bin_edges=nan_fill case
        if self.extrapolation_treatment == 'zero_fill':
            fill_val = 0

        orig_axis_in_fin = orig_spectrum.spectral_axis.to(fin_spec_axis.unit)

        out_flux_arr = np.interp(fin_spec_axis.value, orig_axis_in_fin.value,
                                 orig_spectrum.flux.value, left=fill_val, right=fill_val)
        out_flux = Quantity(out_flux_arr, unit=orig_spectrum.flux.unit)

        new_unc = None
        if orig_spectrum.uncertainty is not None:
            out_unc_arr = np.interp(fin_spec_axis.value, orig_axis_in_fin.value,
                                    orig_spectrum.uncertainty.array,
                                    left=fill_val, right=fill_val)
            new_unc = orig_spectrum.uncertainty.__class__(array=out_unc_arr,
                                                          unit=orig_spectrum.unit)

        return Spectrum1D(spectral_axis=fin_spec_axis,
                          flux=out_flux,
                          uncertainty=new_unc)

class SplineInterpolatedResampler(ResamplerBase):
    """
    This resample algorithim uses a cubic spline interpolator. Any uncertainty
    is also interpolated using an identical spline.

    Parameters
    ----------
    extrapolation_treatment : str
        What to do when resampling off the edge of the spectrum.  Can be
        ``'nan_fill'`` to have points beyond the edges by set to NaN, or
        ``'zero_fill'`` to be set to zero.

    Examples
    --------

    To resample an input spectrum to a user specified spectral axis grid using
    a cubic spline interpolator:

    >>> import numpy as np
    >>> import astropy.units as u
    >>> from specutils import Spectrum1D
    >>> from specutils.manipulation import SplineInterpolatedResampler
    >>> input_spectra = Spectrum1D(
    ...     flux=np.array([1, 3, 7, 6, 20]) * u.mJy,
    ...     spectral_axis=np.array([2, 4, 12, 16, 20]) * u.nm)
    >>> resample_grid = [1, 5, 9, 13, 14, 17, 21, 22, 23] * u.nm
    >>> fluxc_resample = SplineInterpolatedResampler()
    >>> output_spectrum1D = fluxc_resample(input_spectra, resample_grid) # doctest: +IGNORE_OUTPUT

    """
    def __init__(self, bin_edges='nan_fill'):
        super().__init__(bin_edges)

    def resample1d(self, orig_spectrum, fin_spec_axis):
        """
        Call interpolation, repackage new spectra

        Parameters
        ----------
        orig_spectrum : `~specutils.Spectrum1D`
            The original 1D spectrum.
        fin_spec_axis : Quantity
            The desired spectral axis array.

        Returns
        -------
        resample_spectrum : `~specutils.Spectrum1D`
            An output spectrum containing the resampled `~specutils.Spectrum1D`
        """
        orig_axis_in_new = orig_spectrum.spectral_axis.to(fin_spec_axis.unit)
        flux_spline = CubicSpline(orig_axis_in_new.value, orig_spectrum.flux.value,
                                   extrapolate=self.extrapolation_treatment != 'nan_fill')
        out_flux_val = flux_spline(fin_spec_axis.value)

        new_unc = None
        if orig_spectrum.uncertainty is not None:
            unc_spline = CubicSpline(orig_axis_in_new.value, orig_spectrum.uncertainty.array,
                                       extrapolate=self.extrapolation_treatment != 'nan_fill')
            out_unc_val = unc_spline(fin_spec_axis.value)
            new_unc = orig_spectrum.uncertainty.__class__(array=out_unc_val, unit=orig_spectrum.unit)

        if self.extrapolation_treatment == 'zero_fill':
            origedges = orig_spectrum.spectral_axis.bin_edges
            off_edges = (fin_spec_axis < origedges[0]) | (origedges[-1] < fin_spec_axis)
            out_flux_val[off_edges] = 0
            if new_unc is not None:
                new_unc.array[off_edges] = 0

        return Spectrum1D(spectral_axis=fin_spec_axis,
                          flux=out_flux_val*orig_spectrum.flux.unit,
                          uncertainty=new_unc)

import numpy,from scipy,import astropy,from astropy


'''

该文件用于ROI_summer 数据集上,192 大小对256大小的sar-optical图像匹配,
统计匹配位置误差,匹配FLAG(<5 PIXELS), 匹配关键点数量, 关键点误差(<= 2 PIXELS)
related results are saved in eval_results_XX.txt

'''

import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import cv2
import numpy as np
import math
import os
import time
from glob import glob

class models():
    def __init__(self,dataset,weights_list):
        self.dataset = dataset
        self.weights_list - weights_list

def model_choose(model_path):

    if 'agl_cls_cspdense' in model_path:

        from old.networks_angle_cls_csp_dense import EmbeddingNet, LDMNet
        embede_model = EmbeddingNet()
        embede_model.cuda()
        ldm_model = LDMNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using agl_cls_cspdense net')
        return embede_model,ldm_model

    elif 'HardNet' in model_path:

        from old.networks_l2_sos import EmbeddingNet, LDMNet
        embede_model = EmbeddingNet()
        embede_model.cuda()
        ldm_model = LDMNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using Hardnet net')

        return embede_model,ldm_model

    elif 'classification' in model_path:
        from old.networks_angle_classification import EmbeddingNet, LDMNet
        embede_model = EmbeddingNet()
        embede_model.cuda()
        ldm_model = LDMNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using classification simple net')

        return embede_model,ldm_model

    elif 'TFeat' in model_path:
        from old.networks_tfeat import TNet_Rocket, LDMNet
        embede_model = TNet_Rocket()
        embede_model.cuda()
        ldm_model = LDMNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using Tfeat net')

        return embede_model,ldm_model

    elif 'MatchNet' in model_path:
        from old.network_matchnet_512 import FeatureNet, CLSNet
        embede_model = FeatureNet()
        embede_model.cuda()
        ldm_model = CLSNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using matchnet net...')

        return embede_model,ldm_model

    elif '_l2_cspdense' in model_path:
        from old.networks_l2_cls_csp_dense import EmbeddingNet, LDMNet
        embede_model  = EmbeddingNet()
        embede_model.cuda()
        ldm_model = LDMNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using l2_cspdense net...')
        return embede_model,ldm_model

    elif 'entropy_cspdense' in model_path:
        from old.networks_entropyloss_cls_csp_dense import EmbeddingNet, CLSNet
        embede_model  = EmbeddingNet()
        embede_model.cuda()
        ldm_model = CLSNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using entropy_csp_dense net...')
        return embede_model,ldm_model

    elif 'agl_dense' in model_path:
        from old.networks_angle_cls_dense import EmbeddingNet,LDMNet
        embede_model  = EmbeddingNet()
        embede_model.cuda()
        ldm_model = LDMNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using agl_dense net...')
        return embede_model,ldm_model

    elif 'cspdense64_joint' in model_path:
        from old.networks_l2_cls_csp_dense import EmbeddingNet
        from joint_model.jointNet import JointNet,LDMNet
        embede_model  = JointNet(EmbeddingNet())
        embede_model.cuda()
        ldm_model = LDMNet(embede_model).cuda().float()
        # ldm_model = LDMNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using cspdense64_joint net...')
        global joint_flag
        joint_flag = True
        return embede_model,ldm_model

    elif 'arcl2_cspdense' in model_path:
        from networks_arcl2_csp_dense import EmbeddingNet,LDMNet
        embede_model  = EmbeddingNet()
        embede_model.cuda()
        ldm_model = LDMNet(embede_model)
        ldm_model.cuda().eval()
        load_model(ldm_model,model_path)
        print('using arcl2 cspdense net...')
        return embede_model,ldm_model

def load_model(model,model_path):
    checkpoint = torch.load(model_path,map_location='cuda:0')
    model.load_state_dict(checkpoint['model_state_dict'])
    return

def kpts2descriptors(kpts,img,model,batch_size=128,patch_size=64,use_gpu=True):
    descrs = []
    length = len(kpts)
    shards = int(np.ceil(length/batch_size))
    h,w = img.shape[:2]
    sx,sy,ex,ey = 0,0,0,0
    ltoffsets = patch_size//2
    rboffsets = patch_size-ltoffsets
    for i in range(shards):
        patches = []
        batch_kpts = kpts[i*batch_size:min((i+1)*batch_size,length)]
        for kp in batch_kpts:
            x, y = kp.pt
            x = int(x)
            y = int(y)
            if x <= ltoffsets:
                sx = 0
                ex = sx + patch_size
            else:
                ex = min(x+rboffsets,w)
                sx = ex-patch_size
            if y <= ltoffsets:
                sy = 0
                ey = sy + patch_size
            else:
                ey = min(y+rboffsets,h)
                sy = ey-patch_size
            patch = img[int(sy):int(ey),int(sx):int(ex)]
            assert patch.shape[0] == patch_size and patch.shape[1] == patch_size,str(patch.shape)+'  '+str(x)+' '+str(y)
            patches.append(patch)

        patches = torch.from_numpy(np.asarray(patches)).float()
        patches = torch.unsqueeze(patches, 1)
        if use_gpu:
            patches = patches.cuda()

        descrs.append(model(patches).detach().cpu().numpy())

    descrs = np.concatenate(descrs,axis=0)
    # print(len(descrs))
    return descrs

def matches2offsets(matches,queryKp,trainKp):
    offsets = []
    for m in matches:
        train_x,train_y = trainKp[m.trainIdx].pt
        train_x,train_y = round(train_x),round(train_y)
        query_x,query_y = queryKp[m.queryIdx].pt
        query_x,query_y = round(query_x),round(query_y)
        offset_x = query_x - train_x
        offset_y = query_y - train_y
        if 0 < offset_x < 289 and 0 < offset_y < 289:
            offsets.append([offset_x,offset_y])
    return offsets

def matched_kpt_summary(offsets,gt_point):
    # 统计匹配上的关键点数量和kpt平均误差
    matched_cnt = 0
    matched_diff = 0.
    diff = np.array(offsets) - np.array(gt_point)
    distance = np.sqrt(diff[:,0]*diff[:,0]+diff[:,1]*diff[:,1])
    matched_cnt = np.sum(distance <= 2)
    mask = (distance <= 2).reshape(-1,1)
    if matched_cnt != 0:
        matched_diff = np.mean(distance[distance <= 2])
        mask = np.concatenate([mask,mask],1)
        diff = abs(diff[mask].reshape(-1,2))
    else:
        matched_diff = 0
        diff = np.zeros([1,2])
    # print(distance[distance <= 2])
    # print('match cnt%f, matched diff%f'%(matched_cnt,matched_diff))
    return matched_cnt, matched_diff, [np.mean(diff[:,0]),np.mean(diff[:,1])]

def match_images(img1,img2,model,fp_detector,MIN_MATCH_COUNT=4,homo=True,thresh=1.20,knn=1,gt_point=[0,0]):
    normalize_img1 = cv2.normalize(img1,dst=None,alpha=450,beta=10,norm_type=cv2.NORM_MINMAX)
    normalize_img2 = cv2.normalize(img2,dst=None,alpha=450,beta=10,norm_type=cv2.NORM_MINMAX)
    # normalize_img1 = img1
    # normalize_img2 = img2
    kp1 = fp_detector.detect(normalize_img1, None)
    kp2 = fp_detector.detect(normalize_img2, None)
    desc_tfeat1 = kpts2descriptors(kp1,img1,model)
    desc_tfeat2 = kpts2descriptors(kp2,img2,model)
    #print('query.shape:',desc_tfeat1.shape)
    #print('train.shape:',desc_tfeat2.shape)
    bf = cv2.BFMatcher(cv2.NORM_L2)

    matches = bf.knnMatch(desc_tfeat1,desc_tfeat2, k=knn)
    good = []
    for m in matches:
        for mm in m:
            if mm.distance < thresh:
                good.append(mm)
    #print('num good matches:',len(good))

    # src_pts = np.float32([ kp1[m.queryIdx].pt for m in good]).reshape(-1,1,2)
    # dst_pts = np.float32([ kp2[m.trainIdx].pt for m in good]).reshape(-1,1,2)
    # offsets_k,flag = matches2offsets_v2(src_pts,dst_pts)

    offsets = matches2offsets(good,kp1,kp2)
    matched_cnt, matched_diff = 0. ,0.
    diff_xy = [0,0]
    length = len(offsets)
    if length >= 2:
        matched_cnt, matched_diff, diff_xy = matched_kpt_summary(offsets,gt_point)

    if length == 0:
        offsets = [144,144]
    elif length <= 2:
        offsets = offsets[0]
    else:
        #offsets = kmeans(offsets)
        offsets = find_most_common(offsets)

    good_temp = []
    px,py = 0,0
    for m in good:
        qx,qy = kp1[m.queryIdx].pt
        #qx,qy = round(qx),round(qy)
        tx,ty = kp2[m.trainIdx].pt
        #tx,ty = round(tx),round(ty)
        if abs(qx-tx - offsets[0]) < 1 and abs(qy-ty- offsets[1]) < 1:
            good_temp.append(m)
            px += (qx-tx)
            py += (qy-ty)
    good = good_temp

    match_img = cv2.drawMatches(img1,kp1,img2,kp2,good,None,(0,0,255),flags=2)
    bk_gd = np.ones((match_img.shape))
    bk_gd[0:img1.shape[1],0:img1.shape[1]] = cv2.cvtColor(img1,cv2.COLOR_GRAY2BGR)
    bk_gd[0:img2.shape[1],img1.shape[1]:(img1.shape[1]+img2.shape[1])] = cv2.cvtColor(img2,cv2.COLOR_GRAY2BGR)
    bk_gd = cv2.rectangle(bk_gd,(offsets[0],offsets[1]),(int(offsets[0]+img2.shape[1]),int(offsets[1]+img2.shape[1])),(0,0,255),3)

    return offsets,[match_img,bk_gd],[matched_cnt,matched_diff,diff_xy]

def find_most_common(offsets):
    kernel = np.zeros((7,7))
    for i in range(7):
        for j in range(7):
            kernel[i,j] = 1 - math.sqrt((i-3)**2+(j-3)**2)/5

    # cv2.imshow('kernel',kernel)
    # cv2.waitKey(0)
    array = np.zeros((289,289))
    for x,y in offsets:
        ksy,key,ksx,kex = 0,7,0,7
        asy = y-3
        aey = y + 4
        asx = x-3
        aex = x+4
        if asy < 0:
            asy = 0
            ksy = 3-y
        if aey > 289:
            key = 7+289-aey
            aey = 289
        if asx < 0:
            asx = 0
            ksx = 3-x
        if aex > 289:
            kex = 7+289-aex
            aex = 289
        #print(sy,ey,sx,ex)

        array[asy:aey,asx:aex] += kernel[ksy:key,ksx:kex]
    index = array.reshape(-1).argmax()
    offsets = [index%289,index//289]
    return offsets

def validate(model,fp_detector,data_path,save_path):

    match_count = 0
    match_error = 0
    error_x = 0
    error_y = 0
    kpt_cnt, kpt_diff = [],[]
    cnt = 1
    opt_imgs_path = glob(data_path+'/opt_*')
    for i,opt_path in enumerate(opt_imgs_path):
        # try:
        x,y = opt_path.split('/')[-1].split('.')[0].split('_')[2:]
        x,y = float(x),float(y)
        sar_path = opt_path.replace('opt','sar')
        opt_img = cv2.imread(opt_path,0)
        sar_img = cv2.imread(sar_path,0)
        offsets,result_img, matched_info = match_images(opt_img,sar_img,model,fp_detector,homo=False,thresh=7,knn=2,gt_point=[x,y]) # use opt as queryImg, sar as trainImg

        distance = math.sqrt(math.pow(offsets[0]-x,2) + math.pow(offsets[1]-y,2))
        match_flag =  distance < 5
        kpt_cnt.append(matched_info[0])
        kpt_diff.append(matched_info[1])
        if match_flag:
            error_x += abs(offsets[0] - x)
            error_y += abs(offsets[1] - y)
            match_error += distance
            match_count += 1

        cv2.imwrite('./%d_%.2f_%d_kset.png'%(i,distance,match_flag),result_img[0])
        cv2.imwrite('./%d_%.2f_%d_kset_circle.png'%(i,distance,match_flag),result_img[1])
        cnt += 1

    kpt_avg_cnt = np.mean(np.array(kpt_cnt))
    kpt_avg_diff = np.mean(np.array(kpt_diff))

    return match_count,match_error/match_count, error_x/match_count, error_y/match_count,[kpt_avg_cnt,kpt_avg_diff]

if __name__ == '__main__':

    model_path = 'weights_arcl2_cspdense_Rocket/thetabeta_ap_66_684_0.363_28.287.tar'
    embede_model,_ = model_choose(model_path)

    save_path='best_results_diffsize_rocket/{}'.format(model_path.split('/')[-2])
    # fp_detector = cv2.xfeatures2d.SIFT_create(6000)
    # fp_detector = cv2.xfeatures2d.HarrisLaplaceFeatureDetector_create()
    fp_detector = cv2.FastFeatureDetector_create()
    #fp_detector = cv2.BRISK.create()
    start = time.time()

    data_path = 'test_image'
    mc,me,mex,mey,kpt_info = validate(embede_model,fp_detector,data_path,save_path)
    # print('{} test finished!'.format(model_path.split('/')[0].split('_')[1]))
    save_txt = '{}.txt'.format(model_path.split('/')[-2])
    save_data1 = 'matched key point average count: {}, key point average diff:{}\n'.format(kpt_info[0],kpt_info[1])
    save_data2 = 'match_count:%d, location match error:%f, x_error:%f, y_error:%f'%(mc,me,mex,mey)
    print('time consumed: %.2fs'%(time.time()-start))
    print(save_data1)
    print(save_data2)


    # -*- coding: utf-8 -*-
# mpc_nbody/mpc_nbody/parse_input.py

'''
----------------------------------------------------------------------------
mpc_nbody's module for parsing OrbFit + ele220 elements

Mar 2020
Mike Alexandersen & Matthew Payne & Matthew Holman

This module provides functionalities to
(a) read an OrbFit .fel/.eq file with heliocentric ecliptic cartesian els
(b) read ele220 element strings
(c) convert the above to barycentric equatorial cartesian elements

This is meant to prepare the elements for input into the n-body integrator
----------------------------------------------------------------------------
'''

# Import third-party packages
# -----------------------------------------------------------------------------
import os, sys
import numpy as np
from astropy.time import Time
import getpass

if getpass.getuser() in ['matthewjohnpayne']:  # Payne's dev laptop set up differently ...:
    sys.path.append('/Users/matthewjohnpayne/Envs/mpcvenv/')
import mpcpp.MPC_library as mpc

# Import neighbouring packages
# -----------------------------------------------------------------------------

# Default for caching stuff using lru_cache
# -----------------------------------------------------------------------------

# Constants and stuff
# -----------------------------------------------------------------------------
DATA_PATH = os.path.realpath(os.path.dirname(__file__))
au_km = 149597870.700  # This is now a definition

# Data classes/methods
# -----------------------------------------------------------------------------

class ParseElements():
    '''
    Class for parsing elements and returning them in the correct format.
    '''

    def __init__(self, input_file=None, filetype=None, save_parsed=False ):

        # The variables that will be used to hold the elements
        # - They get populated by *parse_orbfit* & *make_bary_equatorial*
        self.helio_ecl_vec_EXISTS   = False
        self.helio_ecl_vec          = None
        self.helio_ecl_cov_EXISTS   = False
        self.helio_ecl_cov          = None
        self.bary_eq_vec_EXISTS     = False
        self.bary_eq_vec            = None
        self.bary_eq_cov_EXISTS     = False
        self.bary_eq_cov            = None

        # If input filename provided, process it:
        if isinstance(input_file, str) & isinstance(filetype, str):
            if filetype == 'ele220':
                self.parse_ele220(input_file)
            if (filetype == 'fel') | (filetype == 'eq'):
                self.parse_orbfit(input_file)
            self.make_bary_equatorial()
            if save_parsed:
                self.save_elements()
        else:
            print("Keywords 'input_file' and/or 'filetype' missing; "
                  "initiating empty object.")

    def save_elements(self, output_file='holman_ic'):
        """
        Save the barycentric equatorial cartesian elements to file.

        Inputs:
        -------
        output_file : string, filename to write elements to.

        The file is overwritten if it already exists.
        """
        self.tstart = self.time.tdb.jd
        outfile = open(output_file, 'w')
        outfile.write(f"tstart {self.tstart:}\n")
        outfile.write("tstep +20.0\n")
        outfile.write("trange 600.\n")
        outfile.write("geocentric 0\n")
        outfile.write("state\n")

        # For whatever reason, we are writing this over two lines
        # - perhaps to compare against JPL?
        for n,coeff in enumerate(self.bary_eq_vec):
            suffix = '\n' if n in [2,5] else ''
            outfile.write(f"{coeff: 18.15e} " + suffix)

    def parse_ele220(self, ele220file=None):
        '''
        Parse a file containing a single ele220 line.
        Currently returns junk data.
        NOT ACTUALLY IMPLEMENTED YET!!!
        '''
        if ele220file is None:
            raise TypeError("Required argument 'ele220file'"
                            " (pos 1) not found")

        # make fake data & set appropriate variables
        self._get_and_set_junk_data()

    def parse_orbfit(self, felfile):
        '''
        Parse a file containing OrbFit elements for a single object & epoch.
        Currently returns junk data.

        Inputs:
        -------
        felfile : string, filename of fel/eq formatted OrbFit output

        Populates:
        --------
        self.helio_ecl_vec_EXISTS   : Boolean
        self.helio_ecl_vec          : 1D np.ndarray
        self.helio_ecl_cov_EXISTS   : Boolean
        self.helio_ecl_cov          : 1D np.ndarray
        self.time                   : astropy Time object
        '''

        # Read the contents of the orbfit output "fel" file
        obj = {}
        with open(felfile,'r') as fh:
            el = fh.readlines()
        cart_head = '! Cartesian position and velocity vectors\n'

        # Only do this if the file actually has cartesian coordinates.
        if el.count(cart_head) > 0:
            # get Cartesian Elements out of the file contents
            carLoc = len(el) - 1 - list(reversed(el)).index(cart_head)
            carEls = el[carLoc:carLoc + 25]

            # Form an array of the heliocentric ecliptic cartesian coefficients
            (_, car_x, car_y, car_z, car_dx, car_dy, car_dz
                       ) = carEls[1].split()
            self.helio_ecl_vec = np.array([ float(car_x), float(car_y),  float(car_z), \
                                            float(car_dx), float(car_dy), float(car_dz)]
                                            )
            self.helio_ecl_vec_EXISTS = True

            # Using Astropy.time for time conversion,
            # because life's too short for timezones and time scales.
            _, mjd_tdt, _ = carEls[2].split()
            self.time = Time(float(mjd_tdt), format='mjd', scale='tt')

            # Parse carEls (the contents of the orbfit file) to get
            # the cartesian covariance matrix
            self.helio_ecl_cov_EXISTS, self.helio_ecl_cov = _parse_Covariance_List(carEls)

        else:
            raise TypeError("There does not seem to be any valid elements "
                            f"in the input file {felfile:}")

    def make_bary_equatorial(self):
        '''
        Transform heliocentric-ecliptic coordinates into
        barycentric equatorial coordinates

        requires:
        ----------
        self.helio_ecl_vec_EXISTS   : Boolean
        self.helio_ecl_vec          : 1D np.ndarray
        self.helio_ecl_cov_EXISTS   : Boolean
        self.helio_ecl_cov          : 2D np.ndarray

        populates:
        ----------
        self.bary_eq_vec_EXISTS     = Boolean
        self.bary_eq_vec            = 1D np.ndarray
        self.bary_eq_cov_EXISTS     = Boolean
        self.bary_eq_cov            = 2D np.ndarray
        '''
        if self.helio_ecl_vec_EXISTS :
            # Transform the helio-ecl-coords to bary-eq-coords
            # NB 2-step transformation for the vector (posn,vel)
            self.bary_eq_vec   = equatorial_helio2bary(
                                    ecliptic_to_equatorial(self.helio_ecl_vec),
                                    self.time.tdb.jd
                                )
            # Set boolean as well (not sure if we'll really use these ...)
            self.bary_eq_vec_EXISTS = True

        if self.helio_ecl_cov_EXISTS:
            # Only need to do a rotation for the CoV
            self.bary_eq_cov = ecliptic_to_equatorial(self.helio_ecl_cov)

            # Set booleans as well (not sure if we'll really use these ...)
            self.bary_eq_cov_EXISTS = True

        if not self.helio_ecl_vec_EXISTS and not self.helio_ecl_cov_EXISTS:
            raise TypeError("There does not seem to be any valid helio_ecl to transform into bary_eq")

        return True

    def _get_and_set_junk_data(self, BaryEqDirect=False ):
        """Just make some junk data for saving."""
        self.time                           = Time(2458849.5, format='jd', scale='tdb')
        v   = np.array( [3., 2., 1., 0.3, 0.2, 0.1] )
        CoV = 0.01 * np.ones((6,6))

        # Default is to make helio-ecl, then calc bary-eq from that
        if not BaryEqDirect:
            self.helio_ecl_vec              = v
            self.helio_ecl_vec_EXISTS       = True

            self.helio_ecl_cov              = CoV
            self.helio_ecl_cov_EXISTS       = True

            self.make_bary_equatorial()

        # Alternative is to directly set bary-eq
        else:
            self.bary_eq_vec                = v
            self.bary_eq_vec_EXISTS         = True

            self.bary_eq_cov                = CoV
            self.bary_eq_cov_EXISTS         = True

# Functions
# -----------------------------------------------------------------------------

def ecliptic_to_equatorial(input, backwards=False):
    '''
    Rotates a cartesian vector or Cov-Matrix from mean ecliptic to mean equatorial.

    Backwards=True converts backwards, from equatorial to ecliptic.

    inputs:
    -------
    input : 1-D or 2-D arrays
     - If 1-D, then len(input) must be 3 or 6
     - If 2-D, then input.shape must be (6,6)

    output:
    -------
    output : np.ndarray
     - same shape as input
    '''

    # Ensure we have an array
    input = np.atleast_1d(input)

    # The rotation matricees we may use
    direction = -1 if backwards else +1
    R3 = mpc.rotate_matrix(mpc.Constants.ecl * direction)
    R6 = np.block( [ [R3, np.zeros((3,3))],[np.zeros((3,3)),R3] ])

    # Vector input => Single rotation operation
    if   input.ndim == 1 and input.shape[0] in [3,6]:
        R      = R6 if input.shape[0] == 6 else R3
        output = R @ input

    # Matrix (CoV) input => R & R.T
    elif input.ndim == 2 and input.shape == (6,6):
        R = R6
        output = R @ input @ R.T

    # Unknown input
    else:
        sys.exit(f'Does not compute: input.ndim=={input.ndim} , input.shape={input.shape}')

    assert output.shape == input.shape
    return output

def equatorial_helio2bary(input_xyz, jd_tdb, backwards=False):
    '''
    Convert from heliocentric to barycentic cartesian coordinates.
    backwards=True converts backwards, from bary to helio.
    input:
        input_xyz - np.ndarray length 3 or 6
        backwards - boolean
    output:
        output_xyz  - np.ndarray
                    - same shape as input_xyz

    input_xyz MUST BE EQUATORIAL!!!
    '''
    direction = -1 if backwards else +1

    # Ensure we have an array of the correct shape to work with
    input_xyz = np.atleast_1d(input_xyz)
    assert input_xyz.ndim == 1
    assert input_xyz.shape[0] in [3,6]

    # Position & Motion of the barycenter w.r.t. the heliocenter (and vice-versa)
    delta, delta_vel = mpc.jpl_kernel[0, 10].compute_and_differentiate(jd_tdb)

    # Work out whether we need xyz or xyzuvw
    delta = delta if input_xyz.shape[0] == 3 else np.block([delta,delta_vel])

    # Shift vectors & return
    return input_xyz + delta * direction / au_km

def _old_parse_Covariance_List(Els):
    '''
    Convenience function for reading and splitting the covariance
    lines of an OrbFit file.
    Not intended for user usage.
    '''
    ElCov  = []
    covErr = ""
    for El in Els:
        if El[:4] == ' COV':
            ElCov.append(El)
    if len(ElCov) == 7:
        _, c11, c12, c13 = ElCov[0].split()
        _, c14, c15, c16 = ElCov[1].split()
        _, c22, c23, c24 = ElCov[2].split()
        _, c25, c26, c33 = ElCov[3].split()
        _, c34, c35, c36 = ElCov[4].split()
        _, c44, c45, c46 = ElCov[5].split()
        _, c55, c56, c66 = ElCov[6].split()
    if len(ElCov) != 7:
        c11, c12, c13, c14, c15, c16, c22 = "", "", "", "", "", "", ""
        c23, c24, c25, c26, c33, c34, c35 = "", "", "", "", "", "", ""
        c36, c44, c45, c46, c55, c56, c66 = "", "", "", "", "", "", ""
        covErr = ' Empty covariance Matrix for '
    return (covErr, c11, c12, c13, c14, c15, c16, c22, c23, c24, c25, c26,
            c33, c34, c35, c36, c44, c45, c46, c55, c56, c66)

def _parse_Covariance_List(Els):
    '''
    Convenience function for reading and splitting the covariance
    lines of an OrbFit file.
    Not intended for user usage.
    # MJP : 20200901 : Suggest to just make & return the required matrix
    '''
    # Set-up array of zeroes
    CoV        = np.zeros( (6,6) )
    CoV_EXISTS = False

    # Populate triangle directly
    ElCov=[]
    for El in Els:
        if El[:4] == ' COV':
            ElCov.append(El)
    if len(ElCov) == 7:
        _, CoV[0,0],CoV[0,1],CoV[0,2] = ElCov[0].split() # c11, c12, c13
        _, CoV[0,3],CoV[0,4],CoV[0,5] = ElCov[1].split() # c14, c15, c16
        _, CoV[1,1],CoV[1,2],CoV[1,3] = ElCov[2].split() # c22, c23, c24
        _, CoV[1,4],CoV[1,5],CoV[2,2] = ElCov[3].split() # c25, c26, c33
        _, CoV[2,3],CoV[2,4],CoV[2,5] = ElCov[4].split() # c34, c35, c36
        _, CoV[3,3],CoV[3,4],CoV[3,5] = ElCov[5].split() # c44, c45, c46
        _, CoV[4,4],CoV[4,5],CoV[5,5] = ElCov[6].split() # c55, c56, c66

        # Populate the symmetric part
        for i in range(1,6):
            for j in range(i):
                # # MA: Killed totally annoying and unneccessary print
                #print(f'Setting Cov[{i,j}] = CoV{[j,i]}')
                CoV[i,j]=CoV[j,i]

        # Set boolean
        CoV_EXISTS = True
    return CoV_EXISTS, CoV