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

# # Lab 2: Astronomical Imaging I
#
# *Gather round, for I shall tell a tale old as time.
# Long, long ago, astronomers gazed at the heavens, and with naught but their eyes, recorded the positions of the stars they saw there.
# Then, with the telescope, long, heavy contraptions which reached for the skies like an outstretched arm, (thanks, f/20 focal lengths)
# they gazed through eyepieces in freezing domes, falling, on occasion, to their deaths from cages suspended at prime foci.*
#
# *Next came glass -- emulsion plates -- which upon extended exposure revealed upon their ghostly frames the splotches of stars and soon, galaxies and nebulae.
# Many a grad student, of course, spent their nights twiddling thumbs over micrometer dials to keep these stars in place. Manual guiding... this story teller shudders to imagine it.
# And yet, even then, with permanent record, no measure could be made that weren't 'tween the eyes of an observer, peering at the glass through magnifying eyepiece, assigning grades of brightness and, amazingly, pretty much nailing the spectral classification of stars by their absorption features*.
#
# *And after this painstaking work, came the earliest CCDs, fractured by detector failures, riddled with readout noise, consumed by cosmic ray hits, laid low by low quantum efficiencies.*
#
# *Now... we use Python.*
#

# <hr>
# Astronomical images are one of the basic building blocks of astronomical research. While we now obtain data in myriad ways, from gravitational waves to neutrinos, from spectra to polarimetry, the basic tenant of astronomy (and it's most recongizable public impact) is the images we make of the sky.
#
# That is why the first science topic we'll be tackling is imaging. And along the way, we'll learn how `astropy` makes our lives *so* much easier when dealing with images, and how to use robust functions to perform the analyses one might want to carry out on images (after we're done admiring their beauty).
#
# # If the glove `FITS`
#
# Many of you are probably familiar with the `FITS` standard. It stands for `Flexible Image Transport System`, and for all intents and purposes, it acts as a container (much like a zip file). Within it, one can store several kinds of information: images (2d arrays of values), as the name implies, headers (dictionary like objects with metadata), and sometimes tabular data as well. A `FITS` file can contain any number of *extensions*, which just refer to "slots" where stuff is stored. Every "slot" has a `header` attribute and a `data` attribute. Thus, you could store a hundred astronomical images in 1 file, each with a different extension and a different header describing the image in that slot.
#
# ```{tip}
# Contextual clues are important when dealing with `FITS` files. Files that are explicitly single images almost always have the image stored in the 0th extension, while files that are explicitly table data tend to have the table stored in the 1st extension (with the 0th empty).
# ```
#
# The `FITS` standard is pretty old, and may be retired soon, but almost all ground based telescopes still use it for their native image output, so it behooves us to know how to get image data out of `FITS` files.
#
# ## Problem 1: Loading a FITS file
#
# Write a function which takes as its argument a `string` filepath to a `FITS` file, and should have an optional argument to set the extension (default 0). It should then load the given extension of that `fits` file using a [context manager](https://docs.astropy.org/en/stable/io/fits/index.html#opening-a-fits-file), and return a tuple containing the header and data of that extension.
#
# The function should have documentation that describes the inputs and outputs. Documentation is **incredibly important** when writing code, both in-line and (# comments) and at the top of functions, methods, and classes.
#
# In this class, we'll be using [Sphinx-compatible](https://www.sphinx-doc.org/en/master/) documentation written in the [Numpy/Scipy style](https://numpydoc.readthedocs.io/en/latest/format.html). There are several reasons to do this.
#
# 1. It is a user-friendly and write-friendly, readable documentation format that is easy to add to your functions.
# 2. It can be rendered by Sphinx, the most popular automatic documentation renderer. If you've ever read the online documentation pages for functions in, e.g., numpy and scipy, those pages were rendered automatically based on the docstrings of the functions in question. This is possible with tools like Sphinx, but the documentation must be formatted correctly for this to work.
#
# In the cell below, I've provided a dummy function which takes in any number of inputs (mininum 3) and chooses a random input to return. The formatting of the documentation is shown there (as well as in the link above).
#
#

# In[5]:

import numpy as np

def random_return(a,b,c,*args):
    '''
    A function which requires three inputs which are floats, accepts any number of additional inputs (of any kind), and returns one randomly chosen input.

    Parameters
    ----------
    a: int
        description of this integer input
    b: int
        description of this integer input
    c: int
        description of this integer input
    *args: tuple
        any additional arguments get stored here

    Returns
    -------
    choice
        The randomly selected input (type not specified)
    '''
    full_input_list = [a,b,c] + list(args)
    choice = np.random.choice(full_input_list)
    return choice

# In[11]:

random_return(1,5,4,6,4,21,6)

# When our function has been imported in some code, we can use the `help` command to see the documentation at any time:

# In[9]:

help(random_return)

# You will also notice my use of `*args` in this function. This allows me to enter additional function arguments. Similar is `**kwargs`, which allows additional arguments tied to an input keyword. The former gets stored in a tuple, while the latter gets stored in a dictionary. We'll be using these a lot in class --- you can refresh or learn the basics in Section 2.8 of the chapter on functional programming [here](https://prappleizer.github.io/Tutorials/FunctionalProgramming/FunctionalProgramming_web.html).
#
# So as a brief overview of documentation, it contains
# 1. A brief summary of the function
# 2. The word Parameters with the next line having underlines of the same length
# 3. arguments, which are followed by a colon and the data type(s). On the next line, indented, descriptions of the arguments
# 4. The word Returns, with the same underline scheme
# 5. The returned objects, labeled the same way as the input.
#
# Also above, we saw how to format when no data type is specified. There, additional inputs could've been *any* data type, so we can't be sure what the output will be. The main thing we didn't cover is optional arguments. Those are set like this:
# ```
# a: int, optional
#     Description of the thing. (default 5)
# ```
# So we mark it as optional, and then give the default for it.
#
# With that, you're ready to write and document your code below. **All functions you write in this class should have documentation**.

# In[10]:

from astropy.io import fits
import os

# your code
def load_fits(...):
    '''

    '''
    pass

# ## Problem 2: Data I/O
# ### Problem 2.1
# Using the function you created above, read in the header and data of the file `antenna_Rband.fits` which came with the lab assignment.
# ```{note}
# While this may seem small, the fact that your read-in is now one line instead of ~5 does improve your efficiency! But only if you use your function enough times to overcome the initial time spent writing it...
# We also have the flexibility to take our function and give it more features over time, which we will do in this class.
#
# ```

# ### Problem 2.2
# Next, we need to plot the image data. There are several operations that we almost always perform when plotting astronomical data, as well as several user-preferences for how we "by default" plot images before we begin tweaking things. If you spend any time as an astronomer, you will plot quite literally *thousands* of images --- why set all these settings every time, when we can write a handy function to do it for us?
#
# Write a function which takes in as input arguments
# - an image (2D array or masked array)
#
# as well as the following optional arguments (so set a default)
# - figsize (default (15,13) )
# - cmap (default 'gray_r')
# - scale (default 0.5)
# - \*\*kwargs (see [here](https://prappleizer.github.io/Tutorials/FunctionalProgramming/FunctionalProgramming_web.html))
#
# Inside the function, create figure and axes objects using `plt.subplots()`. When working in notebooks, it is often useful to set the `figsize` argument of subplots to a nice large-ish value, such as `(15,13)`, which will make the image fill most of the notebook. Since *your* function has set figsize as an argument, you can feed `figsize` directly into the `subplots` call, so that a user of the function can leave the default or set their own.
#
# Next, use `ax.imshow()` to actually plot the image. You'll want to save the output of this, e.g., `im = ax.imshow(...)`. In this plotting call, set `imshow`'s argument `origin='lower'`. We *always* want to do this when dealing with imaging, as we want (0,0) to be a coordinate.
# ```{note}
# By default, matplotlib uses a "matrix-style" plotting, where 0 of the y axis is in the *top* left corner, and 0 of the x axis is in the *bottom* left corner.
# ```
# Also within the call to `imshow()`, feed in the cmap from your function (i.e., `cmap=cmap`). The other critical `imshow()` arguments are `vmin` and `vmax`, which sets the saturation points (and thus the contrast) of the image.
#
# We haven't set `vmin` and `vmax` as arguments of our outer function, but because of `kwargs`, we can still create a default here that can be overriden from outside.
#
# As a default, within your function, calculate the mean and standard deviation of the image. Set some temporary variables with the quantities `mu - scale*sigma` and `mu + scale*sigma` (where here `mu` is the calculated mean and `sigma` is the calculated std dev, and `scale` was the optional input). Next, check the kwargs dictionary (which will exist in your function because we added the packing argument `**kwargs` to our function. IF `vmin` and `vmax` are in this dictionary, plug those into your imshow command. Otherwise, use the values determined by the calculation above. Bonus point for accomodating either no vmin/vmax entered, just vmin or vmax, or both (using the calculated values when things aren't provided).
#
# Your function should **return** the created `fig` and `ax` objects so the user can continue to tweak them.
#
# Run your function and test its outputs. Once you're satisfied it's working, use it to plot the provided data. Find either a vmin/vmax pair, or a choice of `scale` which makes the image look pretty!

# In[ ]:

import matplotlib.pyplot as plt
import numpy as np

# Your code
def implot(...):
    '''
    WRITE DOCSTRING HERE
    '''
    pass #replace with your code

# In[42]:

# I've included my output below for your comparison. Overwrite it with your own!

# ### Problem 2.3
#
# Copy your function down, we're going to keep messing with it.
#
# So far, we've made it so that with a simple function call and, potentially, with just the input of an image array, we get out a nice plot with a scaling, colormap, and origin selection. In this section, we are going to allow (optionally) for a colorbar to be added to the figure. We're also going to add in the ability for the figure to be plotted in celestial coordinates (i.e., RA and DEC) instead of pixel units, if information about the image (via the world coordinate system, WCS) exists in the image header.
#
# ```{note}
# Generally, WCS information is present in the headers of *published* images like the one here, but *not* present in raw data gathered at the telescope. This is because images need to undergo a process known as *plate solving* to determine both the direction (coordinates) toward the image, as well as the pixel scale (i.e., how many arcseconds per pixel in the image).
# ```
#
# Add three new optional arguments to your function.
# - colorbar = False
# - header = None
# - wcs = None
#
# Let's start with the colorbar. At the end of your plotting commands, check if `colorbar=True`, and if so, create a colorbar via `plt.colorbar()`, setting the `mappable` argument to whatever you saved the output of `ax.imshow()` into above. Also set the `ax` argument to be your ax; this will tell `matplotlib` to steal a bit of space from that axis to make room for the colorbar.
#
#
#

