"""
Defines comparer functions used by FormulaGrader and its subclasses.

Simple Comparer Functions
=========================

A comparer function must have signature
`comparer_func(comparer_params_eval, student_eval, utils)` and should return
True, False, 'partial', or a dictionary with required key 'grade_decimal' and
optional key 'msg'. When `FormulaGrader` (or its subclasses) call your custom
comparer function, `comparer_func`'s argument values are:

- `comparer_params_eval`: The `comparer_params` list, numerically evaluated
  according to variable and function sampling.
- `student_eval`: The student's input, numerically evaluated according to
  variable and function sampling.
- `utils`: A convenience object that may be helpful when writing custom
  comparer functions. It has attributes:

    - `utils.tolerance`: The tolerance specified in grader configuration,
      `0.01%` by default
    - `utils.within_tolerance(x, y)`: checks that `y` is within specified
      tolerance of `x`. Can handle scalars, vectors, and matrices.
      If tolerance was specified as a percentage, then checks that
      `|x-y| < tolerance * x`.

    Comparer functions used inside `MatrixGrader` have the following additional
    `utils` method:

    - `utils.validate_shape(student_eval, shape)`: Checks that `student_eval`
      has specified `shape`, where `shape` is a Numpy shape tuple.

A comparer function must return either:

  - a boolean, or
  - a dictionary with keys:
      - `'grade_decimal'`: number between 0 and 1 (required)
      - `'ok'`: `True` or `False` or `'partial'` (optional, inferred from
        grade_decimal by default)
      - `'msg'`: a feedback message (optional, defaults to `''`)

NOTE: doctests in this module show how the comparer function would be used
      inside a grader

Correlated Comparers
====================
See ./baseclasses.py and ./linear_comparer.py for examples.
"""
from __future__ import print_function, division, absolute_import, unicode_literals

from numbers import Number
import numpy as np

import six
from voluptuous import Schema, Required, Any, Range, All

from mitxgraders.exceptions import InputTypeError, StudentFacingError
from mitxgraders.helpers.validatorfuncs import is_callable, Nullable, text_string
from mitxgraders.helpers.calc.mathfuncs import is_nearly_zero
from mitxgraders.helpers.calc.math_array import are_same_length_vectors, is_vector
from mitxgraders.comparers.baseclasses import Comparer, CorrelatedComparer

def identity_transform(x):
    """
    Returns the input.

    Note: used instead of lambdas because it prints a nice name.
    """
    return x

class EqualityComparer(Comparer):
    """
    This comparer checks for equality between the student and instructor results,
    up to a desired tolerance. If desired, a transforming function can be applied
    before the comparison is carried out.

    comparer_params: ['expect']

    By default, equality_comparer is used in FormulaGrader, NumericalGrader and MatrixGrader.
    >>> from mitxgraders import *
    >>> equality_comparer == EqualityComparer()
    True
    >>> grader = FormulaGrader(
    ...     answers='2*x',
    ...     variables=['x']
    ... )
    >>> grader(None, 'x')['ok']
    False
    >>> grader(None, 'x*2')['ok']
    True

    The following example applies cosine to the expected answer and student input
    before comparison:
    >>> import numpy as np
    >>> grader = FormulaGrader(
    ...     answers={
    ...         'comparer': EqualityComparer(transform=np.cos),
    ...         'comparer_params': ['x']
    ...     },
    ...     variables=['x']
    ... )
    >>> grader(None, 'x')['ok']
    True
    >>> grader(None, '-x')['ok']
    True
    >>> grader(None, 'x + 2*pi')['ok']
    True
    >>> grader(None, 'x + pi')['ok']
    False

    The following example takes the norm of the expected answer and student input
    before comparison. Note the different method of changing the comparer.
    >>> MatrixGrader.set_default_comparer(EqualityComparer(transform=np.linalg.norm))
    >>> grader = MatrixGrader(
    ...     answers='[1, 0, 0]'
    ... )
    >>> grader(None, '[1, 0, 0]')['ok']
    True
    >>> grader(None, '[0, 1, 0]')['ok']
    True
    >>> grader(None, '[1/sqrt(2), 0, 1/sqrt(2)]')['ok']
    True
    >>> MatrixGrader.reset_default_comparer()

    FormulaGrader and MatrixGrader both have set_default_comparer() and
    reset_default_comparer() methods.

    """
    schema_config = Schema({
        Required('transform', default=None): All(
            Nullable(is_callable),
            # if f is None, coerce to identity function
            lambda f: identity_transform if f is None else f
        )
    })

    @staticmethod
    def validate(expected_eval, student_eval, utils):
        if hasattr(utils, 'validate_shape'):
            # in numpy, scalars have empty tuples as their shapes
            shape = tuple() if isinstance(expected_eval, Number) else expected_eval.shape
            utils.validate_shape(student_eval, shape)

    def __call__(self, comparer_params_eval, student_eval, utils):
        expected_eval = comparer_params_eval[0]
        self.validate(expected_eval, student_eval, utils)

        transform = self.config['transform']
        expected_eval = transform(expected_eval)
        student_eval = transform(student_eval)

        return utils.within_tolerance(expected_eval, student_eval)

equality_comparer = EqualityComparer()

