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

# # Lab 3: Building a Photometric Pipeline
#
# In this lab, we'll be using classes and functions to build a pipeline which will automatically extract the fluxes of stars in an image. We're all familiar with aperture photometry, but in this case, we're going to take the additional step of convolving our image with a PSF.
#
# Our pipeline will be split into several steps.
#
# 1. Reading in an image
# 2. Finding the local peaks in the image (the stars)
# 3. Calculating the centroid of each peak region
# 4. Convolving with the PSF and extracting flux.
#
#

# ```{warning}
# When you modify a class and re-run it, you need to re-instantiate (or make a new instance) of the class for it to reflect the changes. This is a common mistake in notebooks, where a class might get changed but your instance of it didn't get re-created. Watch out for this during the lab -- the safest bet is to just copy-paste your instantiation code (`a = MyObj()` type thing) into each problem so you sure you're working with the latest version of your class.
#
# You'll be copying and pasting (and overwriting) your class in this assignment quite a bit. That's not super realistic to reality (where you'd pick one cell and keep modifying it there), but for grading purposes we ask you follow along so we can give credit for each method of the class you complete and demonstrate along the way.
# ```

# ## Assignment Rubric
# ````{panels}
# The points for this assignment are as follows:
# - Problem 1 (4 pts)
# - Problem 2 (8 pts)
# - Problem 3 (5 pts)
# - Problem 4 (3 pts)
# ````

# ## Problem 1
# In this problem, we'll load in the image(s) for use and begin constructing our `PSFPhot` class.
#
# ### Problem 1.1
# Create a script (`.py` file) in this directory and copy into it your `load_fits()` and final `implot()` functions from last lab. Then import them here.
#
# Use your `load_fits()` to read in in the image `2020-04-15-0001.fits` and plot it with `implot()`.
#
# This image was taken of the M81/M82 field using the Dragonfly Telephoto Array (in a narrowband configuration).

# In[ ]:

# import your functions here

# In[1]:

# Use them to load and show the image we'll be working with

# ```{hint}
# M82 should be visible left of center, M81 on the left hand edge. Two features caused by detector amp glow are also visible at the bottom of the image.
# ```

# ### Problem 1.2
#
# Finish the syntax for creating the `PSFPhot` class below. The init for the class should ask for three string paths: one to an image, one to a dark file, and one to a flat file. Within `__init__()`, use your `load_fits()` function to get the just the science image and header into the class. Store them as class attributes `self.data_init` and `self.data_header`.
#
# Don't forget to add a docstring to your init! It doesn't return anything, but sets the initial attributes.

# In[ ]:

class PSFPhot():

# ### Problem 1.3
#
# Add a method to your class called `dark_subtract()` which takes an image and a string-path to a dark file. Use `load_fits()` to read it in the dark, and subtract it from the image, and return the output.
#
# Add another method called `flat_field()` which takes in an image and a string-path to a flatfield file. Use `load_fits()` to read it in, normalize it to its maximum value, and divide the input image by it.
#
# Finally, within your init function, set a new class attribute called `self.data_calibrated`, and set it equal to the output of your calibration methods as run on the set of image paths that have been input.
#
# ```{hint}
# This step can be done in one line with nested functions if you prefer.
# ``
#
# Finally Use `implot()` to plot both the `data_init` and `data_calibrated` and show that your calibrations worked. You should see the amp glow regions dissapear, and the image sky should look nice and uniform.

# In[527]:

class PSFPhot():

# ```{note}
# Running the dark/flat methods within init means they'll happen right away when someone instantiates an object of this class. It's a matter of code design and personal preference to do this, versus having the user run those methods manually as part of the UX. I choose to wrap it into the init because it's a basic, common thing all inputs will need done to them.
# ```

# ### Problem 1.4
#
# The final step in setting up our image for proper analysis is the subtraction of the globally varying sky background. Performing that fit (generally a low order 2D polynomial of some kind) is beyond the scope of this lab, but we encourage you to look into this step more if you are interested.
#
# Instead, we're going to return to the `sep` package and take advantage of its background estimation feature. Recall it reads in the base image array and a separate mask. We want the user of this class to be able to supply a mask, so we won't run this step automatically in `__init__()`.
#
# Using the same stucture as in Lab 2, we'll calculate the background using `sep`.
#
# Add a method to your class that is called `subtract_background`, whose only argument (besides `self`) should be an optional mask (default `None`).
#
# Inside, use `sep` to measure the background and access, e.g., the `bkg.back()` method, and subtract this background-image from `self.data_calibrated`. You should set `mask=mask` in the sep call so your method input gets passed through. For convenient access, store the output of `bkg.back()` in a class attribute called `self.background`, and then set `self.image` equal to the background subtracted image. This method is known as a *setter*, because we're setting a class attribute but not returning anything.
#
# It can be handy here to add a print statement at the end of the method saying something like `"Background estimated; output saved to attribute 'image' "`. Since this is the main image we'll be using from now on, I've elected to give it the short name `image`. You may rather use something like `data_bg_sub`, for example
#
# ```{warning}
# Don't forget about the C order switch with sep.
# ```

# In[528]:

class PSFPhot():

# Our class usage will now look something like
# ```
# pipe = PSFPhot(im_path,dark_path,flat_path)
# pipe.subtract_background()
# ```
#
# or, if we create a mask (as we will below),
# ```
# pipe = PSFPhot(im_path,dark_path,flat_path)
# mask = # Something here
# pipe.subtract_background(mask=mask)
# ```
#

# ## Problem 2
#
# Our goal is to estimate the PSF of the above image, then measure fluxes of the stars and galaxies here accounting for the PSF.
#
# To start this process, we need to locate the stars in this image. We saw how to segment an image using `sep` last time, but in this lab we are going to carry out this step ourselves, using two methods.
#
# ### Problem 2.1
#
# Before we do this, we want to take the step of masking out several regions of the image which may register as peaks but which are not nicely isolated stars. In particular, the two galaxies need to be masked out when both estimating the background and when looking for point sources.
#
#
# Create a mask of the dimensions of the image data, containing `False` everywhere except where we want to mask out the galaxies (rough rectangular regions are fine for this masking). It's easiest to make the full array `False` first, then set the regions we want to mask to `True` via image indexing.
#
# Add a method to your class called `set_image_mask()` which overwrites `self.image` with a numpy masked array containing what used to be `self.image` as the data, and a mask that the user inputs.
#
# Finally, instantiate your class into a variable and check that the masking and subtraction are both working. You should plot your image by accessing it (it should be a class attribute by now). It should look like mine below.
# ```{warning}
# If we have the memory to, it's often worth avoiding overwriting class attributes. Here, we could set the masked image as `self.image_masked`, for example. I've elected not to here for two reasons. One: simplicity. Two: the unmasked version of the image is easily accessible still via `self.image.data`, the data attribute of numpy masked arrays.
# ```

# In[535]:

# Solution
mask = #define mask

class PSFPhot():

# In[536]:

# instantiate your class, run the subtraction, input the mask, see that it all works.

#
# ```{note}
# In our use case for this lab, the mask we will use when background subtracting is the same as the one we'll use for blocking out regions from being included in peak finding and such. But our code is flexible enough to, e.g., mask when background-subtracting but not when peak finding, or vice versa. Our mask setter only affects the image data, and the mask argument in the subtract background method only affects background subtraction.
# ```
#
# ```{tip}
# As it stands, `set_image_mask()` can only be run *after* `subtract_background()`, because that method sets `self.image` for the first time. We can use the `hasattr(self,'image')` check to see, when `set_image_mask()` is run, if that previous method was already run and `self.image` exists. For now, it's ok to assume the user will use the class in order, but a more robust code would either raise an exception or otherwise handle this case in some way.
# ```

# Plot the `.background` attribute of your class. It seems our image has a spatially varying background that includes a gradient across the image.

# In[ ]:

#plot the background

# ### Problem 2.2
#
# Now that we have the appropriate image regions masked, we can move on to the peak finder.
#
# The "fast" or "efficient" method of doing this involves some scipy filtering operations. But for our purposes, the "slow way" (iterating over the image pixels) takes ~few seconds to run, and is worth doing to build intuition.
#
# Add a method to your class to find peaks in an image by looping over each pixel and checking its neighbors, with a "peak" being defined as a region of higher flux than all adjacent pixels (i.e., the 8 surrounding pixels). In order to not pick up random noise pixels, also take an input called `threshold`. Within your algorithm, don't return any pixels which are "peaks" but for which the pixel value is below this threshold.
#
# ```{hint}
# :class: dropdown
# This algorithm needs to avoid the edges of the image (since indexing "i+1" won't work there). Creating a 1 pixel "buffer" in your looping will prevent this.
# ```

# In[70]:

# your code here

# The looping solution is slow, and will not scale well if we have to run on many images, but for one image is okay.
#
# There are several solutions which generally involve either **filtering** the image or **cross correlating** the image with a template. Here's one such solution.

# In[399]:

from scipy.ndimage import maximum_filter

def findpeaks_maxfilter(image, threshold):
    '''
    Algorithm for finding peaks (above a threshold) in an image

    Parameters
    ----------
    image: array_like
        2D array containing the image of interest.
    threshold: float
        minimum pixel value for inclusion in search

    Returns
    -------
    peaks: array_like
        array containing the x and y coordinates of peak regions.
    '''
    neighborhood = np.ones((3,3),dtype=bool) # just 3x3 True, defining the neighborhood over which to filter
    # find local maximum for each pixel
    amax = maximum_filter(image, footprint=neighborhood) #max filter will set each 9-square region in the image to the max in that region.

    peaks = np.where((image == amax) & (image >= threshold)) #find the pixels unaffected by the max filter.
    peaks = np.array([peaks[0],peaks[1]]).T
    return peaks

# Let's take a moment to understand how this algorithm works. The key is in the `maximum_filter()` step. Filtering a 2D image is a process carried out in fourier space, which is what allows scipy to carry it out quickly. But what is maximum filtering?
#
# ```{admonition} Definition
# Maximum Filtering is the process by which all pixels in local neighborhoods within an array are raised to the maximum value of any pixel in that neighborhood
# ```
#
# Let's look at a 1D case. Below, I define a 1D array that has some peaks in it.

# In[110]:

array_1dpeaks = np.array([1,1,2,1,2,1,2,1,2,3,4,3,2,1,2,4,5,6,8,6,5,4,3,2,4,3,4,2,1,0.5,1,2,1,2,3,2])

# Our data looks like this:

# In[122]:

fig, ax = plt.subplots(figsize=(6,6))
ax.plot(array_1dpeaks,lw=1.5,alpha=0.9)
ax.plot(array_1dpeaks,'o',color='C0')

# Let's now run the maximum filter on this data and plot it's output. I'm going to pick a neighborhood of 3, which means +/- 1 pixel around each location.

# In[112]:

mf = maximum_filter(array_1dpeaks, footprint=np.ones(3,dtype=bool))

# In[120]:

fig, ax = plt.subplots(figsize=(6,6))
ax.plot(array_1dpeaks,lw=1.5,alpha=0.9,label='original array')
ax.plot(array_1dpeaks,'o',color='C0')
ax.plot(mf,label='max filtered array')
eq, = np.where(array_1dpeaks==mf)
ax.plot(np.arange(len(mf))[eq],mf[eq],'o',label='local peaks')
ax.legend();

# What the filtering has done is for every 3 pixel neighborhood across this array, it's raised the value of all three pixels to the maximum value across the three. So we see that anywhere the three pixels were, e.g., (1,2,1), they are all now 2. What you should notice looking at this plot is that the max filtering has also identified true peaks in our data! Notice that the only spots where the orange curve (the max filtered version of the data) is equal to the original array is exactly at locations that are local maxima. This is because when applying max filtering to an array, the only values *unchanged* by the filtering are those that *are* local maxima (in the neighborhood defined).
#
# And thus, we have our peaks! All we need to do is find out `where()` the max filtered array equals the original array. Of course, we can also put in a threshold (in this example, maybe 2.5) to ensure low level noise doesn't enter in. This is why the `findpeaks_maxfilter()` function has a threshold option as well.
#
# ```{note}
# You may notice that the first index in the array is marked as a peak, despite not being one. Edges are always troublesome with these algorithms, and they normally have multiple options for how edges are handled, e.g, interpolating a constant value, reflecting over the edge, etc. For our lab, we're not going to worry about this.
# ```
#