# In[46]:

# Your code

# ```{tip}
# When I do this, the colorbar is matching the figure height, rather than the axis height. If that bugs you like it bugs me, check out [this solution](https://stackoverflow.com/a/33505522) from StackOverflow, which you can adapt within your function to make the cbar match the axis height.
# ```
#
# In order to plot in RA and DEC coordinates, we need to first have an `astropy` `WCS` object associated with the image in question. You can import `WCS` from `astropy.wcs`. WCS objects are created from the headers of plate-solved fits files. In our function, we allow the user to input either a header or a WCS object directly. More on WCS can be found in the lecture notes, or [here](https://docs.astropy.org/en/stable/wcs/index.html).
#
# Within your function, check if a wcs is input -- if it is, we're good to go and can safely ignore `header` (even if it is provided). If instead only `header` is provided, use the `WCS()` function to create a new wcs object from that header.
# ```{tip}
# You'll want to do this at the very top of your function.
# ```
#
# We now need to move our `fig, ax = ....` creation line into an if-statement. If we are using WCS "stuff", you'll need to set a `projection` for your plot that uses the wcs. This is accomplished as follows:
#
# ```
# fig, ax = plt.subplots(...,subplot_kw={'projection':wcs})
# ```
# where `wcs` is whatever you've named the output of `WCS(header)` or is the WCS input directly into the function.
#
#

# In[4]:

from astropy.wcs import WCS

# your code

# ```{warning}
# In this case, we will get an error from our function that happens because of some distortion coefficient nonsense between astropy and drizzled HST images. Since it's not pertinent to our lab, I'm providing below a snippet of code you should use to fix your header before running `WCS(header)`.
# ```

# In[203]:

def strip_SIP(header):
    A_prefixes = [i for i in header.keys() if i.startswith('A_')]
    B_prefixes = [i for i in header.keys() if i.startswith('B_')]
    for a,b in zip(A_prefixes,B_prefixes):
        del header[a]
        del header[b]
    return header

# If this worked correctly, when you add the header you read in from our image, you should now see the axes of your plot change from pixels to `pos.eq.ra` and `pos.eq.dec`. We're now looking at actual on-sky coordinates! Yay!
#
# ### Problem 2.4
# Within the if-blocks of your function that sets the `ax` to be a wcs projection, set the $x$ and $y$ labels to read "Right Ascension \[hms\]" and "Declination \[degrees\]" in fontsize 15.
#
# Lastly, to polish things off, use `ax.tick_params()` to set inward, larger ticks, and increase the axis tick label size to 15.
#
# ```{note}
# You'll notice (especially with larger ticks) that the are not perpendicular to the axes spines. This is because this particular image has been rotated with respect to the standard celestial coordinate axes. This can be seen more clearly if you add the following to your function:
# `ax.coords.grid(color='gray', alpha=0.5, linestyle='solid')`. Try doing that, adding an optional keyword to your function called 'grid' and enabling this command if it is set.
# ```
#
# ```{tip}
# To check if a condition is true (e.g, `if condition == True:`), you can just type `if condition:`
# ```
#
# It's taken us some time, but this image could now be placed in a scientific publication. And, since we have a handy function for it, we can make images that look this nice on the fly with a short one line command, yet still have a lot of flexibility over many important inputs. And of course, the figure and axes objects are returned, so one could run this function and then continue tweaking the plot after the fact.
#
# ```{note}
# As a final note, I want to draw your attention to the fact that once you use the `wcs` projection on some plot, it's no longer a normal ax object, it's a wcsax object. This makes changing certain elements of those axes a little more involved than for a standard one. I use [this page](https://docs.astropy.org/en/stable/visualization/wcsaxes/ticks_labels_grid.html) and the others linked there when doing so.
# ```

# ## Problem 3: Cutouts and Aperture Photometry

# When working with astronomical images, like the one we've been using in this lab, it is often advantageous to be working with a cutout -- a chunk of the image centered on a certain coordinate, and of a certain size. For example, if there is an HII region or star cluster of interest in the above image, we may like to zoom in on it to examine it more closely.
#
# Now that we've switched over to using celestial coordinates instead of pixels as our projection frame, zooming in on a region is not as simple as slicing our array, e.g., `image[500:700,200:550]`. On the plus side, the framework we'll learn here is very robust, and will allow for all sorts of useful measurements.
#
#
# To make a cutout, we'll need the `Cutout2D` module within `astropy`, which I'll import below. To provide the position of the cutout, we need to be able to feed in astronomical coordinates. For this, we'll use `SkyCoord`, a module in `astropy.coordinates`. Finally, we'll need to integrate the `units` module in `astropy` to successfully create coordinate objects.

# In[100]:

from astropy.nddata import Cutout2D
from astropy.coordinates import SkyCoord
import astropy.units as u

# Let's start with a `SkyCoord` object. There are several coordinate systems used in astronomy, e.g., Galactic coordinates ($b$, $l$), Equatorial ($RA$, $DEC$). The most common (especially in any extragalactic setting) is RA and DEC (as you can see in the image you've plotted already).
#
# The [documentation](https://docs.astropy.org/en/stable/api/astropy.coordinates.SkyCoord.html) for `SkyCoord` is solid, and worth reading.
#
# The general way we create these objects is, e.g.,

# In[103]:

coord = SkyCoord('12:01:53.6 -18:53:11',unit=(u.hourangle,u.deg))

# ```{tip}
# You can input various types of strings and formattings for the coordinates, just be sure to specify the units as shown above. The documentation shows the various methods.
# ```
#
# ### Problem 3.1
#
# In this case, the coordinates I set above are for NGC 4039, which is the smaller of the two galaxies in the image we're using.
#
# ```{note}
# If at any point you're trying to make a coordinate object for a well known galaxy/object, try, e.g., `coord = SkyCoord.from_name('NGC 4038')`, and ususally that will work!
# ```
#
#
# In the cell below, use the coordinate we've created, plus a size (use 1x1 arcminutes), and the wcs object for our image, to create a `Cutout2D` object.

# In[106]:

#cutout = # your code here

# Now, use your fancy new function to plot the new cutout.
# ```{note}
# Cutout objects contain the image and their own wcs, accessible via, e.g., `cutout.data` and `cutout.wcs`.
# ```
#

# In[ ]:

# plot it

# ### Problem 3.2
#
# We're now going to do some aperture photometry.
# ```{admonition} Definition
# Aperture Photometry is the process of defining a region on an image (generally, but not always circular), and measuring the sum pixel value within that region. The region is known as an aperture, and the "collapsing" of the 2D spatial information about counts in each pixel into a single number is known as photometry.
# ```
#
# Below, I provide a new coordinate, this time centered on the region between the two galaxies. Make a new cutout of that region (again, 1x1 arcmin), and plot it.

# In[120]:

new_coord = SkyCoord('12:01:55.0 -18:52:45',unit=(u.hourangle,u.deg))

# In[121]:

# your code

# In[5]:

#plot it

# In this region, there are a lot of blobby looking roughly circular sources --- Some of the larger ones are HII star forming regions, the smaller ones are likely stars. Later in this lab, we'll use multi-wavelength data to try to suss out what is what.
#
# Often, for calibration purposes, we'd need to create apertures around all those sources in the image. We definitely don't want to do that by hand! Instead, we're going to use the `sep` package.
#
# ```{note}
# Simply `pip install sep` inside your `a330` environment terminal to get it installed, you should then be able to import it.
# ```
#
#

# In[123]:

import sep

# There are three steps to performing aperture photometry with `sep`, which are detailed in [it's documentation](https://sep.readthedocs.io/en/v0.4.x/).
#
# - Estimate the background
# - Find Sources
# - Perform Aperture Photometry
#
# Using the instructions presented in the documentation linked, measure the background of the cutout image, and run the source extractor on it. Don't forget to subtract the background before running the extractor!
#
# To do this, write a function that takes as input the data (in this case, a `cutout.data` object and a threshold scale (to be multiplied by the `globalrms`, and performs these steps, returning the `objects` array. Don't forget to document it!
#
# ```{hint}
# :class: dropdown
# I used a threshold value of 2.0, and got ~100 objects.
# ```
#
# ```{warning}
# When I ran this, I got an error that my "array was not C-contiguous." I found the solution to this issue in [this stackoverflow post](https://stackoverflow.com/a/26782930). Loosely, `sep` uses C-bindings to actually run the heavy lifting code in C rather than Python (it's faster). This means input arrays must be arranged in physical memory the way C is used to. This particular array was not, but it is easy to order it this way.
# ```

# In[173]:

# Your code

def run_sep(...):
    '''
    DOCSTRING HERE
    '''
    pass

# Run your function and store the output in a variable called `objects`. You should now have an object array containing many detected sources.

# In[ ]:

# run func

# The positions of the determined sources are stored in the output structured array and can be indexed with, e.g., `objects['x']` and `objects['y']`. Replot your image of the cutout, but now circle the positions of all your detected sources. Do they line up with actual point sources in the image?
#
# ```{hint}
# :class: dropdown
# Try using ax.plot, setting your symbol to 'o', your color to 'None', your marker edge color (mec) to some color, and the marker size (ms) to some largeish value, for a quick way to circle objects.
# ```

# In[ ]:

# Your code

# In[167]:

# Here's my image, for comparison

# It looks like we've done a pretty adequate job finding the sources in this field -- there'a a few that got missed, and a few we might want to remove, but overall, this is pretty solid.
#
# We need to start working with these objects we've found. While `sep` can perform aperture photometry itself (reading the docs, you can see it is simple to feed in the objects and a pixel radius), we're going to be a bit more careful about things. To better visualize and work with this data, I'd like it to be a `pandas` `DataFrame`. We're going to be using those a lot in this course. Converting a `numpy` structured array to a dataframe is simple: I provide the code below. Run it, and then simply type `df` in a new cell to see a nicely formatted pandas table of our objects.

# In[171]:

import pandas as pd
df = pd.DataFrame(objects)

# In[ ]:

df #run this

# With `sep` providing the first pass, we can cull a few objets from our sample that very clearly are not star forming regions. Using your DataFrame, plot the flux of each detected object using the `flux` column. You should have at least a few that are huge outliers, with dramatically more flux the rest.

# In[6]:

# plot

# Write a function called `remove_outliers` that reads in a dataframe and a flux-min and flux-max. It should filter the dataframe to only include fluxes between the input values, and return the new dataframe. Then use this function on your data, choosing an appropriate cutoff.

# In[182]:

# Your Code

def remove_outliers(...):
    '''
    '''
    pass

# In[183]:

# run the function