class MatrixEntryComparer(CorrelatedComparer):
    """
    Default comparer for MatrixGrader. Compares student and instructor array
    evaluations entry-by-entry for equality. Note that despite the name, this comparer
    works equally well on vectors/matrices/tensors.

    Configuration
    =============
        transform (None | function): same as EqualityComparer (default None)

        entry_partial_credit ('proportional' | number): Determines how partial credit
            is awarded. If set to 'proportional', then credit is proportional to
            the number of correct array entries. If a numeric value betweem 0 and 1
            is provided, this flat rate of partial credit is provided as long as
            some but not all entries are correct. Default is the numeric value 0
            (no partial credit).

        entry_partial_msg (str): A text string message shown when partial credit
            is awarded. The string may optionally contain the formatting key {error_indices},
            which will be replaced with a diagram showing the correct/incorrect entries.
            To show no message, use the the empty string.
            Default value is:
            "Some array entries are incorrect, marked below:\n{error_locations}"
    """

    default_msg = "Some array entries are incorrect, marked below:\n{error_locations}"
    schema_config = EqualityComparer.schema_config.extend({
        Required('entry_partial_credit', default=0): Any(All(float, Range(0, 1)), 0, 1, 'proportional'),
        Required('entry_partial_msg', default=default_msg): text_string
    })

    @staticmethod
    def format_message_with_locations(format_string, locs):
        """
        Returns format_string with {error_locations} replaced by a diagram showing
        correct/incorrect entries.

        Arguments:
            format_string: a string that may contain {error_locations} formatting key.
            locs: a boolean array with False values indicating incorrect entries
        """
        # Not the most elegant way to do these replacements, but this was what
        # I came up with to minimize the amount of extra u prefixes in Python 2

        # These are the edX colors, at least as of July 2019
        bad_str = '\u2717'
        good_str = '\u2713'
        matrix_as_text = six.text_type(locs).replace("  ", " ").replace("[ ", "[")
        matrix_as_text = matrix_as_text.replace("True", good_str).replace("False", bad_str)
        matrix_as_text = matrix_as_text.replace('\n', '
')
        formatted_locs = '
{mat}
'.format(mat=matrix_as_text)
        return format_string.format(error_locations=formatted_locs)

    @staticmethod
    def validate(expected_evals, student_evals, utils):
        for x, y in zip(expected_evals, student_evals):
            EqualityComparer.validate(x, y, utils)

    def __call__(self, comparer_params_evals, student_evals, utils):
        expected_evals = [params[0] for params in comparer_params_evals]
        self.validate(expected_evals, student_evals, utils)

        transform = self.config['transform']
        expected_evals = [transform(x) for x in expected_evals]
        student_evals = [transform(x) for x in student_evals]
        vec_within_tol = np.vectorize(utils.within_tolerance)
        # comparisons_by_eval is a boolean array of entry-by-entry comparisons,
        # one for each comparison. Its numpy shape is (n_evals, *eval_shape)
        comparisons_by_eval = vec_within_tol(expected_evals, student_evals)
        comparisons_summary = np.all(comparisons_by_eval, axis=0)

        num_entries = comparisons_summary.size
        percent_correct = np.sum(comparisons_summary).item()/num_entries
        msg = self.format_message_with_locations(self.config['entry_partial_msg'], comparisons_summary)
        partial_credit = self.config['entry_partial_credit']

        if percent_correct == 1:
            return True
        elif percent_correct == 0:
            return {'ok': False, 'grade_decimal': 0, 'msg': msg}
        elif partial_credit == 'proportional':
            return {'ok': 'partial', 'grade_decimal': percent_correct, 'msg': msg}
        else:
            return {'ok': 'partial', 'grade_decimal': partial_credit, 'msg': msg}

def between_comparer(comparer_params_eval, student_eval, utils):
    """
    Used to check that input is real and between two parameters.

    comparer_params: ['start', 'stop']

    Example:
    >>> from mitxgraders import NumericalGrader
    >>> grader = NumericalGrader(
    ...     answers={
    ...         'comparer': between_comparer,
    ...         'comparer_params': ['1e6', '1e9']
    ...     }
    ... )
    >>> grader(None, '2.5e8')['ok']
    True
    >>> grader(None, '0.001e8')['ok']
    False
    >>> grader(None, '5e7')['ok']
    True

    Input must be real:
    >>> try:
    ...     grader(None, '5e8+2e6*i')['ok']
    ... except InputTypeError as error:
    ...     print(error)
    Input must be real.
    """
    start, stop = comparer_params_eval

    if not np.isreal(student_eval):
        raise InputTypeError("Input must be real.")

    return start <= student_eval <= stop

def congruence_comparer(comparer_params_eval, student_eval, utils):
    """
    Compares the student input to a target, moduli a given modulus.
    Will often set modulus to 2*pi in order to compare angles.

    comparer_params: [target, modulus]

    Example usage:
    >>> from mitxgraders import FormulaGrader
    >>> grader = FormulaGrader(
    ...     answers={
    ...         'comparer': congruence_comparer,
    ...         'comparer_params': [
    ...             'b^2/a', # target
    ...             'c'      # modulus
    ...         ]
    ...     },
    ...     variables=['a', 'b', 'c']
    ... )
    >>> grader(None, 'b^2/a')['ok']
    True
    >>> grader(None, 'b^2/a + 1.5*c')['ok']
    False
    >>> grader(None, 'b^2/a + 2*c  ')['ok']
    True
    """
    expected, modulus = comparer_params_eval

    expected_reduced = expected % modulus
    input_reduced = student_eval % modulus
    return utils.within_tolerance(expected_reduced, input_reduced)

def eigenvector_comparer(comparer_params_eval, student_eval, utils):
    """
    Used to check that a student's answer is an eigenvector of a matrix
    with a given eigenvalue. Ignores scaling of the eigenvector.

    comparer_params: [matrix, eigenvalue]

    Example Usage:
    >>> from mitxgraders import MatrixGrader
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[[1, x], [x, -1]]',    # matrix
    ...             'sqrt(1+x^2)'           # eigenvalue
    ...         ],
    ...         'comparer': eigenvector_comparer
    ...     },
    ...     variables=['x']
    ... )
    >>> grader(None, '[1+sqrt(1+x^2), x]')['ok']
    True
    >>> grader(None, '2*[1+sqrt(1+x^2), x]')['ok']
    True
    >>> grader(None, '[1+sqrt(1+x^2), 1]')['ok']
    False
    >>> grader(None, '[0, 0]') == {
    ...     'ok': False,
    ...     'msg': 'Eigenvectors must be nonzero.',
    ...     'grade_decimal': 0
    ... }
    True

    """

    matrix, eigenvalue = comparer_params_eval

    # matrix is square with shape (n, n); student input should have shape (n, )
    expected_input_shape = (matrix.shape[0], )
    utils.validate_shape(student_eval, expected_input_shape)

    expected = eigenvalue * student_eval
    actual = matrix * student_eval

    if utils.within_tolerance(0, np.linalg.norm(student_eval)):
        return {
            'ok': False,
            'grade_decimal': 0,
            'msg': 'Eigenvectors must be nonzero.'
        }

    return utils.within_tolerance(actual, expected)

def vector_span_comparer(comparer_params_eval, student_eval, utils):
    """
    Check whether student's answer is nonzero and in the span of some given
    vectors.

    comparer_params: A list of vectors

    Usage
    =====

    Use a single vector as comparer_params to test whether student input is
    parallel to a particular vector:
    >>> from mitxgraders import MatrixGrader
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[3, x, 1 + i]',
    ...         ],
    ...         'comparer': vector_span_comparer
    ...     },
    ...     variables=['x'],
    ... )
    >>> grader(None, '[3, x, 1 + i]')['ok']
    True
    >>> grader(None, '[9, 3*x, 3 + 3*i]')['ok']
    True
    >>> grader(None, '[9, 3*x, 3 - 3*i]')['ok']
    False

    Complex scale factors work, too:
    >>> grader(None, '(4 + 2*i)*[3, x, 1 + i]')['ok']
    True

    Student input should be nonzero:
    >>> result = grader(None, '[0, 0, 0]')
    >>> expected = {
    ...     'ok': False,
    ...     'grade_decimal': 0.0,
    ...     'msg': 'Input should be a nonzero vector.'
    ... }
    >>> result == expected
    True

    Input shape is validated:
    >>> try:
    ...     grader(None, '5')
    ... except InputTypeError as error:
    ...     print(error)
    Expected answer to be a vector, but input is a scalar

    Multiple vectors can be provided:
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[1, 1, 0]',    # v0
    ...             '[0, 1, 2]'     # v1
    ...         ],
    ...         'comparer': vector_span_comparer
    ...     },
    ... )

    The vector 2*v0 + 3i*v1 = [2, 2+3i, 6i] is in the span of v0 and v1:
    >>> grader(None, '[2, 2 + 3*i, 6*i]')['ok']
    True

    The comparer_params should be list of equal-length vectors:
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[1, 1, 0]',
    ...             '5'
    ...         ],
    ...         'comparer': vector_span_comparer
    ...     },
    ... )
    >>> try:
    ...     grader(None, '[1, 2, 3]')               # doctest: +ELLIPSIS
    ... except StudentFacingError as error:
    ...     print(error)
    Problem Configuration Error: ...to equal-length vectors
    """

    # Validate the comparer params
    if not are_same_length_vectors(comparer_params_eval):
        raise StudentFacingError('Problem Configuration Error: comparer_params '
            'should be a list of strings that evaluate to equal-length vectors')

    # Validate student input shape
    utils.validate_shape(student_eval, comparer_params_eval[0].shape)

    if utils.within_tolerance(0, np.linalg.norm(student_eval)):
        return {
            'ok': False,
            'grade_decimal': 0,
            'msg': 'Input should be a nonzero vector.'
        }

    # Use ordinary least squares to find an approximation to student_eval
    # that lies within the span of given vectors, then check that the
    # residual-sum is small in comparison to student input.
    column_vectors = np.array(comparer_params_eval).transpose()
    # rcond=-1 uses machine precision for testing singular values
    # In numpy 1.14+, use rcond=None fo this behavior. (we use 1.6)
    ols = np.linalg.lstsq(column_vectors, student_eval, rcond=-1)
    error = np.sqrt(ols[1])

    # Check that error is nearly zero, using student_eval as a reference
    # when tolerance is specified as a percentage
    return is_nearly_zero(error, utils.tolerance, reference=student_eval)

def vector_phase_comparer(comparer_params_eval, student_eval, utils):
    """
    Check that student input equals a given input (to within tolerance), up to
    an overall phase factor.

    comparer_params: [target_vector]

    Usage
    =====

    >>> from mitxgraders import MatrixGrader
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[1, exp(-i*phi)]',
    ...         ],
    ...         'comparer': vector_phase_comparer
    ...     },
    ...     variables=['phi'],
    ... )

    >>> grader(None, '[1, exp(-i*phi)]')['ok']
    True
    >>> grader(None, '[exp(i*phi/2), exp(-i*phi/2)]')['ok']
    True
    >>> grader(None, '[i, exp(i*(pi/2 - phi))]')['ok']
    True

    >>> grader(None, '[1, exp(+i*phi)]')['ok']
    False
    >>> grader(None, '[2, 2*exp(-i*phi)]')['ok']
    False

    The comparer_params should be list with a single vector:
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[1, 1, 0]',
    ...             '[0, 1, 1]'
    ...         ],
    ...         'comparer': vector_phase_comparer
    ...     },
    ... )
    >>> try:
    ...     grader(None, '[1, 2, 3]')               # doctest: +ELLIPSIS
    ... except StudentFacingError as error:
    ...     print(error)
    Problem Configuration Error: ...to a single vector.
    """
    # Validate that author comparer_params evaluate to a single vector
    if not len(comparer_params_eval) == 1 and is_vector(comparer_params_eval[0]):
        raise StudentFacingError('Problem Configuration Error: comparer_params '
            'should be a list of strings that evaluate to a single vector.')

    # We'll check that student input is in the span as target vector and that
    # it has the same magnitude

    in_span = vector_span_comparer(comparer_params_eval, student_eval, utils)

    expected_mag = np.linalg.norm(comparer_params_eval[0])
    student_mag = np.linalg.norm(student_eval)
    same_magnitude = utils.within_tolerance(expected_mag, student_mag)

    return in_span and same_magnitude


#!/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


""" Plot the spectral sequence """

import matplotlib.pyplot as plt
plt.rc("font", family="serif")
plt.rc("text", usetex=True)
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
import numpy as np
from astropy.table import Table
from astropy.cosmology import Planck15
from astropy.time import Time
import glob
import sys
sys.path.append("/Users/annaho/Github/Spectra")
sys.path.append("/Users/annaho/Dropbox/Projects/Research/ZTF18abukavn/code")
from load_lc import get_lc
from normalize import smooth_spec
from measure_snr import get_snr

def get_files():
    files = np.array(glob.glob(
    "/Users/annaho/Dropbox/Projects/Research/ZTF18abukavn/data/spec/ZTF18abukavn/*.ascii"))
    dt = np.zeros(len(files))
    tels = []
    cols = np.array([""]*len(dt), dtype='U10')

    # Read in all of the files, pull out the corresponding dates, and sort by date
    t0 = 2458370.6634 # in JD
    for ii,f in enumerate(files):
        tel = f.split("_")[2]
        tels.append(tel)
        alldat = open(f).readlines()
        if tel == 'LT':
            for line in alldat:
                if 'DATE-OBS' in line:
                    obsdate = line[13:36]
                    t = Time(obsdate, format='isot').jd
                    dt[ii] = t-t0
            cols[ii] = 'magenta'
        elif tel == 'P200':
            for line in alldat:
                if 'UT shutter open' in line:
                    obsdate = line[12:35]
                    t = Time(obsdate, format='isot').jd
                    dt[ii] = t-t0
            cols[ii] = 'lightblue'
        elif tel == 'Keck1':
            for line in alldat:
                if 'DATE_BEG' in line:
                    obsdate = line[13:32]
                    t = Time(obsdate, format='isot').jd
                    dt[ii] = t-t0
            cols[ii] = 'red'
        elif tel == 'DCT':
            obsdate = '2018-09-14T00:00:00' # temporary
            t = Time(obsdate, format='isot').jd
            dt[ii] = t-t0
            cols[ii] = 'yellow'
        elif tel == 'NOT':
            obsdate = '2018-09-17T00:00:00' # temporary
            t = Time(obsdate, format='isot').jd
            dt[ii] = t-t0
            cols[ii] = 'green'
        elif tel == 'P60':
            for line in alldat:
                if 'MJD_OBS' in line:
                    obsdate = float(line[11:25])
                    t = Time(obsdate, format='mjd').jd
                    dt[ii] = t-t0
            cols[ii] = 'black'
        elif tel == 'omr.ascii':
            # first Xinglong spectrum
            t = Time('2018-09-21T11:15:10.0').jd
            dt[ii] = t-t0
        elif tel == 'Bfosc.ascii':
            # second Xinglong spectrum
            t = Time('2018-09-25T11:16:43.0').jd
            dt[ii] = t-t0
        else:
            print("couldn't find telescope")
            print(tel)
    order = np.argsort(dt)
    files_sorted = files[order]
    dt_sorted = dt[order]
    tel_sorted = np.array(tels)[order]
    cols = cols[order]
    return files_sorted, dt_sorted, tel_sorted

def get_res(tel):
    """ Here, this means the width of a line in Angstroms """
    if tel == 'LT':
        res = 18 # Angstrom, res at central wavelength
        res = 30 # add a couple of Ang?
    elif tel == 'P200':
        res = 10 # determined by eye from the spectrum
        # basically, width of a galaxy emission line is 10 AA
        # and each pixel is 1 AA
    elif tel == 'Keck1':
        res = 7*2 # determined by eye from spectrum
        # width of a line is around 7 pixels
        # and each pixel is 2 Angstroms
    elif tel == 'NOT':
        # width of a line is around 8 pixels
        # and each pixel is around 2.63 Ang
        res = 8*2.63
    elif tel == 'DCT':
        # width of a line is around 7 pixels
        # and each pixel is 2.2 Ang
        res = 7*2.2
    elif tel == 'P60':
        res = 20
    elif 'ascii' in tel:
        # Xinglong spectrum
        res = 26
    else:
        res = 1
    return res

def load_spec(f, tel):
    """ load data from spec file """
    dat = np.loadtxt(f)
    wl = dat[:,0]
    flux = dat[:,1]
    if tel == 'Keck':
        eflux = dat[:,3]
    else:
        # need to estimate uncertainty from scatter
        eflux = np.array([get_snr(wl, flux, 6000, 6200)]*len(wl))
    ivar = 1/eflux**2
    return wl, flux, ivar

def plot_spec(ax, x, y, tel, epoch):
    """ plot the spectrum """
    choose_x = np.logical_and(x >= 3200, x<= 9300)
    choose = choose_x
    ax.plot(
            x[choose], y[choose], c='lightgrey',
            drawstyle='steps-mid', lw=0.5, alpha=0.4)
    return ax

def plot_smoothed_spec(ax, x, y, ivar, tel, epoch, ls='-', lw=0.5, c='black', label=None):
    """ plot the smoothed spectrum """
    res = get_res(tel)
    choose_x = np.logical_and(x >= 3200, x<= 9300)
    choose = choose_x
    smoothed = smooth_spec(x, y, ivar, res*3)
    ax.plot(
            x[choose], smoothed[choose], c=c,
            drawstyle='steps-mid', lw=lw, ls=ls, alpha=1.0, label=label)
    dt_str = r"+%s\,d" %str(np.round(epoch, 1))
    ax.text(
            x[choose][-1]+100, smoothed[choose][-1],  s=dt_str,
            horizontalalignment='left', verticalalignment='center',
            fontsize=14)
    return ax

def choose_lines(z, dt):
    """ choose galaxy emission lines given the epoch """
    balmer = np.array([6564.61, 4862.68, 4341.68, 4102.89, 3970.072])
    oiii = np.array([4363, 4932.6, 4960.295, 5008.24]) # O III
    oii = np.array([3727.092, 3729.875])
    nii = np.array([6549.86])
    oi = np.array([6302.046, 6365.536])
    gal_wl = np.hstack((balmer, oiii, oii)) * (z+1)
    return gal_wl

def plot_lines(ax, z, tel, dt):
    """ Plot galaxy emission lines for a particular redshift """
    res = get_res(tel)
    gal_wl = choose_lines(z, dt)
    for val in gal_wl:
        ax.axvspan(
                val-res/2, val+res/2, ls='--', color='grey', lw=0.5, alpha=0.5)

def clip_lines(wl, flux, z, tel, dt):
    res = get_res(tel)
    gal_wl = choose_lines(z, dt)
    for line in gal_wl:
        choose = np.logical_and(wl >= line-res/2, wl <= line+res/2)
        flux = np.interp(wl, wl[~choose], flux[~choose]) # interp over features
    return wl, flux

def get_tellurics():
    start = np.array([7594, 6853])
    end = np.array([7678, 6950])
    return start, end

def clip_tellurics(wl, flux):
    start, end = get_tellurics()
    for ii,beg in enumerate(start):
        choose = np.logical_and(wl >= beg, wl <= end[ii])
        flux = np.interp(wl, wl[~choose], flux[~choose])
    return wl, flux

def plot_tellurics():
    col = 'pink'
    plt.axvspan(7594, 7678, ls='--', color=col, lw=0.5, alpha=0.5)
    plt.axvspan(6853, 6950, ls='--', color=col, lw=0.5, alpha=0.5)

def fluxcal(wl, flux, dt_spec):
    """ Flux-calibrate to R-band light curve """
    # get r-band LC
    dt, filt, mag, emag = get_lc()
    det = np.logical_and(mag<99, ~np.isnan(mag))
    nondet = np.logical_or(mag==99, np.isnan(mag))
    choose = np.logical_and(det, filt=='r')
    # interpolate to this epoch
    rval = np.interp(dt_spec, dt[choose], mag[choose])
    # TEMP: r is roughly 658nm +/- 138nm
    # TEMP: assume AB mag
    lam = 6580 # in angstroms
    c = 3E18 # angstrom/s
    fnu = 1E-23 * 3631 * 10**(rval/(-2.5)) # erg/s/cm2/Hz
    flam = fnu * (c/lam**2) # should be erg/s/cm2/AA
    # scale factor
    flam_meas = np.interp(lam, wl, flux)
    #scale = (flam/flam_meas)/1E-15
    return wl, flux

if __name__=="__main__":
    z = 0.03154

    files, epochs, tels = get_files()
    start = 0
    end = 19
    files = files[start:end]
    epochs = epochs[start:end]
    tels = tels[start:end]
    nfiles = len(files)

    fig,axarr = plt.subplots(
            1, 2, figsize=(10,10), sharex=True)

    for ii,f in enumerate(files):
        if ii < nfiles/2:
            ax = axarr[0]
        else:
            ax = axarr[1]
        tel = tels[ii]
        dt = epochs[ii]
        wl, flux, ivar = load_spec(f, tel)
        print(tel)
        wl, flux = clip_lines(wl, flux, z, tel, dt)
        wl, flux = clip_tellurics(wl, flux)
        wl, flux = fluxcal(wl, flux, dt)
        if ii < nfiles/2:
            scale = (flux[wl > 3800][0])/2
            plot_spec(ax, wl, flux/scale+nfiles/2-ii%(nfiles/2), tel, dt)
            plot_smoothed_spec(
                ax, wl, flux/scale+nfiles/2-ii%(nfiles/2), ivar, tel, dt)
        else:
            scale = (flux[wl > 4600][0])/2
            plot_spec(ax, wl, flux/scale+nfiles/2-ii%(nfiles/2), tel, dt)
            plot_smoothed_spec(
                ax, wl, flux/scale+nfiles/2-ii%(nfiles/2), ivar, tel, dt)
        ax.tick_params(axis='both', labelsize=14)
    axarr[0].set_ylabel(
            r"Scaled $F_{\lambda}$ + constant",
            fontsize=16)
    axarr[0].set_xlabel(r"Observed Wavelength (\AA)", fontsize=16)
    axarr[1].set_xlabel(r"Observed Wavelength (\AA)", fontsize=16)
    axarr[1].get_yaxis().set_ticks([])
    plt.xlim(3000, 11000)
    #plt.xlim(4900, 5200)
    plt.subplots_adjust(wspace=0)
    axarr[0].set_ylim(0,12)
    axarr[1].set_ylim(1,12)

    #plt.tight_layout()
    plt.savefig("spec_sequence.eps", dpi=300, bbox_inches='tight')
    #plt.show()
    #plt.close()


# Contains the functions to load the state dict into the model and data preprocessing
from download import download_model
import torch.nn.functional as F
from pathlib import Path
import torch.nn as nn
import numpy as np
import torchtext
import requests
import pickle
import spacy
import torch
import os
import re

nlp = spacy.load("en_core_web_sm")

class CNN(nn.Module):

    def __init__(self, vocab_size, embedding_dim,
                 n_filters, filter_sizes, output_dim,
                 dropout, pad_idx):

        super().__init__()

        self.embedding = nn.Embedding(
            vocab_size, embedding_dim, padding_idx=pad_idx)

        '''
        ModuleList means an arbirtary sized list of filter sizes can be provided
        and the list comprehension will create conv layers for each of the filters
        '''
        self.convs = nn.ModuleList([
            nn.Conv2d(in_channels=1, out_channels=n_filters,
                      kernel_size=(fs, embedding_dim)) for fs in filter_sizes])

        self.fc = nn.Linear(len(filter_sizes) * n_filters, output_dim)

        self.dropout = nn.Dropout(dropout)

    def forward(self, text):
        '''
        In PyTorch RNNs want the input with batch dim second, CNNs want the batch dim first
        we permute the input to make it the right shape for the CNN
        '''
        text = text.permute(1, 0)

        # Text passed through embedding layer to get embeddings
        embedded = self.embedding(text)

        '''
        A conv layer wants the second dim of the input to be a channel dim
        text does not have a channel dim, so the tensor is unsqueezed to create one
        '''
        embedded = embedded.unsqueeze(1)

        # Iterates through the list of conv layers applying each conv layer to get list of conv outputs
        conved = [F.relu(conv(embedded)).squeeze(3) for conv in self.convs]

        '''
        Conv outputs are passed through a max pooling that takes the maximum value over a dimension
        the idea being that the "maximum value" is the most important feature for determining the sentiment
        which corresponds to the most important n-gram in the review
        '''
        pooled = [F.max_pool1d(conv, conv.shape[2]).squeeze(2)
                  for conv in conved]

        '''
        The model has 100 filters of 3 different sizes, therefore 300 n-grams that could be important
        which we concatenate into a single vector and pass through a dropout layer and finally a linear layer
        (NOTE: dropout is set to 0 during inference time)
        '''
        cat = self.dropout(torch.cat(pooled, dim=1))

        # passed through linear layer to make predictions
        return self.fc(cat)

s3_model_url = 'https://sent-model.s3.eu-west-2.amazonaws.com/conv-sentiment_model1.pt'
path_to_model = download_model(s3_model_url, model_name="conv-sentiment_model1.pt")

model = CNN(25002, 300, 100, [3, 4, 5], 1, 0.55, 1)
model.load_state_dict(torch.load(path_to_model, map_location='cpu'))

s3_word_dict_url = 'https://sent-model.s3.eu-west-2.amazonaws.com/word_dict.pkl'
path_to_dict = download_model(s3_word_dict_url, model_name="word_dict.pkl")

with open(path_to_dict, 'rb') as f:
    TEXT = pickle.load(f)

def predict_sentiment(sentence, model=model, min_len=5):

    model.eval()

    sentence = sentence.lower()
    # Remove punctuation
    sentence = re.sub(r'[^\w\s]', '', sentence)

    tokenized = [tok.text for tok in nlp.tokenizer(sentence)]

    '''
    If the length of the sentence is shorter than the length of the largest filter
    then the sentence must be padded to the length of the largest filter
    '''
    if len(tokenized) < min_len:
        tokenized += ['<pad>'] * (min_len - len(tokenized))
    indexed = [TEXT.stoi[t] for t in tokenized]
    tensor = torch.LongTensor(indexed)
    tensor = tensor.unsqueeze(1)
    prediction = torch.sigmoid(model(tensor))

    probs = [{'name': index, 'prob': prediction.item()}
             for index in np.argsort(prediction.item())]

    return (sentence, probs)


r"""
Unit testing for Sage objects
"""

# ****************************************************************************
#       Copyright (C) 2009 Nicolas M. Thiery <nthiery at users.sf.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#                  https://www.gnu.org/licenses/
# ****************************************************************************

import unittest
import sys
import traceback

class TestSuite(object):
    """
    Test suites for Sage objects.

    EXAMPLES::

        sage: TestSuite(ZZ).run()

    No output means that all tests passed. Which tests?
    In practice this calls all the methods ``._test_*`` of this
    object, in alphabetic order::

        sage: TestSuite(1).run(verbose = True)
        running ._test_category() . . . pass
        running ._test_eq() . . . pass
        running ._test_new() . . . pass
        running ._test_nonzero_equal() . . . pass
        running ._test_not_implemented_methods() . . . pass
        running ._test_pickling() . . . pass

    Those methods are typically implemented by abstract
    super classes, in particular via categories, in order to
    enforce standard behavior and API, or provide mathematical
    sanity checks. For example if ``self`` is in the category of
    finite semigroups, this checks that the multiplication is
    associative (at least on some elements)::

        sage: S = FiniteSemigroups().example(alphabet = ('a', 'b'))
        sage: TestSuite(S).run(verbose = True)
        running ._test_an_element() . . . pass
        running ._test_associativity() . . . pass
        running ._test_cardinality() . . . pass
        running ._test_category() . . . pass
        running ._test_construction() . . . pass
        running ._test_elements() . . .
          Running the test suite of self.an_element()
          running ._test_category() . . . pass
          running ._test_eq() . . . pass
          running ._test_new() . . . pass
          running ._test_not_implemented_methods() . . . pass
          running ._test_pickling() . . . pass
          pass
        running ._test_elements_eq_reflexive() . . . pass
        running ._test_elements_eq_symmetric() . . . pass
        running ._test_elements_eq_transitive() . . . pass
        running ._test_elements_neq() . . . pass
        running ._test_enumerated_set_contains() . . . pass
        running ._test_enumerated_set_iter_cardinality() . . . pass
        running ._test_enumerated_set_iter_list() . . . pass
        running ._test_eq() . . . pass
        running ._test_new() . . . pass
        running ._test_not_implemented_methods() . . . pass
        running ._test_pickling() . . . pass
        running ._test_some_elements() . . . pass

    The different test methods can be called independently::

        sage: S._test_associativity()

    Debugging tip: in case of failure of some test, use ``%pdb on`` to
    turn on automatic debugging on error. Run the failing test
    independently: the debugger will stop right where the first
    assertion fails. Then, introspection can be used to analyse what
    exactly the problem is. See also the ``catch = False`` option to
    :meth:`.run`.

    When meaningful, one can further customize on which elements
    the tests are run. Here, we use it to *prove* that the
    multiplication is indeed associative, by running the test on
    all the elements::

        sage: S._test_associativity(elements = S)

    Adding a new test boils down to adding a new method in the class
    of the object or any super class (e.g. in a category). This method
    should use the utility :meth:`._tester` to handle standard options
    and report test failures. See the code of
    :meth:`._test_an_element` for an example. Note: Python's testunit
    convention is to look for methods called ``.test*``; we use instead
    ``._test_*`` so as not to pollute the object's interface.

    Eventually, every implementation of a :class:`SageObject` should
    run a :class:`TestSuite` on one of its instances in its doctest
    (replacing the current ``loads(dumps(x))`` tests).

    Finally, running ``TestSuite`` on a standard Python object does
    some basic sanity checks::

        sage: TestSuite(int(1)).run(verbose = True)
        running ._test_new() . . . pass
        running ._test_pickling() . . . pass

    TODO:

     - Allow for customized behavior in case of failing assertion
       (warning, error, statistic accounting).
       This involves reimplementing the methods fail / failIf / ...
       of unittest.TestCase in InstanceTester

     - Don't catch the exceptions if ``TestSuite(..).run()`` is called
       under the debugger, or with ``%pdb`` on (how to detect this? see
       ``get_ipython()``, ``IPython.Magic.shell.call_pdb``, ...)
       In the mean time, see the ``catch=False`` option.

     - Run the tests according to the inheritance order, from most
       generic to most specific, rather than alphabetically. Then, the
       first failure will be the most relevant, the others being
       usually consequences.

     - Improve integration with doctests (statistics on failing/passing tests)

     - Add proper support for nested testsuites.

     - Integration with unittest:
       Make TestSuite inherit from unittest.TestSuite?
       Make ``.run(...)`` accept a result object

     - Add some standard option ``proof = True``, asking for the
       test method to choose appropriately the elements so as to
       prove the desired property. The test method may assume that
       a parent implements properly all the super categories. For
       example, the ``_test_commutative`` method of the category
       ``CommutativeSemigroups()`` may just check that the
       provided generators commute, implicitly assuming that
       generators indeed generate the semigroup (as required by
       ``Semigroups()``).
    """

    def __init__(self, instance):
        """
        TESTS::

            sage: TestSuite(ZZ)
            Test suite for Integer Ring
        """
        from sage.structure.sage_object import SageObject
        if not isinstance(instance, (SageObject, PythonObjectWithTests)):
            instance = PythonObjectWithTests(instance)
        self._instance = instance

    def __repr__(self):
        """
        TESTS::

            sage: TestSuite(ZZ)
            Test suite for Integer Ring
        """
        return "Test suite for %s" % self._instance

    def run(self, category=None, skip=[], catch=True, raise_on_failure=False,
            **options):
        """
        Run all the tests from this test suite:

        INPUT:

         - ``category``         - a category; reserved for future use
         - ``skip``             - a string or list (or iterable) of strings
         - ``raise_on_failure`` - a boolean (default: False)
         - ``catch``            - a boolean (default: True)

        All other options are passed down to the individual tests.

        EXAMPLES::

            sage: TestSuite(ZZ).run()

        We now use the ``verbose`` option::

            sage: TestSuite(1).run(verbose = True)
            running ._test_category() . . . pass
            running ._test_eq() . . . pass
            running ._test_new() . . . pass
            running ._test_nonzero_equal() . . . pass
            running ._test_not_implemented_methods() . . . pass
            running ._test_pickling() . . . pass

        Some tests may be skipped using the ``skip`` option::

            sage: TestSuite(1).run(verbose = True, skip ="_test_pickling")
            running ._test_category() . . . pass
            running ._test_eq() . . . pass
            running ._test_new() . . . pass
            running ._test_nonzero_equal() . . . pass
            running ._test_not_implemented_methods() . . . pass
            sage: TestSuite(1).run(verbose = True, skip =["_test_pickling", "_test_category"])
            running ._test_eq() . . . pass
            running ._test_new() . . . pass
            running ._test_nonzero_equal() . . . pass
            running ._test_not_implemented_methods() . . . pass

        We now show (and test) some standard error reports::

            sage: class Blah(SageObject):
            ....:     def _test_a(self, tester): pass
            ....:     def _test_b(self, tester): tester.fail()
            ....:     def _test_c(self, tester): pass
            ....:     def _test_d(self, tester): tester.fail()

            sage: TestSuite(Blah()).run()
            Failure in _test_b:
            Traceback (most recent call last):
              ...
            AssertionError: None
            ------------------------------------------------------------
            Failure in _test_d:
            Traceback (most recent call last):
              ...
            AssertionError: None
            ------------------------------------------------------------
            Failure in _test_pickling:
            Traceback (most recent call last):
              ...
            ...PicklingError: Can't pickle <class '__main__.Blah'>: attribute
            lookup ...Blah... failed
            ------------------------------------------------------------
            The following tests failed: _test_b, _test_d, _test_pickling

            sage: TestSuite(Blah()).run(verbose = True)
            running ._test_a() . . . pass
            running ._test_b() . . . fail
            Traceback (most recent call last):
              ...
            AssertionError: None
            ------------------------------------------------------------
            running ._test_c() . . . pass
            running ._test_category() . . . pass
            running ._test_d() . . . fail
            Traceback (most recent call last):
              ...
            AssertionError: None
            ------------------------------------------------------------
            running ._test_new() . . . pass
            running ._test_not_implemented_methods() . . . pass
            running ._test_pickling() . . . fail
            Traceback (most recent call last):
              ...
            ...PicklingError: Can't pickle <class '__main__.Blah'>: attribute
            lookup ...Blah... failed
            ------------------------------------------------------------
            The following tests failed: _test_b, _test_d, _test_pickling

            File "/opt/sage/local/lib/python/site-packages/sage/misc/sage_unittest.py", line 183, in run
            test_method(tester = tester)

        The ``catch=False`` option prevents ``TestSuite`` from
        catching exceptions::

            sage: TestSuite(Blah()).run(catch=False)
            Traceback (most recent call last):
              ...
              File ..., in _test_b
                def _test_b(self, tester): tester.fail()
              ...
            AssertionError: None

        In conjunction with ``%pdb on``, this allows for the debugger
        to jump directly to the first failure location.
        """
        if isinstance(skip, str):
            skip = [skip]
        else:
            skip = tuple(skip)

        # The class of exceptions that will be caught and reported;
        # other exceptions will get through. () catches nothing.
        catch_exception = Exception if catch else ()

        tester = instance_tester(self._instance, **options)
        failed = []
        for method_name in dir(self._instance):
            if method_name[0:6] == "_test_" and method_name not in skip:
                # TODO: improve pretty printing
                # could use the doc string of the test method?
                tester.info(tester._prefix + "running .%s() . . ." % method_name, newline=False)
                test_method = getattr(self._instance, method_name)
                try:
                    test_method(tester=tester)
                    tester.info(" pass")
                except catch_exception as e:
                    failed.append(method_name)
                    if isinstance(e, TestSuiteFailure):
                        # The failure occurred in a nested testsuite
                        # which has already reported the details of
                        # that failure
                        if not tester._verbose:
                            print(tester._prefix + "Failure in {}".format(method_name))
                    else:
                        if tester._verbose:
                            tester.info(" fail")
                        else:
                            print(tester._prefix + "Failure in {}:".format(method_name))
                        s = traceback.format_exc()
                        print(tester._prefix + s.strip().replace("\n", "\n" + tester._prefix))
                        print(tester._prefix + "-" * 60)
        if failed:
            print(tester._prefix + "The following tests failed: {}".format(", ".join(failed)))
            if raise_on_failure:
                raise TestSuiteFailure

class TestSuiteFailure(AssertionError):
    pass

def instance_tester(instance, tester=None, **options):
    """
    Return a gadget attached to ``instance`` providing testing utilities.

    EXAMPLES::

        sage: from sage.misc.sage_unittest import instance_tester
        sage: tester = instance_tester(ZZ)

        sage: tester.assertTrue(1 == 1)
        sage: tester.assertTrue(1 == 0)
        Traceback (most recent call last):
        ...
        AssertionError: False is not true
        sage: tester.assertTrue(1 == 0, "this is expected to fail")
        Traceback (most recent call last):
        ...
        AssertionError: this is expected to fail

        sage: tester.assertEqual(1, 1)
        sage: tester.assertEqual(1, 0)
        Traceback (most recent call last):
        ...
        AssertionError: 1 != 0

    The available assertion testing facilities are the same as in
    :class:`unittest.TestCase` [UNITTEST]_, which see (actually, by a slight
    abuse, tester is currently an instance of this class).

    TESTS::

        sage: instance_tester(ZZ, tester = tester) is tester
        True
    """
    if tester is None:
        return InstanceTester(instance, **options)
    else:
        assert not options
        assert tester._instance is instance
        return tester

class InstanceTester(unittest.TestCase):
    """
    A gadget attached to an instance providing it with testing utilities.

    EXAMPLES::

        sage: from sage.misc.sage_unittest import InstanceTester
        sage: InstanceTester(instance = ZZ, verbose = True, elements = [1,2,3])
        Testing utilities for Integer Ring

    This is used by ``SageObject._tester``, which see::

        sage: QQ._tester()
        Testing utilities for Rational Field
    """

    # On Python 3 this attribute defaults to True, causing the AssertionErrors
    # output by failed test cases to produce longer error messages than the
    # default error messages on Python 2.  So for backwards compatibility of
    # existing test cases we disable these "long messages" (which don't gain us
    # all that much anyways)
    longMessage = False

    def __init__(self, instance, elements=None, verbose=False, prefix="",
                 max_runs=4096, max_samples=None, **options):
        """
        A gadget attached to an instance providing it with testing utilities.

        EXAMPLES::

            sage: from sage.misc.sage_unittest import InstanceTester
            sage: InstanceTester(instance = ZZ, verbose = True, elements = [1,2,3])
            Testing utilities for Integer Ring

        This is used by ``SageObject._tester``, for example::

            sage: QQ._tester()
            Testing utilities for Rational Field
        """
        unittest.TestCase.__init__(self)
        self._instance = instance
        self._verbose = verbose
        self._elements = elements
        self._prefix = prefix
        self._max_runs = max_runs
        self._max_samples = max_samples

    def runTest(self):
        """
        Trivial implementation of :meth:`unittest.TestCase.runTest` to
        please the super class :class:`TestCase`. That's the price to
        pay for abusively inheriting from it.

        EXAMPLES::

            sage: from sage.misc.sage_unittest import InstanceTester
            sage: tester = InstanceTester(ZZ, verbose = True)
            sage: tester.runTest()
        """
        pass

    def info(self, message, newline=True):
        """
        Display user information

        EXAMPLES::

            sage: from sage.misc.sage_unittest import InstanceTester
            sage: tester = InstanceTester(ZZ, verbose = True)

            sage: tester.info("hello"); tester.info("world")
            hello
            world

            sage: tester = InstanceTester(ZZ, verbose = False)
            sage: tester.info("hello"); tester.info("world")

            sage: tester = InstanceTester(ZZ, verbose = True)
            sage: tester.info("hello", newline = False); tester.info(" world")
            hello world
        """
        if self._verbose:
            if newline:
                sys.stdout.write(message + "\n")
            else:
                sys.stdout.write(message)
            sys.stdout.flush()

    def __repr__(self):
        """
        EXAMPLES::

            sage: from sage.misc.sage_unittest import InstanceTester
            sage: InstanceTester(ZZ, verbose = True)
            Testing utilities for Integer Ring

        """
        return "Testing utilities for %s" % self._instance

    def some_elements(self, S=None, repeat=None):
        """
        Return a list (or iterable) of elements of the instance on which
        the tests should be run.

        This is only meaningful for container objects like parents.

        INPUT:

        - ``S`` -- a set of elements to select from.  By default this
          will use the elements passed to this tester at creation
          time, or the result of :meth:`.some_elements` if no elements
          were specified.

        - ``repeat`` -- integer (default: None).  If given, instead returns
          a list of tuples of length ``repeat`` from ``S``.

        OUTPUT:

        A list of at most ``self._max_runs`` elements of ``S^r``,
        or a sample of at most ``self._max_samples`` if that is not ``None``.

        EXAMPLES:

        By default, this calls :meth:`.some_elements` on the instance::

            sage: from sage.misc.sage_unittest import InstanceTester
            sage: class MyParent(Parent):
            ....:     def some_elements(self):
            ....:         return [1,2,3,4,5]
            ...
            sage: tester = InstanceTester(MyParent())
            sage: list(tester.some_elements())
            [1, 2, 3, 4, 5]

            sage: tester = InstanceTester(MyParent(), max_runs=3)
            sage: list(tester.some_elements())
            [1, 2, 3]

            sage: tester = InstanceTester(MyParent(), max_runs=7)
            sage: list(tester.some_elements())
            [1, 2, 3, 4, 5]

            sage: tester = InstanceTester(MyParent(), elements=[1,3,5])
            sage: list(tester.some_elements())
            [1, 3, 5]

            sage: tester = InstanceTester(MyParent(), elements=[1,3,5], max_runs=2)
            sage: list(tester.some_elements())
            [1, 3]

            sage: tester = InstanceTester(FiniteEnumeratedSet(['a','b','c','d']), max_runs=3)
            sage: tester.some_elements()
            ['a', 'b', 'c']

            sage: tester = InstanceTester(FiniteEnumeratedSet([]))
            sage: list(tester.some_elements())
            []

            sage: tester = InstanceTester(ZZ)
            sage: ZZ.some_elements()             # yikes, shamelessly trivial ...
            <generator object ..._some_elements_from_iterator at 0x...>
            sage: list(tester.some_elements())
            [0, 1, -1, 2, -2, ..., 49, -49, 50]

            sage: tester = InstanceTester(ZZ, elements = ZZ, max_runs=5)
            sage: list(tester.some_elements())
            [0, 1, -1, 2, -2]

            sage: tester = InstanceTester(ZZ, elements = srange(100), max_runs=5)
            sage: list(tester.some_elements())
            [0, 1, 2, 3, 4]

            sage: tester = InstanceTester(ZZ, elements = srange(3), max_runs=5)
            sage: list(tester.some_elements())
            [0, 1, 2]

        The ``repeat`` keyword can give pairs or triples from ``S``::

            sage: list(tester.some_elements(repeat=2))
            [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1)]

        You can use ``max_samples`` to sample at random, instead of in order::

            sage: tester = InstanceTester(ZZ, elements = srange(8), max_samples = 4)
            sage: all(t in srange(8) for t in tester.some_elements())
            True
            sage: all(s in srange(8) and t in srange(8) for s,t in tester.some_elements(repeat=2))
            True

        Test for :trac:`15919`, :trac:`16244`::

            sage: Z = IntegerModRing(25) # random.sample, which was used pre #16244, has a threshold at 21!
            sage: Z[1]                   # since #8389, indexed access is used for ring extensions
            Traceback (most recent call last):
            ...
            ValueError: variable name '1' does not start with a letter
            sage: tester = InstanceTester(Z, elements=Z, max_runs=5)
            sage: list(tester.some_elements())
            [0, 1, 2, 3, 4]

            sage: C = cartesian_product([Z]*4)
            sage: len(C)
            390625
            sage: tester = InstanceTester(C, elements = C, max_runs=4)
            sage: list(tester.some_elements())
            [(0, 0, 0, 0), (0, 0, 0, 1), (0, 0, 0, 2), (0, 0, 0, 3)]
        """
        S = S or self._elements or self._instance.some_elements()
        from sage.misc.misc import some_tuples
        return list(some_tuples(S, repeat, self._max_runs, self._max_samples))

class PythonObjectWithTests(object):
    """
    Utility class for running basis tests on a plain Python object
    (that is not in SageObject). More test methods can be added here.

    EXAMPLES::

        sage: TestSuite("bla").run()
    """
    def __init__(self, instance):
        """
        EXAMPLES::

            sage: from sage.misc.sage_unittest import PythonObjectWithTests
            sage: x = PythonObjectWithTests(int(1)); x
            <sage.misc.sage_unittest.PythonObjectWithTests object at ...>
            sage: TestSuite(x).run()
        """
        self._instance = instance

    def _test_pickling(self, **options):
        """
        Checks that the instance in self can be pickled and unpickled properly.

        EXAMPLES::

            sage: from sage.misc.sage_unittest import PythonObjectWithTests
            sage: PythonObjectWithTests(int(1))._test_pickling()

        .. SEEALSO::

            :func:`dumps`, :func:`loads`
        """
        tester = instance_tester(self, **options)
        from sage.misc.all import loads, dumps
        tester.assertEqual(loads(dumps(self._instance)), self._instance)

    def _test_new(self, **options):
        """
        Check that ``cls.__new__(cls)`` does not crash Python, with
        ``cls`` either the tested instance (if it's a type) or the type
        of the instance.

        It is perfectly legal for ``__new__`` to raise ordinary
        exceptions.

        EXAMPLES::

            sage: TestSuite(int(1)).run(verbose=True)
            running ._test_new() . . . pass
            running ._test_pickling() . . . pass
            sage: TestSuite(int).run(verbose=True)
            running ._test_new() . . . pass
            running ._test_pickling() . . . pass
        """
        cls = self._instance
        if not isinstance(cls, type):
            cls = type(cls)
        try:
            cls.__new__(cls)
        except Exception:
            pass


#!/usr/bin/env python2
from PIL import Image
import numpy as np

im = Image.open('pic.jpg')

w, h = im.size
print( "w:", w ,"h:", h)
if w != 240 or h != 240:
	raise ValueError('Image must be same dimensions as display \
		({0}x{1}).' .format(240, 240))

with open('/dev/fb1', 'wb') as f:
	img = np.asarray(im)
	pix = np.zeros((240,240,2), dtype = np.uint8)
	pix[...,[1]] = np.add(np.bitwise_and(img[...,[0]],0xF8),np.right_shift(img[...,[1]],5))
	pix[...,[0]] = np.add(np.bitwise_and(np.left_shift(img[...,[1]],3),0xE0),np.right_shift(img[...,[2]],3))
	pix = pix.flatten().tobytes()
	f.write(pix)