# ### Problem 2.3
#
# How fast is max filtering over your looping algorithm?
#
# Use the `%%timeit` magic command in your notebook to test how fast the two algorithms are respectively. If you were working with a sample of 1000 images, how long would the looping algorithm take compared to the max-filtering case?
#
# ```{hint}
# Check out the "Timing Code" page in the "Quick Tips" sidebar tab of our website to see examples of using the `timeit` module.
# ```

# In[ ]:

# time your function

# In[ ]:

# time the max filter function

# How many times faster is the max filtering?

# In[ ]:

# Calculate

# How long would it take to run each on 1000 images (extrapolating from your results)?

# In[ ]:

# Calculate

# For the rest of the lab, you can either leave in your peak finder, or replace it with the max-filtering version (turning the function above into a method).
#
# Either way, add a peak-finding method into your class, which should be runnable with no inputs, but has an optional threshold argument to set a minimum flux value for inclusion.

# In[ ]:

# your class here

# ### Problem 2.4
#
# Run your peakfinder on your masked image, and assemble the list of peaks you'll use moving forward. Your peakfinder method should save a class attribute `self.peak_locations` which contains the ($x,y$) pairs of points associated with your peaks.
#
#
# Plot up the original image, but use our "source circling" technique from last lab to circle all the peaks found in the masked image. I show mine below.
#
# ```{hint}
# Recall that `implot()` returns a set of `fig, ax`, so you can then plot the circles onto the same axis as the image.
# ```

# In[549]:

# re-instantiate your object, run the subtraction, then your new peak finder here.

# ### Problem 2.5
#
# This should look pretty good -- most of what's circled above is clearly a star/point source in the image. However, one problem with this method is that single hot pixels are going to be registered as peaks, even via the clever algorithm. We need a way to eliminate these from our sample before moving on.
#
# Adjust your peak-finder method to add in something that checks that not only is there a true peak, but that at least 4 of the pixels around the peak are also elevated in flux (I used 0.5 times the peak flux). The easiest way is to loop over the peaks after they're found and institute the check --- there are far fewer peaks than pixels, so this doesn't significantly affect the runtime. But feel free to find a cleverer solution!
#
# ```{warning}
# Be careful with transpositions in this problem. when you plot coordinates, you plot(x,y), but when you index the image, you index image[y,x]. Always be tracking which is which!
# ```

# In[619]:

class PSFPhot():

# Re-find your peaks using your newer, better algorithm, and plot them below as before.

# In[ ]:

# your code here

# Notice that we've decreased our total number of peaks. But you should find that now, everything currently circled looks like a bright, "resolved" star. (resolved insofar as the PSF is spreading the light of the star over multiple pixels).
#
# ### Problem 2.6
#
# In your the image above, you should see that ~8-10 stars look like they are circled by several very closely overlapping circles all targeting the same star. Infer (or investigate and determine) why this has happened, and write your answer below.
#
#
#

# *answer here*

# ## Problem 3
#
# We now have a function that can return the peaks in a given image. Our next step is going to be to estimate the exact center of those peaks (stars) using their **centroid**.
#
# ```{admonition} Definition
# The centroid is the light-weighted-mean of a set of pixels. It is not always the maximum-valued pixel, and is determined to sub-pixel accuracy.
# ```
#
# Many of you have seen the centroid formula, but as a reminder, it looks like this (in 1D):
#
#
# $$
# x_{\rm com} = \frac{\sum{x_i \hat{f}_i}}{\sum \hat{f}_i},
# $$
#
# where $x_i$ are the positions and $\hat{f}_i$ are the fluxes at those positions.
#
# In 2D, when working with images, the $x$ and $y$ centers of mass are independent, and the 2D centroid is just the location ($x_{\rm com}$, $y_{\rm com}$).
#
# ### Problem 3.1
#
# Add a method to your class called `centroid()` which should read in an $x$ and $y$ peak location (these will be returned by your peak finder), then creates a cutout by indexing `self.image` in a window of N (user-settable), and determine the centroid of this window. The $x,y$ location of this centroid should be returned.
#
# One subtlety --- We want to use a window size greater than 1 pixel on either side of each peak. But because our peak finder is likely to have found peaks near the edge of the detector (both because it needs only 1 pixel-thick borders and because it handles edges), if we write a centroid function that, e.g., uses a 10x10 pixel window, we'll end up trying to index over the edge of the original image whenever there's a star near the edge. Because of this, your function should raise an exception if a peak position is entered whose distance from an edge is less than half the window size.

# In[622]:

class PSFPhot():
    # copy down and add to your class

# ### Problem 3.2
# Use your `centroid()` function to confirm that the algorithm is working by testing it on a few individual peaks from your peak list, and make a plot showing the window region and the determined centroid (along with the location of the input peak). I'm leaving behind a demo of what I mean below. The blue point is the pixel with the peak flux, while the crosshairs center on the determined centroid
#
# ```{Note}
# It's usually ok if the centroid is not at what appears to be the very center of the light distribution of the star. Often due to seeing and detector effects, along with tricks of the stretch you happen to be using, the centroid doesn't look perfectly centered. But it shouldn't be super off, either -- see below.
# ```
# In my solution below, I plot 25 stars from my peak finder, with peaks and centroids shown. You should get similar looking results as you query stars in frame.
#

# In[624]:

# You need not plot 25, but check a handful of stars to make sure you're in business.

# ### Problem 3.3
#
# If you recall from above, you determined why the peak algorithm occasionally marked a bunch of pixels as peaks within the same star, which shouldn't happen. It should be clear from your answer that these stars will be unusable for the purposes of measuring the PSF of stars in the image. We thus need to remove these ~8-10 stars from our sample.
#
# Write a method `check_stars()` which takes in the list of centroids, and identifies these cases. The easiest way to do this is to simply iterate through the list of centroids and compute the distance to all other centroids in the list. Any subgroups with small distances should be removed. (Say, 5-10 pixels).
#
#
# This method should return the final list of centroids for use. Plot them over the data to confirm these "stacked" peak cases have been removed.
#

# In[ ]:

# your code here

# ## Problem 4
#
# Armed with a dark-subtracted, flat-fielded, background-subtracted image, as well as with a list of centroids corresponding to stars in our image, we are ready to estimate the PSF.
#
# There are two main functional forms typically used to fit star profiles: 2D Gaussians, and Moffat profiles (which combines the shapes of a Gaussian and Lorentzian to best match both the inner and outer regions of the PSF).
#
# We're going to use the [`Gaussian2D`](https://docs.astropy.org/en/stable/api/astropy.modeling.functional_models.Gaussian2D.html) class from `astropy` to do this:

# In[433]:

from astropy.modeling.functional_models import Gaussian2D

# For each star, a Gaussian2D profile (normalized) will be used "as the PSF". The parameters we need to know for this profile are $x,y$, for which we'll use the centroids we calculated earlier, the amplitude (set by the normalization), and $\sigma_x,\sigma_y$, the standard deviations in the two axes. For this lab, we're going to assume our stars are circular ($\sigma_x=\sigma_y$). This is a strictly incorrect, but not a bad assumption for most cases. All other optional arguments we won't need, primarily due to the assumption of circularity.
#
# ```{note}
# We are going to make a point estimate of the "size" of the stars in our image, which constrains us from using a more fancy model for the PSF. An example of a more sophisticated setup would be *fitting* a Gaussian or Moffat profile to every star, and in a Bayesian framework marginalizing over the stars to determine the best-fit PSF (including ellipticity, etc) for the image, or, even fancier, interpolating a PSF model which varies over the detector.
# ```
#
# PSF photometry works by multiplying the *data* (say, a cutout around a star) by the *estimated PSF* during the fluxing stage. Instead of picking a radius and performing aperture photometry (which includes fully all pixels within the aperture and throws out all pixels beyond), this method attempts to weight each pixel fractionally by how likely it is to be stellar flux, with the weighting coming from the PSF of the detector. This means further pixels may still be included, but will contribute less than pixels near the center of the star.
#
# The formula for measuring the PSF flux of a star is
#
# $$
# f_{\rm PSF} = \frac{\sum \hat{f_i} p_i}{\sum p_i^2},
# $$
#
# where $\hat{f_i}$ are the fluxes in your image and $p_i$ is your PSF estimate. This formula should be reminiscent of the centroiding formula; it's a similar weighting scheme.

# `Gaussian2D` is a class, but we want to interact with it pretty simply, and have simplified inputs. I've made a quick wrapper function below which allows us to enter $\sigma_x$ and $\sigma_y$ and then $x,y$ grids created via `np.meshgrid()`, and creates the Gaussian and evaluates it on our grid.

# In[490]:

def eval_gauss(x_arr,y_arr,sigma_x,sigma_y,mu_x,mu_y):

    g = Gaussian2D.evaluate(x=x_arr,y=y_arr,amplitude=1,theta=0,x_mean=mu_x,
                   y_mean=mu_y,
                   x_stddev=sigma_x,
                   y_stddev=sigma_y)
    g/=np.sum(g)
    return g

# In[491]:

xx, yy = np.meshgrid(np.arange(loc[0]-10,loc[0]+10),
                     np.arange(loc[1]-10,loc[1]+10))
model = eval_gauss(x_arr=xx,y_arr=yy,sigma=3,mu_x=testx,mu_y=testy)

# In[485]:

fig, ax = plt.subplots(figsize=(8,8))
ax.imshow(model,origin='lower')

# As we can see, I now have a model for the PSF which I can easily create for given inputs. We're going to do this for cutouts around each star, and instead of a random $\sigma$, we're going to estimated it using the second moment (moment of inertia) of the star itself.
#
# The formula for this (from Markevich et al. 1989) is
#
# $$
# \sigma_x = \left[\frac{\sum x_i^2 \hat{f}_i}{\sum \hat{f}_i} - x_{\rm com}^2\right]^{1/2}
# $$
#
# $$
# \sigma_y = \left[\frac{\sum y_i^2 \hat{f}_i}{\sum \hat{f}_i} - y_{\rm com}^2\right]^{1/2}
# $$
#
# In this case, we'll need to use `meshgrid()` directly within our second moment function, as you can see it depends on the difference between the pixels and the centroid.
#
# ### Problem 4.1
#
# Add a method to your class called `second_moment` which reads in an image cutout, the meshgrid input (xx,yy) (or constructs it), and finally, a centroid (x and y). Inside, use the formulas above to determine sigma_x and sigma_y, and return them

# In[494]:

# Below, I show that a functional form of my second moment code returns 2.92 and 2.98 when run on the example psf I made above (which had a true sigma of 3).

# In[495]:

second_moment(model,xx,yy,testx,testy)

# Within 10%, this calculation, which requires no fitting, tells us $\sigma$ (under the assumption the distribution is Gaussian). Thus, by running out image cutouts through this function, we can derive for ourselves a good $\sigma$ to choose in our Gaussian model of the PSF.

# ### Problem 4.2
#
# We now have everything we need to run our pipeline.
#
# Right now, our final three methods (centroider, star checker, and width-finder) all *return* their values. I had you do it that way for testing purposes. To finalize our "object oriented pipeline" though, we should make a wrapper method that will bring them together and handle the outputs.
#
#
# Add a new, final method called `psf_photometry`. When the user runs this method, it should first feed the peaks into the centroid code one by one, assembling a set of centroids. It should then construct cutouts of `self.image` around each peak (or centroid), and feed those, plus the centroids, and appropriate meshgrids into the second moment function to save a pair sigma_x and sigma_y for each star as well. And finally, it should use the eval_gauss function I've provided above to create a reasonable PSF model, then carry out the PSF photometry (using the equation above).
#
# At the end of it all, you can save the centroids, widths, and psf-fluxes all to class attributes, where they will remain accessible. You can also find a nice way to return them for extra credit (see below).