# In[7]:

# plot again

# Re-plot the set of sources you have now, over the data.

# In[8]:

# plot

# You should see that some of the circles which were over bright parts of the galaxy are no longer here.
# ```{note}
# You may find that there are some visible sources which, despite tinkering, don't get caught by `sep`. That's okay -- if we really wanted them, we could just go in and put down apertures by hand. Generally when performing a step like this in the field, we have a pre-determined set of points to use, because we have measured fluxes for them, and wish to flux calibrate our data.
# ```
#
# ## Problem 4: Continuum Subtraction
#
# At this point, we have the tools necessary to make flux measurements in apertures (at least, via `sep` -- we will also learn how to do this with the `photutils` package). However, $R$ band photometry of HII regions is not exceedingly interesting. More interesting is the flux in $H\alpha$, an emission line caused by Hydrogen recombination. This line (at 6563 Angstrom) is located *within* the $R$ band, and is imaged using a narrow filter compared ot $R$.
#
# The flux measured by an $H\alpha$ filter contains both the flux from the emission line as well as the underlying continuum --- essentially, the starlight from the galaxy. If we can get a clean measure of the flux in $H\alpha$ (sans this continuum), we can make a direct estimate of the star formation rate of the system.
#
# ### Problem 4.1
# To do this, we need to access an $H\alpha$ image, and then subtract off the continuum present (which we'll infer from our $R$ band image). Located in the lab directory is a file called `antenna_Haband.fits`. Use your fits loader function from above to read in this new image, create an equivalent sized and centered cutout to the $R$ band data, and plot it, with the same sources found in the $R$ band circled.

# In[9]:

# Your Code

# There's a few interesting things to note right away here. Looking just north of the center of the image, there's now a large amount of flux in $H\alpha$ coming from some blobs which do not appear in the $R$ band. This is likely due to the fact that by zero-ing in on the ionized gas, we can see the large envelopes of gas being lit up by the star formation in this region.
# ```{note}
# The active merger scenario between these two galaxies is triggering a lot of star formation.
# ```
#
# ### Problem 4.2
#
# To get a better idea of how the $H\alpha$ flux compares to the $R$ band distribution, use the `plt.contour()` tool to measure contours of $H\alpha$, drawing them over the
# $H\alpha$ image and tweaking the levels until you think you are well tracing the distribution. Then, plot those contours over the $R$ band data instead.
#
# ```{hint}
# It is often beneficial to use `np.logspace()` when defining contour levels, as it allows you to cover large dynamic range with fewer contours. You may also find it helpful to set your contour `alpha` to something < 1, to better see both the image underneath and the contours.
# ```

# In[10]:

# Your Code

# To help guide you, I've shown what I got for my plot below --- you need not emulate it exactly. The blue box shown in the image will be useful to you in the next part of the problem.

# In[267]:

# My solution

# We can now see things a lot more clearly! It is obvious that the $H\alpha$ contours trace some of the sources we were seeing in the $R$ band, but that the gas is more extended around these clusters of sources, as we might expect.
#
# ### Problem 4.3
# We must now attempt a continuum subtraction of the data. We can't simply subtract the $R$ band from the $H\alpha$ band, because the $R$ band filter is much wider, and thus for similar exposure times, collects many more photons than the narrower $H\alpha$ filter. Ideally, one could use a set of foreground stars (which are pure continuum sources in both images) to measure fluxes in both, and find a scaling constant. Here, all of the sources we see in this image are most likely actually HII regions. This means if we use them to scale between our images, we will likely oversubtract true flux.
#
# Instead, what I'm going to do here (which may be slightly sketchy), is pick a "blank" region of continuum emission from the $R$ band which has no $H\alpha$ contours, and assert that this patch must have the same flux in both images. In the picture above, I indicated a blue rectangular patch on the sky. This is what we'll be using to measure our continuum-to-narrowband ratio.
#
# To make this patch, we're going to use `SkyRectangularAperture`, from `photutils`. You can pip install `photutils` in your `a330` environment if you don't have it. Then do the following imports:

# In[268]:

from photutils.aperture import SkyRectangularAperture
from photutils.aperture import aperture_photometry

# As shown in [it's documentation](https://photutils.readthedocs.io/en/stable/api/photutils.aperture.SkyRectangularAperture.html), we need to feed in a `SkyCoord` position, as well as a width and a height.
# I've provided the new coordinate to use below.
#
# Create a rectangle of your own, measuring $0.27''$ by $0.2''$. You can now use the `aperture_photometry()` function to measure the flux in your aperture, applied to a certain image. Using the [documentation](https://photutils.readthedocs.io/en/stable/api/photutils.aperture.aperture_photometry.html#photutils.aperture.aperture_photometry) as needed, find the flux in this box for both the $R$ band and $H\alpha$ band data, and then determine the ratio between those values. Don't forget about the `wcs`!

# In[11]:

patch_cent = SkyCoord('12:01:53.7 -18:52:52',unit=(u.hourangle,u.deg))

# Your code

# We can now use this ratio to perform our subtraction.
#
# ### Problem 4.4
# Now, fill in the function below, which should read in your rectangular patch, the $R$ band image, and the $H\alpha$ image. Within the function, copy in the code that determines the ratio, and then use the ratio you found to scale your $R$ band image, then subtract it from the $H\alpha$ image. The function should return this new image array.
#
# Plot up your continuum subtracted image using your `implot()` function, adding back in the apertures we found earlier and the contours we made from the full $H\alpha$ image.

# In[12]:

# your code
def continuum_subtract(...):
    '''
    DOCSTRING HERE
    '''
    pass

# What do you notice about the distribution of apertures with respect to the distribution of $H\alpha$ gas? Is there a strong alignment of the $R$ band sources and the gas? What does this imply about most of the $R$ band sources?
#
#

# *Answer here*

# ### Problem 4.5
#
# Lastly, let's load up the $B$ band image. As you probably know, the $B$ band traces bluer light, and thus will more preferentially see young, hot stars (whereas the $R$ band traces the main sequence and turn off stellar distribution).
#
# Load up the `antenna_Bband.fits` image, make a cutout, and plot it below. Use your `load_fits()` function and your `implot()` function. Make a new set of countours from your *continuum subtracted* $H\alpha$ data, and overplot that onto the $B$ band data. What do you see?
#

# In[13]:

# Your code

# In my own image, there are some clusters of $B$-band flux that align well with the $H\alpha$ contours, and some $B$ band sources out on their own. It is unsurprising that there is a correspondance between them; excess in $B$ light implies the prescence of UV radiation as well, from young O/B stars. It is this radiation responsible for ionizing the gas that is shining in $H\alpha$, so the $H\alpha$-emitting gas should be loosely clustered around sources bright in the $B$ band (the experiment would be even cleaner if we used, say, *GALEX* FUV data).

# ### Bonus Question (up to 3 points)
#
# If you want to take your analysis farther, try measuring some fluxes off of your continuum subtracted $H\alpha$ data, placing several manual apertures down over the regions of highest $H\alpha$ concentration (which also align with $B$ band concentrations).
#
# The measures you get from this will be in counts on the detector. We need to convert these counts into flux units (e.g., erg s$^{-1}$ cm$^{-2}$). To do this,
# - pull the 'PHOTFLAM' keyword from the header of your $H\alpha$ data
# - also pull the 'EXPTIME' value from the header.
#
# Start by taking your fluxes in counts, multiplying by the `PHOTFLAM` value, and dividing by the `EXPTIME` value. This puts you in erg s$^{-1}$ cm$^{-2}$ *per Angstrom*. Thus, what we have is technically a *flux density*. To convert to a true flux, we'll simply integrate over the bandpass... but for this example, let's assume a constant flux across the bandpass, which for the ACS 658N filter ($H\alpha$), is 136.27 angstroms.
#
# Once you have fluxes, use the distance to NGC 4038/9 to determine the luminosity in erg/s of your collective set of sources, and then use the $SFR(H\alpha)$ calibration of [Kennicut & Bell](https://iopscience.iop.org/article/10.1086/319025/fulltext/52481.text.html) to convert to an SFR. How does your answer compare to, e.g.,
# - The SFR of the Milky Way?
# - The SFR of the Orion Nebula?
# - The SFR of the Tarantula Nebula?

import numpy,import astropy,from astropy


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

#
# #  Tutorial 1: "What" models
# __Content creators:__ Matt Laporte, Byron Galbraith, Konrad Kording
#
# __Content reviewers:__ Dalin Guo, Aishwarya Balwani, Madineh Sarvestani, Maryam Vaziri-Pashkam, Michael Waskom
#
# We would like to acknowledge [Steinmetz _et al._ (2019)](https://www.nature.com/articles/s41586-019-1787-x) for sharing their data, a subset of which is used here.
#

# ___
# # Tutorial Objectives
# This is tutorial 1 of a 3-part series on different flavors of models used to understand neural data. In this tutorial we will explore 'What' models, used to describe the data. To understand what our data looks like, we will visualize it in different ways. Then we will compare it to simple mathematical models. Specifically, we will:
#
# - Load a dataset with spiking activity from hundreds of neurons and understand how it is organized
# - Make plots to visualize characteristics of the spiking activity across the population
# - Compute the distribution of "inter-spike intervals" (ISIs) for a single neuron
# - Consider several formal models of this distribution's shape and fit them to the data "by hand"

# In[ ]:

#@title Video 1: "What" Models
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='KgqR_jbjMQg', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video

# # Setup
#
#

# Python requires you to explictly "import" libraries before their functions are available to use. We will always specify our imports at the beginning of each notebook or script.

# In[ ]:

import numpy as np
import matplotlib.pyplot as plt

# Tutorial notebooks typically begin with several set-up steps that are hidden from view by default.
#
# **Important:** Even though the code is hidden, you still need to run it so that the rest of the notebook can work properly. Step through each cell, either by pressing the play button in the upper-left-hand corner or with a keyboard shortcut (`Cmd-Return` on a Mac, `Ctrl-Enter` otherwise). A number will appear inside the brackets (e.g. `[3]`) to tell you that the cell was executed and what order that happened in.
#
# If you are curious to see what is going on inside each cell, you can double click to expand. Once expanded, double-click the white space to the right of the editor to collapse again.

# In[ ]:

#@title Figure Settings
import ipywidgets as widgets #interactive display

get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")

# In[ ]:

#@title Helper functions

#@markdown Most of the tutorials make use of helper functions
#@markdown to simplify the code that you need to write. They are defined here.

# Please don't edit these, or worry about understanding them now!

