#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jan 08 9:02 2021

@author: MCR

All functions associated with the determining the centroid positions of the
order 1 and 2 SOSS spectra trace.
"""

import numpy as np
import emcee
import warnings
from SOSS.trace import tracepol as tp
from SOSS.extract.empirical_trace import plotting as plotting

warnings.simplefilter(action='ignore', category=FutureWarning)

# hack to get around the fact that relative paths are constantly messing up atm
path = '/Users/michaelradica/Documents/GitHub/jwst-mtl/SOSS/'

def determine_stack_dimensions(stack, header=None, verbose=False):
    ''' Determine the size of the stack array. Will be called by
    get_uncontam_centroids and make_trace_mask.

    Parameters
    ----------
    stack : array of floats (2D)
        Data frame. Assumes DMS orientation.
        This array could be a native pixel size SOSS subarray or FF.
        It could also be a 2D trace reference file in which case padding exists
        around the edges, and the pixels may be oversampled by some integer factor.
    header : fits header
        Header associated to the stack array.
        If the header is None then some assumptions will be made regarding the stack array.
        If the header is passed, then specific keywords will be read in it to assess what
        the stack array is. This ensures that a 2D Trace Reference file will be digested
        properly.

    Returns
    -------
    dimx, dimy : The dimensions of the stack array.
    xos, yos : The oversampling factor (integer) of the stack array.
    xnative, ynative : The dimensions of the stack image, expressed in native pixels.
    padding : the size of padding all around the image, in units of native pixels.
    working_pixel_bool : a 2D array of the same size as stack with boolean values of
        False where pixels are not light sensitive (the reference pixels). True elsewhere.
    '''

    # Dimensions of the subarray.
    dimy, dimx = np.shape(stack)

    # Determine what is the input stack based either on its dimensions or on
    # the header if passed. Construct a mask of working pixels in case the
    # stack contains reference pixels.
    if header is None:
        # No header passed - Assume stack is valid SOSS subarray or FF, i.e.
        # 2048x256 or 2040x252 (working pixels) or multiple of if oversampled
        # 2048x96 or 2040x96 (working pixels) or multiple of
        # 2048x2048 or 2040x2040 (working pixels) or multiple of
        if (dimx % 2048) == 0:
            # stack is a multiple of native pixels.
            xnative = 2048
            # The x-axis oversampling is:
            xos = int(dimx / 2048)
        elif (dimx % 2040) == 0:
            # stack is a multiple of native *** working *** pixels.
            xnative = 2040
            # The y-axis oversampling is:
            xos = int(dimx / 2040)
        else:
            # stack x dimension has unrecognized size.
            print('Stack X dimension has unrecognized size of {:}. Accepts 2048, 2040 or multiple of.'.format(dimx))
            sys.exit()
        # Check if the Y axis is consistent with the X axis.
        acceptable_ydim = [96, 256, 252, 2040, 2048]
        yaxis_consistent = False
        for accdim in acceptable_ydim:
            if dimy / (accdim*xos) == 1:
                # Found the acceptable dimension
                yos = np.copy(xos)
                ynative = np.copy(accdim)
                yaxis_consistent = True
        if yaxis_consistent is False:
            # stack y dimension is inconsistent with the x dimension.
            print('Stack Y dimension ({:}) is inconsistent with X dimension ({:}) for acceptable SOSS arrays'.format(dimy,dimx))
            sys.exit()
        # Construct a boolean mask (true or false) of working pixels
        working_pixel_bool = np.full((dimy, dimx), True)

        # For dimensions where reference pixels would have been included in
        # stack, mask those reference pixels out.
        # Sizes 96, 252 and 2040 should not contain any reference pixel.
        if xnative == 2048:
            # Mask out the left and right columns of reference pixels
            working_pixel_bool[:, 0:xos * 4] = False
            working_pixel_bool[:, -xos * 4:] = False
        if ynative == 2048:
            # Mask out the top and bottom rows of reference pixels
            working_pixel_bool[0:yos * 4, :] = False
            working_pixel_bool[-yos * 4:, :] = False
        if ynative == 256:
            # Mask the top rows of reference pixels
            working_pixel_bool[-yos * 4:, :] = False

        # Initialize padding to zero in this case because it is not a 2D Trace ref file
        padding = int(0)

    else:
        # header was passed
        # Read in the relevant keywords
        xos, yos = int(header['OVERSAMP']), int(header['OVERSAMP'])
        padding = int(header['PADDING'])
        # The 2D Trace profile is for FULL FRAME so 2048x2048
        xnative, ynative = int(2048), int(2048)
        # Check that the stack respects its intended format
        if dimx != ((xnative+2*padding)*xos):
            # Problem
            print('The header passed is inconsistent with the X dimension of the stack.')
            sys.exit()
        if dimy != ((ynative+2*padding)*yos):
            # Problem
            print('The header passed is inconsistent with the Y dimension of the stack.')
            sys.exit()
        # Construct a mask of working pixels. The 2D Trace REFERENCE file does
        # not contain any reference pixel. So all are True.
        working_pixel_bool = np.full((dimy, dimx), True)

    # For debugging purposes...
    if verbose is True:
        print('dimx={:}, dimy={:}, xos={:}, yos={:}, xnative={:}, ynative={:}'.format(dimx, dimy, xos, yos, xnative, ynative))

    return(dimx, dimy, xos, yos, xnative, ynative, padding, working_pixel_bool)

def _do_emcee(xref, yref, xdat, ydat, showprogress=False):
    '''Calls the emcee package to preform an MCMC determination of the best
    fitting rotation angle and offsets to map the reference centroids onto the
    data for the first order.

    Parameters
    ----------
    xref, yref : array of float
        X and Y trace centroids respectively to be used as a reference point,
        for example: as returned by get_om_centroids.
    xdat, ydat : array of float
        X and Y trace centroids determined from the data, for example: as
        returned by get_uncontam_centroids.
    showprogress: bool
        If True, show the emcee progress bar.

    Returns
    -------
    sampler : emcee EnsembleSampler object
        MCMC fitting results.
    '''

    # Set up the MCMC run.
    initial = np.array([0, 0, 0])  # Initial guess parameters
    pos = initial + 0.5*np.random.randn(32, 3)
    nwalkers, ndim = pos.shape

    sampler = emcee.EnsembleSampler(nwalkers, ndim, _log_probability,
                                    args=[xref, yref, xdat, ydat])
    # Run the MCMC for 5000 steps - it has generally converged
    # within ~3000 steps in trial runs.
    sampler.run_mcmc(pos, 5000, progress=showprogress)

    return sampler

def get_centerofmass_centroids(stack, header=None, badpix=None, tracemask=None,
                               verbose=False):
    '''Determine the x, y positions of the trace centroids from an
    exposure using a center-of-mass analysis. Works for either order if there
    is no contamination, or for order 1 on a detector where the two orders
    are overlapping.
    This is an adaptation of Loïc's get_order1_centroids which can better
    deal with a bright second order.

    Parameters
    ----------
    stack : array of floats (2D)
        Data frame. Assumes DMS orientation.
        This array could be a native pixel size SOSS subarray or FF.
        It could also be a 2D trace reference file in which case padding exists
        around the edges, and the pixels may be oversampled by some integer factor.
    header : fits header
        Header associated to the stack array.
        If the header is None then some assumptions will be made regarding the stack array.
        If the header is passed, then specific keywords will be read in it to assess what
        the stack array is. This ensures that a 2D Trace Reference file will be digested
        properly.
    badpix : array of floats (2D) with anything different than zero meaning a bad pixel
        Optional input bad pixel mask to apply to the stack. Should be of
        the same dimensions as the stack.
    tracemask : array of floats (2D) with anything different than zero meaning a
        masked out pixel. The spirit is to have zeros along one spectral order with
        a certain width.
    specpix_bounds : native spectral pixel bounds to consider in fitting the trace. Most
        likely used for the 2nd and 3rd orders, not for the 1st order.

    Returns
    -------
    tracexbest : np.array
        Best estimate data x centroid.
    traceybest : np.array
        Best estimate data y centroids.
    '''

    # Call the script that determines the dimensions of the stack. It handles
    # regular science images of various subaaray sizes, with or without the
    # reference pixels, oversampled or not. It also handles the 2D Trace
    # Reference File.
    dimx, dimy, xos, yos, xnative, ynative, padding, working_pixel_bool = \
        determine_stack_dimensions(stack, header=header)

    # Make a numpy mask array of the working pixels
    working_pixel_mask = np.ma.array(np.ones((dimy, dimx)), mask=np.invert(working_pixel_bool))
    # Fill the working pixel mask with NaN
    working_pixel_mask = np.ma.filled(working_pixel_mask, np.nan)

    # Check for the optional input badpix and create a bad pixel numpy mask
    if badpix is not None:
        # 1) Check the dimension is the same as stack
        # TODO:
        # 2) Create the numpy.ma array with it
        # The bad pixel mask has values of 'one' for valid pixels.
        badpix_mask = np.ma.array(np.ones((dimy, dimx)), mask=(badpix != 0))
    else:
        # Create a mask with all valid pixels (all ones)
        badpix_mask = np.ma.array(np.ones((dimy, dimx)))
    # Fill the bad pixels with NaN
    badpix_mask = np.ma.filled(badpix_mask, np.nan)

    # Check for the optional input tracemask and create a trace numpy mask
    if tracemask is not None:
        # 1) Check the dimension is the same as stack
        # TODO:
        # 2) Create the numpy.ma array with it
        # The trace mask has values of 'one' for valid pixels.
        trace_mask = np.ma.array(np.ones((dimy, dimx)), mask=(tracemask == 0))
    else:
        # Create a mask with all pixels in the trace (all ones)
        trace_mask = np.ma.array(np.ones((dimy, dimx)))
    # Fill the trace mask with NaN
    trace_mask = np.ma.filled(trace_mask, np.nan)

    # Multiply working pixel mask, bad pixel mask and trace mask
    # The stack image with embedded numpy mask is stackm
    stackm = stack * badpix_mask * working_pixel_mask * trace_mask

    # Identify the floor level of all 2040 working cols to subtract it first.
    floorlevel = np.nanpercentile(stackm, 10, axis=0)
    backsub = stackm - floorlevel
    # Find centroid - first pass, use all pixels in the column.
    # Normalize each column
    norm = backsub / np.nanmax(backsub, axis=0)
    # Create 2D Array of pixel positions
    rows = (np.ones((dimx, dimy)) * np.arange(dimy)).T
    # CoM analysis to find centroid
    com = (np.nansum(norm * rows, axis=0) / np.nansum(norm, axis=0)).data
    # Adopt these trace values as best
    tracex_best = np.arange(dimx)
    tracey_best = np.copy(com)
    # Second pass, find centroid on a subset of pixels
    # from an area around the centroid determined earlier.
    tracex = np.arange(dimx)
    tracey = np.zeros(dimx)*np.nan
    row = np.arange(dimy)
    w = 30 * yos
    for i in range(dimx):
        miny = np.int(np.nanmax([np.around(tracey_best[i] - w), 0]))
        maxy = np.int(np.nanmax([np.around(tracey_best[i] + w), dimy - 1]))
        val = backsub[miny:maxy, i] / np.nanmax(backsub[:, i])
        ind = np.where(np.isfinite(val))
        thisrow = (row[miny:maxy])[ind]
        thisval = val[ind]
        com = np.sum(thisrow * thisval) / np.sum(thisval)
        # Ensure that the centroid position is not getting too close to an edge
        # such that it is biased.
        if (not np.isfinite(com)) or (com <= 5*yos) or (com >= (ynative-6)*yos):
            continue
        # For a bright second order, it is likely that the centroid at this
        # point will be somewhere in between the first and second order.
        # If this is the case (i.e. the pixel value of the centroid is very low
        # compared to the column average), restrict the range of pixels
        # considered to be above the current centroid.
        if backsub[int(com)][i] < np.nanmean(backsub[(int(com) - w):(int(com) + w), i]):
            miny = np.int(np.nanmax([np.around(com), 0]))
            maxy = np.int(np.nanmin([np.around(com + 2*w), dimy - 1]))
            val = backsub[miny:maxy, i] / np.nanmax(backsub[:, i])
            ind = np.where(np.isfinite(val))
            thisrow = (row[miny:maxy])[ind]
            thisval = val[ind]
            com = np.sum(thisrow * thisval) / np.sum(thisval)
        tracey[i] = com
    # Adopt these trace values as best.
    tracex_best = np.copy(tracex)
    tracey_best = np.copy(tracey)

    # Third pass - fine tuning.
    tracex = np.arange(dimx)
    tracey = np.zeros(dimx) * np.nan
    row = np.arange(dimy)
    w = 16 * yos
    for i in range(len(tracex_best)):
        miny = np.int(np.nanmax([np.around(tracey_best[i] - w), 0]))
        maxy = np.int(np.nanmax([np.around(tracey_best[i] + w), dimy - 1]))
        val = backsub[miny:maxy, i] / np.nanmax(backsub[:, i])
        ind = np.where(np.isfinite(val))
        thisrow = (row[miny:maxy])[ind]
        thisval = val[ind]
        com = np.sum(thisrow * thisval) / np.sum(thisval)
        tracex[i] = np.copy(tracex_best[i])
        tracey[i] = np.copy(com)
    # Update with the best estimates
    tracex_best = np.copy(tracex)
    tracey_best = np.copy(tracey)

    if verbose is True:
        plt.figure()
        plt.plot(tracex_best, tracey_best)

    # Final pass : Fitting a polynomial to the measured (noisy) positions
    if padding == 0:
        # Only use the non NaN pixels.
        induse = np.isfinite(tracex_best) & np.isfinite(tracey_best)
    else:
        # Important steps in the case of the 2D Trace reference file.
        # Mask out the padded pixels from the fit so it is rigorously the
        # same as for regular science images.
        induse = np.isfinite(tracex_best) & np.isfinite(tracey_best) & \
                 (tracex_best >= xos*padding) & (tracex_best < (dimx-xos*padding))

    # Use a *** fixed *** polynomial order of 11 to keep results consistent
    # from data set to data set. Any systematics would remain fixed.
    polyorder = 11
    param = np.polyfit(tracex_best[induse], tracey_best[induse], polyorder)
    tracey_best = np.polyval(param, tracex_best)

    if verbose is True:
        plt.plot(tracex_best, tracey_best, color='r')
        plt.show()

    return tracex_best, tracey_best

def get_contam_centroids(clear, ref_centroids=None, doplot=False,
                         return_orders=[1, 2, 3], bound=True,
                         showprogress=False):
    '''Get the trace centroids for all orders when there is contamination on
    the detector. Fits the first order centroids using the uncontaminated
    method, and determines the second/third order centroids via the
    well-calibrated relationship between all orders.

    Parameters
    ----------
    clear : np.ndarray (2D)
        CLEAR SOSS exposure data frame.
    ref_centroids : np.ndarray (2D)
        Centroids relative to which to determine rotations parameters.
        Must contain lists of x and y centroids for each order to be returned.
        If None, uses the trace table centroids as a reference.
    doplot : bool
        Whether to plot the results of the reference centroids fit to the
        first order.
    return_orders : list
        Orders for which centroid x and y positions will be returned.
    bound : bool
        If True, only returns centroids that fall on the detector after
        polynomial fitting.
    showprogress : bool
        If True, show the emcee progress bar.

    Returns
    -------
    trans_cen : dict
        Dictionary containing x and y trace centroids for each order in
        return_orders.
    rot_pars : tuple
        Tuple containing the required parameters to transform the reference
        centroids to match the dataL: rotation angle, x offset and y offset.

    Raises
    ------
    ValueError
        If ref_centroids does not match return_orders.
    '''

    # Determine reference centroids for all orders.
    ref_cen = {}
    trans_cen = {}
    # If provided as input.
    if ref_centroids is not None:
        # Ensure length is correct.
        if len(ref_centroids) != 2*len(return_orders):
            raise ValueError('Insufficient reference centroids provided.')
        # Repackage into a dictionary
        for i, order in enumerate(return_orders):
            ref_cen['order '+str(order)] = [ref_centroids[2*i], ref_centroids[2*i+1]]
    # Or from the trace table (currently the optics model).
    else:
        for order in return_orders:
            # Extend centroids off of the detector to compensate for shifts.
            xom, yom, tp = get_om_centroids(np.arange(2148)-50, order=order)
            ref_cen['order '+str(order)] = [xom, yom]

    # Get the order 1 centroids from the data.
    xdat_o1, ydat_o1 = get_uncontam_centroids(clear, np.arange(2048), fit=True)
    trans_cen['order 1'] = [xdat_o1, ydat_o1]

    # Fit the reference centroids to the data for order 1.
    fit = _do_emcee(ref_cen['order 1'][0], ref_cen['order 1'][1], xdat_o1,
                    ydat_o1, showprogress=showprogress)
    # Plot MCMC results if requested.
    if doplot is True:
        plotting._plot_corner(fit)
    # Get fitted rotation parameters.
    flat_samples = fit.get_chain(discard=500, thin=15, flat=True)
    ang = np.percentile(flat_samples[:, 0], 50)
    xshift = np.percentile(flat_samples[:, 1], 50)
    yshift = np.percentile(flat_samples[:, 2], 50)
    rot_params = (ang, xshift, yshift)

    # Get rotated centroids for all other orders.
    for order in return_orders:
        if order == 1:
            continue
        xtrans, ytrans = rot_centroids(ang, xshift, yshift,
                                       ref_cen['order '+str(order)][0],
                                       ref_cen['order '+str(order)][1])
        # Ensure that the centroids cover the whole detector.
        pp = np.polyfit(xtrans, ytrans, 5)
        ytrans = np.polyval(pp, np.arange(2048))
        if bound is True:
            inds = np.where((ytrans >= 0) & (ytrans < 256))[0]
            trans_cen['order '+str(order)] = [np.arange(2048)[inds], ytrans[inds]]
        else:
            trans_cen['order '+str(order)] = [np.arange(2048), ytrans]

    return trans_cen, rot_params

# Needs to be updated whenever we decide on how we will interact with
# new reference files
def get_om_centroids(atthesex=None, order=1):
    '''Get trace profile centroids from the NIRISS SOSS optics model.
    These centroids include the standard rotation of 1.489 deg about
    (1514, 486) to transform from the optics model into the CV3 coordinate
    system.

    Parameters
    ----------
    atthesex : list of floats
        Pixel x values at which to evaluate the centroid position.
    order : int
        Diffraction order for which to return the optics model solution.

    Returns
    -------
    xOM : list of floats
        Optics model x centroids.
    yOM : list of floats
        Optics model y centroids.
    tp2 : list of floats
        trace polynomial coefficients.
    '''

    if atthesex is None:
        atthesex = np.linspace(0, 2047, 2048)

    # Derive the trace polynomials.
    tp2 = tp.get_tracepars(filename=path+'/trace/NIRISS_GR700_trace.csv')

    # Evaluate the trace polynomials at the desired coordinates.
    w = tp.specpix_to_wavelength(atthesex, tp2, order, frame='dms', oversample=2)[0]
    xOM, yOM, mas = tp.wavelength_to_pix(w, tp2, order, frame='dms', oversample=2)

    return xOM, yOM[::-1], tp2

def get_uncontam_centroids(stack, atthesex=np.arange(2048), fit=True):
    '''Determine the x, y positions of the trace centroids from an
    exposure using a center-of-mass analysis. Works for either order if there
    is no contamination, or for order 1 on a detector where the two orders
    are overlapping.
    This is an adaptation of Loïc's get_order1_centroids which can better
    deal with a bright second order.

    Parameters
    ----------
    stack : array of floats (2D)
        Data frame.
    atthesex : list of floats
        Pixel x values at which to extract the trace centroids.
    fit : bool
        If True, fits a 5th order polynomial to the extracted y-centroids,
        and returns the evaluation of this polynomial at atthesex. If False,
        a y-centroid may not be located for each x-pixel in atthesex.

    Returns
    -------
    tracexbest : np.array
        Best estimate data x centroid.
    traceybest : np.array
        Best estimate data y centroids.
    '''

    # Dimensions of the subarray.
    dimx = len(atthesex)
    dimy = np.shape(stack)[0]

    # Identify the floor level of all 2040 working cols to subtract it first.
    floorlevel = np.nanpercentile(stack, 10, axis=0)
    backsub = stack - floorlevel

    # Find centroid - first pass, use all pixels in the column.
    # Normalize each column
    norm = backsub[:, 4:2044] / np.nanmax(backsub[:, 4:2044], axis=0)
    # Create 2D Array of pixel positions
    rows = (np.ones((2040, 256)) * np.arange(256)).T
    # Mask any nan values
    norm_mask = np.ma.masked_invalid(norm)
    # CoM analysis to find centroid
    cx = (np.nansum(norm_mask * rows, axis=0) / np.nansum(norm, axis=0)).data

    # Adopt these trace values as best
    tracex_best = np.arange(2040)+4
    tracey_best = cx

    # Second pass, find centroid on a subset of pixels
    # from an area around the centroid determined earlier.
    tracex = []
    tracey = []
    row = np.arange(dimy)
    w = 30
    for i in range(dimx - 8):
        miny = np.int(np.nanmax([np.around(tracey_best[i] - w), 0]))
        maxy = np.int(np.nanmax([np.around(tracey_best[i] + w), dimy - 1]))
        val = backsub[miny:maxy, i + 4] / np.nanmax(backsub[:, i + 4])
        ind = np.where(np.isfinite(val))
        thisrow = (row[miny:maxy])[ind]
        thisval = val[ind]
        cx = np.sum(thisrow * thisval) / np.sum(thisval)
        # Ensure that the centroid position is not getting too close to an edge
        # such that it is biased.
        if not np.isfinite(cx) or cx <= 5 or cx >= 250:
            continue
        # For a bright second order, it is likely that the centroid at this
        # point will be somewhere in between the first and second order.
        # If this is the case (i.e. the pixel value of the centroid is very low
        # compared to the column average), restrict the range of pixels
        # considered to be above the current centroid.
        if backsub[int(cx)][i+4] < np.nanmean(backsub[(int(cx) - w):(int(cx)+w), i+4]):
            miny = np.int(np.nanmax([np.around(cx), 0]))
            maxy = np.int(np.nanmin([np.around(cx + 2*w), dimy - 1]))
            val = backsub[miny:maxy, i + 4] / np.nanmax(backsub[:, i + 4])
            ind = np.where(np.isfinite(val))
            thisrow = (row[miny:maxy])[ind]
            thisval = val[ind]
            cx = np.sum(thisrow * thisval) / np.sum(thisval)

        tracex.append(i + 4)
        tracey.append(cx)

    # Adopt these trace values as best.
    tracex_best = np.array(tracex) * 1
    tracey_best = np.array(tracey) * 1

    # Third pass - fine tuning.
    tracex = []
    tracey = []
    row = np.arange(dimy)
    w = 16
    for i in range(len(tracex_best)):
        miny = np.int(np.nanmax([np.around(tracey_best[i] - w), 0]))
        maxy = np.int(np.nanmax([np.around(tracey_best[i] + w), dimy - 1]))
        val = backsub[miny:maxy, i + 4] / np.nanmax(backsub[:, i + 4])
        ind = np.where(np.isfinite(val))
        thisrow = (row[miny:maxy])[ind]
        thisval = val[ind]
        cx = np.sum(thisrow * thisval) / np.sum(thisval)

        tracex.append(tracex_best[i])
        tracey.append(cx)

    tracex_best = np.array(tracex)
    tracey_best = np.array(tracey)

    if fit is True:
        # Fit a polynomial to centroids to ensure there is a centroid at each x
        p_o1 = np.polyfit(tracex_best, tracey_best, 5)
        tracey_best = np.polyval(p_o1, atthesex)
        tracex_best = atthesex

    return tracex_best, tracey_best

def _log_likelihood(theta, xmod, ymod, xdat, ydat):
    '''Definition of the log likelihood. Called by _do_emcee.
    xmod/ymod should extend past the edges of the SUBSTRIP256 detector.
    '''
    ang, xshift, yshift = theta
    # Calculate rotated model
    modelx, modely = rot_centroids(ang, xshift, yshift,
                                   xmod, ymod, bound=True)
    # Interpolate rotated model onto same x scale as data
    modely = np.interp(xdat, modelx, modely)

    return -0.5 * np.sum((ydat - modely)**2 - 0.5 * np.log(2 * np.pi * 1))

def _log_prior(theta):
    '''Definition of the priors. Called by _do_emcee.
    Angle within +/- 5 deg (one motor step is 0.15deg).
    X-shift within +/- 100 pixels, TA shoukd be accurate to within 1 pixel.
    Y-shift to within +/- 50 pixels.
    '''
    ang, xshift, yshift = theta

    if -5 <= ang < 5 and -100 <= xshift < 100 and -50 <= yshift < 50:
        return -1
    else:
        return -np.inf

def _log_probability(theta, xmod, ymod, xdat, ydat):
    '''Definition of the final probability. Called by _do_emcee.
    '''
    lp = _log_prior(theta)
    if not np.isfinite(lp):
        return -np.inf

    return lp + _log_likelihood(theta, xmod, ymod, xdat, ydat)

def rot_centroids(ang, xshift, yshift, xpix, ypix, bound=True, atthesex=None,
                  cenx=1024, ceny=50):
    '''Apply a rotation and shift to the trace centroids positions. This
    assumes that the trace centroids are already in the CV3 coordinate system.

    Parameters
    ----------
    ang : float
        The rotation angle in degrees CCW.
    xshift : float
        Offset in the X direction to be rigidly applied after rotation.
    yshift : float
        Offset in the Y direction to be rigidly applied after rotation.
    xpix : float or np.array of float
        Centroid pixel X values.
    ypix : float or np.array of float]
        Centroid pixel Y values.
    bound : bool
        Whether to trim rotated solutions to fit within the subarray256.
    atthesex : list of float
        Pixel values at which to calculate rotated centroids.

    Returns
    -------
    rot_xpix : np.array of float
        xval after the application of the rotation and translation
        transformations.
    rot_ypix : np.array of float
        yval after the application of the rotation and translation
        transformations.
    '''

    # Convert to numpy arrays
    xpix = np.atleast_1d(xpix)
    ypix = np.atleast_1d(ypix)
    # Required rotation in the detector frame to match the data.
    t = np.deg2rad(ang)
    R = np.array([[np.cos(t), -np.sin(t)], [np.sin(t), np.cos(t)]])

    # Rotation center set to o1 trace centroid halfway along spectral axis.
    points1 = np.array([xpix - cenx, ypix - ceny])
    rot_pix = R @ points1

    rot_pix[0] += cenx
    rot_pix[1] += ceny

    # Apply the offsets
    rot_pix[0] += xshift
    rot_pix[1] += yshift

    if xpix.size >= 10:
        if atthesex is None:
            # Ensure that there are no jumps of >1 pixel.
            min = int(round(np.min(rot_pix[0]), 0))
            max = int(round(np.max(rot_pix[0]), 0))
            # Same range as rotated pixels but with step of 1 pixel.
            atthesex = np.linspace(min, max, max-min+1)
        # Polynomial fit to ensure a centroid at each pixel in atthesex
        pp = np.polyfit(rot_pix[0], rot_pix[1], 5)
        # Warn user if atthesex extends beyond polynomial domain.
        if np.max(atthesex) > np.max(rot_pix[0])+25 or np.min(atthesex) < np.min(rot_pix[0])-25:
            warnings.warn('atthesex extends beyond rot_xpix. Use results with caution.')
        rot_xpix = atthesex
        rot_ypix = np.polyval(pp, rot_xpix)
    else:
        # If too few pixels for fitting, keep rot_pix.
        if atthesex is not None:
            print('Too few pixels for polynomial fitting. Ignoring atthesex.')
        rot_xpix = rot_pix[0]
        rot_ypix = rot_pix[1]

    # Check to ensure all points are on the subarray.
    if bound is True:
        inds = [(rot_ypix >= 0) & (rot_ypix < 256) & (rot_xpix >= 0) &
                (rot_xpix < 2048)]
        rot_xpix = rot_xpix[inds]
        rot_ypix = rot_ypix[inds]

    return rot_xpix, rot_ypix


import math
import numpy as np
from tinygrad.tensor import Tensor
from tinygrad.utils import fetch
from tinygrad.nn import BatchNorm2D

USE_TORCH = False

def fake_torch_load(b0):
  import io
  import pickle
  import struct

  # convert it to a file
  fb0 = io.BytesIO(b0)

  # skip three junk pickles
  pickle.load(fb0)
  pickle.load(fb0)
  pickle.load(fb0)

  key_prelookup = {}

  class HackTensor:
    def __new__(cls, *args):
      #print(args)
      ident, storage_type, obj_key, location, obj_size, view_metadata = args[0]
      assert ident == 'storage'

      ret = np.zeros(obj_size, dtype=storage_type)
      key_prelookup[obj_key] = (storage_type, obj_size, ret, args[2], args[3])
      return ret

  class MyPickle(pickle.Unpickler):
    def find_class(self, module, name):
      #print(module, name)
      if name == 'FloatStorage':
        return np.float32
      if name == 'LongStorage':
        return np.int64
      if module == "torch._utils" or module == "torch":
        return HackTensor
      else:
        return pickle.Unpickler.find_class(self, module, name)

    def persistent_load(self, pid):
      return pid

  ret = MyPickle(fb0).load()

  # create key_lookup
  key_lookup = pickle.load(fb0)
  key_real = [None] * len(key_lookup)
  for k,v in key_prelookup.items():
    key_real[key_lookup.index(k)] = v

  # read in the actual data
  for storage_type, obj_size, np_array, np_shape, np_strides in key_real:
    ll = struct.unpack("Q", fb0.read(8))[0]
    assert ll == obj_size
    bytes_size = {np.float32: 4, np.int64: 8}[storage_type]
    mydat = fb0.read(ll * bytes_size)
    np_array[:] = np.frombuffer(mydat, storage_type)
    np_array.shape = np_shape

    # numpy stores its strides in bytes
    real_strides = tuple([x*bytes_size for x in np_strides])
    np_array.strides = real_strides

  return ret

class MBConvBlock:
  def __init__(self, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio):
    oup = expand_ratio * input_filters
    if expand_ratio != 1:
      self._expand_conv = Tensor.zeros(oup, input_filters, 1, 1)
      self._bn0 = BatchNorm2D(oup)
    else:
      self._expand_conv = None

    self.strides = strides
    if strides == (2,2):
      self.pad = [(kernel_size-1)//2-1, (kernel_size-1)//2]*2
    else:
      self.pad = [(kernel_size-1)//2]*4

    self._depthwise_conv = Tensor.zeros(oup, 1, kernel_size, kernel_size)
    self._bn1 = BatchNorm2D(oup)

    num_squeezed_channels = max(1, int(input_filters * se_ratio))
    self._se_reduce = Tensor.zeros(num_squeezed_channels, oup, 1, 1)
    self._se_reduce_bias = Tensor.zeros(num_squeezed_channels)
    self._se_expand = Tensor.zeros(oup, num_squeezed_channels, 1, 1)
    self._se_expand_bias = Tensor.zeros(oup)

    self._project_conv = Tensor.zeros(output_filters, oup, 1, 1)
    self._bn2 = BatchNorm2D(output_filters)

  def __call__(self, inputs):
    x = inputs
    if self._expand_conv:
      x = self._bn0(x.conv2d(self._expand_conv)).swish()
    x = x.pad2d(padding=self.pad)
    x = x.conv2d(self._depthwise_conv, stride=self.strides, groups=self._depthwise_conv.shape[0])
    x = self._bn1(x).swish()

    # has_se
    x_squeezed = x.avg_pool2d(kernel_size=x.shape[2:4])
    x_squeezed = x_squeezed.conv2d(self._se_reduce).add(self._se_reduce_bias.reshape(shape=[1, -1, 1, 1])).swish()
    x_squeezed = x_squeezed.conv2d(self._se_expand).add(self._se_expand_bias.reshape(shape=[1, -1, 1, 1]))
    x = x.mul(x_squeezed.sigmoid())

    x = self._bn2(x.conv2d(self._project_conv))
    if x.shape == inputs.shape:
      x = x.add(inputs)
    return x

class EfficientNet:
  def __init__(self, number=0):
    self.number = number
    global_params = [
      # width, depth
      (1.0, 1.0), # b0
      (1.0, 1.1), # b1
      (1.1, 1.2), # b2
      (1.2, 1.4), # b3
      (1.4, 1.8), # b4
      (1.6, 2.2), # b5
      (1.8, 2.6), # b6
      (2.0, 3.1), # b7
      (2.2, 3.6), # b8
      (4.3, 5.3), # l2
    ][number]

    def round_filters(filters):
      multiplier = global_params[0]
      divisor = 8
      filters *= multiplier
      new_filters = max(divisor, int(filters + divisor / 2) // divisor * divisor)
      if new_filters < 0.9 * filters: # prevent rounding by more than 10%
        new_filters += divisor
      return int(new_filters)

    def round_repeats(repeats):
      return int(math.ceil(global_params[1] * repeats))

    out_channels = round_filters(32)
    self._conv_stem = Tensor.zeros(out_channels, 3, 3, 3)
    self._bn0 = BatchNorm2D(out_channels)
    blocks_args = [
      [1, 3, (1,1), 1, 32, 16, 0.25],
      [2, 3, (2,2), 6, 16, 24, 0.25],
      [2, 5, (2,2), 6, 24, 40, 0.25],
      [3, 3, (2,2), 6, 40, 80, 0.25],
      [3, 5, (1,1), 6, 80, 112, 0.25],
      [4, 5, (2,2), 6, 112, 192, 0.25],
      [1, 3, (1,1), 6, 192, 320, 0.25],
    ]
    self._blocks = []
    # num_repeats, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio
    for b in blocks_args:
      args = b[1:]
      args[3] = round_filters(args[3])
      args[4] = round_filters(args[4])
      for n in range(round_repeats(b[0])):
        self._blocks.append(MBConvBlock(*args))
        args[3] = args[4]
        args[1] = (1,1)

    in_channels = round_filters(320)
    out_channels = round_filters(1280)
    self._conv_head = Tensor.zeros(out_channels, in_channels, 1, 1)
    self._bn1 = BatchNorm2D(out_channels)
    self._fc = Tensor.zeros(out_channels, 1000)
    self._fc_bias = Tensor.zeros(1000)

  def forward(self, x):
    x = x.pad2d(padding=(0,1,0,1))
    x = self._bn0(x.conv2d(self._conv_stem, stride=2)).swish()
    for block in self._blocks:
      #print(x.shape)
      x = block(x)
    x = self._bn1(x.conv2d(self._conv_head)).swish()
    x = x.avg_pool2d(kernel_size=x.shape[2:4])
    x = x.reshape(shape=(-1, x.shape[1]))
    #x = x.dropout(0.2)
    return x.dot(self._fc).add(self._fc_bias.reshape(shape=[1,-1]))

  def load_weights_from_torch(self, gpu):
    # load b0
    # https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/utils.py#L551
    if self.number == 0:
      b0 = fetch("https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth")
    elif self.number == 2:
      b0 = fetch("https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b2-8bb594d6.pth")
    elif self.number == 4:
      b0 = fetch("https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b4-6ed6700e.pth")
    elif self.number == 7:
      b0 = fetch("https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b7-dcc49843.pth")
    else:
      raise Exception("no pretrained weights")

    if USE_TORCH:
      import io
      import torch
      b0 = torch.load(io.BytesIO(b0))
    else:
      b0 = fake_torch_load(b0)

    for k,v in b0.items():
      if '_blocks.' in k:
        k = "%s[%s].%s" % tuple(k.split(".", 2))
      mk = "self."+k
      #print(k, v.shape)
      try:
        mv = eval(mk)
      except AttributeError:
        try:
          mv = eval(mk.replace(".weight", ""))
        except AttributeError:
          mv = eval(mk.replace(".bias", "_bias"))
      vnp = v.numpy().astype(np.float32) if USE_TORCH else v
      mv.data[:] = vnp if k != '_fc.weight' else vnp.T
      if gpu:
        mv.cuda_()


#-------------------------------------------------------------------------------
#
#  Swarm MIO_SHA_2* coefficients loaders
#
# Author: Martin Paces <martin.paces@eox.at>
#
#-------------------------------------------------------------------------------
# Copyright (C) 2018 EOX IT Services GmbH
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies of this Software or works derived from this Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#-------------------------------------------------------------------------------

from io import open
from numpy import arange
from .model_mio import (
    DipoleMIOPrimaryGeomagneticModel, DipoleMIOGeomagneticModel,
    MIO_EARTH_RADIUS,
)
from .coefficients_mio import SparseSHCoefficientsMIO
from .parser_mio import parse_swarm_mio_file

def load_model_swarm_mio_internal(path):
    """ Load internal (secondary field) model from a Swarm MIO_SHA_2* product.
    """
    coefficients, params = load_coeff_swarm_mio_internal(path)
    return _create_mio_model(coefficients, params)

def load_model_swarm_mio_external(path, above_ionosphere=None):
    """ Load external (primary field) model from a Swarm MIO_SHA_2* product.
    """
    with open(path, encoding="ascii") as file_in:
        params = parse_swarm_mio_file(file_in)

    if above_ionosphere is None:
        return _create_composed_mio_model(
            _get_coeff_swarm_mio_external(params, False),
            _get_coeff_swarm_mio_external(params, True),
            params
        )
    else:
        return _create_mio_model(
            _get_coeff_swarm_mio_external(params, above_ionosphere), params
        )

def _create_composed_mio_model(coefficients_below_ionosphere,
                               coefficients_above_ionosphere, params):
    return DipoleMIOPrimaryGeomagneticModel(
        _create_mio_model(coefficients_below_ionosphere, params),
        _create_mio_model(coefficients_above_ionosphere, params),
        height=params["height"],
    )

def _create_mio_model(coefficients, params):
    return DipoleMIOGeomagneticModel(
        coefficients, north_pole=(params["lat_NGP"], params["lon_NGP"]),
        wolf_ratio=params["wolf_ratio"], height=params["height"],
    )

def load_coeff_swarm_mio_internal(path):
    """ Load internal model coefficients and other parameters
    from a Swarm MIO_SHA_2* product file.
    """
    with open(path, encoding="ascii") as file_in:
        data = parse_swarm_mio_file(file_in)

    return SparseSHCoefficientsMIO(
        data["nm"], data["gh"],
        ps_extent=(data["pmin"], data["pmax"], data["smin"], data["smax"]),
        is_internal=True,
    ), data

def load_coeff_swarm_mio_external(path, above_ionosphere=True):
    """ Load external model coefficients from a Swarm MIO_SHA_2* product file.
    Use the `above_ionosphere` to pick the right model variant.
    """
    with open(path, encoding="ascii") as file_in:
        data = parse_swarm_mio_file(file_in)

    return _get_coeff_swarm_mio_external(data, above_ionosphere), data

def _get_coeff_swarm_mio_external(data, above_ionosphere):
    """ Create coefficient object for the given source data. """
    indices = data["nm"]
    coefficients = data["qs"]
    if above_ionosphere:
        is_internal = True
        coefficients = convert_external_mio_coeff(
            data["degree_max"], indices, coefficients, data["height"]
        )
    else:
        is_internal = False

    return SparseSHCoefficientsMIO(
        indices, coefficients,
        ps_extent=(data["pmin"], data["pmax"], data["smin"], data["smax"]),
        is_internal=is_internal,
    )

def convert_external_mio_coeff(degree, indices, coefficients, height):
    """ Convert external coefficients to internal ones. """
    nrrad = -(1.0 + height/MIO_EARTH_RADIUS)
    order = arange(degree + 1, dtype='float')
    scale = (order/(order + 1)) * nrrad**(2*order + 1)
    return (scale[indices[:, 0]] * coefficients.transpose()).transpose()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : Dengpan Fu (fdpan@mail.ustc.edu.cn)

import os, sys
import numpy as np
import time, pickle
import scipy.io as sio
import torch
from torch.nn import functional as F

def load_data(path, dtype='torch', mat_type=False):
    if not mat_type:
        with open(path, 'rb') as f:
            data = pickle.load(f)
    else:
        data = sio.loadmat(path)
    q_feat, q_pid, q_cam = data['q_feat'], data['q_id'], data['q_cam']
    g_feat, g_pid, g_cam = data['g_feat'], data['g_id'], data['g_cam']
    if dtype == 'torch':
        q_feat, g_feat = torch.from_numpy(q_feat), torch.from_numpy(g_feat)
    out = [q_feat, q_pid, q_cam, g_feat, g_pid, g_cam]
    return out

def pairwise_distance(x, y, dist_type='cosine'):
    """ Calculate pairwise distance """
    if dist_type == 'euclidean':
        m, n = x.size(0), y.size(0)
        x = x.view(m, -1)
        y = y.view(n, -1)
        dist = torch.pow(x, 2).sum(1).unsqueeze(1).expand(m, n) + \
               torch.pow(y, 2).sum(1).unsqueeze(1).expand(n, m).t()
        dist.addmm_(1, -2, x, y.t())
    elif dist_type == 'cosine':
        x = F.normalize(x)
        y = F.normalize(y)
        dist = torch.mm(x, y.t()).mul_(-1).add(1)
    else:
        raise TypeError("Unknown dist_type={}.".format(dist_type))
    return dist

def mah_dist(x, y, M=None):
    """ Calculate Mahalanobis Distance if Matrix `M` provided,
        Or calculate pairwise Euclidean distance
    """
    if M is None:
        return pairwise_distance(x, y, 'euclidean')
    u = (x.matmul(M)*x).sum(1)
    v = (y.matmul(M)*y).sum(1)
    uv = x.matmul(M).matmul(y.t())
    return u.view(-1, 1) + v.view(1, -1) - 2 * uv

def tensor2numpy(x):
    if isinstance(x, torch.Tensor):
        x = x.cpu().numpy()
    if not isinstance(x, np.ndarray):
        x = np.array(x)
    return x

def numpy2tensor(x):
    if isinstance(x, np.ndarray):
        x = torch.from_numpy(x)
    if not isinstance(x, torch.Tensor):
        x = torch.Tensor(x)
    return x

def print_scores(mAP, cmc_scores, p_str=''):
    if p_str:
        print('{:<15}:'.format(p_str), end='')
    print(('[mAP: {:5.2%}], [cmc1: {:5.2%}], [cmc5: {:5.2%}],'
        ' [cmc10: {:5.2%}]').format(mAP, *cmc_scores[[0, 4, 9]]))

def compute_mAP(index, good_index, junk_index):
    ap = 0
    cmc = torch.IntTensor(len(index)).zero_()
    if good_index.size==0:   # if empty
        cmc[0] = -1
        return ap,cmc

    # remove junk_index
    mask = np.in1d(index, junk_index, invert=True)
    index = index[mask]

    # find good_index index
    ngood = len(good_index)
    mask = np.in1d(index, good_index)
    rows_good = np.where(mask)[0]

    cmc[rows_good[0]:] = 1
    d_recall = 1.0/ngood
    precision = (np.arange(len(rows_good), dtype=np.float) + 1) / (rows_good + 1)
    if rows_good[0] == 0:
        old_precision = np.ones(len(rows_good))
        old_precision[1:] = np.arange(1, len(rows_good), dtype=np.float) / rows_good[1:]
    else:
        old_precision = np.arange(len(rows_good), dtype=np.float) / rows_good
    ap = np.sum((precision + old_precision) / 2. * d_recall)

    return ap, cmc

def compute_score(dist, q_id, q_cam, g_id, g_cam, verbose=True, out_aps=False):
    dist = tensor2numpy(dist)
    q_id, q_cam = tensor2numpy(q_id), tensor2numpy(q_cam)
    g_id, g_cam = tensor2numpy(g_id), tensor2numpy(g_cam)
    t1 = time.time()
    cmc = torch.IntTensor(len(g_id)).zero_()
    ap = 0.0
    aps = []
    for i in range(len(q_id)):
        index = dist[i].argsort()
        ql, qc, gl, gc = q_id[i], q_cam[i], g_id, g_cam

        good_index = np.where((gl==ql) & (gc!=qc))[0]
        junk_index = np.where(((gl==ql) & (gc==qc)) | (gl==-1))[0]

        ap_tmp, cmc_tmp = compute_mAP(index, good_index, junk_index)
        aps.append(ap_tmp)
        if cmc_tmp[0]==-1:
            continue
        cmc = cmc + cmc_tmp
        ap += ap_tmp
        if verbose and i % 500 == 0:
            print("Precessing [{:d}/{:d}], Using: {:.3f}s ...".format(
                i+1, len(q_id), time.time() - t1))

    cmc = cmc.float()
    cmc = cmc/len(q_id) #average cmc
    ap = ap/len(q_id)
    if not out_aps:
        return ap, cmc
    else:
        return ap, cmc, aps


# Authors: Sai Nudurupati & Erkan Istanbulluoglu, 21May15
# Edited: 15Jul16 - to conform to Landlab version 1.
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

from landlab import load_params
from landlab.plot import imshow_grid
from landlab.components import (PrecipitationDistribution, Radiation,
                                PotentialEvapotranspiration, SoilMoisture,
                                Vegetation, VegCA)

GRASS = 0
SHRUB = 1
TREE = 2
BARE = 3
SHRUBSEEDLING = 4
TREESEEDLING = 5

def compose_veg_grid(grid, percent_bare=0.4, percent_grass=0.2,
                     percent_shrub=0.2, percent_tree=0.2):
    """Compose spatially distribute PFT."""
    no_cells = grid.number_of_cells
    shrub_point = int(percent_bare * no_cells)
    tree_point = int((percent_bare + percent_shrub) * no_cells)
    grass_point = int((1 - percent_grass) * no_cells)

    veg_grid = np.full(grid.number_of_cells, BARE, dtype=int)
    veg_grid[shrub_point:tree_point] = SHRUB
    veg_grid[tree_point:grass_point] = TREE
    veg_grid[grass_point:] = GRASS

    np.random.shuffle(veg_grid)
    return veg_grid

def initialize(data, grid, grid1):
    """Initialize random plant type field.

    Plant types are defined as the following:

    *  GRASS = 0
    *  SHRUB = 1
    *  TREE = 2
    *  BARE = 3
    *  SHRUBSEEDLING = 4
    *  TREESEEDLING = 5
    """
    grid1.at_cell['vegetation__plant_functional_type'] = compose_veg_grid(
        grid1, percent_bare=data['percent_bare_initial'],
        percent_grass=data['percent_grass_initial'],
        percent_shrub=data['percent_shrub_initial'],
        percent_tree=data['percent_tree_initial'])

    # Assign plant type for representative ecohydrologic simulations
    grid.at_cell['vegetation__plant_functional_type'] = np.arange(6)
    grid1.at_node['topographic__elevation'] = np.full(grid1.number_of_nodes,
                                                      1700.)
    grid.at_node['topographic__elevation'] = np.full(grid.number_of_nodes,
                                                     1700.)
    precip_dry = PrecipitationDistribution(
        mean_storm_duration=data['mean_storm_dry'],
        mean_interstorm_duration=data['mean_interstorm_dry'],
        mean_storm_depth=data['mean_storm_depth_dry'])
    precip_wet = PrecipitationDistribution(
        mean_storm_duration=data['mean_storm_wet'],
        mean_interstorm_duration=data['mean_interstorm_wet'],
        mean_storm_depth=data['mean_storm_depth_wet'])

    radiation = Radiation(grid)
    pet_tree = PotentialEvapotranspiration(grid, method=data['PET_method'],
                                           MeanTmaxF=data['MeanTmaxF_tree'],
                                           delta_d=data['DeltaD'])
    pet_shrub = PotentialEvapotranspiration(grid, method=data['PET_method'],
                                            MeanTmaxF=data['MeanTmaxF_shrub'],
                                            delta_d=data['DeltaD'])
    pet_grass = PotentialEvapotranspiration(grid, method=data['PET_method'],
                                            MeanTmaxF=data['MeanTmaxF_grass'],
                                            delta_d=data['DeltaD'])
    soil_moisture = SoilMoisture(grid, **data) # Soil Moisture object
    vegetation = Vegetation(grid, **data) # Vegetation object
    vegca = VegCA(grid1, **data) # Cellular automaton object

    # Initializing inputs for Soil Moisture object
    grid.at_cell['vegetation__live_leaf_area_index'] = (
        1.6 * np.ones(grid.number_of_cells))
    grid.at_cell['soil_moisture__initial_saturation_fraction'] = (
        0.59 * np.ones(grid.number_of_cells))

    return (precip_dry, precip_wet, radiation, pet_tree, pet_shrub,
            pet_grass, soil_moisture, vegetation, vegca)

def empty_arrays(n, grid, grid1):
    precip = np.empty(n) # Record precipitation
    inter_storm_dt = np.empty(n) # Record inter storm duration
    storm_dt = np.empty(n) # Record storm duration
    time_elapsed = np.empty(n) # To record time elapsed from the start of simulation

    # Cumulative Water Stress
    veg_type = np.empty([n / 55, grid1.number_of_cells], dtype=int)
    daily_pet = np.zeros([365, grid.number_of_cells])
    rad_factor = np.empty([365, grid.number_of_cells])
    EP30 = np.empty([365, grid.number_of_cells])

    # 30 day average PET to determine season
    pet_threshold = 0  # Initializing pet_threshold to ETThresholddown
    return (precip, inter_storm_dt, storm_dt, time_elapsed, veg_type,
            daily_pet, rad_factor, EP30, pet_threshold)

def create_pet_lookup(radiation, pet_tree, pet_shrub, pet_grass, daily_pet,
                      rad_factor, EP30, grid):
    for i in range(0, 365):
        pet_tree.update(float(i) / 365.25)
        pet_shrub.update(float(i) / 365.25)
        pet_grass.update(float(i) / 365.25)
        daily_pet[i] = [pet_grass._PET_value, pet_shrub._PET_value,
                   pet_tree._PET_value, 0., pet_shrub._PET_value,
                   pet_tree._PET_value]
        radiation.update(float(i) / 365.25)
        rad_factor[i] = grid.at_cell['radiation__ratio_to_flat_surface']

        if i < 30:
            if i == 0:
                EP30[0] = daily_pet[0]
            else:
                EP30[i] = np.mean(daily_pet[:i], axis=0)
        else:
            EP30[i] = np.mean(daily_pet[i - 30:i], axis=0)

def save(sim, inter_storm_dt, storm_dt, precip, veg_type, yrs,
         walltime, time_elapsed):
    np.save(sim + '_Tb', inter_storm_dt)
    np.save(sim + '_Tr', storm_dt)
    np.save(sim + '_P', precip)
    np.save(sim + '_VegType', veg_type)
    np.save(sim + '_Years', yrs)
    np.save(sim + '_Time_Consumed_minutes', walltime)
    np.save(sim + '_CurrentTime', time_elapsed)

def plot(sim, grid, veg_type, yrs, yr_step=10):
    pic = 0
    years = range(0, yrs)
    cmap = mpl.colors.ListedColormap(
        ['green', 'red', 'black', 'white', 'red', 'black'])
    bounds = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5]
    norm = mpl.colors.BoundaryNorm(bounds, cmap.N)
    print 'Plotting cellular field of Plant Functional Type'
    print 'Green - Grass; Red - Shrubs; Black - Trees; White - Bare'

    # Plot images to make gif.
    for year in range(0, yrs, yr_step):
        filename = 'year_' + "%05d" % year
        pic += 1
        plt.figure(pic, figsize=(10, 8))
        imshow_grid(grid, veg_type[year], values_at='cell', cmap=cmap,
                    grid_units=('m', 'm'), norm=norm, limits=[0, 5],
                    allow_colorbar=False)
        plt.title(filename, weight='bold', fontsize=22)
        plt.xlabel('X (m)', weight='bold', fontsize=18)
        plt.ylabel('Y (m)', weight='bold', fontsize=18)
        plt.xticks(fontsize=14, weight='bold')
        plt.yticks(fontsize=14, weight='bold')
        plt.savefig(sim + '_' + filename)

    grass_cov = np.empty(yrs)
    shrub_cov = np.empty(yrs)
    tree_cov = np.empty(yrs)
    grid_size = float(veg_type.shape[1])

    for x in range(0, yrs):
        grass_cov[x] = (veg_type[x][veg_type[x] == GRASS].size / grid_size) * 100
        shrub_cov[x] = ((veg_type[x][veg_type[x] == SHRUB].size / grid_size) *
                        100 + (veg_type[x][veg_type[x] == SHRUBSEEDLING].size /
                        grid_size) * 100)
        tree_cov[x] = ((veg_type[x][veg_type[x] == TREE].size / grid_size) *
                       100 + (veg_type[x][veg_type[x] == TREESEEDLING].size /
                       grid_size) * 100)

    pic += 1
    plt.figure(pic, figsize=(10, 8))
    plt.plot(years, grass_cov, '-g', label='Grass', linewidth=4)
    plt.hold(True)
    plt.plot(years, shrub_cov, '-r', label='Shrub', linewidth=4)
    plt.hold(True)
    plt.plot(years, tree_cov, '-k', label='Tree', linewidth=4)
    plt.ylabel('% Area Covered by Plant Type', weight='bold', fontsize=18)
    plt.xlabel('Time in years', weight='bold', fontsize=18)
    plt.xticks(fontsize=12, weight='bold')
    plt.yticks(fontsize=12, weight='bold')
    plt.legend(loc=0, prop={'size': 16, 'weight': 'bold'})
    plt.savefig(sim + '_percent_cover')


"""Data input for A.S.A.P model."""
from __future__ import print_function, division, unicode_literals

import os
import pickle

import numpy as np

import emcee

from astropy.table import Table, Column
from astropy.cosmology import FlatLambdaCDM

from . import ensemble

__all__ = ["load_obs", "load_um", "save_pickle", "load_pickle", "load_npz_results",
           "save_results_to_npz"]

def load_dsigma(cfg, verbose=False):
    """Read in the observed weak lensing sigma profiles.

    Parameters
    ----------
    cfg : dict
        Configuration parameters for the observations.
    verbose : boolen
        Blah, blah, blah

    Return
    ------
    """
    # Now the DeltaSigma data are stored in a numpy array
    # Mass bin is defined by: min_logm1, min_logm2, max_logm1, and max_logm2
    wl_dsigma = np.load(cfg['dsigma'])

    cfg['wl_n_bin'] = len(wl_dsigma)
    if verbose:
        if cfg['wl_n_bin'] > 1:
            print("# There are %d DSigma profiles in this sample" %
                  cfg['wl_n_bin'])
        else:
            print("# There is 1 DSigma profile in this sample")

    return wl_dsigma

def load_obs(cfg, verbose=True):
    """Load the observed data.

    Parameters
    ----------
    cfg : dict
        Configuration parameters for the observations.
    verbose : boolen
        Blah, blah, blah

    Return
    ------
        Dictionary that contains all the observations.

    """
    # Galaxy catalog.
    mass = Table.read(cfg['galaxy'])
    minn = np.array(mass[cfg['minn_col']])
    mtot = np.array(mass[cfg['mtot_col']])

    # Observed DeltaSigma profiles
    wl_dsigma = load_dsigma(cfg, verbose=verbose)
    cfg['dsigma_n_data'] = len(wl_dsigma[0]['dsigma']) * cfg['wl_n_bin']

    # Stellar mass functions
    if os.path.splitext(cfg['smf_inn'])[-1] == '.npy':
        smf_inn = np.load(cfg['smf_inn'])
    else:
        smf_inn = Table.read(cfg['smf_inn'])

    if os.path.splitext(cfg['smf_tot'])[-1] == '.npy':
        smf_tot = np.load(cfg['smf_tot'])
    else:
        smf_tot = Table.read(cfg['smf_tot'])

    # This is for a specific format of SMF
    cfg['smf_inn_min'] = np.min(smf_inn['logm_0'])
    cfg['smf_inn_max'] = np.max(smf_inn['logm_1'])
    cfg['smf_inn_nbin'] = len(smf_inn)

    cfg['smf_tot_min'] = np.min(smf_tot['logm_0'])
    cfg['smf_tot_max'] = np.max(smf_tot['logm_1'])
    cfg['smf_tot_nbin'] = len(smf_tot)

    cfg['ngal_use'] = ((mtot >= cfg['smf_tot_min']) &
                       (minn >= cfg['smf_inn_min'])).sum()

    cfg['min_mtot'] = cfg['smf_tot_min'] - 0.1

    cfg['smf_n_data'] = cfg['smf_tot_nbin'] + cfg['smf_inn_nbin']

    # Covariance of the SMF
    if cfg['smf_cov'] is not None:
        smf_cov = np.load(cfg['smf_cov'])
        assert cfg['smf_n_data'] == len(smf_cov)
    else:
        smf_cov = None

    if verbose:
        print("# SMF for total stellar mass: ")
        print("  %7.4f -- %7.4f in %d bins" % (cfg['smf_tot_min'],
                                               cfg['smf_tot_max'],
                                               cfg['smf_tot_nbin']))
        print("# SMF for inner stellar mass: ")
        print("  %7.4f -- %7.4f in %d bins" % (cfg['smf_inn_min'],
                                               cfg['smf_inn_max'],
                                               cfg['smf_inn_nbin']))

    logms_inn = minn[mtot >= cfg['smf_tot_min']]
    logms_tot = mtot[mtot >= cfg['smf_tot_min']]

    if os.path.isfile(cfg['smf_full']):
        smf_full = Table.read(cfg['smf_full'])
        smf_full[smf_full['smf'] <= 0]['smf'] = 1E-8
        smf_full[smf_full['smf_low'] <= 0]['smf_low'] = 1E-9
        smf_full[smf_full['smf_upp'] <= 0]['smf_upp'] = 1E-7
        smf_full = smf_full
    else:
        smf_full = None

    if verbose:
        print("# For inner stellar mass: ")
        print("    %d bins at %5.2f < logMinn < %5.2f" %
              (cfg['smf_inn_nbin'], cfg['smf_inn_min'],
               cfg['smf_inn_max']))
        print("# For total stellar mass: ")
        print("    %d bins at %5.2f < logMtot < %5.2f" %
              (cfg['smf_tot_nbin'], cfg['smf_tot_min'],
               cfg['smf_tot_max']))

    # Redshift range and observed volume
    cosmo = FlatLambdaCDM(H0=cfg['h0'] * 100, Om0=cfg['omega_m'])
    cfg['volume'] = (
        (cosmo.comoving_volume(np.nanmax(mass[cfg['z_col']])) -
         cosmo.comoving_volume(np.nanmin(mass[cfg['z_col']]))) *
        (cfg['area'] / 41254.0)).value

    if verbose:
        print("# The volume of the HSC data is %15.2f Mpc^3" % cfg['volume'])

    return {'mass': mass, 'minn': minn, 'mtot': mtot,
            'logms_inn': logms_inn, 'logms_tot': logms_tot,
            'wl_dsigma': wl_dsigma,
            'smf_inn': smf_inn, 'smf_tot': smf_tot,
            'smf_full': smf_full, 'smf_cov': smf_cov}, cfg

def load_um(cfg, verbose=True):
    """Load the UniverseMachine data.

    Parameters
    ----------
    cfg : dict
        Configuration parameters for the UniverseMachine model.
    verbose : boolen
        Blah, blah, blah

    Return
    ------
        Dictionary that contains all the UniverseMachine data.

    """
    # Mock galaxy catalog
    um_mock = Table(np.load(cfg['galaxy']))

    # Only select the useful columns
    cols_use = ['halo_id', 'upid', 'sm', 'icl', 'x', 'y', 'z',
                'mtot_galaxy', 'mstar_mhalo', 'logms_gal',
                'logms_icl', 'logms_tot', 'logms_halo',
                'logmh_vir', 'logmh_peak', 'logmh_host']
    um_mock_use = um_mock[cols_use]

    # Value added a few useful columns
    um_mock_use.add_column(Column(data=(um_mock_use['mtot_galaxy'] /
                                        um_mock_use['mstar_mhalo']),
                                  name='frac_cen_tot'))
    um_mock_use.add_column(Column(data=(um_mock_use['sm'] /
                                        um_mock_use['mtot_galaxy']),
                                  name='frac_ins_cen'))
    um_mock_use.add_column(Column(data=(um_mock_use['icl'] /
                                        um_mock_use['mtot_galaxy']),
                                  name='frac_exs_cen'))
    um_mock_use = um_mock_use.as_array()

    # Load the pre-compute lensing pairs
    um_mass_encl = np.load(cfg['dsigma'])
    assert len(um_mock_use) == len(um_mass_encl)

    # Mask for central galaxies
    mask_central = (um_mock_use['upid'] == -1)
    if verbose:
        print("# %d out of %d galaxies are central" % (
            mask_central.sum(), len(um_mock_use)))

    # Mask for massive enough halo
    mask_mass = (um_mock_use[cfg['halo_col']] >= cfg['min_mvir'])

    return {'um_mock': um_mock_use[mask_mass],
            'um_mass_encl': um_mass_encl[mask_mass, :],
            'mask_central': mask_central[mask_mass]}

def save_pickle(pickle_file, data):
    """Save some data in pickle format.

    Parameters
    ----------
    pickle_file : string
        Name of the output Pickle file.
    data : array or dict
        Some data
    """
    pickle.dump(data, open(pickle_file, 'wb'))
    pickle_file.close()

    return

def load_pickle(pickle_file):
    """Load the pickled pickle.

    Parameters
    ----------
    pickle_file : string
        Name of the output Pickle file.
    """
    data = pickle.load(open(pickle_file, 'rb'))
    pickle_file.close()

    return data

def load_npz_results(mcmc_file):
    """Retrieve the MCMC results from .npz file."""
    mcmc_data = np.load(mcmc_file)

    return (mcmc_data['samples'], mcmc_data['chains'],
            mcmc_data['lnprob'], mcmc_data['best'],
            mcmc_data['position'], mcmc_data['acceptance'])

def save_results_to_npz(mcmc_results, mcmc_sampler, mcmc_file,
                        mcmc_ndims, verbose=True, frac=0.1, tol=20, c=5):
    """Save the MCMC run results."""
    mcmc_position, mcmc_lnprob, _ = mcmc_results

    mcmc_samples = mcmc_sampler.chain[:, :, :].reshape((-1, mcmc_ndims))
    mcmc_chains = mcmc_sampler.chain
    mcmc_lnprob = mcmc_sampler.lnprobability

    mcmc_params_stats = ensemble.mcmc_samples_stats(mcmc_samples)

    # Best parameter using the best log(prob)
    mcmc_best = mcmc_sampler.flatchain[mcmc_sampler.flatlnprobability.argmax()]

    # Best parameters using the mean of the last few samples
    _, n_step, n_dim = mcmc_chains.shape
    mcmc_mean = np.nanmean(
        mcmc_chains[:, -int(n_step * frac):, :].reshape([-1, n_dim]), axis=0)

    # Auto-correlation time
    try:
        tau = mcmc_sampler.get_autocorr_time(quiet=False, tol=tol, c=c)
        print("# Current autocorrelation time is", tau)
    except emcee.autocorr.AutocorrError:
        print("# The chain is shorter than {} x tau right now...".format(tol))
        tau = None

    np.savez(mcmc_file,
             samples=mcmc_samples, lnprob=np.array(mcmc_lnprob),
             best=np.array(mcmc_best), mean=np.asarray(mcmc_mean),
             chains=mcmc_chains, tau=tau,
             position=np.asarray(mcmc_position),
             acceptance=np.array(mcmc_sampler.acceptance_fraction))

    if verbose:
        print("#------------------------------------------------------")
        print("#  Mean acceptance fraction",
              np.mean(mcmc_sampler.acceptance_fraction))
        print("#------------------------------------------------------")
        print("#  Best ln(Probability): %11.5f" % np.max(mcmc_lnprob))
        print(mcmc_best)
        print("#------------------------------------------------------")
        print("#  Best parameters (mean):")
        print(mcmc_mean)
        print("#------------------------------------------------------")
        for param_stats in mcmc_params_stats:
            print(param_stats)
        print("#------------------------------------------------------")