# In[ ]:

# I know this was long! Congrats for getting here!

# ## Bonus Problem (+1 Extra Credit)
#
# A convenient way to return the set of measurements carried out by our pipeline is via a `flux table`, like the one returned by `aperture_photometry` from `astropy`. For our purposes, since we'll be spending a lot of time with data frames, let's go with that.
#
# Modify your final `psf_photometry` method to assemble a nice `DataFrame` which contains as colummns `peak_x`, `peak_y`, `centroid_x`, `centroid_y`, `width_x`, `width_y`, and `psf_flux`. Each row should be the relevant values for a given star that was originally found by your peakfinder (but which made it through our sample cleaning stage).
#
# ```{tip}
# Inside your final method, you probably have lists or arrays containing the output of your centroid and width methods as they looped over the stars (and your find peak function returns a list of positions outright). So if you set up an empty DataFrame, you can set, e.g., `df['peak_x'] = peaks[:,0]` or something like `df['width_x'] = out_widths`, to set up these columns. As long as you never messed with the order of the lists, it should turn into the desired output frame.
# ```

# In[ ]:


""" The script reads a file like
`data/sample_D3.5Mpc_Emin56EeV_1000nuclei.txt` obtained with
`psample.py` and then creates a sample of "UHECRs" coming from a
particular source, using "nucleus_NNNEeV.txt.xz" obtained with
galback.py and located in a directory that corresponds to the
selected Nside and the GMF model (e.g., in data/jf/512/).

Input:
    * source_id='CenA' -- the "name" of an UHECR source
    * Nini=1000 -- the number of nuclei in the initial sample
      (obtained with `psample.py`)
    * Emin=56 [EeV]
    * Nside=512 -- the HEALPix parameter
    * source_vicinity_radius = 1 [deg]
    * a file `data/sample_D..._sorted.txt` prepared with `psample.py`
    * GMF = 'PTKN11'|'JF12ST' -- the GMF model

Output:
    * a file like
      `src_sample_CenA_D3.5_Emin56_N1000_R1_Nside512.txt.xz`
      to be placed in `data/GMF/sources/` where GMF=[jf|pt]
      depending on the chosen GMF model.

An output file contains 8 columns:
    * arrival direction at Earth (lat_ini, lon_ini), degrees;
    * arrival direction at the boundary of the Galaxy (lat_res,
      lon_res); these will be similar to each other because of the
      small radius of the vicinity of the direction to the given source;
    * angular separation between (lat_res, lon_res) and the direction to
      the source;
    * (Z, E) pair;
    * number of the cell that contains (lat_res, lon_res) on the healpix
      grid.

Notation (lat_ini, lon_ini) and (lat_res, lon_res) originates from the
fact that we perform backtracking of (anti-)nuclei, so that their
initial position is at Earth. Perhaps, they should be renamed to, e.g.,
(lat_earth, lon_earth) and (lat_gal, lon_gal).

The script reads a file like `data/sample_D3.5Mpc..._sorted.txt`
obtained with `psample.py`, thus *Nini* and *Emin* must be exactly
the same as used in `psample.py` for the distance to the given
source (see *source_id*).

It then searches for a sample of Nini events in a given
`source_vicinity_radius` from files `nucleus_NNNEeV.txt.xz` obtained
with `galback.py`.  An output file is called `src_sample_ID_....txt`,
where ID is the *source_id*.  It must be placed in `data/sources/`.  The
sample can then be used to plot a figure and/or create a mix of samples
with an isotropic background.

*NB: an output sample does NOT necessarily contain Nini nuclei since
there may be less "suitable" (Z,E) pairs coming form the vicinity of the
source on a given healpix' grid! Use greater Nside if necessary or
increase the radius of the vicinity*

`convert_src_sample.py` can be used to convert data from
`data/GMF/sources/src_sample_...NsideXXX.txt.xz` file to data (and file)
calculated for another Nside.

An output file is called src_sample_ID....txt.xz and is placed
in data/GMF/sources/
"""

import sys
import numpy as np
from PyAstronomy.pyasl import getAngDist
import argparse

import lzma
import healpy as hp
import os

#______________________________________________________________________
#source_id = 'NGC253'
source_id = 'CenA'
#source_id = 'M82'
#source_id = 'M87'
#source_id = 'FornaxA'
#______________________________________________________________________

# Total number of EECRs to be found in the given vicinity of the source.
# This is the same number as used in psample.py
Nini = 10000

# Minimum energy used to create a sample, EeV:
Emin = '56'

# Healpix grid used in simulations:
#Nside = 512
Nside = 256

# Radius of the source neighborhood, deg
source_vicinity_radius = 1

shiftA = '1'

# Model of the Galactic Magnetic Field
#GMF = 'jf'
GMF = 'pt'

#______________________________________________________________________
# No settings below

source_data = {
    # Name : [source_lon, source_lat, D_src]
    'M82': [141.4095,40.5670,'3.5'],
    'CenA': [309.5159,19.4173,'3.5'],
    'NGC253': [97.3638,-87.9645,'3.5'],
    'NGC6946': [95.71873,11.6729,'6.0'],
    'M87': [283.7777,74.4912, '18.5'],
    'FornaxA': [240.1627,-56.6898,'20.0']
}

def get_source_data(source_id):
    if source_id in source_data:
        return tuple(source_data[source_id])
    else:
        raise ValueError('Unknown source!')

cline_parser = argparse.ArgumentParser(description='E,Z pair sampling',
                                       formatter_class=argparse.ArgumentDefaultsHelpFormatter)
def add_arg(*pargs, **kwargs):
    cline_parser.add_argument(*pargs, **kwargs)

add_arg('--source_id', type=str, help='one of ' + ' , '.join(source_data.keys()), default=source_id)
add_arg('--GMF', type=str, help='pt or jf', default=GMF)
add_arg('--Emin', type=str, help='Emin in EeV', default=Emin)
add_arg('--Nini', type=int, help='Number of cosmic ray nuclei to form a sample', default=Nini)
add_arg('--Nside', type=int, help='Nside for output healpix grid', default=Nside)
add_arg('--shiftA', type=str, help='A factor to shift and atomic mass by', default=shiftA)

args = cline_parser.parse_args()
source_id = args.source_id
GMF = args.GMF
Emin = args.Emin
Nini = args.Nini
Nside = args.Nside
shiftA = float(args.shiftA)

source_lon, source_lat, D_src = get_source_data(source_id)

#______________________________________________________________________
# A file in spectra_1s/

gmf_dir = GMF + '/'

# File with the spectrum that was "propagated" with CRPropa
infile = ('data/sample_D'+D_src+'_Emin'+Emin+'_'
        + str(Nini))
if shiftA != 1.0:
    infile += '_shift' + args.shiftA

infile += 'nuclei_sorted.txt'

outfile = ('src_sample_' + source_id + '_D' + D_src
        + '_Emin' + Emin
        + '_N' + str(Nini) + '_R'
        + str(source_vicinity_radius) + '_Nside' + str(Nside))
if shiftA != 1.0:
    outfile += '_shift' + args.shiftA

# A template for the output file
# lat, lon (ini); lat, lon (res); angsep; Z, E, # in the healpix map
outdata = np.zeros([Nini,8],dtype=float)

# Read the file and extract data
data = np.loadtxt(infile,dtype=int)

# Z, E/EeV, multiplicity of the nuclei
Z = data[:,0]
E = data[:,1]
M = data[:,2]

# Number of nucleus_energy.txt.xz files to be proceeded
N_files = len(Z)
#N_files = 2

# A list of nuclei to prepare the name of a file to read on the fly
nuclei = ['proton','helium','lithium','beryllium','boron','carbon',
        'nitrogen','oxygen','fluorine','neon','sodium','magnesium',
        'aluminium','silicon','phosphorus','sufur','chlorine','argon',
        'potassium','calcium','scandium','titanium','vanadium',
        'chromium','manganese','iron']

# A counter of EECRs added to the output file (result)
k = 0

# A counter for missing (Z,E) pairs
err_no = 0

# Let us check first that all necessary files exist
for i in np.arange(N_files):
    nucleus = nuclei[Z[i]-1]
    filename = nucleus + '_{:03d}EeV.txt.xz'.format(E[i])
    file2load = ('data/'+gmf_dir+str(Nside)+'/'+filename)

    print(nucleus + ' (Z = ' +str(Z[i])+ ') ' + str(E[i]) + ' EeV')

    # Check file
    if os.path.isfile(file2load):
        print('OK')
    else:
        # This is to be copy-pasted in run_galback.py or
        # a similar script
        print(('-----> MISSING: '+ nucleus
                + '\t(Z,E) = {:2d}\t{:3d}\t{:3d} '
                + 'nuclei needed')
                .format(Z[i],E[i],M[i]))
        err_no += M[i]

if err_no>0:
    print('Files for some (Z,E) are missing.')
    print('{:} nuclei will not be included in the output'
            .format(err_no))
    # Normally, we should exit at this point but let us allow
    # going forward in case the number of missing nuclei is much
    # less than the size of the sample.
    #sys.exit()

print('-------------------------------------\n')

# _____________________________________________________________________
# A cycle over all nuclei listed in infile
for i in np.arange(N_files):
    nucleus = nuclei[Z[i]-1]
    filename = nucleus + '_{:03d}EeV.txt.xz'.format(E[i])

    print(nucleus + ' (Z = ' +str(Z[i])+ ') ' + str(E[i]) + ' EeV')

    try:
        file2load = ('data/'+gmf_dir+str(Nside)+'/'+filename)
        with lzma.open(file2load,'rt') as f:
            summary = np.genfromtxt(f,dtype=float)

        # Coordinates observed at Earth
        #lat_earth_deg = summary[:,0]
        #lon_earth_deg = summary[:,1]
        # Coordinates at the boundary of the Galaxy
        lat_gal_deg = summary[:,2]
        lon_gal_deg = summary[:,3]
        # We need angular separtion between the source and
        # arrival directions at the Galaxy boundary
        ang_sep = getAngDist(source_lon, source_lat,
                lon_gal_deg, lat_gal_deg)

        # Find EECRs that fit in the given source vicinity at
        # the Galaxy boundary, if any
        close_dirs = np.asarray(
                np.where(ang_sep<=source_vicinity_radius+0.01))
        N_close = np.size(close_dirs)

        if N_close:
            print('{:4d} close EECRs'.format(N_close))
            print('To be selected: ' + str(M[i]) + '\n')

            close_dirs = np.reshape(close_dirs,N_close)

            # This is the 1st version: arrival directions are chosen
            # w/0 replacement. Due to this, we sometimes lack a number
            # of arrival directions/nuclei in the output sample
            #sample_size = np.min([M[i],N_close])
            #sample = np.random.choice(close_dirs,size=sample_size,
            #        replace=False)

            # Version 2. Allow replacement in case there are not enough
            # close arrival directions. This will give an output closer
            # to the given spectrum
            sample_size = M[i]
            if M[i]<=N_close:
                sample = np.random.choice(close_dirs,size=sample_size,
                        replace=False)
            else:
                sample = np.random.choice(close_dirs,size=sample_size,
                        replace=True)

            # OK, this is for the output file
            for j in sample:
                try:
                    outdata[k,0:4] = summary[j,0:4]
                    outdata[k,4] = ang_sep[j]
                    outdata[k,5] = Z[i]
                    outdata[k,6] = E[i]
                    outdata[k,7] = j
                    k += 1
                except IndexError as er:  # temporary fix or rare error TODO: find out the reason of
                    print('skipping sample', j, ':', er, file=sys.stderr)

            summary = []

        else:
            print('No close EECRs')
            print('To be selected: ' + str(M[i]) + '\n')

    # File is not available
    except IOError:
        print('\n-------> ' + filename + ' file not found!\n')
        #sys.exit()

#______________________________________________________________________
# Write down the output file and xz it