def restrict_spike_times(spike_times, interval):
  """Given a spike_time dataset, restrict to spikes within given interval.

  Args:
    spike_times (sequence of np.ndarray): List or array of arrays,
      each inner array has spike times for a single neuron.
    interval (tuple): Min, max time values; keep min <= t < max.

  Returns:
    np.ndarray: like `spike_times`, but only within `interval`
  """
  interval_spike_times = []
  for spikes in spike_times:
    interval_mask = (spikes >= interval[0]) & (spikes < interval[1])
    interval_spike_times.append(spikes[interval_mask])
  return np.array(interval_spike_times, object)

# In[ ]:

#@title Data retrieval
#@markdown This cell downloads the example dataset that we will use in this tutorial.
import io
import requests
r = requests.get('https://osf.io/sy5xt/download')
if r.status_code != 200:
  print('Failed to download data')
else:
  spike_times = np.load(io.BytesIO(r.content), allow_pickle=True)['spike_times']

# ---
#
# # Section 1: Exploring the Steinmetz dataset
#
# In this tutorial we will explore the structure of a neuroscience dataset.
#
# We consider a subset of data from a study of [Steinmetz _et al._ (2019)](https://www.nature.com/articles/s41586-019-1787-x). In this study, Neuropixels probes were implanted in the brains of mice. Electrical potentials were measured by hundreds of electrodes along the length of each probe. Each electrode's measurements captured local variations in the electric field due to nearby spiking neurons. A spike sorting algorithm was used to infer spike times and cluster spikes according to common origin: a single cluster of sorted spikes is causally attributed to a single neuron.
#
# In particular, a single recording session of spike times and neuron assignments was loaded and assigned to `spike_times` in the preceding setup.
#
# Typically a dataset comes with some information about its structure. However, this information may be incomplete. You might also apply some transformations or "pre-processing" to create a working representation of the data of interest, which might go partly undocumented depending on the circumstances. In any case it is important to be able to use the available tools to investigate unfamiliar aspects of a data structure.
#
# Let's see what our data looks like...

# ## Section 1.1: Warming up with `spike_times`

# What is the Python type of our variable?

# In[ ]:

type(spike_times)

# You should see `numpy.ndarray`, which means that it's a normal NumPy array.
#
# If you see an error message, it probably means that you did not execute the set-up cells at the top of the notebook. So go ahead and make sure to do that.
#
# Once everything is running properly, we can ask the next question about the dataset: what's its shape?

# In[ ]:

spike_times.shape

# There are 734 entries in one dimension, and no other dimensions. What is the Python type of the first entry, and what is *its* shape?

# In[ ]:

idx = 0
print(
  type(spike_times[idx]),
  spike_times[idx].shape,
  sep="\n",
)

# It's also a NumPy array with a 1D shape! Why didn't this show up as a second dimension in the shape of `spike_times`? That is, why not `spike_times.shape == (734, 826)`?
#
# To investigate, let's check another entry.

# In[ ]:

idx = 321
print(
  type(spike_times[idx]),
  spike_times[idx].shape,
  sep="\n",
)

# It's also a 1D NumPy array, but it has a different shape. Checking the NumPy types of the values in these arrays, and their first few elements, we see they are composed of floating point numbers (not another level of `np.ndarray`):

# In[ ]:

i_neurons = [0, 321]
i_print = slice(0, 5)

for i in i_neurons:
  print(
    "Neuron {}:".format(i),
    spike_times[i].dtype,
    spike_times[i][i_print],
    "\n",
    sep="\n"
  )

# Note that this time we've checked the NumPy `dtype` rather than the Python variable type. These two arrays contain floating point numbers ("floats") with 32 bits of precision.
#
# The basic picture is coming together:
# - `spike_times` is 1D, its entries are NumPy arrays, and its length is the number of neurons (734): by indexing it, we select a subset of neurons.
# - An array in `spike_times` is also 1D and corresponds to a single neuron; its entries are floating point numbers, and its length is the number of spikes attributed to that neuron. By indexing it, we select a subset of spike times for that neuron.
#
# Visually, you can think of the data structure as looking something like this:
#
# ```
# | . . . . . |
# | . . . . . . . . |
# | . . . |
# | . . . . . . . |
# ```
#
# Before moving on, we'll calculate and store the number of neurons in the dataset and the number of spikes per neuron:

# In[ ]:

n_neurons = len(spike_times)
total_spikes_per_neuron = [len(spike_times_i) for spike_times_i in spike_times]

print(f"Number of neurons: {n_neurons}")
print(f"Number of spikes for first five neurons: {total_spikes_per_neuron[:5]}")

# In[ ]:

#@title Video 2: Exploring the dataset
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='oHwYWUI_o1U', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video

# ## Section 1.2: Getting warmer: counting and plotting total spike counts
#
# As we've seen, the number of spikes over the entire recording is variable between neurons. More generally, some neurons tend to spike more than others in a given period. Lets explore what the distribution of spiking looks like across all the neurons in the dataset.

# Are most neurons "loud" or "quiet", compared to the average? To see, we'll define bins of constant width in terms of total spikes and count the neurons that fall in each bin. This is known as a "histogram".
#
# You can plot a histogram with the matplotlib function `plt.hist`. If you just need to compute it, you can use the numpy function `np.histogram` instead.

# In[ ]:

plt.hist(total_spikes_per_neuron, bins=50, histtype="stepfilled")
plt.xlabel("Total spikes per neuron")
plt.ylabel("Number of neurons");

# Let's see what percentage of neurons have a below-average spike count:

# In[ ]:

mean_spike_count = np.mean(total_spikes_per_neuron)
frac_below_mean = (total_spikes_per_neuron < mean_spike_count).mean()
print(f"{frac_below_mean:2.1%} of neurons are below the mean")

# We can also see this by adding the average spike count to the histogram plot:

# In[ ]:

plt.hist(total_spikes_per_neuron, bins=50, histtype="stepfilled")
plt.xlabel("Total spikes per neuron")
plt.ylabel("Number of neurons")
plt.axvline(mean_spike_count, color="orange", label="Mean neuron")
plt.legend();

# This shows that the majority of neurons are relatively "quiet" compared to the mean, while a small number of neurons are exceptionally "loud": they must have spiked more often to reach a large count.
#
# ### Exercise 1: Comparing mean and median neurons
#
# If the mean neuron is more active than 68% of the population, what does that imply about the relationship between the mean neuron and the median neuron?
#
# *Exercise objective:* Reproduce the plot above, but add the median neuron.
#

# In[ ]:

# To complete the exercise, fill in the missing parts (...) and uncomment the code

median_spike_count = ...  # Hint: Try the function np.median

# plt.hist(..., bins=50, histtype="stepfilled")
# plt.axvline(..., color="limegreen", label="Median neuron")
# plt.axvline(mean_spike_count, color="orange", label="Mean neuron")
# plt.xlabel("Total spikes per neuron")
# plt.ylabel("Number of neurons")
# plt.legend()

# [*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D1_ModelTypes/solutions/W1D1_Tutorial1_Solution_b3411d5d.py)
#
# *Example output:*
#
#
#
#

#
# *Bonus:* The median is the 50th percentile. What about other percentiles? Can you show the interquartile range on the histogram?

# ---
#
# # Section 2: Visualizing neuronal spiking activity

# ## Section 2.1: Getting a subset of the data
#
# Now we'll visualize trains of spikes. Because the recordings are long, we will first define a short time interval and restrict the visualization to only the spikes in this interval. We defined a utility function, `restrict_spike_times`, to do this for you. If you call `help()` on the function, it will tell you a little bit about itself:

# In[ ]:

help(restrict_spike_times)

# In[ ]:

t_interval = (5, 15)  # units are seconds after start of recording
interval_spike_times = restrict_spike_times(spike_times, t_interval)

# Is this a representative interval? What fraction of the total spikes fall in this interval?

# In[ ]:

original_counts = sum([len(spikes) for spikes in spike_times])
interval_counts = sum([len(spikes) for spikes in interval_spike_times])
frac_interval_spikes = interval_counts / original_counts
print(f"{frac_interval_spikes:.2%} of the total spikes are in the interval")

# How does this compare to the ratio between the interval duration and the experiment duration? (What fraction of the total time is in this interval?)
#
# We can approximate the experiment duration by taking the minimum and maximum spike time in the whole dataset. To do that, we "concatenate" all of the neurons into one array and then use `np.ptp` ("peak-to-peak") to get the difference between the maximum and minimum value:

# In[ ]:

spike_times_flat = np.concatenate(spike_times)
experiment_duration = np.ptp(spike_times_flat)
interval_duration = t_interval[1] - t_interval[0]

frac_interval_time = interval_duration / experiment_duration
print(f"{frac_interval_time:.2%} of the total time is in the interval")

# These two values—the fraction of total spikes and the fraction of total time—are similar. This suggests the average spike rate of the neuronal population is not very different in this interval compared to the entire recording.
#
# ## Section 2.2: Plotting spike trains and rasters
#
# Now that we have a representative subset, we're ready to plot the spikes, using the matplotlib `plt.eventplot` function. Let's look at a single neuron first:

# In[ ]:

neuron_idx = 1
plt.eventplot(interval_spike_times[neuron_idx], color=".2")
plt.xlabel("Time (s)")
plt.yticks([]);

# We can also plot multiple neurons. Here are three:

# In[ ]:

neuron_idx = [1, 11, 51]
plt.eventplot(interval_spike_times[neuron_idx], color=".2")
plt.xlabel("Time (s)")
plt.yticks([]);

# This makes a "raster" plot, where the spikes from each neuron appear in a different row.
#
# Plotting a large number of neurons can give you a sense for the characteristics in the population. Let's show every 5th neuron that was recorded:

# In[ ]:

neuron_idx = np.arange(0, len(spike_times), 5)
plt.eventplot(interval_spike_times[neuron_idx], color=".2")
plt.xlabel("Time (s)")
plt.yticks([]);

# *Question*: How does the information in this plot relate to the histogram of total spike counts that you saw above?

# In[ ]:

#@title Video 3: Visualizing activity
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='QGA5FCW7kkA', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video

# ---
#
# # Section 3: Inter-spike intervals and their distributions

# Given the ordered arrays of spike times for each neuron in `spike_times`, which we've just visualized, what can we ask next?
#
# Scientific questions are informed by existing models. So, what knowledge do we already have that can inform questions about this data?
#
# We know that there are physical constraints on neuron spiking. Spiking costs energy, which the neuron's cellular machinery can only obtain at a finite rate. Therefore neurons should have a refractory period: they can only fire as quickly as their metabolic processes can support, and there is a minimum delay between consecutive spikes of the same neuron.
#
# More generally, we can ask "how long does a neuron wait to spike again?" or "what is the longest a neuron will wait?" Can we transform spike times into something else, to address questions like these more directly?
#
# We can consider the inter-spike times (or interspike intervals: ISIs). These are simply the time differences between consecutive spikes of the same neuron.
#
# ### Exercise 2: Plot the distribution of ISIs for a single neuron
#
# *Exercise objective:* make a histogram, like we did for spike counts, to show the distribution of ISIs for one of the neurons in the dataset.
#
# Do this in three steps:
#
# 1. Extract the spike times for one of the neurons
# 2. Compute the ISIs (the amount of time between spikes, or equivalently, the difference between adjacent spike times)
# 3. Plot a histogram with the array of individual ISIs

# In[ ]:

def compute_single_neuron_isis(spike_times, neuron_idx):
  """Compute a vector of ISIs for a single neuron given spike times.

  Args:
    spike_times (list of 1D arrays): Spike time dataset, with the first
      dimension corresponding to different neurons.
    neuron_idx (int): Index of the unit to compute ISIs for.

  Returns:
    isis (1D array): Duration of time between each spike from one neuron.
  """
  #############################################################################
  # Students: Fill in missing code (...) and comment or remove the next line
  raise NotImplementedError("Exercise: compute single neuron ISIs")
  #############################################################################

  # Extract the spike times for the specified neuron
  single_neuron_spikes = ...

  # Compute the ISIs for this set of spikes
  # Hint: the function np.diff computes discrete differences along an array
  isis = ...

  return isis

# Uncomment the following lines when you are ready to test your function
# single_neuron_isis = compute_single_neuron_isis(spike_times, neuron_idx=283)
# plt.hist(single_neuron_isis, bins=50, histtype="stepfilled")
# plt.axvline(single_neuron_isis.mean(), color="orange", label="Mean ISI")
# plt.xlabel("ISI duration (s)")
# plt.ylabel("Number of spikes")
# plt.legend()

# [*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D1_ModelTypes/solutions/W1D1_Tutorial1_Solution_4792dbfa.py)
#
# *Example output:*
#
#
#
#

# ---
#
# In general, the shorter ISIs are predominant, with counts decreasing rapidly (and smoothly, more or less) with increasing ISI. However, counts also rapidly decrease to zero with _decreasing_ ISI, below the maximum of the distribution (8-11 ms). The absence of these very low ISIs agrees with the refractory period hypothesis: the neuron cannot fire quickly enough to populate this region of the ISI distribution.
#
# Check the distributions of some other neurons. To resolve various features of the distributions, you might need to play with the value of `n_bins`. Using too few bins might smooth over interesting details, but if you use too many bins, the random variability will start to dominate.
#
# You might also want to restrict the range to see the shape of the distribution when focusing on relatively short or long ISIs. *Hint:* `plt.hist` takes a `range` argument

# ---
#
# # Section 4: What is the functional form of an ISI distribution?

# In[ ]:

#@title Video 4: ISI distribution
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='DHhM80MOTe8', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video

# The ISI histograms seem to follow continuous, monotonically decreasing functions above their maxima. The function is clearly non-linear. Could it belong to a single family of functions?
#
# To motivate the idea of using a mathematical function to explain physiological phenomena, let's define a few different function forms that we might expect the relationship to follow: exponential, inverse, and linear.

# In[ ]:

def exponential(xs, scale, rate, x0):
  """A simple parametrized exponential function, applied element-wise.

  Args:
    xs (np.ndarray or float): Input(s) to the function.
    scale (float): Linear scaling factor.
    rate (float): Exponential growth (positive) or decay (negative) rate.
    x0 (float): Horizontal offset.

  """
  ys = scale * np.exp(rate * (xs - x0))
  return ys

def inverse(xs, scale, x0):
  """A simple parametrized inverse function (`1/x`), applied element-wise.

  Args:
    xs (np.ndarray or float): Input(s) to the function.
    scale (float): Linear scaling factor.
    x0 (float): Horizontal offset.

  """
  ys = scale / (xs - x0)
  return ys

def linear(xs, slope, y0):
  """A simple linear function, applied element-wise.

  Args:
    xs (np.ndarray or float): Input(s) to the function.
    slope (float): Slope of the line.
    y0 (float): y-intercept of the line.

  """
  ys = slope * xs + y0
  return ys

# ### Interactive Demo: ISI functions explorer
#
# Here is an interactive demo where you can vary the parameters of these functions and see how well the resulting outputs correspond to the data. Adjust the parameters by moving the sliders and see how close you can get the lines to follow the falling curve of the histogram. This will give you a taste of what you're trying to do when you *fit a model* to data.
#
# "Interactive demo" cells have hidden code that defines an interface where you can play with the parameters of some function using sliders. You don't need to worry about how the code works – but you do need to **run the cell** to enable the sliders.
#

# In[ ]:

#@title

#@markdown Be sure to run this cell to enable the demo
# Don't worry about understanding this code! It's to setup an interactive plot.
single_neuron_idx = 283
single_neuron_spikes = spike_times[single_neuron_idx]
single_neuron_isis = np.diff(single_neuron_spikes)

counts, edges = np.histogram(
  single_neuron_isis,
  bins=50,
  range=(0, single_neuron_isis.max())
)

functions = dict(
  exponential=exponential,
  inverse=inverse,
  linear=linear,
)

colors = dict(
  exponential="C1",
  inverse="C2",
  linear="C4",
)

@widgets.interact(
  exp_scale=widgets.FloatSlider(1000, min=0, max=20000, step=250),
  exp_rate=widgets.FloatSlider(-10, min=-200, max=50, step=1),
  exp_x0=widgets.FloatSlider(0.1, min=-0.5, max=0.5, step=0.005),
  inv_scale=widgets.FloatSlider(1000, min=0, max=3e2, step=10),
  inv_x0=widgets.FloatSlider(0, min=-0.2, max=0.2, step=0.01),
  lin_slope=widgets.FloatSlider(-1e5, min=-6e5, max=1e5, step=10000),
  lin_y0=widgets.FloatSlider(10000, min=0, max=4e4, step=1000),
)
def fit_plot(
  exp_scale=1000, exp_rate=-10, exp_x0=0.1,
  inv_scale=1000, inv_x0=0,
  lin_slope=-1e5, lin_y0=2000,
):
  """Helper function for plotting function fits with interactive sliders."""
  func_params = dict(
    exponential=(exp_scale, exp_rate, exp_x0),
    inverse=(inv_scale, inv_x0),
    linear=(lin_slope, lin_y0),
  )
  f, ax = plt.subplots()
  ax.fill_between(edges[:-1], counts, step="post", alpha=.5)
  xs = np.linspace(1e-10, edges.max())
  for name, function in functions.items():
    ys = function(xs, *func_params[name])
    ax.plot(xs, ys, lw=3, color=colors[name], label=name);
  ax.set(
      xlim=(edges.min(), edges.max()),
      ylim=(0, counts.max() * 1.1),
      xlabel="ISI (s)",
      ylabel="Number of spikes",
  )
  ax.legend()

# In[ ]:

#@title Video 5: Fitting models by hand
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='uW2HDk_4-wk', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video

# # Summary
#
# In this tutorial, we loaded some neural data and poked at it to understand how the dataset is organized. Then we made some basic plots to visualize (1) the average level of activity across the population and (2) the distribution of ISIs for an individual neuron. In the very last bit, we started to think about using mathematical formalisms to understand or explain some physiological phenomenon. All of this only allowed us to understand "What" the data looks like.
#
# This is the first step towards developing models that can tell us something about the brain. That's what we'll focus on in the next two tutorials.


# python src/chapter7/chapter7note.py
# python3 src/chapter7/chapter7note.py
'''
Class Chapter7_1

Class Chapter7_2

Class Chapter7_3

Class Chapter7_4
'''

from __future__ import division, absolute_import, print_function

import sys as _sys
import math as _math
import random as _random
import time as _time
from random import randint as _randint
from copy import copy as _copy, deepcopy as _deepcopy
from numpy import arange as _arange

if __name__ == '__main__':
    import quicksort
    import stooge
else:
    from . import quicksort
    from . import stooge

class Chapter7_1:
    def note(self):
        '''
        Summary
        ====
        Print chapter7.1 note

        Example
        ====
        >>> Chapter7_1().note()
        '''
        print('chapter7.1 note as follow')
        print('第7章 快速排序')
        print('快速排序是一种排序算法，对包含n个数的输入数组进行排序，最坏情况的运行时间为Θ(n^2)')
        print('虽然这个最坏情况运行时间比较差，但是快速排序通常是用于排序最佳的实用选择，这是因为其平均性能相当好')
        print('快速排序期望的运行时间为Θ(nlgn),且Θ(nlgn)记号中隐含的常数因子很小')
        print('快速排序能够进行就地排序，在虚存坏境中也能很好地工作')
        print('7.1 快速排序的描述')
        print('像合并排序一样，快速排序也是基于分治模式的')
        print(' 1.分解:数组A[p..r]被划分成两个(可能为空的)子数组A[p..q-1]和A[q+1..r]')
        print('  使得A[p..q-1]中的每个元素都小于等于A(q),而且，小于等于A[q+1..r]')
        print('  下标q也在这个划分过程中进行计算')
        print(' 2.解决:通过递归调用快速排序，对子数组A[p..q-1]和A[q+1..r]排序')
        print(' 3.合并:因为这两个子数组是就地排序的(不开辟新的数组),将他们合并不需要任何操作，整个数组A[p..r]已经排好序')
        print('子数组快速排序伪代码')
        print('QUICKSORT(A,p,r)')
        print(' 1. if q < r')
        print(' 2.   q <- PARTITION(A,p,r)')
        print(' 3.       QUICKSORT(A,p,q-1)')
        print(' 3.       QUICKSORT(A,q+1,r)')
        print('排序完整的数组A，调用QUICKSORT(A,0,len(A))即可')
        print('快速排序算法的关键是PARTITION过程，它对子数组A[q..r]进行就地重排')
        print('PARTITION(A,p,r)')
        print(' 1. x <- A[r]')
        print(' 2. i <- p-1')
        print(' 3. for j <- p to r-1')
        print(' 4.  if A[j] <= x')
        print(' 5.      i <- i+1')
        print(' 6.      exchange A[i] <-> A[j]')
        print(' 7. exchange A[i+1] <-> A[r]')
        print(' 8. return i + 1')
        A = [8, 9, 6, 7, 4, 5, 2, 3, 1]
        print('数组A', _deepcopy(A), '的快速排序过程为:',
            quicksort.quicksort(A))
        A = [13, 19, 9, 5, 12, 8, 7, 4, 11, 2, 6, 21]
        print('练习7.1-1 数组A', _deepcopy(A),
            '的一步partition过程得到middle索引为：',
            quicksort.partition(A, 0, len(A) - 1))
        A = [11, 11, 11, 11, 11]
        print('练习7.1-2 数组A', _deepcopy(A),
            '的一步partition过程得到middle索引为：',
            quicksort.partition(A, 0, len(A) - 1))
        print('练习7.1-3 就一个长度为n的for循环，且一定会执行，所以时间复杂度为Θ(n)，然后用确界的夹逼定义证明')
        print('练习7.1-4 不等号方向改变即可')
        # python src/chapter7/chapter7note.py
        # python3 src/chapter7/chapter7note.py