header = ('#   lat_ini    lon_ini    lat_res    lon_res     angsep   '
          'Z   E   cell_no\n')

with open('data/' + gmf_dir + 'sources/' + outfile + '.txt','w') as d:
    d.write(header)
    for i in np.arange(Nini):
        d.write('{:11.5f}{:11.5f}{:11.5f}{:11.5f}{:11.5f}{:4d}{:5d}{:9d}\n'.
                format(outdata[i,0],outdata[i,1],
                   outdata[i,2],outdata[i,3],
                   outdata[i,4],int(outdata[i,5]),
                   int(outdata[i,6]),int(outdata[i,7])))

os.system('xz data/' + gmf_dir + 'sources/' + outfile + '.txt')

#______________________________________________________________________


from init_hosts import load_all_hosts, write_hosts_to_script
import time, subprocess
import numpy as np

UPDATE_TIME_INTERVAL = 8
PROBABLITY_TO_KILL = 0.75
LOCAL_HOST_ID = 1
KILL_WORKER_ENABLE = 0

def load_curr_hosts():
    child = subprocess.Popen(['./discover_hosts.sh'], stdout=subprocess.PIPE)
    stdout, _ = child.communicate()
    _curr_hosts = str(stdout, encoding='utf-8').split("\n")[:-1]
    curr_hosts = {}
    # print(curr_hosts)
    for i in _curr_hosts:
        curr_hosts[int(i[-3])-1] = (i[:-2], int(i[-1]))

    print("current working hosts: ", curr_hosts)
    return curr_hosts

def mutate(curr_hosts, all_hosts, action):
    if action == -1:
        # remove workers
        num_to_remove = np.random.randint(1, len(curr_hosts))
        ids = np.random.choice(list(curr_hosts.keys()), num_to_remove, replace=False)
        print("remove %d hosts: " %num_to_remove, ids)
    if action == 1:
        num_to_add = np.random.randint(1, len(all_hosts)-len(curr_hosts)+1)
        ids = np.random.choice(list(set(all_hosts.keys()).difference(set(curr_hosts.keys()))), num_to_add, replace=False)
        print("add %d hosts: " %num_to_add, ids)

    return ids

def kill_worker(host):
    subprocess.run(['ssh', 'root@'+host, '"pkill -f python'])

# def update_hosts():
if __name__ == "__main__":
    all_hosts = load_all_hosts('./hosts.json')
    # print(all_hosts)
    curr_hosts = load_curr_hosts()
    # print(curr_hosts)
    # i = 10
    while(i>0):
    # while(1):
        time.sleep(UPDATE_TIME_INTERVAL)
        # i -= 1
        if KILL_WORKER_ENABLE:
            # only kill a worker if some probablity requirement is met
            p = np.random.uniform()
            if p >= PROBABLITY_TO_KILL:
                candidants = list(curr_hosts.keys())
                if LOCAL_HOST_ID in candidants:
                    candidants.remove(LOCAL_HOST_ID)
                worker_to_kill = np.random.choice(candidants)
                host = curr_hosts[worker_to_kill][0]
                kill_worker(host)

        curr_hosts = load_curr_hosts()
        # print(curr_hosts)
        if (len(curr_hosts) == len(all_hosts)):
            # in this case, mutation must remove workers
            ids = mutate(curr_hosts, all_hosts, action=-1)
            write_hosts_to_script(ids, all_hosts, mode="update", action=-1)
            # print(load_curr_hosts())
            continue
        elif len(curr_hosts) == 1:
            # in this case, mutation must add workers
            ids = mutate(curr_hosts, all_hosts, action=1)
            write_hosts_to_script(ids, all_hosts, mode="update", action=1)
            # print(load_curr_hosts())
            continue
        else:
            # now mutation can either add or remove workers
            # action=1 implies addition, action=-1 implies removal
            action = np.random.choice([1, -1], 1)
            ids = mutate(curr_hosts, all_hosts, action)
            write_hosts_to_script(ids, all_hosts, mode="update", action=action)
            # print(load_curr_hosts())
            continue


from pathlib import Path
import logging
import tensorrt as trt

from numba.cgutils import printf

LOGGER = logging.getLogger(__name__)
# LOGGER = trt.Logger(trt.Logger.VERBOSE)

class SSD:
    PLUGIN_PATH = None
    ENGINE_PATH = None
    MODEL_PATH = None
    INPUT_SHAPE = ()
    OUTPUT_NAME = None

    @classmethod
    def add_plugin(cls, ssd_graph):
        raise NotImplementedError

    @classmethod
    def build_engine(cls, trt_logger, batch_size, calib_dataset=Path.home() / 'VOCdevkit' / 'VOC2007' / 'JPEGImages'):
        import graphsurgeon as gs
        import uff
        from . import calibrator

        # convert TensorFlow ssd_graph into UFF
        dynamic_graph = gs.DynamicGraph(str(cls.MODEL_PATH))
        dynamic_graph = cls.add_plugin(dynamic_graph)
        uff_model = uff.from_tensorflow(
            dynamic_graph.as_graph_def(),
            [cls.OUTPUT_NAME],
            quiet=True
        )

        with trt.Builder(trt_logger) as builder, builder.create_network() as network, trt.UffParser() as parser:
            builder.max_workspace_size = 1 << 30
            builder.max_batch_size = batch_size
            LOGGER.info('Building engine with batch size: %d', batch_size)
            LOGGER.info('This may take a while...')

            if builder.platform_has_fast_fp16:
                builder.fp16_mode = True
            if builder.platform_has_fast_int8:
                builder.int8_mode = True
                builder.int8_calibrator = calibrator.SSDEntropyCalibrator(cls.INPUT_SHAPE,
                                                                          data_dir=calib_dataset,
                                                                          cache_file=Path(__file__).parent /
                                                                          f'{cls.__name__}_calib_cache')
            LOGGER.info('This may take a while.1.')
            parser.register_input('Input', cls.INPUT_SHAPE)
            parser.register_output('MarkOutput_0')
            parser.parse_buffer(uff_model, network)
            engine = builder.build_cuda_engine(network)
            LOGGER.info('This may take a while.3')

            if engine is None:
                LOGGER.critical('Failed to build engine')
                return None

            LOGGER.info("Completed creating engine")
            with open(cls.ENGINE_PATH, 'wb') as engine_file:
                engine_file.write(engine.serialize())
            return engine

class SSDMobileNetV1(SSD):
    ENGINE_PATH = Path(__file__).parent / 'ssd_mobilenet_v1_coco.trt'
    MODEL_PATH = Path(__file__).parent / 'ssd_mobilenet_v1_coco.pb'
    INPUT_SHAPE = (3, 300, 300)
    OUTPUT_NAME = 'NMS'
    NMS_THRESH = 0.5
    TOPK = 100

    @classmethod
    def add_plugin(cls, ssd_graph):
        import tensorflow as tf
        import graphsurgeon as gs

        all_assert_nodes = ssd_graph.find_nodes_by_op("Assert")
        ssd_graph.remove(all_assert_nodes, remove_exclusive_dependencies=True)
        all_identity_nodes = ssd_graph.find_nodes_by_op("Identity")
        ssd_graph.forward_inputs(all_identity_nodes)

        Input = gs.create_plugin_node(
            name="Input",
            op="Placeholder",
            dtype=tf.float32,
            shape=[1, *cls.INPUT_SHAPE]
        )

        PriorBox = gs.create_plugin_node(
            name="MultipleGridAnchorGenerator",
            op="GridAnchor_TRT",
            minSize=0.2,
            maxSize=0.95,
            aspectRatios=[1.0, 2.0, 0.5, 3.0, 0.33],
            variance=[0.1, 0.1, 0.2, 0.2],
            featureMapShapes=[19, 10, 5, 3, 2, 1],
            numLayers=6
        )

        NMS = gs.create_plugin_node(
            name="NMS",
            op="NMS_TRT",
            shareLocation=1,
            varianceEncodedInTarget=0,
            backgroundLabelId=0,
            confidenceThreshold=1e-8,
            nmsThreshold=cls.NMS_THRESH,
            topK=100,
            keepTopK=100,
            numClasses=91,
            inputOrder=[0, 2, 1],
            confSigmoid=1,
            isNormalized=1
        )

        concat_priorbox = gs.create_node(
            "concat_priorbox",
            op="ConcatV2",
            dtype=tf.float32,
            axis=2
        )

        concat_box_loc = gs.create_plugin_node(
            "concat_box_loc",
            op="FlattenConcat_TRT",
            dtype=tf.float32,
            axis=1,
            ignoreBatch=0
        )

        concat_box_conf = gs.create_plugin_node(
            "concat_box_conf",
            op="FlattenConcat_TRT",
            dtype=tf.float32,
            axis=1,
            ignoreBatch=0
        )
        # Create a mapping of namespace names -> plugin nodes.
        namespace_plugin_map = {
            "MultipleGridAnchorGenerator": PriorBox,
            "Postprocessor": NMS,
            "Preprocessor": Input,
            "ToFloat": Input,
            "image_tensor": Input,
            "MultipleGridAnchorGenerator/Concatenate": concat_priorbox,
            "concat": concat_box_loc,
            "concat_1": concat_box_conf
        }

        # Create a new ssd_graph by collapsing namespaces
        ssd_graph.collapse_namespaces(namespace_plugin_map)
        # Remove the outputs, so we just have a single output node (NMS).
        # If remove_exclusive_dependencies is True, the whole ssd_graph will be removed!
        ssd_graph.remove(ssd_graph.graph_outputs, remove_exclusive_dependencies=False)
        ssd_graph.find_nodes_by_op("NMS_TRT")[0].input.remove("Input")
        ssd_graph.find_nodes_by_name("Input")[0].input.remove("image_tensor:0")

        return ssd_graph

class SSDMobileNetV2(SSD):
    ENGINE_PATH = Path(__file__).parent / 'ssd_mobilenet_v2_coco.trt'
    MODEL_PATH = Path(__file__).parent / 'ssd_mobilenet_v2_coco.pb'
    # ENGINE_PATH = Path(__file__).parent / 'frozen_inference_graph.trt'
    # MODEL_PATH = Path(__file__).parent / 'frozen_inference_graph.pb'

    INPUT_SHAPE = (3, 300, 300)
    OUTPUT_NAME = 'NMS'
    NMS_THRESH = 0.5
    TOPK = 100

    @classmethod
    def add_plugin(cls, ssd_graph):
        import tensorflow as tf
        import graphsurgeon as gs

        all_assert_nodes = ssd_graph.find_nodes_by_op("Assert")
        ssd_graph.remove(all_assert_nodes, remove_exclusive_dependencies=True)
        all_identity_nodes = ssd_graph.find_nodes_by_op("Identity")
        ssd_graph.forward_inputs(all_identity_nodes)

        Input = gs.create_plugin_node(
            name="Input",
            op="Placeholder",
            dtype=tf.float32,
            shape=[1, *cls.INPUT_SHAPE]
        )

        PriorBox = gs.create_plugin_node(
            name="GridAnchor",
            op="GridAnchor_TRT",
            minSize=0.2,
            maxSize=0.95,
            aspectRatios=[1.0, 2.0, 0.5, 3.0, 0.33],
            variance=[0.1, 0.1, 0.2, 0.2],
            featureMapShapes=[19, 10, 5, 3, 2, 1],
            numLayers=6
        )

        NMS = gs.create_plugin_node(
            name="NMS",
            op="NMS_TRT",
            shareLocation=1,
            varianceEncodedInTarget=0,
            backgroundLabelId=0,
            confidenceThreshold=1e-8,
            nmsThreshold=cls.NMS_THRESH,
            topK=100,
            keepTopK=100,
            numClasses=91,
            inputOrder=[1, 0, 2],
            confSigmoid=1,
            isNormalized=1
        )

        concat_priorbox = gs.create_node(
            "concat_priorbox",
            op="ConcatV2",
            dtype=tf.float32,
            axis=2
        )

        concat_box_loc = gs.create_plugin_node(
            "concat_box_loc",
            op="FlattenConcat_TRT",
            dtype=tf.float32,
            axis=1,
            ignoreBatch=0
        )

        concat_box_conf = gs.create_plugin_node(
            "concat_box_conf",
            op="FlattenConcat_TRT",
            dtype=tf.float32,
            axis=1,
            ignoreBatch=0
        )

        namespace_plugin_map = {
            "MultipleGridAnchorGenerator": PriorBox,
            "Postprocessor": NMS,
            "Preprocessor": Input,
            "ToFloat": Input,
            "image_tensor": Input,
            "Concatenate": concat_priorbox,
            "concat": concat_box_loc,
            "concat_1": concat_box_conf
        }

        # Create a new ssd_graph by collapsing namespaces
        ssd_graph.collapse_namespaces(namespace_plugin_map)
        # Remove the outputs, so we just have a single output node (NMS).
        # If remove_exclusive_dependencies is True, the whole ssd_graph will be removed!
        ssd_graph.remove(ssd_graph.graph_outputs, remove_exclusive_dependencies=False)
        ssd_graph.find_nodes_by_op("NMS_TRT")[0].input.remove("Input")

        return ssd_graph