class Chapter7_2:
    def note(self):
        '''
        Summary
        ====
        Print chapter7.2 note

        Example
        ====
        >>> Chapter7_2().note()
        '''
        print('chapter7.2 note as follow')
        print('7.2 快速排序的性能')
        print('快速排序的运行时间与划分是否对称有关，而后者由与选择了哪个元素来进行划分有关')
        print('如果划分是对称的，那么快速排序算法从渐进上与合并算法一样快，否则就和插入排序一样慢')
        print('快速情况的最坏情况划分行为发生在划分过程产生的两个区域分别包含n-1个元素和1个0元素的时候')
        print('假设每次划分都出现了这种不对称划分，划分的时间代价为Θ(n),故算法的运行时间可以递归地写为')
        print('T(n)=T(n-1)+T(0)+Θ(n),递归式的解为T(n)=Θ(n^2)')
        print('快速排序的最坏情况并不比插入排序的最坏情况更好')
        print('另外，当一个已经排序好时，快速排序运行时间Θ(n^2)，插入排序运行时间Θ(n)')
        print('快速排序最佳情况是是其中一个字问题的大小为[n/2],另一个问题的大小为[n/2]-1')
        print('在这种情况下，快速排序的运行时间要快的多，T(n)<=T(n/2)+Θ(n)')
        print('根据主定理，以上递归式的解为O(nlgn)')
        print('平衡的划分')
        print('快速排序的平均运行时间与其最佳运行时间很接近')
        print('练习7.2-1 T(1)=T(0)+Θ(n),T(2)=T(1)+Θ(n),T(n)=T(n-1)+Θ(n)')
        print(' 从第一个式子加到第n个式子，T(n)=Θ(n^2)')
        print('练习7.2-2 数组A中的每个元素都相同时，也属于元素已经排序好的情况，',
            '所以调用PARTITION子程序每次都会得到最差的分配，所以最坏情况运行时间为T(n)=Θ(n^2)')
        print('练习7.2-3 根据书中的描写降序排序好的元素，会导致每次分配都得到最差的情况，又递归式和主定理得T(n)=Θ(n^2)')
        print('练习7.2-4 对已经排序好的支票对于快速排序来说属于最差情况输入，运行时间为O(n^2)')
        print(' 而对于插入排序来说却是最优输入，运行时间为O(n)')
        print('练习7.2-5 [滑稽]但是平均的运行时间仍然为O(nlgn)')
        print('练习7.2-6 略')
        # python src/chapter7/chapter7note.py
        # python3 src/chapter7/chapter7note.py

class Chapter7_3:
    '''
    chapter7.3 content : note, function, etc..

    See Also
    ========
    Chapter7_1 Chapter7_2 Chapter7_4
    '''
    def note(self):
        '''
        Summary
        ====
        Print chapter7.3 note

        Example
        ====
        >>> Chapter7_3().note()
        '''
        print('chapter7.3 note as follow')
        print('7.3 快速排序的随机化版本')
        print('在探讨快速排序的平均性态过程中，假定输入数据的所有排列都是等可能的')
        print('但在工程中，这个假设就不会总是成立')
        print('虽然第五章介绍过一些随机算法，但是如果采用一种不同的，称为随机取样的随机化技术的话，可以使分析更加简单')
        print('在这种方法中，不是时钟采用A[r]作为主元，而是从子数组A[p..r]中随机选择一个元素')
        print('然后将这个随机元素与A[r]交换作为主元')
        print('因为主元元素是随机选择的，在期望的平均情况下，对输入数组的划分比较对称')
        A = [8, 7, 6, 5, 4, 3, 2, 1]
        print('数组[8, 7, 6, 5, 4, 3, 2, 1]的随机化快速排序：',
            quicksort.randomized_quicksort(A))
        print('练习7.3-1:大部分时候输入的待排序序列我们是不知道的，而对于快速排序来讲，一个平均的输入才能反映其算法性能，最坏情况出现的概率比较小')
        print('练习7.3-2:最佳情况调用Θ(n)次，最坏情况调用Θ(n^2)次')
        # python src/chapter7/chapter7note.py
        # python3 src/chapter7/chapter7note.py

class Chapter7_4:
    '''
    chapter7.4 content : note, function, etc..

    See Also
    ========
    Chapter7_1 Chapter7_2 Chapter7_3
    '''
    def note(self):
        '''
        Summary
        ====
        Print chapter7.4 note

        Example
        ====
        ``Chapter7_4().note()``
        '''
        print('chapter7.4 note as follow')
        print('7.4 快速排序分析')
        print('7.4.1 最坏情况分析')
        print('如果快速排序中每一层递归上所做的都是最坏情况划分，则运行时间为Θ(n^2)')
        print('7.4.2 期望的运行时间')
        print('RANDOMZIED-QUICKSORT的平均情况运行时间为O(nlgn)')
        print('运行时间和比较')
        print('quicksort的运行时间是由花在过程PARTITION上的时间所决定的。')
        print('每当PARTITION过程被调用时，就要选出一个主元元素，后续对QUICKSORT和PARTITION的各次递归调用中，都不会包含该元素')
        print('于是，在快速排序算法的整个执行过程中，最多只可能调用PARTITION过程n次，调用一次PARTITION的时间为O(1)在加上一段时间')
        print('引理7.1 设当QUICKSORT在一个包含n个元素的数组上运行时，PARTITION在第四行所做的比较次数为X,那么QUICKSORT的运行时间为O(n+X)')
        print('练习7.4-1 递归式子T(n)=max(T(q)+T(n-q-1)+Θ(n))中，T(n)=Ω(n^2)')
        print('练习7.4-2 快速排序的最佳情况运行时间为Ω(nlgn)')
        print('练习7.4-3 略')
        print('练习7.4-4 RANDOMIZED-QUICKSORT算法期望的运行时间为Ω(nlgn)')
        print('练习7.4-5 对插入排序来说，当其输入已经是几乎排好序的，运行时间是很快的')
        print(' 当在一个长度小于k的子数组上调用快速排序时，让它不做任何排序就返回。',
            '当顶层的快速排序调用返回后，对整个数组运行插入排序来完成排序过程。',
            '这一排序算法的期望运行时间为O(nk+nlg(n/k))')
        print('练习7.4-6 PARTITION过程做这样的修改，从数组A中随机地选出三个元素，并围绕这三个数的中数(即这三个元素的中间值)进行划分',
            '求出以a的函数形式表示的、最坏情况中a:(1-a)划分的近似概率')
        A = [13, 19, 9, 5, 12, 8, 7, 4, 11, 2, 6, 21]
        print('思考题7-1：数组A', _deepcopy(A), '的HOARE-PARTITION算法过程为:',
            quicksort.hoare_partition(A, 0, len(A) - 1))
        print('数组A', _deepcopy(A), '的HOARE-QUICKSORT的过程为：', quicksort.hoare_quicksort(A))
        print('思考题7-2:对快速排序算法的另一种分析')
        print(' 着重关注每一次QUICKSORT递归调用的期望运行时间，而不是执行的比较次数')
        print(' a) 给定一个大小为n的数组，任何特定元素被选为主元的概率为1/n')
        print('思考题7-3 Stooge排序')
        A = [8, 7, 56, 43, 21]
        print('数组A', _deepcopy(A), '的Stooge排序结果为:', stooge.stoogesort(A), A)
        print('思考题7-4 快速排序的堆栈深度')
        print(' 7.1中的快速排序算法包含有两个对其自身的递归调用,但是第二个递归不是必须的')
        A = [8, 7, 56, 43, 21]
        print('数组A', _deepcopy(A), '的尾递归快速排序结果为:', quicksort.morequicksort(A))
        print('思考题7-5 \"三数取中\"划分 也就是主元素RANDOMIZED-QUICKSORT的RANDOMIZED-PARTITION过程')
        print(' 三数取中方法仅仅影响其运行时间Ω(nlgn)中的常数因子')
        print('思考题7-6 对区间的模糊排序:算法的目标是对这些区间进行模糊排序')
        print('模糊排序算法的期望运行时间为Θ(nlgn),但当所有区间都重叠时，期望的运行时间为Θ(n)')
        # python src/chapter7/chapter7note.py
        # python3 src/chapter7/chapter7note.py

chapter7_1 = Chapter7_1()
chapter7_2 = Chapter7_2()
chapter7_3 = Chapter7_3()
chapter7_4 = Chapter7_4()

def printchapter7note():
    '''
    print chapter7 note.
    '''
    print('Run main : single chapter seven!')
    chapter7_1.note()
    chapter7_2.note()
    chapter7_3.note()
    chapter7_4.note()

# python src/chapter7/chapter7note.py
# python3 src/chapter7/chapter7note.py
if __name__ == '__main__':
    printchapter7note()
else:
    pass


# coding=utf-8
import numpy as np
import torch
import os
import torch.nn.functional as F
import cv2
import random