class SSDInceptionV2(SSD):
    ENGINE_PATH = Path(__file__).parent / 'ssd_inception_v2_coco.trt'
    MODEL_PATH = Path(__file__).parent / 'ssd_inception_v2_coco.pb'
    INPUT_SHAPE = (3, 300, 300)
    OUTPUT_NAME = 'NMS'
    NMS_THRESH = 0.5
    TOPK = 100

    @classmethod
    def add_plugin(cls, ssd_graph):
        import tensorflow as tf
        import graphsurgeon as gs

        all_assert_nodes = ssd_graph.find_nodes_by_op("Assert")
        ssd_graph.remove(all_assert_nodes, remove_exclusive_dependencies=True)
        all_identity_nodes = ssd_graph.find_nodes_by_op("Identity")
        ssd_graph.forward_inputs(all_identity_nodes)

        # Create TRT plugin nodes to replace unsupported ops in Tensorflow ssd_graph
        Input = gs.create_plugin_node(
            name="Input",
            op="Placeholder",
            dtype=tf.float32,
            shape=[1, *cls.INPUT_SHAPE]
        )

        PriorBox = gs.create_plugin_node(
            name="GridAnchor",
            op="GridAnchor_TRT",
            minSize=0.2,
            maxSize=0.95,
            aspectRatios=[1.0, 2.0, 0.5, 3.0, 0.33],
            variance=[0.1, 0.1, 0.2, 0.2],
            featureMapShapes=[19, 10, 5, 3, 2, 1],
            numLayers=6
        )

        NMS = gs.create_plugin_node(
            name="NMS",
            op="NMS_TRT",
            shareLocation=1,
            varianceEncodedInTarget=0,
            backgroundLabelId=0,
            confidenceThreshold=1e-8,
            nmsThreshold=cls.NMS_THRESH,
            topK=100,
            keepTopK=100,
            numClasses=91,
            inputOrder=[0, 2, 1],
            confSigmoid=1,
            isNormalized=1
        )

        concat_priorbox = gs.create_node(
            "concat_priorbox",
            op="ConcatV2",
            dtype=tf.float32,
            axis=2
        )

        concat_box_loc = gs.create_plugin_node(
            "concat_box_loc",
            op="FlattenConcat_TRT",
            dtype=tf.float32,
            axis=1,
            ignoreBatch=0
        )

        concat_box_conf = gs.create_plugin_node(
            "concat_box_conf",
            op="FlattenConcat_TRT",
            dtype=tf.float32,
            axis=1,
            ignoreBatch=0
        )

        # Create a mapping of namespace names -> plugin nodes.
        namespace_plugin_map = {
            "MultipleGridAnchorGenerator": PriorBox,
            "Postprocessor": NMS,
            "Preprocessor": Input,
            "ToFloat": Input,
            "image_tensor": Input,
            "MultipleGridAnchorGenerator/Concatenate": concat_priorbox,
            "concat": concat_box_loc,
            "concat_1": concat_box_conf
        }

        # Create a new ssd_graph by collapsing namespaces
        ssd_graph.collapse_namespaces(namespace_plugin_map)
        # Remove the outputs, so we just have a single output node (NMS).
        # If remove_exclusive_dependencies is True, the whole ssd_graph will be removed!
        ssd_graph.remove(ssd_graph.graph_outputs, remove_exclusive_dependencies=False)
        return ssd_graph


import torchvision.models
import torch.nn as nn
from torchvision.transforms import ToTensor, ToPILImage
import torch
from PIL import Image
import os
import torch.optim as optim
import random
import torchvision.utils as vutils
import torch.nn.functional as F
from pidmodel import *
import glob
from PIL import ImageEnhance
import time
import numpy as np

def define_path():
    gtpath = '/home/lir0b/Code/perceptual-reflection-removal/synthetic'
    simupath = '/home/lir0b/Code/TransparenceDetection/src/pid/exp/synthetic_multi'
    ckptpath = '/home/lir0b/Code/TransparenceDetection/src/pid/ckpt'
    if not os.path.exists(gtpath) and not os.path.exists(simupath) and not os.path.exists(ckptpath):
        gtpath = '/ibex/scratch/lir0b/data/synthetic'
        simupath = '/ibex/scratch/lir0b/data/synthetic_multi'
        ckptpath = '/ibex/scratch/lir0b/TransparenceDetection/src/pid/ckpt'
    if not os.path.exists(gtpath) and not os.path.exists(simupath) and not os.path.exists(ckptpath):
        gtpath = '/ibex/scratch/lir0b/data/synthetic'
        simupath = '/ibex/scratch/lir0b/data/synthetic_multi'
        ckptpath = '/ibex/scratch/lir0b/TransparenceDetection/src/pid/ckpt'
    if not os.path.exists(gtpath) and not os.path.exists(simupath) and not os.path.exists(ckptpath):
        gtpath = '/ibex/scratch/qius0a/data/synthetic'
        simupath = '/ibex/scratch/qius0a/data/synthetic_multi'
        ckptpath = '/ibex/scratch/qius0a/TransparenceDetection/src/pid/ckpt'
    if not os.path.exists(gtpath) and not os.path.exists(simupath) and not os.path.exists(ckptpath):
        gtpath = '/ibex/scratch/qius0a/data/synthetic'
        simupath = '/ibex/scratch/qius0a/data/synthetic_multi'
        ckptpath = '/ibex/scratch/qius0a/TransparenceDetection/src/pid/ckpt'
    return gtpath, simupath, ckptpath

def train_sepmodel_singlenet():
    use_CUDA = True
    from dataloader import DataLoader
    from pprint import pprint
    gtpath, simupath, ckptpath = define_path()

    polar_name = 'SepNet'
    psf_name = 'PSFNet'
    discriminator_name = 'netD'
    refD_name = 'refD'
    transD_name = 'transD'

    eps = 1e-5
    dl = DataLoader(simupath, gtpath)
    Epochs = 100
    t = 0
    numsample = 100
    rsz = (320, 320)
    save_epoch = True
    continue_train = False
    save_image = True
    device = torch.device("cuda:0" if use_CUDA else "cpu")
    sepnet = Sepmodelv2().to(device)
    optimizerG = optim.Adam(sepnet.parameters(), lr=0.0002, betas=(0.5, 0.999))
    refD = RefDiscriminator()
    transD = TransDiscriminator()
    refD.to(device)
    transD.to(device)
    refD.apply(weights_init)
    transD.apply(weights_init)

    optimizer_refD = optim.Adam(refD.parameters(), lr=0.0002, betas=(0.5, 0.999))
    optimizer_transD = optim.Adam(transD.parameters(), lr=0.0002, betas=(0.5, 0.999))

    weights = [1.0, 0, 0, 3, 1]
    #weights = [1.33, 0, 0, 0, 0]
    brightness = 1
    if continue_train:
        savefoldername = '1.00_0.00_0.00_1581451445'
        ckptpath = os.path.join(ckptpath, savefoldername)
    else:
        ckptpath = os.path.join(ckptpath, '%.2f_%.2f_%.2f_%d' % (weights[0], weights[1], weights[2], int(time.time())))
    if not os.path.exists(ckptpath):
        os.makedirs(ckptpath)
    for i in range(Epochs):
        savepathi = os.path.join(ckptpath, str(i), polar_name)
        if os.path.exists(savepathi) and continue_train:
            savepathi_1 = os.path.join(ckptpath, str(i + 1), polar_name)
            if not os.path.exists(savepathi_1):
                sepnet.load_state_dict(torch.load(os.path.join(ckptpath, str(i), polar_name)))
                #psfnet.load_state_dict(torch.load(os.path.join(ckptpath, str(i), psfname)))
                refD.load_state_dict(torch.load(os.path.join(ckptpath, str(i), refD_name)))
                transD.load_state_dict(torch.load(os.path.join(ckptpath, str(i), transD_name)))
            continue
        for idx in range(numsample):
            print(idx)
            #idx = random.randint(0, numsample) % len(dl.simpair)
            #idx = random.randint(0, numsample) % len(dl.simpair)
            ref_id, trans_id = dl.simpair[idx]
            theta_pol_filelist = dl.get_theta_pol_filelist(ref_id, trans_id)
            ref_trans_pair = dl.get_ref_trans_pair(ref_id, trans_id)

            Ir = preprocess_image(ref_trans_pair[0], rsz, brightness)
            if Ir is None:
                continue
            Ir_cuda = to_CUDA(Ir)

            It = preprocess_image(ref_trans_pair[1], rsz, brightness)
            if It is None:
                continue
            It_cuda = to_CUDA(It)

            target = torch.cat((Ir_cuda, It_cuda), dim=1).cuda()
            theta_idx = 0
            t = 0

            for idxt, thetapath in enumerate(theta_pol_filelist):
                for idxp, polpath in enumerate(thetapath):
                    img = preprocess_image(polpath, rsz, brightness)
                    if img is None:
                        continue
                    img_cuda = to_CUDA(img)

                    if False:
                        sepnet.zero_grad()
                        [Ir_pred, It_pred, I_obs] = sepnet([img_cuda])
                        allloss = sepnet.loss([Ir_cuda, It_cuda], [Ir_pred, It_pred], img_cuda, I_obs, weights)
                        [loss_sep, loss_vgg_sep, loss_exclu, loss_backward] = allloss
                        total_loss = loss_sep + loss_vgg_sep + loss_exclu + loss_backward

                    if True:
                        sepnet.zero_grad()
                        [Ir_pred, It_pred] = sepnet(img_cuda)
                        allloss = sepnet.loss([Ir_pred, It_pred], [Ir_cuda, It_cuda])
                        [loss_ref, loss_trans, total_loss] = allloss

                        loss_ref_adv = refD.gloss(refD(img_cuda, Ir_pred))
                        loss_trans_adv = transD.gloss(transD(img_cuda, It_pred))
                        total_loss = total_loss+loss_ref_adv*1000000+loss_trans_adv*1000000

                        #It_pred = It_res_pred + img_cuda
                        #[loss_sep, loss_vgg_sep, loss_exclu, loss_backward] = allloss

                    total_loss.backward(retain_graph=True)
                    optimizerG.step()
                    G_x = total_loss.item()

                    # calc discriminator loss
                    refD.zero_grad()
                    transD.zero_grad()
                    isref_pred = refD(img_cuda, Ir_pred)
                    istrans_pred = transD(img_cuda, It_pred)
                    isref_gt = refD(img_cuda, Ir_cuda)
                    istrans_gt = transD(img_cuda, It_cuda)

                    refDloss = refD.dloss(isref_pred, isref_gt)
                    refDloss.backward(retain_graph=True)
                    optimizer_refD.step()

                    transDloss = transD.dloss(istrans_pred, istrans_gt)
                    transDloss.backward(retain_graph=True)
                    optimizer_transD.step()

                    #isreal = refD(target)
                    #IrIt = torch.cat([Ir_pred, It_pred], dim=1)
                    #ispred = netD(IrIt)
                    #Dloss = netD.loss(isreal, ispred)
                    #Dloss.backward()
                    #D_x = Dloss.mean().item()
                    #optimizerD.step()

                    #netD.zero_grad()
                    #isreal = netD(target)
                    #IrIt = torch.cat([Ir_pred, It_pred], dim=1)
                    #ispred = netD(IrIt)
                    #Dloss = netD.loss(isreal, ispred)
                    #Dloss.backward()
                    #D_x = Dloss.mean().item()
                    #optimizerD.step()

                    if idx % 10 == 0:
                        ps = '%d_%d_%d_%d tot:%.6f, refl:%.6f, transl:%.6f, refD:%.6f, transD:%.6f, refG:%.6f transG:%.6f\n'
                        print(ps % (i, idx, idxt, idxp,
                                    total_loss.item(),
                                    loss_ref.item(),
                                    loss_trans.item(),
                                    refDloss.item(),
                                    transDloss.item(),
                                    loss_ref_adv.item(),
                                    loss_trans_adv.item()))

                    if idx % 10 == 0 and False:
                        print('%d_%d_%d A:%.2f Lsep:%.6f Lvgg:%.6f Lexc:%.6f Lb:%.6f \n' % (i, idxt, t, total_loss.item(),
                                                                                loss_sep.item(),
                                                                                loss_vgg_sep.item(),
                                                                                loss_exclu.item(),
                                                                                loss_backward.item()))
                        print('%d L:%.2f D:%.2f' % (i, G_x, D_x))

                    if save_image and idx % 10 == 0:
                        subfoldername = str(ref_id)+'_'+str(trans_id)
                        if not os.path.exists(os.path.join(ckptpath, str(i))):
                            os.makedirs(os.path.join(ckptpath, str(i)))
                        if not os.path.exists(os.path.join(ckptpath, str(i), subfoldername)):
                            os.makedirs(os.path.join(ckptpath, str(i), subfoldername))

                        Irpath = os.path.join(ckptpath, str(i), subfoldername, str(ref_id) + '_r.jpg')
                        if not os.path.exists(Irpath):
                            vutils.save_image(Ir_cuda.clone(), Irpath, normalize=False)
                        Itpath = os.path.join(ckptpath, str(i), subfoldername, str(trans_id) + '_t.jpg')
                        if not os.path.exists(Itpath):
                            vutils.save_image(It_cuda.clone(), Itpath, normalize=False)

                        Itpath = os.path.join(ckptpath, str(i), subfoldername, str(trans_id) + '_' + str(idxt) + '_res_pred_t.jpg')
                        if not os.path.exists(Itpath):
                            pass
                            #vutils.save_image(It_res_pred.clone(), Itpath, normalize=False)

                        Irpath_pred = os.path.join(ckptpath, str(i), subfoldername, str(ref_id) + '_' + str(idxt) + '_r_pred.jpg')
                        if not os.path.exists(Irpath_pred):
                            vutils.save_image(Ir_pred.clone(), Irpath_pred, normalize=False)
                        Itpath_pred = os.path.join(ckptpath, str(i), subfoldername, str(trans_id) + '_' + str(idxt) + '_t_pred.jpg')
                        if not os.path.exists(Itpath_pred):
                            vutils.save_image(It_pred.clone(), Itpath_pred, normalize=False)

                        #ori_path = os.path.join(ckptpath, str(i), str(ref_id)+'_'+str(trans_id) + '_' + str(idxt) + '_ori.jpg')
                        #vutils.save_image(img_cuda, ori_path, normalize=True)
                        I_obs_path = os.path.join(ckptpath, str(i), subfoldername, str(ref_id)+'_'+str(trans_id) + '_' + str(idxt) + '_obs.jpg')
                        vutils.save_image(img_cuda.clone(), I_obs_path, normalize=False)

        if save_epoch:
            savepathi = os.path.join(ckptpath, str(i))
            if not os.path.exists(savepathi):
                os.makedirs(savepathi)
            torch.save(sepnet.state_dict(), os.path.join(savepathi, polar_name))
            torch.save(refD.state_dict(), os.path.join(savepathi, refD_name))
            torch.save(transD.state_dict(), os.path.join(savepathi, transD_name))


import tensorflow as tf
import numpy as np

FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_integer('image_size', 224,
                            """Provide square images of this size.""")
tf.app.flags.DEFINE_integer('num_preprocess_threads', 12,
                            """Number of preprocessing threads per tower. """
                            """Please make this a multiple of 4.""")
tf.app.flags.DEFINE_integer('num_readers', 12,
                            """Number of parallel readers during train.""")
tf.app.flags.DEFINE_integer('input_queue_memory_factor', 12,
                            """Size of the queue of preprocessed images. """
                            """Default is ideal but try smaller values, e.g. """
                            """4, 2 or 1, if host memory is constrained. See """
                            """comments in code for more details.""")
def parse_example_proto(examples_serialized):
    feature_map={
        'image/encoded': tf.FixedLenFeature([], dtype=tf.string,
                                            default_value=''),
        'bboxes': tf.VarLenFeature(dtype=tf.float32)
        }
    features=tf.parse_single_example(examples_serialized, feature_map)
    bboxes = tf.sparse_tensor_to_dense(features['bboxes'])
    r = 5*tf.random_uniform((1,), minval=0, maxval=tf.size(bboxes, out_type=tf.int32)/5, dtype=tf.int32) # r = 8*tf.random_uniform((1,), minval=0, maxval=tf.size(bboxes, out_type=tf.int32)/8, dtype=tf.int32)
    bbox = tf.gather_nd(bboxes, [r,r+1,r+2,r+3,r+4]) # bbox = tf.gather_nd(bboxes, [r,r+1,r+2,r+3,r+4,r+5,r+6,r+7])

    return features['image/encoded'], bbox

def eval_image(image, height, width):
    image = tf.image.central_crop(image, central_fraction=0.875)
    image = tf.expand_dims(image, 0)
    image = tf.image.resize_bilinear(image, [height, width],
                                     align_corners=False)
    image = tf.squeeze(image, [0])

    return image

def distort_color(image, thread_id):
    color_ordering = thread_id % 2
    if color_ordering == 0:
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
        image = tf.image.random_hue(image, max_delta=0.2)
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
    elif color_ordering == 1:
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
        image = tf.image.random_hue(image, max_delta=0.2)
    image = tf.clip_by_value(image, 0.0, 1.0)

    return image

def distort_image(image, height, width, thread_id):
    distorted_image = tf.image.random_flip_left_right(image)
    distorted_image = distort_color(distorted_image, thread_id)
    return distorted_image

def image_preprocessing(image_buffer, train, thread_id=0):
    height = FLAGS.image_size
    width = FLAGS.image_size
    image = tf.image.decode_png(image_buffer, channels=3)
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    image = tf.image.resize_images(image, [height,width])
    if train:
        image = distort_image(image, height, width, thread_id)
    #else:
    #    image = eval_image(image, height, width)
    image = tf.subtract(image, 0.5)
    image = tf.multiply(image, 2.0)
    return image

def batch_inputs(data_files, train, num_epochs, batch_size,
                 num_preprocess_threads, num_readers):
    print(train)
    if train:
        filename_queue = tf.train.string_input_producer(data_files,
                                                        num_epochs,
                                                        shuffle=True,
                                                        capacity=16)
    else:
        filename_queue = tf.train.string_input_producer(data_files,
                                                        num_epochs,
                                                        shuffle=False,
                                                        capacity=1)

    examples_per_shard = 1024
    min_queue_examples = examples_per_shard * FLAGS.input_queue_memory_factor
    if train:
        print('pass')
        examples_queue = tf.RandomShuffleQueue(
            capacity=min_queue_examples+3*batch_size,
            min_after_dequeue=min_queue_examples,
            dtypes=[tf.string])
    else:
        examples_queue = tf.FIFOQueue(
            capacity=examples_per_shard + 3 * batch_size,
            dtypes=[tf.string])

    if num_readers > 1:
        enqueue_ops = []
        for _ in range(num_readers):
            reader = tf.TFRecordReader()
            _, value = reader.read(filename_queue)
            enqueue_ops.append(examples_queue.enqueue([value]))
        tf.train.queue_runner.add_queue_runner(
            tf.train.queue_runner.QueueRunner(examples_queue,enqueue_ops))
        examples_serialized = examples_queue.dequeue()
    else:
        reader = tf.TFRecordReader()
        _, examples_serialized = reader.read(filename_queue)

    images_and_bboxes=[]
    for thread_id in range(num_preprocess_threads):
        image_buffer, bbox = parse_example_proto(examples_serialized)
        image = image_preprocessing(image_buffer, train, thread_id)
        images_and_bboxes.append([image, bbox])

    images, bboxes = tf.train.batch_join(
        images_and_bboxes,
        batch_size=batch_size,
        capacity=2*num_preprocess_threads*batch_size)

    height = FLAGS.image_size
    width = FLAGS.image_size
    depth = 3

    images = tf.cast(images, tf.float32)
    images = tf.reshape(images, shape=[batch_size, height, width, depth])

    return images, bboxes

def distorted_inputs(data_files, num_epochs, train=True, batch_size=None):
    with tf.device('/cpu:0'):
        print(train)
        images, bboxes = batch_inputs(
            data_files, train, num_epochs, batch_size,
            num_preprocess_threads=FLAGS.num_preprocess_threads,
            num_readers=FLAGS.num_readers)

    return images, bboxes

def inputs(data_files, num_epochs=1, train=False, batch_size=1):
    with tf.device('/cpu:0'):
        print(train)
        images, bboxes = batch_inputs(
            data_files, train, num_epochs, batch_size,
            num_preprocess_threads=FLAGS.num_preprocess_threads,
            num_readers=1)

    return images, bboxes


# run in pydro38_test env
import os, sys, glob, logging, configparser
from datetime import datetime
import numpy as np
from scipy.interpolate import griddata
import pyproj
from pyproj import Transformer
from pyproj.crs import CompoundCRS
import gdal
from gdal import ogr
import rasterio
gdal.UseExceptions()

LOGGER = logging.getLogger('vyperdatum')

def load_config(config_filename = 'vyperdatum.config'):
    """
    Load the provided configuration file.
    """
    if not os.path.isfile(config_filename):
        raise ValueError(f'file not found: {config_filename}')
    config = {}
    config_file = configparser.ConfigParser()
    config_file.read(config_filename)
    sections = config_file.sections()
    for section in sections:
        config_file_section = config_file[section]
        for key in config_file_section:
            config[key] = config_file_section[key]

    if 'loggername' in config:
        global LOGGER
        LOGGER = logging.getLogger(config['loggername'])
    if 'inpath' in config:
        global inpath
        tmp = config['inpath']
        inpath = fr'{tmp}'
    if 'outpath' in config:
        global outpath
        tmp = config['outpath']
        outpath = fr'{tmp}'
    if 'vdatum_directory' in config:
        global VDATUM_DIRECTORY
        tmp = config['vdatum_directory']
        VDATUM_DIRECTORY = fr'{tmp}'
    else:
        raise ValueError('No VDatum path defined.')

def check_gdal_version():
    """
    Check the version of gdal imported to ensure it meets the requirements
    of this script

    Returns
    -------
    None.

    """
    version = gdal.VersionInfo()
    major = int(version[0])
    minor = int(version[1:3])
    bug = int(version[3:5])
    if major == 3 and minor >= 1:
        pass
    else:
        msg = f'The version of GDAL must be >= 3.1.  Version found: {version}'
        LOGGER.error(msg)
        raise ValueError(msg)

def update_vdatum_data_directory():
    global VDATUM_DIRECTORY
    orig_proj_paths = pyproj.datadir.get_data_dir()
    if VDATUM_DIRECTORY not in orig_proj_paths:
        pyproj.datadir.append_data_dir(VDATUM_DIRECTORY)