def class_weights():
    """
    COCO train2014每个样本类的频率
    是用于处理样本不均衡.
    “样本偏斜是指数据集中正负类样本数量不均，比如正类样本有10000个，负类样本只有100个，
    这就可能使得超平面被“推向”负类（因为负类数量少，分布得不够广），影响结果的准确性。
    """
    weights = 1 / torch.FloatTensor(
        [187437, 4955, 30920, 6033, 3838, 4332, 3160, 7051, 7677, 9167, 1316, 1372, 833, 6757, 7355, 3302, 3776, 4671,
         6769, 5706, 3908, 903, 3686, 3596, 6200, 7920, 8779, 4505, 4272, 1862, 4698, 1962, 4403, 6659, 2402, 2689,
         4012, 4175, 3411, 17048, 5637, 14553, 3923, 5539, 4289, 10084, 7018, 4314, 3099, 4638, 4939, 5543, 2038, 4004,
         5053, 4578, 27292, 4113, 5931, 2905, 11174, 2873, 4036, 3415, 1517, 4122, 1980, 4464, 1190, 2302, 156, 3933,
         1877, 17630, 4337, 4624, 1075, 3468, 135, 1380])
    weights /= weights.sum()
    # tensor([1.4458e-04, 5.4690e-03, 8.7642e-04, 4.4918e-03, 7.0606e-03, 6.2555e-03,
    #         8.5756e-03, 3.8433e-03, 3.5299e-03, 2.9561e-03, 2.0592e-02, 1.9751e-02,
    #         3.2532e-02, 4.0105e-03, 3.6844e-03, 8.2068e-03, 7.1766e-03, 5.8015e-03,
    #         4.0034e-03, 4.7492e-03, 6.9342e-03, 3.0010e-02, 7.3518e-03, 7.5358e-03,
    #         4.3708e-03, 3.4216e-03, 3.0868e-03, 6.0153e-03, 6.3433e-03, 1.4554e-02,
    #         5.7682e-03, 1.3812e-02, 6.1546e-03, 4.0695e-03, 1.1282e-02, 1.0078e-02,
    #         6.7544e-03, 6.4907e-03, 7.9445e-03, 1.5896e-03, 4.8073e-03, 1.8621e-03,
    #         6.9077e-03, 4.8924e-03, 6.3182e-03, 2.6873e-03, 3.8613e-03, 6.2816e-03,
    #         8.7444e-03, 5.8428e-03, 5.4867e-03, 4.8888e-03, 1.3297e-02, 6.7679e-03,
    #         5.3629e-03, 5.9193e-03, 9.9292e-04, 6.5886e-03, 4.5690e-03, 9.3283e-03,
    #         2.4252e-03, 9.4322e-03, 6.7143e-03, 7.9352e-03, 1.7863e-02, 6.5742e-03,
    #         1.3686e-02, 6.0705e-03, 2.2772e-02, 1.1772e-02, 1.7371e-01, 6.8901e-03,
    #         1.4437e-02, 1.5371e-03, 6.2483e-03, 5.8605e-03, 2.5208e-02, 7.8139e-03,
    #         2.0073e-01, 1.9637e-02])

    return weights

def xyxy2xywh(x):  # Convert bounding box format from [x1, y1, x2, y2] to [x, y, w, h]
    y = torch.zeros(x.shape) if x.dtype is torch.float32 else np.zeros(x.shape)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2
    y[:, 2] = x[:, 2] - x[:, 0]
    y[:, 3] = x[:, 3] - x[:, 1]
    return y

def xywh2xyxy(x):  # Convert bounding box format from [x, y, w, h] to [x1, y1, x2, y2]
    y = torch.zeros(x.shape) if x.dtype is torch.float32 else np.zeros(x.shape)
    y[:, 0] = (x[:, 0] - x[:, 2] / 2)
    y[:, 1] = (x[:, 1] - x[:, 3] / 2)
    y[:, 2] = (x[:, 0] + x[:, 2] / 2)
    y[:, 3] = (x[:, 1] + x[:, 3] / 2)
    return y

def bbox_iou(box1, box2, x1y1x2y2=True):
    """
    Returns the IoU of two bounding boxes
    """
    if x1y1x2y2:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
    else:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2

    # get the coordinates of the intersection rectangle
    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)
    # Intersection area
    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, 0) * torch.clamp(inter_rect_y2 - inter_rect_y1, 0)
    # Union Area
    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)

    return inter_area / (b1_area + b2_area - inter_area + 1e-16)

def build_targets(pred_boxes, pred_conf, pred_cls, target, anchor_wh, nA, nC, nG, batch_report):
    """return tx, ty, tw, th, tconf, tcls, nCorrect, nT:number of targets """
    nB = len(target)  # number of images in batch
    nT = [len(x) for x in target]  # targets per image
    tx = torch.zeros(nB, nA, nG, nG)  # nB:batch size(4)
    ty = torch.zeros(nB, nA, nG, nG)  # nA:number of anchors(3),
    tw = torch.zeros(nB, nA, nG, nG)  # nG:number of grid points(13) = img_dim/stride
    th = torch.zeros(nB, nA, nG, nG)
    tconf = torch.ByteTensor(nB, nA, nG, nG).fill_(0)  # 在函数后面加 _  是改变自身的意思
    tcls = torch.ByteTensor(nB, nA, nG, nG, nC).fill_(0)  # nC = number of classes
    TP = torch.ByteTensor(nB, max(nT)).fill_(0)
    FP = torch.ByteTensor(nB, max(nT)).fill_(0)
    FN = torch.ByteTensor(nB, max(nT)).fill_(0)

    TC = torch.ShortTensor(nB, max(nT)).fill_(-1)  # target category  目标类别

    for b in range(nB):
        nTb = nT[b]  # number of targets
        if nTb == 0:
            continue
        t = target[b]
        if batch_report:
            FN[b, :nTb] = 1

        # 转换为相对于框的位置
        TC[b, :nTb], gx, gy, gw, gh = t[:, 0].long(), t[:, 1] * nG, t[:, 2] * nG, t[:, 3] * nG, t[:, 4] * nG
        # 获取网格框索引并防止溢出（即13个锚点上的13.01）
        '''
        clamp表示夹紧，夹住的意思，torch.clamp(input,min,max,out=None)-> Tensor
        将input中的元素限制在[min,max]范围内并返回一个Tensor
        '''
        gi = torch.clamp(gx.long(), min=0, max=nG - 1)
        gj = torch.clamp(gy.long(), min=0, max=nG - 1)

        # iou of targets-anchors (using wh only)
        box1 = t[:, 3:5] * nG
        # box2 = anchor_grid_wh[:, gj, gi]
        box2 = anchor_wh.unsqueeze(1).repeat(1, nTb, 1)

        # torch.prod(input): 返回所有元素的乘积
        inter_area = torch.min(box1, box2).prod(2)
        iou_anch = inter_area / (gw * gh + box2.prod(2) - inter_area + 1e-16)

        # Sekect best iou_pred and anchor
        iou_anch_best, a = iou_anch.max(0)  # best anchor [0-2] for each target

        # Select best unique target-anchor combinations
        if nTb > 1:
            iou_order = np.argsort(-iou_anch_best)  # best to worst

            # Unique anchor selection(slower but retains original order)
            u = torch.cat((gi, gj, a), 0).view(3, -1).numpy()
            _, first_unique = np.unique(u[:, iou_order], axis=1, return_index=True)  # 第一个独特的指数

            i = iou_order[first_unique]
            # 最佳anchor必须与目标共享重要的共性（iou）
            i = i[iou_anch_best[i] > 0.10]
            if len(i) == 0:
                continue

            a, gj, gi, t = a[i], gj[i], gi[i], t[i]
            if len(t.shape) == 1:
                t = t.view(1, 5)
        else:
            if iou_anch_best < 0.10:
                continue
            i = 0

        tc, gx, gy, gw, gh = t[:, 0].long(), t[:, 1] * nG, t[:, 2] * nG, t[:, 3] * nG, t[:, 4] * nG

        # Coordinates  坐标
        # b : number of images in batch
        # a : anchor
        tx[b, a, gj, gi] = gx - gi.float()
        ty[b, a, gj, gi] = gy - gj.float()

        # Width and height(yolo method)
        tw[b, a, gj, gi] = torch.log(gw / anchor_wh[a, 0])
        th[b, a, gj, gi] = torch.log(gh / anchor_wh[a, 1])

        # One-hot encoding of label
        tcls[b, a, gj, gi, tc] = 1
        tconf[b, a, gj, gi] = 1

        if batch_report:
            # predicted classes and confidence
            tb = torch.cat((gx - gw / 2, gy - gh / 2, gx + gw / 2, gy + gh / 2)).view(4, -1).t()  # target boxes
            pcls = torch.argmax(pred_cls[b, a, gj, gi], 1).cpu()
            pconf = torch.sigmoid(pred_conf[b, a, gj, gi]).cpu()
            iou_pred = bbox_iou(tb, pred_boxes[b, a, gj, gi].cpu())

            TP[b, i] = (pconf > 0.5) & (iou_pred > 0.5) & (pcls == tc)
            FP[b, i] = (pconf > 0.5) & (TP[b, i] == 0)
            FN[b, i] = pconf <= 0.5
    return tx, ty, tw, th, tconf, tcls, TP, FP, FN, TC