def get_gtx_grid_list():
    """
    Built a dicionary of all available VDatum grids.
    """
    global VDATUM_DIRECTORY
    search_path = os.path.join(VDATUM_DIRECTORY, '*/*.gtx')
    gtx_list = glob.glob(search_path)
    if len(gtx_list) == 0:
        raise ValueError(f'No GTX files found in the provided VDatum directory: {VDATUM_DIRECTORY}')
    grids = {}
    for gtx in gtx_list:
        gtx_path, gtx_file = os.path.split(gtx)
        gtx_path, gtx_folder = os.path.split(gtx_path)
        gtx_name = '/'.join([gtx_folder, gtx_file])
        gtx_subpath = os.path.join(gtx_folder, gtx_file)
        grids[gtx_name] = gtx_subpath
    return grids

def get_vdatum_region_polygons():
    """
    Get a list of all the kml files that are in the VDatum directory.

    Returns
    -------
    List

    """
    global VDATUM_DIRECTORY
    search_path = os.path.join(VDATUM_DIRECTORY, '*/*.kml')
    kml_list = glob.glob(search_path)
    if len(kml_list) == 0:
        raise ValueError(f'No kml files found in the provided VDatum directory: {VDATUM_DIRECTORY}')
    geom = {}
    for kml in kml_list:
        kml_path, kml_file = os.path.split(kml)
        root_dir, kml_name = os.path.split(kml_path)
        geom[kml_name] = kml
    return geom

def get_interesecting_vdatum_regions(datafilepath):
    """
    Find the vdatum regions that intersect the given data.
    """
    dataset = gdal.Open(datafilepath)
    is_raster = dataset.RasterCount > 0
    if is_raster:
        # get raster bounds
        crs = pyproj.CRS.from_wkt(dataset.GetProjectionRef())
        transform = dataset.GetGeoTransform()
        pixelWidth = transform[1]
        pixelHeight = transform[5]
        cols = dataset.RasterXSize
        rows = dataset.RasterYSize
        x0 = transform[0]
        y1 = transform[3]
        x1 = x0+cols*pixelWidth
        y0 = y1-rows*pixelHeight
        if crs.is_projected:
            unproject = pyproj.Proj(proj='utm', zone = 19, ellps = 'WGS84')
            ul = unproject(x0,x1, inverse = True)
            ur = unproject(x1, y1, inverse = True)
            lr = unproject(x1, y0, inverse = True)
            ll = unproject(x0, y0, inverse = True)
        else:
            ul = (x0, y1)
            ur = (x1, y1)
            lr = (x1, y0)
            ll = (x0, y0)
        # build polygon from raster
        ring = ogr.Geometry(ogr.wkbLinearRing)
        ring.AddPoint(ul[0], ul[1])
        ring.AddPoint(ur[0], ur[1])
        ring.AddPoint(lr[0], lr[1])
        ring.AddPoint(ll[0], ll[1])
        ring.AddPoint(ul[0], ul[1])
        dataGeometry = ogr.Geometry(ogr.wkbPolygon)
        dataGeometry.AddGeometry(ring)
    else:
        raise NotImplementedError('Not handling XYZ data yet')
    dataset = None
    # get all the regions represented by geometry files
    geom_list = get_vdatum_region_polygons()
    # see if the regions intersect with the provided geometries
    intersecting_regions = []
    for region in geom_list:
        vector = ogr.Open(geom_list[region])
        layer_count = vector.GetLayerCount()
        for m in range(layer_count):
            layer = vector.GetLayerByIndex(m)
            feature_count = layer.GetFeatureCount()
            for n in range(feature_count):
                feature = layer.GetFeature(n)
                feature_name = feature.GetField(0)
                if feature_name[:15] == 'valid-transform':
                    valid_vdatum_poly = feature.GetGeometryRef()
                    if dataGeometry.Intersect(valid_vdatum_poly):
                        intersecting_regions.append(region)
        vector = None
    return intersecting_regions

def run_pipeline(xx, yy, zz, incrs, region_name):
    """
    MLLW to NAVD88
    """
    req_hcrs_epsg = 26919
    req_vcrs_epsg = 'mllw'
    out_vcrs_epsg = 5703
    # parse the provided CRS
    cmpd_incrs = CompoundCRS.from_wkt(incrs.to_wkt())
    if len(cmpd_incrs.sub_crs_list) == 2:
        inhcrs, invcrs = cmpd_incrs.sub_crs_list
        assert inhcrs.to_epsg() == req_hcrs_epsg
        assert invcrs.to_epsg() == req_vcrs_epsg
    elif not cmpd_incrs.is_vertical:
        assert incrs.to_epsg() == req_hcrs_epsg
    # build the output crs
    out_cmpd_crs = CompoundCRS(name="NAD83 UTM19 + NAVD88", components=[f"EPSG:{req_hcrs_epsg}", f"EPSG:{out_vcrs_epsg}"])
    # get the transform at the sparse points
    grids = get_gtx_grid_list()
    transformer = Transformer.from_pipeline(f'proj=pipeline \
                                              step inv proj=utm zone=19 \
                                              step inv proj=vgridshift grids={grids[f"{region_name}/mllw.gtx"]} \
                                              step proj=vgridshift grids={grids[f"{region_name}/tss.gtx"]} \
                                              step proj=utm zone=19')

    result = transformer.transform(xx=xx, yy=yy, zz=zz)
    LOGGER.debug(f'Applying pipeline: {transformer}')

    return result, out_cmpd_crs

def get_datum_sep(raster_name, transform_sampling_distance, region_list):
    """
    Use the provided raster and pipeline to get the separation over the raster area.

    Returns
    -------
    sep

    """
    with rasterio.open(raster_name) as raster:
        elev = raster.read(1)
        transform = raster.transform
        crs = raster.crs
        nodata = raster.nodata
        profile = raster.profile
    # get sparse raster x and y positions
    sy, sx = profile['height'], profile['width']
    resy, resx = transform[4], transform[0]
    y0, x0 = transform[5], transform[2]
    y1 = y0 + sy * resy
    x1 = x0 + sx * resx
    nx = np.round(np.abs((x1 - x0) / transform_sampling_distance)).astype(int)
    ny = np.round(np.abs((y1 - y0) / transform_sampling_distance)).astype(int)
    x_sampled = np.linspace(x0, x1, nx)
    y_sampled = np.linspace(y0, y1, ny)
    yy, xx = np.meshgrid(y_sampled, x_sampled, indexing = 'ij')
    zz = np.zeros(yy.shape)
    y,x = np.mgrid[y0:y1:resy, x0:x1:resx]
    sep = np.full(y.shape, np.nan)

    for region in region_list:
        start = datetime.now()
        try:
            result, new_crs = run_pipeline(xx.flatten(), yy.flatten(), zz.flatten(), crs, region)
        except pyproj.ProjError as e:
            print_paths = '\n'.join(pyproj.datadir.get_data_dir().split(';'))
            LOGGER.error(f'Proj pipeline failed. pyproj paths: \n{print_paths}')
            raise e
        dt = datetime.now() - start
        LOGGER.debug(f'Transforming {len(yy.flatten())} points took {dt} seconds for {region}')
        vals = result[2].flatten()
        valid_idx = np.squeeze(np.argwhere(~np.isinf(vals)))

        if len(valid_idx) == 0:
            LOGGER.debug('No valid points found from gridding in {region}. Putting all points through proj pipeline directly.')
        else:
            LOGGER.debug('interpolating to original grid for {region}')
            start = datetime.now()
            points = np.array([result[1][valid_idx],result[0][valid_idx]]).T
            valid_vals = vals[valid_idx]
            try:
                region_sep = griddata(points, valid_vals, (y,x))
                idx = ~np.isnan(region_sep)
                sep[idx] = region_sep[idx]
            except Exception as error:
                msg = f'{error.__class__.__name__}: {error}\n\n'
                LOGGER.error(msg)
                return None, None
            dt = datetime.now() - start
            LOGGER.debug(f'Interpolating {len(y.flatten())} points took {dt} seconds for {region}')
        missing_idx = np.where(np.isnan(sep) & (elev != nodata))
        num_nan = len(missing_idx[0])
        if num_nan > 0:
            missing, new_crs = run_pipeline(x[missing_idx], y[missing_idx], np.zeros(num_nan), crs, region)
            missing_sep = missing[2]
            still_inf_idx = np.where(np.isinf(missing_sep))
            missing_sep[still_inf_idx] = np.nan
            sep[missing_idx] = missing_sep
    return sep, new_crs

def apply_sep(infile, outfile, sep, new_crs, pass_untransformed = True):
    """

    Parameters
    ----------
    outfile : TYPE
        DESCRIPTION.
    sep : TYPE
        DESCRIPTION.

    Returns
    -------
    None.

    """
    with rasterio.open(infile) as raster:
            data = raster.read()
            transform = raster.transform
            nodata = raster.nodata
            profile = raster.profile
            band_names = raster.descriptions

    missing_idx = np.where(np.isnan(sep) & (data[0] != nodata))
    num_nan = len(missing_idx[0])
    LOGGER.debug(f'Found {num_nan} values that are outside the separation model')
    if pass_untransformed:
        z_values = data[0, missing_idx[0], missing_idx[1]]
        # make the new uncertainty
        u_values = 3 - 0.06 * z_values
        u_values[np.where(z_values > 0)] = 3.0
        uncert = data[1]
        uncert[missing_idx] = u_values
        if num_nan > 0:
            LOGGER.debug(f'Inserting {num_nan} values that were not transformed into new file, but applying CATZOC D vertical uncertainty.')
    elev = (data[0] - sep)
    elev[np.where(np.isnan(elev))] = nodata
    if pass_untransformed:
        elev[missing_idx] = z_values
        data[0] = elev
        data[1] = uncert
    else:
        data[:, missing_idx[0], missing_idx[1]] = nodata

    data = np.round(data[:2].astype(np.float), decimals = 2).astype(np.float32)
    band_names = band_names[:2]

    write_gdal_geotiff(outfile, data, new_crs, transform, nodata, band_names)

def write_gdal_geotiff(outfile, data, pyproj_crs, transform, nodata, band_names):

    numlyrs, rows, cols = data.shape
    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(outfile, cols, rows, numlyrs, gdal.GDT_Float32)
    outRaster.SetGeoTransform(transform.to_gdal())
    for lyr_num in range(numlyrs):
        outband = outRaster.GetRasterBand(lyr_num + 1)
        outband.SetNoDataValue(nodata)
        outband.SetDescription(band_names[lyr_num])
        outband.WriteArray(data[lyr_num])
    outRaster.SetProjection(pyproj_crs.to_wkt())
    # outRaster.SetMetadata({'TIFFTAG_IMAGEDESCRIPTION' : TIFFTAG_IMAGEDESCRIPTION,
    #                        'TIFFTAG_ARTIST': TIFFTAG_ARTIST,
    #                        'TIFFTAG_COPYRIGHT': TIFFTAG_COPYRIGHT,
    #                        })
    outband.FlushCache()

def transform_raster(infilepathname, outfilepathname):
    """

    Returns
    -------
    None.

    """
    LOGGER.debug(f'Begin work on {os.path.basename(infilepathname)}')
    region_list = get_interesecting_vdatum_regions(infilepathname)
    if len(region_list) > 0:
        sep,crs = get_datum_sep(infilepathname, 100, region_list)
        apply_sep(infilepathname, outfilepathname, sep, crs)
    else:
        LOGGER.debug(f'no region found for {os.path.basename(datapath)}')

def set_logger(outpath):
    LOGGER.setLevel(logging.DEBUG)
    log_formatter = logging.Formatter('[%(asctime)s] %(name)-9s %(levelname)-8s: %(message)s')

    log_name = f'_to_NAVD88_{datetime.now():%Y%m%d_%H%M%S}.log'
    log_filename = os.path.join(outpath, log_name)
    log_file = logging.FileHandler(log_filename)
    log_file.setFormatter(log_formatter)
    log_file.setLevel(logging.DEBUG)
    LOGGER.addHandler(log_file)

    output_console = logging.StreamHandler(sys.stdout)
    output_console.setFormatter(log_formatter)
    output_console.setLevel(logging.DEBUG)
    LOGGER.addHandler(output_console)