def model_info(model):  # Plots a line-by-line description of a PyTorch model
    n_p = sum(x.numel() for x in model.parameters())  # number parameters
    n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients
    print('\n%5s %50s %9s %12s %20s %12s %12s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
    for i, (name, p) in enumerate(model.named_parameters()):
        name = name.replace('module_list.', '')
        print('%5g %50s %9s %12g %20s %12.3g %12.3g' % (
            i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
    print('Model Summary: %g layers, %g parameters, %g gradients\n' % (i + 1, n_p, n_g))

def load_classes(path):
    """
    Loads class labels at 'path'
    """
    fp = open(path, 'r')
    names = fp.read().split('\n')[:-1]
    return names

def ap_per_class(tp, conf, pred_cls, target_cls):
    """ Compute the average precision, given the recall and precision curves.
    Method originally from https://github.com/rafaelpadilla/Object-Detection-Metrics.
    # Arguments
        tp:    True positives (list).
        conf:  Objectness value from 0-1 (list).
        pred_cls: Predicted object classes (list).
        target_cls: True object classes (list).
    # Returns
        The average precision as computed in py-faster-rcnn.
    """

    # lists/pytorch to numpy
    tp, conf, pred_cls, target_cls = np.array(tp), np.array(conf), np.array(pred_cls), np.array(target_cls)

    # Sort by objectness
    i = np.argsort(-conf)
    tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]

    # Find unique classes
    unique_classes = np.unique(np.concatenate((pred_cls, target_cls), 0))

    # Create Precision-Recall curve and compute AP for each class
    ap, p, r = [], [], []
    for c in unique_classes:
        i = pred_cls == c
        n_gt = sum(target_cls == c)  # Number of ground truth objects
        n_p = sum(i)  # Number of predicted objects

        if (n_p == 0) and (n_gt == 0):
            continue
        elif (n_p == 0) or (n_gt == 0):
            ap.append(0)
            r.append(0)
            p.append(0)
        else:
            # Accumulate FPs and TPs
            fpc = np.cumsum(1 - tp[i])
            tpc = np.cumsum(tp[i])

            # Recall
            recall_curve = tpc / (n_gt + 1e-16)
            r.append(tpc[-1] / (n_gt + 1e-16))

            # Precision
            precision_curve = tpc / (tpc + fpc)
            p.append(tpc[-1] / (tpc[-1] + fpc[-1]))

            # AP from recall-precision curve
            ap.append(compute_ap(recall_curve, precision_curve))

    return np.array(ap), unique_classes.astype('int32'), np.array(r), np.array(p)

def compute_ap(recall, precision):
    """ Compute the average precision, given the recall and precision curves.
    Code originally from https://github.com/rbgirshick/py-faster-rcnn.
    # Arguments
        recall:    The recall curve (list).
        precision: The precision curve (list).
    # Returns
        The average precision as computed in py-faster-rcnn.
    """
    # correct AP calculation
    # first append sentinel values at the end

    mrec = np.concatenate(([0.], recall, [1.]))
    mpre = np.concatenate(([0.], precision, [0.]))

    # compute the precision envelope
    for i in range(mpre.size - 1, 0, -1):
        mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

    # to calculate area under PR curve, look for points
    # where X axis (recall) changes value
    i = np.where(mrec[1:] != mrec[:-1])[0]

    # and sum (\Delta recall) * prec
    ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap

def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
    """
    Removes detections with lower object confidence score than 'conf_thres' and performs
    Non-Maximum Suppression to further filter detections.
    Returns detections with shape:
        (x1, y1, x2, y2, object_conf, class_score, class_pred)
    """

    output = [None for _ in range(len(prediction))]
    for image_i, pred in enumerate(prediction):
        # Filter out confidence scores below threshold
        # Get score and class with highest confidence

        # cross-class NMS (experimental)
        cross_class_nms = False
        if cross_class_nms:
            # thresh = 0.85
            thresh = nms_thres
            a = pred.clone()
            _, indices = torch.sort(-a[:, 4], 0)  # sort best to worst
            a = a[indices]
            radius = 30  # area to search for cross-class ious
            for i in range(len(a)):
                if i >= len(a) - 1:
                    break

                close = (torch.abs(a[i, 0] - a[i + 1:, 0]) < radius) & (torch.abs(a[i, 1] - a[i + 1:, 1]) < radius)
                close = close.nonzero()

                if len(close) > 0:
                    close = close + i + 1
                    iou = bbox_iou(a[i:i + 1, :4], a[close.squeeze(), :4].reshape(-1, 4), x1y1x2y2=False)
                    bad = close[iou > thresh]

                    if len(bad) > 0:
                        mask = torch.ones(len(a)).type(torch.ByteTensor)
                        mask[bad] = 0
                        a = a[mask]
            pred = a

        x, y, w, h = pred[:, 0], pred[:, 1], pred[:, 2], pred[:, 3]
        a = w * h  # area
        ar = w / (h + 1e-16)  # aspect ratio

        log_w, log_h, log_a, log_ar = torch.log(w), torch.log(h), torch.log(a), torch.log(ar)

        # n = len(w)
        # shape_likelihood = np.zeros((n, 60), dtype=np.float32)
        # x = np.concatenate((log_w.reshape(-1, 1), log_h.reshape(-1, 1)), 1)
        # from scipy.stats import multivariate_normal
        # for c in range(60):
        # shape_likelihood[:, c] = multivariate_normal.pdf(x, mean=mat['class_mu'][c, :2], cov=mat['class_cov'][c, :2, :2])

        class_prob, class_pred = torch.max(F.softmax(pred[:, 5:], 1), 1)

        v = ((pred[:, 4] > conf_thres) & (class_prob > .3))
        v = v.nonzero().squeeze()
        if len(v.shape) == 0:
            v = v.unsqueeze(0)

        pred = pred[v]
        class_prob = class_prob[v]
        class_pred = class_pred[v]

        # If none are remaining => process next image
        nP = pred.shape[0]
        if not nP:
            continue

        # From (center x, center y, width, height) to (x1, y1, x2, y2)
        box_corner = pred.new(nP, 4)
        xy = pred[:, 0:2]
        wh = pred[:, 2:4] / 2
        box_corner[:, 0:2] = xy - wh
        box_corner[:, 2:4] = xy + wh
        pred[:, :4] = box_corner

        # Detections ordered as (x1, y1, x2, y2, obj_conf, class_prob, class_pred)
        detections = torch.cat((pred[:, :5], class_prob.float().unsqueeze(1), class_pred.float().unsqueeze(1)), 1)
        # Iterate through all predicted classes
        unique_labels = detections[:, -1].cpu().unique()
        if prediction.is_cuda:
            unique_labels = unique_labels.cuda(prediction.device)

        nms_style = 'OR'  # 'AND' or 'OR' (classical)
        for c in unique_labels:
            # Get the detections with the particular class
            detections_class = detections[detections[:, -1] == c]
            # Sort the detections by maximum objectness confidence
            _, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
            detections_class = detections_class[conf_sort_index]
            # Perform non-maximum suppression
            max_detections = []

            if nms_style == 'OR':  # Classical NMS
                while detections_class.shape[0]:
                    # Get detection with highest confidence and save as max detection
                    max_detections.append(detections_class[0].unsqueeze(0))
                    # Stop if we're at the last detection
                    if len(detections_class) == 1:
                        break
                    # Get the IOUs for all boxes with lower confidence
                    ious = bbox_iou(max_detections[-1], detections_class[1:])

                    # Remove detections with IoU >= NMS threshold
                    detections_class = detections_class[1:][ious < nms_thres]

            elif nms_style == 'AND':  # 'AND'-style NMS, at least two boxes must share commonality to pass, single boxes erased
                while detections_class.shape[0]:
                    if len(detections_class) == 1:
                        break

                    ious = bbox_iou(detections_class[:1], detections_class[1:])

                    if ious.max() > 0.5:
                        max_detections.append(detections_class[0].unsqueeze(0))

                    # Remove detections with IoU >= NMS threshold
                    detections_class = detections_class[1:][ious < nms_thres]

            if len(max_detections) > 0:
                max_detections = torch.cat(max_detections).data
                # Add max detections to outputs
                output[image_i] = max_detections if output[image_i] is None else torch.cat(
                    (output[image_i], max_detections))

    return output

def write_cfg(cfgfile, cfg, precent):
    with open(cfgfile, 'r') as f:
        lines = f.read().split('\n')  # store the lines in a list
        lines = [x for x in lines if len(x) > 0]  # get read of the empty lines
        lines = [x for x in lines if x[0] != '#']  # get rid of comments
        # lines = [x.rstrip().lstrip() for x in lines]  # get rid of fringe whitespaces\

    block = {}
    blocks = []
    # D:/yolotest/cfg/yolov3.cfg
    # prunedcfg = os.path.join('./'.join(cfgfile.split("/")[0:-1]), "prune_" + cfgfile.split("/")[-1])
    if not os.path.exists('sparsity_2_prune_cfg'):
        os.mkdir('sparsity_2_prune_cfg')
    prunedcfg = os.path.join("sparsity_2_prune_cfg/prune_{}_".format(precent) + cfgfile.split("/")[-1])
    for line in lines:
        if line[0] == "[":  # This marks the start of a new block
            if len(block) != 0:  # If block is not empty, implies it is storing values of previous block.
                blocks.append(block)  # add it the blocks list
                block = {}  # re-init the block
            block["type"] = line[1:-1].rstrip()
        else:
            key, value = line.split("=")
            block[key.rstrip()] = value.lstrip()
    blocks.append(block)
    x = 0
    # print(blocks[1])
    for block in blocks:
        if 'batch_normalize' in block:
            block['filters'] = cfg[x]
            x = x + 1
    ##
    with open(prunedcfg, 'w') as f:
        for block in blocks:
            for i in block:
                if i == "type":
                    f.write('\n')
                    f.write("[" + block[i] + "]\n")
                    for j in block:
                        if j != "type":
                            f.write(j + "=" + str(block[j]) + '\n')
    print('save pruned cfg file in %s' % prunedcfg)
    return prunedcfg

def route_problem(model, ind):
    ds = list(model.children())
    dsas = list(ds[0].children())

    # print('-----------',dsas[90])
    sum1 = 0
    # print(dsas[90].named_children())
    for k in range(ind + 1):
        # print('k:',k)
        for i in dsas[k].named_children():
            # print('i:',i)
            if "_".join(i[0].split("_")[0:-1]) == 'conv_with_bn':
                sum1 = sum1 + 1
    # print(sum1)
    return sum1 - 1

def dontprune(model):
    dontprune = []
    nnlist = model.module_list
    for i in range(len(nnlist)):
        for name in nnlist[i].named_children():
            if name[0].split("_")[0] == 'shortcut':
                if 'conv' in list(nnlist[name[1].froms + i].named_children())[0][0]:
                    dontprune.append(name[1].froms + i)
                else:
                    dontprune.append(name[1].froms + i - 1)
                dontprune.append(i - 1)
    return dontprune

def coco_class_count(path='../coco/labels/train2014/'):
    import glob

    nC = 80  # number classes
    x = np.zeros(nC, dtype='int32')
    files = sorted(glob.glob('%s/*.*' % path))
    for i, file in enumerate(files):
        labels = np.loadtxt(file, dtype=np.float32).reshape(-1, 5)
        x += np.bincount(labels[:, 0].astype('int32'), minlength=nC)
        print(i, len(files))

def plot_results():
    # Plot YOLO training results file 'results.txt'
    import glob
    import numpy as np
    import matplotlib.pyplot as plt
    plt.figure(figsize=(16, 8))
    s = ['X', 'Y', 'Width', 'Height', 'Objectness', 'Classification', 'Total Loss', 'Precision', 'Recall', 'mAP']
    files = sorted(glob.glob('results*.txt'))
    for f in files:
        results = np.loadtxt(f, usecols=[2, 3, 4, 5, 6, 7, 8, 17, 18, 16]).T  # column 16 is mAP
        n = results.shape[1]
        for i in range(10):
            plt.subplot(2, 5, i + 1)
            plt.plot(range(1, n), results[i, 1:], marker='.', label=f)
            plt.title(s[i])
            if i == 0:
                plt.legend()

def plot_one_box(x, img, color=None, label=None, line_thickness=None):  # Plots one bounding box on image img
    tl = line_thickness or round(0.002 * max(img.shape[0:2])) + 1  # line thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1)  # filled
        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

# def model_info(model, report='full'):
#     # Plots a line-by-line description of a PyTorch model
#     n_p = sum(x.numel() for x in model.parameters())  # number parameters
#     n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients
#     if report is 'full':
#         print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
#         for i, (name, p) in enumerate(model.named_parameters()):
#             name = name.replace('module_list.', '')
#             print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
#                   (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
#     print('Model Summary: %g layers, %g parameters, %g gradients' % (len(list(model.parameters())), n_p, n_g))