def clear_logger():
    for handler in LOGGER.handlers[:]:
            handler.close()
            LOGGER.removeHandler(handler)

if __name__ == '__main__':
    if len(sys.argv) == 1:
        load_config()
    else:
        config_filename = sys.argv[1]
        load_config(config_filename)
    #set_logger(outpath)
    update_vdatum_data_directory()
    check_gdal_version()
    flist = glob.glob(os.path.join(inpath, '*.tiff'))
    for datapath in flist:
        outfilepath = os.path.join(outpath, os.path.basename(datapath))
        transform_raster(datapath, outfilepath)
    clear_logger()


#!/usr/bin/env python3

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Copyright 2015 California Institute of Technology. ALL RIGHTS RESERVED.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# United States Government Sponsorship acknowledged. This software is subject to
# U.S. export control laws and regulations and has been classified as 'EAR99 NLR'
# (No [Export] License Required except when exporting to an embargoed country,
# end user, or in support of a prohibited end use). By downloading this software,
# the user agrees to comply with all applicable U.S. export laws and regulations.
# The user has the responsibility to obtain export licenses, or other export
# authority as may be required before exporting this software to any 'EAR99'
# embargoed foreign country or citizen of those countries.
#
# Author: Piyush Agram
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

###These are demonstration scripts for the UNAVCO InSAR processing workshop.
###These scripts will not be maintained over the long term and should not
###be mistaken for official Applications within ISCE.
###These scripts are meant to demo the use of ISCE as a modular library

import argparse
import isce
import isceobj
import numpy as np
import shelve
import os
import datetime
from isceobj.Constants import SPEED_OF_LIGHT

def cmdLineParse():
    '''
    Command line parser.
    '''

    parser = argparse.ArgumentParser( description='Create DEM simulation for merged images')
    parser.add_argument('-a','--alks', dest='alks', type=int, default=1,
            help = 'Number of azimuth looks')
    parser.add_argument('-r','--rlks', dest='rlks', type=int, default=1,
            help = 'Number of range looks')
    parser.add_argument('-d', '--dem', dest='dem', type=str, required=True,
            help = 'Input DEM to use')
    parser.add_argument('-m', '--master', dest='master', type=str, required=True,
            help = 'Dir with master frame')
    parser.add_argument('-b', '--bbox', dest='bbox', type=float, nargs=4, required=True,
            help='SNWE bounding box')
    parser.add_argument('-i', '--input', dest='infile', type=str, required=True,
            help='Input file to be geocoded')
    parser.add_argument('-t', '--type', dest='method', type=str, default='nearest',
            help='Interpolation method to use')
    parser.add_argument('-n', '--nativedop', dest='nativedop', action='store_true',
            default=False)
    return parser.parse_args()

def runGeo(frame, demImage, inImage,
        looks=(1,1), doppler=None,nativedop=False,
        bbox=None, method='nearest'):
    from zerodop.geozero import createGeozero
    from isceobj.Planet.Planet import Planet

    #####Run Topo
    planet = Planet(pname='Earth')
    topo = createGeozero()
    topo.configure()

    alooks = looks[0]
    rlooks = looks[1]

    tStart = frame.sensingStart

    topo.slantRangePixelSpacing = frame.getInstrument().getRangePixelSize()
    topo.prf = frame.getInstrument().getPulseRepetitionFrequency()
    topo.radarWavelength = frame.getInstrument().getRadarWavelength()
    topo.orbit = frame.orbit
    topo.width = inImage.getWidth()
    topo.length = inImage.getLength()
    topo.wireInputPort(name='dem', object=demImage)
    topo.wireInputPort(name='planet', object=planet)
    topo.wireInputPort(name='tobegeocoded', object=inImage)
    topo.numberRangeLooks = rlooks
    topo.numberAzimuthLooks = alooks
    topo.lookSide =  frame.instrument.platform.pointingDirection()
    topo.setSensingStart(tStart)
    topo.rangeFirstSample = frame.startingRange
    topo.method=method
    topo.demCropFilename = 'crop.dem'
    topo.geoFilename = inImage.filename + '.geo'

    if inps.nativedop and (doppler is not None):
        try:
            topo.dopplerCentroidCoeffs = [x/topo.prf for x in doppler._coeffs]
        except:
            topo.dopplerCentroidCoeffs = [x/topo.prf for x in doppler]
    else:
        topo.dopplerCentroidCoeffs = [0.]

    topo.snwe = bbox
    topo.geocode()

    print('South: ', topo.minimumGeoLatitude)
    print('North: ', topo.maximumGeoLatitude)
    print('West:  ', topo.minimumGeoLongitude)
    print('East:  ', topo.maximumGeoLongitude)
    return

if __name__ == '__main__':

    #####Parse command line
    inps = cmdLineParse()

    #####Load master metadata
    db = shelve.open( os.path.join(inps.master, 'data'), flag='r')
    frame = db['frame']
    try:
        dop = db['doppler']
    except:
        dop = frame._dopplerVsPixel
    db.close()

    ####Setup dem
    demImage = isceobj.createDemImage()
    demImage.load(inps.dem + '.xml')
    demImage.setAccessMode('read')

    ####Setup input file
    inImage = isceobj.createImage()
    inImage.load(inps.infile + '.xml')
    inImage.setAccessMode('read')

    ####Geocode the image
    runGeo(frame,demImage,inImage, nativedop=inps.nativedop,
            looks=(inps.alks, inps.rlks), doppler=dop,
            bbox=inps.bbox, method=inps.method)


import time
import cv2
import numpy as np
import onnx
import onnxruntime
from imread_from_url import imread_from_url

from .utils import download_gdrive_tar_model, download_image
from .resnet50Classifier import Resnet50Classifier

np.random.seed(2021)
colors = np.random.randint(255, size=(1000, 3), dtype=int)

class Detector1K():

	def __init__(self, detector_model_path, classifier_model_path, detector_threshold = 0.2, classifier_threshold = 0.6):

		self.threshold = detector_threshold

		# Initialize detection model (download if necessary)
		models_gdrive_id = "1mVxOy65EsLhNgtqGrydOrhMkZpaktW0w"
		download_gdrive_tar_model(models_gdrive_id, detector_model_path)
		self.initialize_model(detector_model_path)

		# Initialize classification model (download if necessary)
		self.classifier = Resnet50Classifier(classifier_model_path, classifier_threshold)

		# Set the crop offset
		self.crop_offset = 0

	def __call__(self, image):

		return self.detect_objects(image)

	def initialize_model(self, model_path):

		self.session = onnxruntime.InferenceSession(model_path)

		# Get model info
		self.get_model_input_details()
		self.get_model_output_details()

	def detect_objects(self, image):

		input_tensor = self.prepare_input(image)

		# Perform inference on the image
		outputs = self.inference(input_tensor)

		# Process output data
		detections = self.process_output(outputs)

		# Classify detections
		self.label_detections = self.classify_objects(detections, image)

		return self.label_detections

	def prepare_input(self, img):

		self.img_height, self.img_width, self.img_channels = img.shape

		# Transform the image for inference
		img = cv2.resize(img,(self.input_width, self.input_height))
		img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

		img = img.transpose(2, 0, 1)
		input_tensor = img[np.newaxis,:,:,:].astype(np.float32)

		return input_tensor

	def inference(self, input_tensor):

		start = time.time()
		outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
		# print("Detection time: ", time.time() - start)

		return outputs

	def process_output(self, outputs):

		# Get all output details
		boxes = np.squeeze(outputs[0])
		classes = np.squeeze(outputs[1])
		scores = np.squeeze(outputs[2])
		num_objects = int(outputs[3][0])

		results = []
		for i in range(num_objects):
			if scores[i] >= self.threshold:

				y1 = (self.img_height * boxes[i][0]).astype(int)
				y2 = (self.img_height * boxes[i][2]).astype(int)
				x1 = (self.img_width * boxes[i][1]).astype(int)
				x2 = (self.img_width * boxes[i][3]).astype(int)

				result = {
				  'bounding_box': np.array([x1, y1, x2, y2],dtype=int),
				  'class_id': 0,
				  'label': "",
				  'detection_score': scores[i],
				  'classification_score': 0,
				}
				results.append(result)
		return results

	def classify_objects(self, detections, img):

		label_detections = detections.copy()

		for det_idx, detection in enumerate(detections):

			# Crop image using the bounding boc and pass it to the classifier
			box = detection['bounding_box']
			crop_left = max(box[0]-self.crop_offset,0)
			crop_top = max(box[1]-self.crop_offset,0)
			crop_right = min(box[2]+self.crop_offset,self.img_width)
			crop_bottom = min(box[3]-self.crop_offset,self.img_height)
			crop_img = img[crop_top:crop_bottom,crop_left:crop_right]

			# Classify the object and update the detection data
			label_id, label, score = self.classifier(crop_img)

			if label_id:
				label_detections[det_idx]['class_id'] = label_id
				label_detections[det_idx]['label'] = label
				label_detections[det_idx]['classification_score'] = score

		return label_detections

	def get_model_input_details(self):

		model_inputs = self.session.get_inputs()
		self.input_name = self.session.get_inputs()[0].name

		self.input_shape = self.session.get_inputs()[0].shape
		self.input_height = self.input_shape[2]
		self.input_width = self.input_shape[3]

	def get_model_output_details(self):

		model_outputs = self.session.get_outputs()
		self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]

	def draw_detections(self, image, draw_scores):

		for idx, detection in enumerate(self.label_detections):

			box = detection['bounding_box']
			class_id = detection['class_id']
			label = detection['label']
			class_score = int(100*detection['classification_score'])
			det_score = int(100*detection['detection_score'])

			color = (int(colors[class_id,0]), int(colors[class_id,1]), int(colors[class_id,2]))
			size = min([self.img_height, self.img_width])*0.002
			text_thickness = int(min([self.img_height, self.img_width])*0.004)
			textSize = cv2.getTextSize(text=f'class {class_score}% - det {det_score}%',
				fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=size, thickness=text_thickness)[0][1]*1.6

			cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), color, text_thickness)

			if not detection['class_id']:
				continue
			cv2.putText(image, f'{label}', (box[0], box[1] + int(textSize*0.8)),
						cv2.FONT_HERSHEY_SIMPLEX, size, color, text_thickness, cv2.LINE_AA)
			if not draw_scores:
				continue
			cv2.putText(image, f'class: {class_score}%', (box[0], box[1] + int(textSize*1.8)),
						cv2.FONT_HERSHEY_SIMPLEX, size, color, text_thickness, cv2.LINE_AA)
			cv2.putText(image, f'det: {det_score}%', (box[0], box[1] + int(textSize*2.8)),
						cv2.FONT_HERSHEY_SIMPLEX, size, color, text_thickness, cv2.LINE_AA)

		return image

if __name__ == '__main__':

	draw_scores = False

	detection_model_path = '../models/object_localizer_float32.onnx'
	detection_threshold = 0.3

	classification_model_path = '../models/resnet50_vd_ssld_float32.onnx'
	classification_threshold = 0.5

	# image = imread_from_url("https://upload.wikimedia.org/wikipedia/commons/thumb/5/55/Il_cuore_di_Como.jpg/800px-Il_cuore_di_Como.jpg")
	image = download_image("https://upload.wikimedia.org/wikipedia/commons/5/55/Il_cuore_di_Como.jpg")
	detector = Detector1K(detection_model_path, classification_model_path, detection_threshold, classification_threshold)

	start_time = time.time()
	detections = detector(image)
	detection_img = detector.draw_detections(image, draw_scores)
	# print("total: ",time.time()-start_time)

	cv2.namedWindow("Detections", cv2.WINDOW_NORMAL)
	cv2.imwrite("../output.jpg", detection_img)
	cv2.imshow("Detections", detection_img)
	cv2.waitKey(0)