#! /usr/bin/env python

from sys import argv
import numpy as np
from ase.io import read, write
from ase.build.surface import fcc111

if argv[1] == "-h" or argv[1] == "--help":
    print("""
Python script that generates a molecular junction for transport computation
using DFTB+.

Author: Alessandro Pirrotta
MIT Licence

Example 1:
$ python generate_molecular_junction.py molecule.xyz c r 0 1

using molecule 'molecule.xyz' as device region
it generates a cluster type molecular junction (c) with no pbc,
removes terminal hydrogen (r) on the end atoms 0 and 1.

Example 2:
$ python generate_molecular_junction.py molecule.xyz s nr

generate a supercell type molecular junction with pbc (s),
does not removes terminal hydrogen (nr) and automatically looks for
sulphur atoms as end atoms.
""")
    exit()

def find_nearest(atom, atoms):
    # create empty list
    dists = []
    # define coordinates for atom
    coord = atom.position
    # loop over atoms and append distance to list
    for a in atoms:
        if a.symbol == 'H':
            v = a.position - coord
            dist = np.sqrt(v[0]**2+v[1]**2+v[2]**2)
            dists.append(dist)
            continue
        dists.append(1.7)
    # make list to np.array
    dists = np.array(dists)
    # find index with shortest distance using argmin
    index = np.argmin(dists)
    if dists[index] > 1.7:  # atom closer than C
        print 'found no H bonded to S, deleted something, consider revising'
    return index

if len(argv) > 4:
    # print("reading end atoms from input")
    endatom1 = int(argv[4])
    endatom2 = int(argv[5])

# create gold electrode
a = 4
pbc = bool(True)

b = a
c = 6
lelectrode = fcc111('Au', size=(a, b, c), vacuum=5.0, orthogonal=True)

# Au-S 2.72AA / Au-Au 2.88AA
# distance between two adjacent gold layers
# deltag           2.35558909829
# perpendicular distance between endatom and electrode surface
# delta s_g        2.14441090171

# read the molecule file
# molecule = read('mol.xyz')
molecule = read(argv[1])
# place electrode close to the origin of axis
n = (lelectrode[0].position)
n = - n                                 # put atom 1 at the origin of axis
lelectrode.translate(n)
natoms = len(molecule)

try:
    endatom1
except NameError:
    endatom1 = -1
    for atom in molecule:
        # if your endatoms are S atoms and there are only 2 in the molecule
        if atom.symbol == 'S':
            if endatom1 < 0:
                endatom1 = atom.index
            else:
                endatom2 = atom.index

# remove terminal hydrogens
if len(argv) > 3 and argv[3] == 'r':
    list = []
    for atom in molecule:
        if atom.symbol == 'S':
            list.append(atom.index)
    del molecule[find_nearest(molecule[list[1]], molecule)]
    del molecule[find_nearest(molecule[list[0]], molecule)]
    natoms = len(molecule)

# align the sulphur atoms along the z axis
po = (molecule[endatom1].position)
lo = (molecule[endatom2].position)
v = lo-po
z = [0, 0, 1]
molecule.rotate(v, z)
molecule.rotate('z', np.pi+np.pi/4.0)

# place the molecule on top of a hollow fcc site for axa face
if a == 6:
    q = (lo + [0, 0, 4.5] - lelectrode[(b*a*2)+(6*4)].position)
if a == 3:
    q = (lo + [0, 0, 4.5] - lelectrode[(a*b*2)+(4)].position)
if a == 4:
    q = (lo + [0, 0, 4.5] - lelectrode[(b*a*2)+(9)].position)
q = - q
molecule.translate(q)

# define distance betweet gold layers
deltag = abs(lelectrode[0].position[2] - lelectrode[(a*b)+1].position[2])
# define end to end distance
v = molecule[endatom2].position - molecule[endatom1].position
lenght_molecule = np.sqrt(v[0]**2+v[1]**2+v[2]**2)
molecule.translate([0, 0, -(deltag)])
molecule.extend(lelectrode)
# make a list of indexes relative to gold atoms
au_list = []
for atom in molecule:
    if atom.symbol == 'Au':
        au_list.append(atom.index)

# number of gold atoms on a single layer
ng = int(a*b)
# distance along the z axis between one layer and the adjacent one:
deltag = abs(lelectrode[0].position[2] - lelectrode[(a*b)+1].position[2])
v = molecule[endatom2].position-molecule[natoms+2].position
deltas_g = abs(v[2])

for x in range(c):
    ly = []
    ly = au_list[(x*a*b):(a*b*(x+1))]
    layer = molecule[ly]
    layer.translate([0, 0,
                    -(abs(lenght_molecule)+abs(deltas_g*2)+abs(deltag*(2*x)))])
    molecule.extend(layer)

if argv[2] == "s" or argv[2] == "S":
    molecule.set_pbc(True)
if argv[2] == "c" or argv[2] == "C":
    molecule.set_pbc(False)
n = (molecule[natoms+a*b*c-a*b].position)
n = -n
molecule.translate(n)

if argv[2] == "s" or argv[2] == "S":
    lele = fcc111('Au',
                  size=(a, b, 18),
                  vacuum=deltas_g+lenght_molecule,
                  orthogonal=True)
    cell = lele.get_cell()
    molecule.set_cell([cell[0],
                       cell[1],
                       [0., 0., molecule[len(molecule)-1].position[2]-20.0]])
    molecule.set_pbc((True, True, True))
    write('moljunct_S.gen', molecule)

write('moljunct_C.gen', molecule)

# printing device, source and drain atom indeces
print 1, natoms
print natoms+1, natoms+a*b*6
print natoms+a*b*6+1, len(molecule)

exit()


#!/usr/local/sci/bin/python2.7
#*****************************
#
# convert 3 hourly data to pentads
#
#
#************************************************************************
'''
Author: Robert Dunn
Created: March 2016
Last update: 12 April 2016
Location: /project/hadobs2/hadisdh/marine/PROGS/Build

-----------------------
CODE PURPOSE AND OUTPUT
-----------------------
Takes 3 hourly data and creates pentads.  First summing across the 5 days and then the 8 timestamps

-----------------------
LIST OF MODULES
-----------------------
utils.py

-----------------------
DATA
-----------------------
Input data and output data stored in:
/project/hadobs2/hadisdh/marine/ICOADS.2.5.1/GRIDS2/

Requires 3hrly grids.

-----------------------
HOW TO RUN THE CODE
-----------------------
python2.7 3hrlies_to_pentads.py --suffix relax --period day --start_year YYYY --end_year YYYY

python2.7 3hrlies_to_pentads.py --help
will show all options

-----------------------
OUTPUT
-----------------------
/project/hadobs2/hadisdh/marine/ICOADS.2.5.1/GRIDS2/

Plots to appear in
/project/hadobs2/hadisdh/marine/ICOADS.2.5.1/PLOTS2/

-----------------------
VERSION/RELEASE NOTES
-----------------------

Version 1 (release date)
---------

Enhancements

Changes

Bug fixes

-----------------------
OTHER INFORMATION
-----------------------
'''

import os
import datetime as dt
import numpy as np
import sys
import argparse
import matplotlib
matplotlib.use('Agg')

import calendar
import netCDF4 as ncdf
import gc

import utils
import set_paths_and_vars
defaults = set_paths_and_vars.set()

OBS_ORDER = utils.make_MetVars(defaults.mdi, multiplier = False)

# what size grid (lat/lon)
DELTA_LAT = 1
DELTA_LON = 1
DELTA_HOUR = 3

# set up the grid
grid_lats = np.arange(-90 + DELTA_LAT, 90 + DELTA_LAT, DELTA_LAT)
grid_lons = np.arange(-180 + DELTA_LAT, 180 + DELTA_LON, DELTA_LON)

#************************************************************************
def read_data(settings, suffix, name, year, grid_lats, grid_lons, period, N_OBS_PER_DAY):
    '''
    Read in the data from the netCDF files

    :param Settings settings: object to hold all filepaths etc.
    :param str suffix: used to determine whether using strict or relaxed criteria
    :param str name: variable name
    :param int year: year to read
    :param array grid_lats: latitudes
    :param array grid_lons: longitudes
    :param str period: which period (day/night/all)
    :param int N_OBS_PER_DAY: number of observation times per day

    :returns: var_3hrlys - array of 3hrly data for single variable
    '''

    if suffix == "relax":
        N_OBS_OVER_DAYS = 1
        N_OBS_OVER_PENTAD = 2

    elif suffix == "strict":
        N_OBS_OVER_DAYS = 2
        N_OBS_OVER_PENTAD = 4

    # set up empty data array
    var_3hrlys = np.ma.zeros([utils.days_in_year(year)*N_OBS_PER_DAY, len(grid_lats), len(grid_lons)])
    var_3hrlys.mask = np.zeros([utils.days_in_year(year)*N_OBS_PER_DAY, len(grid_lats), len(grid_lons)])
    var_3hrlys.fill_value = settings.mdi

    year_start = dt.datetime(year, 1, 1, 0, 0)

    for month in np.arange(12) + 1:
        print year, month

        month_start = utils.day_of_year(year, month)
        month_end = month_start + calendar.monthrange(year, month)[1]

        filename = "{}/{}_1x1_3hr_{}{:02d}_{}_{}.nc".format(settings.DATA_LOCATION, settings.OUTROOT, year, month, period, suffix)

        ncdf_file = ncdf.Dataset(filename,'r', format='NETCDF4')

        if month == 12:
            # run to end of year if december
            var_3hrlys[month_start*N_OBS_PER_DAY:, :, :] = ncdf_file.variables[name][:]
        else:
            var_3hrlys[month_start*N_OBS_PER_DAY:month_end*N_OBS_PER_DAY, :, :] = ncdf_file.variables[name][:]

    return var_3hrlys # read_data

#************************************************************************
def process_february(var_3hrlys, doMask = True):
    '''
    Process the set of 3hly obs that fall on 29th Feb.

    Store them, then remove from the main array.

    :param array var_3hrlys: array for 3hrly data for one variable
    :param bool doMask: process the mask (for data yes, for n_obs no)

    :returns: var_3hrlys, incl_feb29th - updated data and the Feb 29th obs.
    '''

    assert var_3hrlys.shape[0] == 366

    # extract 6-day pentad
    incl_feb29th = var_3hrlys[55:61, :, :, :]

    if doMask:
        # remove the data of Feb 29th from array
        # np.ma.delete doesn't exist, so have to copy mask separately
        mask = var_3hrlys.mask

    var_3hrlys = np.delete(var_3hrlys, 59, 0)

    if doMask:
        mask = np.delete(mask, 59, 0)
        var_3hrlys = np.ma.array(var_3hrlys, mask = mask)
        del mask

    return var_3hrlys, incl_feb29th # process_february

#************************************************************************
def do_conversion(suffix = "relax", start_year = defaults.START_YEAR, end_year = defaults.END_YEAR, period = "all", doQC = False, doBC = False):
    '''
    Convert 3 hrlies to pentads 1x1

    First get pentad average of 3hrly values (so values at 0, 3, 6, ... averaged over 5 days)
    Then get average over the pentad.

    :param str suffix: "relax" or "strict" criteria
    :param int start_year: start year to process
    :param int end_year: end year to process
    :param str period: which period to do day/night/all?
    :param bool doQC: incorporate the QC flags or not
    :param bool doBC: work on the bias corrected data

    :returns:
    '''
    settings = set_paths_and_vars.set(doBC = doBC, doQC = doQC)

    # KW Added SUFFIX variable because all hourlies/dailies/monthlies now have suffix 'strict' (4/2 per daily/day-night)
    # or 'relax' (2/1 per daily/day-night)
    if suffix == "relax":
        N_OBS_OVER_DAYS = 1 # at least 1 obs at this 3 hr timestamp from 5 days in pentad
        N_OBS_OVER_PENTAD = 2

    elif suffix == "strict":
        N_OBS_OVER_DAYS = 2
        N_OBS_OVER_PENTAD = 4  # at least 4 timestamps (of 8) in pentad, could be 2 for local 'relax' setting

    N_OBS_PER_DAY = 24/DELTA_HOUR

    for year in np.arange(start_year, end_year + 1):

        all_pentads =  np.ma.zeros([len(OBS_ORDER), 73, len(grid_lats), len(grid_lons)])
        all_pentads.mask =  np.zeros([len(OBS_ORDER), 73, len(grid_lats), len(grid_lons)])

        # read in a years worth of 3hrly data
        for v, var in enumerate(OBS_ORDER):
            # arrays too massive to process all variables at once.
            print var.name

            var_3hrlys = read_data(settings, suffix, var.name, year, grid_lats, grid_lons, period, N_OBS_PER_DAY)

            # reshape to days x 3hrly obs (365(366),8,180,360)
            var_3hrlys = var_3hrlys.reshape(-1, N_OBS_PER_DAY, var_3hrlys.shape[1], var_3hrlys.shape[2])

            # process the leap-year if appropriate
            if calendar.isleap(year):
                var_3hrlys, incl_feb29th  = process_february(var_3hrlys, doMask = True)
            else:
                assert var_3hrlys.shape[0] == 365

            # get pentadly values for each timestep (73,5,8,180,360)
            shape = var_3hrlys.shape
            var_3hrlys = var_3hrlys.reshape(-1, 5, shape[-3], shape[-2], shape[-1]) # n_pentads x days x hrs x lat x lon

            n_days_per_timestamp = np.ma.count(var_3hrlys, axis = 1) # n_pentads x hrs x lat x lon

            # get average at each timestamp across the pentad - so have N_OBS_PER_DAY averaged values per pentad
            if settings.doMedian:
                pentad_3hrly_grid = utils.bn_median(var_3hrlys, axis = 1) # n_pentads x hrs x lat x lon
            else:
                pentad_3hrly_grid = np.ma.mean(var_3hrlys, axis = 1) # n_pentads x hrs x lat x lon

            pentad_3hrly_grid.mask[n_days_per_timestamp < N_OBS_OVER_DAYS] = True # mask where fewer than N_OBS_OVER_DAYS days have values

            # clear up memory
            del var_3hrlys
            gc.collect()

            # the pentad containing feb 29th is the 11th in the year (KW actually its the 12th, so the 11 in array speak which is what you have done)
            if calendar.isleap(year):
                #  overwrite this with the me(di)an of a 6-day pentad
                if settings.doMedian:
                    pentad_3hrly_grid[11, :, :, :] = utils.bn_median(incl_feb29th, axis = 0)
                else:
                    pentad_3hrly_grid[11, :, :, :] = np.ma.mean(incl_feb29th, axis = 0)

                feb_n_days_per_timestamp = np.ma.count(incl_feb29th, axis = 0)
                pentad_3hrly_grid.mask[11, :, :, :][feb_n_days_per_timestamp < N_OBS_OVER_DAYS] = True
                n_days_per_timestamp[11, :, :, :] = feb_n_days_per_timestamp

                print "processed Feb 29th"

            if settings.plots and v == 0:
                import matplotlib.pyplot as plt
                plt.clf()
                plt.hist(n_days_per_timestamp.reshape(-1), bins = np.arange(-1,7), align = "left", log = True, rwidth=0.5)
                plt.axvline(x = N_OBS_OVER_DAYS-0.5, color = "r")
                plt.title("Number of days with obs at each 3hrly timestamp (over entire year)")
                plt.xlabel("Number of days (max = 5)")
                plt.ylabel("Frequency (log scale)")
                plt.savefig(settings.PLOT_LOCATION + "pentads_n_days_{}_{}_{}.png".format(year, period, suffix))

            # get single pentad values
            n_hrs_per_pentad = np.ma.count(pentad_3hrly_grid, axis = 1) # get the number of pentad-hours present in each pentad
            n_grids_per_pentad = np.sum(n_days_per_timestamp, axis = 1) # get the number of 3hrly 1x1 grids included per pentad 1x1

            # get average at each timestamp across the pentad - so have N_OBS_PER_DAY values per pentad
            if settings.doMedian:
                pentad_grid = utils.bn_median(pentad_3hrly_grid, axis = 1)
            else:
                pentad_grid = np.ma.mean(pentad_3hrly_grid, axis = 1)

            if period == "all":
# KW are you sure this should be n_hrs_per_pentad and not n_grids_per_pentad here? I think it should
                pentad_grid.mask[n_hrs_per_pentad < N_OBS_OVER_PENTAD] = True # mask where fewer than N_OBS_OVER_PENTAD hours have values
            else:
# KW are you sure this should be n_hrs_per_pentad and not n_grids_per_pentad here? I think it should
                pentad_grid.mask[n_hrs_per_pentad < (N_OBS_OVER_PENTAD/2.)] = True # mask where fewer than N_OBS_OVER_PENTAD hours have values

            all_pentads[v, :, :, :] = pentad_grid

            # diagnostics plots of obs/grids per pentad
            if settings.plots and v == 0:
                plt.clf()
                plt.hist(n_hrs_per_pentad.reshape(-1), bins = np.arange(-1,10), align = "left", log = True, rwidth=0.5)
                if period == "all":
                    plt.axvline(x = N_OBS_OVER_PENTAD-0.5, color = "r")
                else:
                    plt.axvline(x = (N_OBS_OVER_PENTAD/2.)-0.5, color = "r")
                plt.title("Number of hrs with obs in each pentad (over entire year)")
                plt.xlabel("Number of days (max = 8)")
                plt.ylabel("Frequency (log scale)")
                plt.savefig(settings.PLOT_LOCATION + "pentads_n_hrs_{}_{}_{}.png".format(year, period, suffix))

            # clear up memory
            del pentad_3hrly_grid
            del pentad_grid
            gc.collect()

        # done all main variables.  Now for number of observations
        print "n_obs"
        n_obs = read_data(settings, suffix, "n_obs", year, grid_lats, grid_lons, period, N_OBS_PER_DAY)
	# KW so we've gone from 8*365hrs,lats,lons to 365,8,lats,lons
        n_obs = n_obs.reshape(-1, N_OBS_PER_DAY, n_obs.shape[1], n_obs.shape[2])
        if calendar.isleap(year):
            n_obs, incl_feb29th  = process_february(n_obs, doMask = True)
        else:
            assert n_obs.shape[0] == 365

        shape = n_obs.shape
	# KW so we're now at pentads, 5days, 8hours, lats, lons
        n_obs = n_obs.reshape(-1, 5, shape[-3], shape[-2], shape[-1]) # pentads x days x hours x lat x lon

	# KW This should sum over the 5days leaving pentads, 8hrs, lats, lons
	# n_obs has -1 as missing data!!! So sum will not work properly
	# set up fill_value as -1
	n_obs.fill_value = -1
        n_obs_per_3hrly_pentad = np.ma.sum(n_obs, axis = 1)
        n_obs_per_3hrly_pentad.fill_value = -1

        if calendar.isleap(year):
            n_obs_per_3hrly_pentad[11, :, :, :] = np.ma.sum(incl_feb29th, axis = 0)

        n_obs_per_pentad = np.ma.sum(n_obs_per_3hrly_pentad, axis = 1)

        # and write out
        times = utils.TimeVar("time", "time since 1/1/{} in hours".format(year), "hours", "time")
        times.data = np.arange(0, all_pentads.shape[1]) * 5 * 24

        out_filename = settings.DATA_LOCATION + settings.OUTROOT + "_1x1_pentad_from_3hrly_{}_{}_{}.nc".format(year, period, suffix)

        utils.netcdf_write(out_filename, all_pentads, n_grids_per_pentad, n_obs_per_pentad, OBS_ORDER, grid_lats, grid_lons, times, frequency = "P")

    return # do_conversion

#************************************************************************
if __name__=="__main__":

    import argparse

    # set up keyword arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--suffix', dest='suffix', action='store', default = "relax",
                        help='"relax" or "strict" completeness, default = relax')
    parser.add_argument('--start_year', dest='start_year', action='store', default = defaults.START_YEAR,
                        help='which year to start run, default = 1973')
    parser.add_argument('--end_year', dest='end_year', action='store', default = defaults.END_YEAR,
                        help='which year to end run, default = present')
    parser.add_argument('--period', dest='period', action='store', default = "all",
                        help='which period to run for (day/night/all), default = "all"')
    parser.add_argument('--doQC', dest='doQC', action='store_true', default = False,
                        help='process the QC information, default = False')
    parser.add_argument('--doBC', dest='doBC', action='store_true', default = False,
                        help='process the bias corrected data, default = False')
    args = parser.parse_args()

    do_conversion(suffix = str(args.suffix), start_year = int(args.start_year), end_year = int(args.end_year), period = str(args.period), doQC = args.doQC, doBC = args.doBC)

# END
# ************************************************************************


#!/usr/bin/env python
# encoding: utf-8
"""
Sample for Multi Pixel Gibbs MCMC.
"""

from collections import OrderedDict
# import pprint
import json
import numpy as np

import fsps

from astropy.utils.console import ProgressBar
from astropy.table import Table, hstack, Column

from sedbot.photconv import micro_jy_to_luminosity
from sedbot.chain import MultiPixelChain

MODEL = None

class MultiPixelGibbsBgSampler(object):
    """MH-in-Gibbs sampler for three level hierarchical models of
    pixel parameters, global parameters, and a linearly sampled background
    parameter.

    Parameters
    ----------
    model_initializer : function
        Function that, when called, creates a model instance.
        This is done to enable multi-processing.
    n_cpu : int
        Number of CPUs to run on.
    """
    def __init__(self, model_initializer,
                 n_cpu=1):
        super(MultiPixelGibbsBgSampler, self).__init__()
        self._n_cpu = n_cpu
        self._model_initializer = model_initializer

        # Set up a local model instance
        self._model = self._model_initializer()

        # Set up a pixel compute pool
        self._map = self._init_pixel_compute_pool()

    def _init_pixel_compute_pool(self):
        """Initialize an ipython.parallel pool for processing pixels."""
        # Put the model in the global namespace for parallel processing
        global MODEL
        MODEL = self._model_initializer()

        # Setup up the pool
        if self._n_cpu == 1:
            map_fcn = map  # single processing
        else:
            map_fcn = None  # TODO
        return map_fcn

    def sample(self, n_iter,
               theta_prop=None,
               phi_prop=None,
               theta0=None,
               phi0=None,
               B0=None,
               chain=None,
               use_bkg_sample_errors=False):
        """Sample for `n_iter` Gibbs steps.

        Parameters
        ----------
        n_iter : int
            Number of Gibbs steps (number of full iterations over all
            parameters)
        theta_prop : ndarray
            Standard deviations of Gaussian proposal distribtions for
            theta (pixel) parameters. If any individual proposal sigma is
            zero then that parameter will not be sampled (but instead kept
            constant in the chain).
        phi_prop : ndarray
            Standard deviations of Gaussian proposal distributions for
            phi (global) parameters. Leave as ``None`` to use the model's
            defaults. If any individual proposal sigma is
            zero then that parameter will not be sampled (but instead kept
            constant in the chain).
        theta0 : ndarray
            Initial values for the theta parameters, a ``(n_pix, n_theta)``
            array.
        phi0 : ndarray
            Initial values for the phi parameters, a ``(n_phi)`` array.
        B0 : ndarray
            Initial values for the background, a ``(n_band)`` array.
            Background is in units of flux per arcsec^2.
        chain : :class:`sedbot.chain.MultiPixelChain`
            A previously-built chain, whose last sample will be used
            as the starting points for this sampling (in liue of setting
            `theta0`, `phi0` and `B0` manually).
        use_bkg_sample_errors : bool
            Use sample standard deviation when updating the background.
        """
        global MODEL

        # Keep these for storage in the HDF5 file later
        self._theta_prop = theta_prop
        self._phi_prop = phi_prop

        if chain is not None:
            theta0 = np.empty((self._model.n_pix, self._model.n_theta),
                              dtype=np.float)
            for i, n in enumerate(self._model.theta_params):
                theta0[:, i] = chain[n][-1]

            phi0 = np.empty(self._model.n_phi, dtype=np.float)
            for i, n in enumerate(self._model.phi_params):
                phi0[i] = chain[n][-1]

            B0 = np.empty(self._model.n_bands, dtype=np.float)
            for i, (n, instr) in enumerate(zip(self._model.observed_bands,
                                               self._model.instruments)):
                name = "B__{0}__{1}".format(instr, n)
                B0[i] = chain[name][-1]
        else:
            assert theta0 is not None
            assert phi0 is not None
            assert B0 is not None

        # initialize memory
        i0 = self._init_chains(n_iter, theta0, phi0, B0)

        with ProgressBar(n_iter) as bar:
            for i in xrange(i0, i0 + n_iter):
                # Sample for all pixels
                args = []
                for ipix in xrange(self._model.n_pix):
                    _post0 = self.pix_lnpost[i - 1, ipix]
                    _theta0 = self.theta[i - 1, ipix, :]
                    _phi0 = self.phi[i - 1, :]
                    _B0 = self.B[i - 1, :]
                    args.append((ipix, _post0, _theta0, _phi0, _B0,
                                 theta_prop))
                results = self._map(pixel_mh_sampler, args)
                for ipix, result in enumerate(results):
                    lnpost, theta, blob, theta_accept = result
                    self.pix_lnpost[i, ipix] = lnpost
                    self.theta[i, ipix, :] = theta
                    if blob is not None:
                        for k, v in blob.iteritems():
                            self.blobs[i][k][ipix] = v
                    else:
                        # repeat previous blob values
                        for k in self.blobs.dtype.fields:
                            self.blobs[i][k][ipix] = self.blobs[i - 1][k][ipix]
                    self._theta_n_accept[ipix, :] += theta_accept

                # Update the background
                model_seds = self.blobs[i]['model_sed']
                B_new, global_lnp, pixel_lnp, pixel_blobs \
                    = self._model.update_background(
                        self.theta[i, :, :],
                        self.phi[i - 1, :],
                        model_seds,
                        sample_errors=use_bkg_sample_errors)
                self.B[i, :] = B_new

                # Sample the global parameters
                phi_new, global_lnp, pixel_lnps, pixel_blobs, n_accept \
                    = self._global_mh(i,
                                      global_lnp,
                                      self.phi[i - 1, :],
                                      self.theta[i, :, :],
                                      self.B[i, :],
                                      phi_prop)
                # fill in new values to chain.
                self.phi[i, :] = phi_new
                self.lnpost[i] = global_lnp
                self._phi_n_accept += n_accept
                # Update ln post and blob data for all pixels too
                # it was done for the pixel-only step, but these values have
                # changed given the new background, etc.
                if pixel_lnps is not None:
                    for ipix, lnp in enumerate(pixel_lnps):
                        self.pix_lnpost[i, ipix] = lnp
                if pixel_blobs is not None:
                    for ipix, blobs in enumerate(pixel_blobs):
                        for k, v in blobs.iteritems():
                            self.blobs[i][k][ipix] = v

                bar.update()

    def _global_mh(self, i, lnpost0, phi, theta, B, phi_prop):
        """Perform MH-in-Gibbs at the global level."""
        phi = np.copy(phi)
        lnpost = lnpost0
        lnpost_pixel = None
        blobs = None
        n_accept = np.zeros(phi.shape[0], dtype=np.int)
        for j in xrange(phi.shape[0]):
            # Gaussian proposal for parameter j, only
            # Skip parameter advancement if proposal sigma is zero.
            if phi_prop[j] > 0.:
                phi_new = np.copy(phi)
                phi_new[j] += phi_prop[j] * np.random.randn()
                global_lnp, pixel_lnp, pixel_blobs \
                    = self._model.sample_global(theta,
                                                phi_new,
                                                B)
                ln_r = global_lnp - lnpost0
                reject = True
                if ~np.isfinite(ln_r):
                    pass
                elif ln_r >= 0.:
                    reject = False
                else:
                    x = np.random.random_sample()
                    if x < np.exp(ln_r):
                        reject = False
                if not reject:
                    # adopt new point
                    phi = phi_new
                    lnpost = global_lnp
                    lnpost_pixel = pixel_lnp
                    blobs = pixel_blobs
                    n_accept[j] = 1
        return phi, lnpost, lnpost_pixel, blobs, n_accept

    def _init_chains(self, n_iter, theta0, phi0, B0):
        """Initialize memory

        Note this could be adapted to extend existing chains.
        """
        # Make parameter chains
        self.theta = np.empty((n_iter + 1,
                               self._model.n_pix,
                               self._model.n_theta),
                              dtype=np.float)
        self.theta.fill(np.nan)
        self.theta[0, :, :] = theta0

        self.phi = np.empty((n_iter + 1,
                             self._model.n_phi),
                            dtype=np.float)
        self.phi.fill(np.nan)
        self.phi[0, :] = phi0

        self.B = np.empty((n_iter + 1,
                           self._model.n_bands),
                          dtype=np.float)
        self.B.fill(np.nan)
        self.B[0, :] = B0

        self.lnpost = np.empty(n_iter + 1, dtype=np.float)
        self.lnpost.fill(np.nan)

        self.pix_lnpost = np.empty((n_iter + 1, self._model.n_pix),
                                   dtype=np.float)
        self.pix_lnpost.fill(np.nan)

        # Make a structured array for the complex blob data
        self.blobs = np.empty(n_iter + 1, dtype=self._model.blob_dtype)
        self.blobs.fill(np.nan)

        # Initialize the starting point of the chains
        global_lnp, pixel_lnp, pixel_blobs \
            = self._model.sample_global(theta0, phi0, B0)
        self.lnpost[0] = global_lnp
        for ipix, lnp in enumerate(pixel_lnp):
            self.pix_lnpost[0, ipix] = lnp
        for ipix, blobs in enumerate(pixel_blobs):
            for k, v in blobs.iteritems():
                self.blobs[0][k][ipix] = v

        # Track acceptance of individual parameters for each pixel
        self._theta_n_accept = np.zeros((self._model.n_pix,
                                         self._model.n_theta),
                                        dtype=np.int)
        self._phi_n_accept = np.zeros(self._model.n_phi, dtype=np.int)

        return 1

    @property
    def median_theta_faccept(self):
        """Median acceptance fraction across all pixels for each parameter"""
        n_samples = float(self.theta.shape[0])
        faccept = np.median(self._theta_n_accept / n_samples, axis=0)
        return faccept

    @property
    def phi_faccept(self):
        """Acceptance fraction of phi parameters."""
        n_samples = float(self.theta.shape[0])
        return self._phi_n_accept / n_samples

    def print_faccept(self):
        """Printed report of the acceptance fraction statistics"""
        faccept = self.median_theta_faccept
        d = OrderedDict(zip(self._model.theta_params, faccept))
        # pp = pprint.PrettyPrinter(indent=4)
        # pp.pprint(d)
        print("theta")
        print(json.dumps(d, indent=4))
        print("phi")
        d = OrderedDict(zip(self._model.phi_params, self.phi_faccept))
        print(json.dumps(d, indent=4))

    @property
    def table(self):
        """An :class:`astropy.table.Table` with the chain."""
        msuns = np.array([fsps.get_filter(n).msun_ab
                          for n in self._model.computed_bands])
        theta_f_accept = json.dumps(dict(zip(self._model.theta_params,
                                             self.median_theta_faccept)))
        phi_f_accept = json.dumps(dict(zip(self._model.phi_params,
                                           self.phi_faccept)))
        meta = OrderedDict((
            ('theta_f_accept', theta_f_accept),
            ('phi_f_accept', phi_f_accept),
            ('observed_bands', self._model.observed_bands),
            ('instruments', self._model.instruments),
            ('computed_bands', self._model.computed_bands),
            ('msun_ab', msuns),
            ('band_indices', self._model.band_indices),
            ('theta_params', self._model.theta_params),
            ('phi_params', self._model.phi_params),
            ('theta_proposal_sigma', self._theta_prop),
            ('phi_proposal_sigma', self._theta_prop),
            ('sed', self._model._seds),
            ('sed_err', self._model._errs),
            ('pixels', self._model.pixel_metadata),
            ('area', self._model._areas)))
        # Make tables for individual chains; stack later
        # FIXME should axis order be changed for theta throughout the sampler?
        # or I can just continue to swarp aces here
        theta_table = Table(np.swapaxes(self.theta, 1, 2),
                            names=self._model.theta_params,
                            meta=meta)
        phi_table = Table(self.phi, names=self._model.phi_params)
        background_names = ["B__{0}__{1}".format(n, b)
                            for n, b in zip(self._model.instruments,
                                            self._model.observed_bands)]
        B_table = Table(self.B, names=background_names)
        blob_table = Table(self.blobs)
        tbl = MultiPixelChain(hstack((theta_table,
                                      phi_table,
                                      B_table,
                                      blob_table)))

        # Add M/L computations for each computed band.
        for i, (band_name, msun) in enumerate(zip(self._model.computed_bands,
                                                  msuns)):
            logLsol = micro_jy_to_luminosity(tbl['model_sed'][:, :, i],
                                             msun,
                                             np.atleast_2d(tbl['d']).T)
            ml = tbl['logMstar'] - logLsol
            colname = "logML_{0}".format(band_name)
            tbl.add_column(Column(name=colname, data=ml))

        return tbl

def pixel_mh_sampler(args):
    """Execute MH-in-Gibbs sampling for a single pixel.

    This function is design to be called by a ``map`` function, hence all
    arguments are wrapped in a single tuple, ``args``.

    Parameters
    ----------
    ipix : int
        Pixel index
    post0 : float
        Ln Posterior probability *for the pixel* given initial values of
        theta0 and phi0.
    theta0 : ndarray
        Initial values of theta parameters.
    phi0 : ndarray
        Initial values of the phi parameters.
    B0 : ndarray
        Initial values of the B, background, parameters.
    theta_prop : ndarray
        Array of Gaussian proposal standard deviations. If an individual
        proposal sigma is zero then that value will not be sampled,
        and kept constant.

    Returns
    -------
    post : float
        Final posterior probability *for this pixel*.
    theta : ndarray
        Theta parameters from this Gibbs step.
    blob : dict
        Stellar population metadata for this pixel given the updated `theta`
        parameters.
    n_accept : ndarray
        Array of length ``n_theta`` whose values are 1 for parameters
        that were updated, and zeros for values that remain constant from
        the previous iteration.
    """
    global MODEL
    ipix, post0, theta0, phi0, B0, theta_prop = args
    theta = np.copy(theta0)
    post = post0
    # MH on each parameter
    blob = None
    n_accept = np.zeros(theta.shape[0], dtype=np.int)
    for i in xrange(theta0.shape[0]):
        # Gaussian proposal for parameter i, only
        if theta_prop[i] > 0.:
            theta_new = theta.copy()
            theta_new[i] += theta_prop[i] * np.random.randn()
            lnpost_new, blob_new = MODEL.sample_pixel(theta_new,
                                                      phi0, B0, ipix)
            ln_r = lnpost_new - post
            # print "pix {0:d} param {1:d} ln_r {2:.2f}".format(ipix, i, ln_r),
            reject = True
            if ~np.isfinite(ln_r):
                # print " reject ln_r not finite"
                pass
            elif ln_r >= 0.:
                # print " accepted".format(ipix, i)
                reject = False
            else:
                x = np.random.random_sample()
                if x < np.exp(ln_r):
                    # print " accept {0:.2f} < r".format(x)
                    reject = False
                # else:
                    # print " reject {0:.2f} > r".format(x)
            if not reject:
                # adopt new point
                theta = theta_new
                blob = blob_new
                post = lnpost_new
                n_accept[i] = 1
    return post, theta, blob, n_accept


"""
(c) RIKEN 2015. All rights reserved.
Author: Keitaro Yamashita

This software is released under the new BSD License; see LICENSE.
"""
import numpy
import itertools
from yamtbx.dataproc.xds import get_xdsinp_keyword
from yamtbx.util import safe_float

class XPARM:
    def __init__(self, xparm_file=None):
        if xparm_file is not None:
            self.parse_xparm_file(xparm_file)
        else:
            self.starting_frame = 1
            self.starting_angle = 0.
            self.osc_range = -1.
            self.rotation_axis = numpy.array((1.,0.,0.))
            self.wavelength = -1.
            self.incident_beam = numpy.array((0.,0.,1.))
            self.nx, self.ny = 0, 0
            self.qx, self.qy = 0., 0.
            self.distance = 0.
            self.origin = numpy.array((0., 0.))
            self.X_axis = numpy.array((1., 0., 0.))
            self.Y_axis = numpy.array((0., 1., 0.))
            self.Z_axis = numpy.array((0., 0., 1.))
            self.spacegroup = 1
            self.unit_cell = numpy.array((100., 100., 100., 90., 90., 90.))
            self.a_axis = numpy.array((100., 0., 0.))
            self.b_axis = numpy.array((0., 100., 0.))
            self.c_axis = numpy.array((0., 0., 100.))
    # __init__()

    def space_group_info(self):
        from cctbx import sgtbx
        g = sgtbx.space_group_info(self.spacegroup).group()
        return g.info()

    def parse_xparm_file(self, xparm_file):
        lines = open(xparm_file).readlines()

        is_new_format = "XPARM.XDS" in lines[0]

        if not is_new_format:
            starting_frame, starting_angle, osc_range, rotx, roty, rotz = lines[0].split()
            wavelength, ibeamx, ibeamy, ibeamz = lines[1].split()
            nx, ny, qx, qy = lines[2].split()
            distance, orgx, orgy = lines[3].split()
            Xx, Xy, Xz = lines[4].split()
            Yx, Yy, Yz = lines[5].split()
            Zx, Zy, Zz = lines[6].split()
            spacegroup, a, b, c, alpha, beta, gamma = lines[7].split()
            ax, ay, az = lines[8].split()
            bx, by, bz = lines[9].split()
            cx, cy, cz = lines[10].split()
        else:
            starting_frame, starting_angle, osc_range, rotx, roty, rotz = lines[1].split()
            wavelength, ibeamx, ibeamy, ibeamz = lines[2].split()
            spacegroup, a, b, c, alpha, beta, gamma = lines[3].split()
            ax, ay, az = lines[4][:15], lines[4][15:30], lines[4][30:45]
            bx, by, bz = lines[5][:15], lines[5][15:30], lines[5][30:45]
            cx, cy, cz = lines[6][:15], lines[6][15:30], lines[6][30:45]
            nseg, nx, ny, qx, qy = lines[7].split()
            orgx, orgy, distance = lines[8].split()
            Xx, Xy, Xz = lines[9].split()
            Yx, Yy, Yz = lines[10].split()
            Zx, Zy, Zz = lines[11].split()

        self.starting_frame = int(starting_frame)
        self.starting_angle = float(starting_angle)
        self.osc_range = float(osc_range)
        self.rotation_axis = numpy.array((float(rotx), float(roty), float(rotz)))
        self.wavelength = float(wavelength)
        self.incident_beam = numpy.array((float(ibeamx), float(ibeamy), float(ibeamz)))
        self.nx = float(nx)
        self.ny = float(ny)
        self.qx = float(qx)
        self.qy = float(qy)
        self.distance = float(distance)
        self.origin = numpy.array((float(orgx), float(orgy)))
        self.X_axis = numpy.array((float(Xx), float(Xy), float(Xz)))
        self.Y_axis = numpy.array((float(Yx), float(Yy), float(Yz)))
        self.Z_axis = numpy.array((float(Zx), float(Zy), float(Zz)))
        self.spacegroup = int(spacegroup)
        self.unit_cell = numpy.array((float(a), float(b), float(c), float(alpha), float(beta), float(gamma)))
        self.a_axis = numpy.array(map(safe_float, (ax, ay, az)))
        self.b_axis = numpy.array(map(safe_float, (bx, by, bz)))
        self.c_axis = numpy.array(map(safe_float, (cx, cy, cz)))
    # parse_xparm_file()

    def set_info_from_xdsinp(self, xdsinp):
        # XXX x, y, z axes

        table = [("STARTING_FRAME", "starting_frame", lambda x: int(x)),
                 ("STARTING_ANGLE", "starting_angle", lambda x: float(x)),
                 ("OSCILLATION_RANGE", "osc_range", lambda x: float(x)),
                 ("ROTATION_AXIS", "rotation_axis", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("X-RAY_WAVELENGTH", "wavelength", lambda x: float(x)),
                 ("INCIDENT_BEAM_DIRECTION", "incident_beam", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("NX", "nx", lambda x: int(x)),
                 ("NY", "ny", lambda x: int(x)),
                 ("QX", "qx", lambda x: float(x)),
                 ("QY", "qy", lambda x: float(x)),
                 ("DETECTOR_DISTANCE", "distance", lambda x: float(x)),
                 ("SPACE_GROUP_NUMBER", "spacegroup", lambda x: int(x)),
                 ("UNIT_CELL_CONSTANTS", "unit_cell", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("UNIT_CELL_A-AXIS", "a_axis", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("UNIT_CELL_B-AXIS", "b_axis", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("UNIT_CELL_C-AXIS", "c_axis", lambda x: numpy.array(map(lambda y:float(y), x.split())))
                 ]
        inp = dict(get_xdsinp_keyword(xdsinp)) # I believe dict() removes duplicated parameters and keeps last.

        for k, at, f in table:
            if k in inp and inp[k].strip() != "":
                setattr(self, at, f(inp[k]))
        if "ORGX" in inp:
            self.origin[0] = float(inp["ORGX"])
        if "ORGY" in inp:
            self.origin[1] = float(inp["ORGY"])
    # set_info_from_xdsinp()

    def xparm_str(self, old_format=False):
        assert not old_format # Currently, only new format is supported!
        xparm_str = """ XPARM.XDS    VERSION March 30, 2013
%6d%14.4f%10.4f%10.6f%10.6f%10.6f
%15.6f%15.6f%15.6f%15.6f
%6d%12.6f%12.6f%12.6f%8.3f%8.3f%8.3f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%10d%10d%10d%12.6f%12.6f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%10d%10d%10d%10d%10d
%8.2f%8.2f%8.2f%9.5f%9.5f%9.5f%9.5f%9.5f%9.5f
""" % (self.starting_frame, self.starting_angle, self.osc_range, self.rotation_axis[0], self.rotation_axis[1], self.rotation_axis[2],
       self.wavelength, self.incident_beam[0], self.incident_beam[1], self.incident_beam[2],
       self.spacegroup, self.unit_cell[0], self.unit_cell[1], self.unit_cell[2], self.unit_cell[3], self.unit_cell[4], self.unit_cell[5],
       self.a_axis[0], self.a_axis[1], self.a_axis[2],
       self.b_axis[0], self.b_axis[1], self.b_axis[2],
       self.c_axis[0], self.c_axis[1], self.c_axis[2],
       1, self.nx, self.ny, self.qx, self.qy,
       self.origin[0], self.origin[1], self.distance,
       self.X_axis[0], self.X_axis[1], self.X_axis[2],
       self.Y_axis[0], self.Y_axis[1], self.Y_axis[2],
       self.Z_axis[0], self.Z_axis[1], self.Z_axis[2],
       1, 1, self.nx, 1, self.ny,
       0., 0., 0., 1., 0., 0., 0., 1., 0.,
       )
        return xparm_str
    # xparm_str()

    def crystal_symmetry(self):
        from cctbx import crystal
        return crystal.symmetry(tuple(self.unit_cell),
                                self.spacegroup)
    # crystal_symmetry()

    def update_cell_based_on_axes(self):
        from yamtbx.util.maths import vectors_angle

        a, b, c = map(lambda x: numpy.linalg.norm(x), (self.a_axis, self.b_axis, self.c_axis))
        al = vectors_angle(self.b_axis, self.c_axis)
        be = vectors_angle(self.c_axis, self.a_axis)
        ga = vectors_angle(self.a_axis, self.b_axis)
        al, be, ga = map(lambda x: numpy.rad2deg(x), (al, be, ga))

        self.unit_cell = numpy.array((a,b,c,al,be,ga))
        print
# class XPARM

def get_xparm_from_integrate_lp(lpfile, frame):
    assert 0 < frame

    keys = {"beam direction": "DIRECT BEAM COORDINATES (REC. ANGSTROEM)",
            "beam center": "DETECTOR ORIGIN (PIXELS) AT",
            "distance": "CRYSTAL TO DETECTOR DISTANCE (mm)",
            "rotation axis": "LAB COORDINATES OF ROTATION AXIS",
            "a axis": "COORDINATES OF UNIT CELL A-AXIS",
            "b axis": "COORDINATES OF UNIT CELL B-AXIS",
            "c axis": "COORDINATES OF UNIT CELL C-AXIS",
            "cell": "UNIT CELL PARAMETERS",
            "spacegroup": "SPACE GROUP NUMBER"
            }
    data = {}

    flag_read = False
    for l in open(lpfile):
        if "PROCESSING OF IMAGES" in l:
            flag_read = False
            l = l.strip()
            first, last = map(lambda x:int(x.strip()), l[l.index("PROCESSING OF IMAGES")+len("PROCESSING OF IMAGES"):].split("..."))
            if first <= frame <= last:
                flag_read = True

        if flag_read:
            for key, s in keys.items():
                if s in l:
                    l = l.strip()
                    val = map(lambda x:float(x.strip()), l[l.index(s)+len(s):].split())
                    data[key] = val

    beam = data["beam direction"]
    rotaxis = data["rotation axis"]
    distance = data["distance"][0]
    orgx, orgy = data["beam center"]
    spacegroup = data["spacegroup"][0]
    a, b, c, alpha, beta, gamma = data["cell"]
    aaxis = data["a axis"]
    baxis = data["b axis"]
    caxis = data["c axis"]

    xp = XPARM("XPARM.XDS")

    xp.rotation_axis = numpy.array(rotaxis)
    xp.incident_beam = numpy.array(beam)
    xp.spacegroup = spacegroup
    xp.unit_cell = numpy.array((a, b, c, alpha, beta, gamma))
    xp.a_axis = numpy.array(aaxis)
    xp.b_axis = numpy.array(baxis)
    xp.c_axis = numpy.array(caxis)
    xp.origin = numpy.array((orgx, orgy))
    xp.distance = distance

    return xp.xparm_str()
# get_xparm_from_integrate_lp()

def find_best_ucaxes_transform(xparm1, xparm2):
    """
    Reference: matrix iteration method in whirligig in CrystFEL-0.6.0 src/whirligig.c: gatinator()

    calculates the matrix m, which gives the best result; axes1 and m * axes2
    returns m, axes1, axes2, rms
    """

    a1, b1, c1 = xparm1.a_axis, xparm1.b_axis, xparm1.c_axis
    a2, b2, c2 = xparm2.a_axis, xparm2.b_axis, xparm2.c_axis

    P = numpy.vstack([a2, b2, c2])
    Q = numpy.vstack([a1, b1, c1])

    min_score = numpy.inf
    best_mat = None

    for x in itertools.product(xrange(-1,2), repeat=9):
        #if x != (1,0,0,0,1,0,0,0,1): continue
        m = numpy.array(x).reshape(3,3)
        if numpy.linalg.det(m) != 1: continue
        #print "Trying"
        #print m

        score = numpy.sum((numpy.dot(m, Q) - P)**2)
        #if x == (1,0,0,0,1,0,0,0,1): print score
        if score < min_score:
            min_score = score
            best_mat = m

    return best_mat, Q, P, numpy.sqrt(min_score/3.)
# find_best_ucaxes_transform()


# coding=utf-8
# Copyright 2021 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Functions related to sampling from a model."""

import jax
import jax.nn
import jax.numpy as jnp
from absl import logging
from flax.training import common_utils
from jax import lax, random

LARGE_NEGATIVE = -1e12

def multinomial(rng, logits):
    """Draws samples from a multinomial distribution.

    Args:
      rng: A jax.random.PRNGKey.
      logits: An array of shape (..., num_categories) containing unnormalized
        log-probabilities.

    Returns:
      An array of shape (...) containing sampled category indices.
    """
    probs = jax.nn.softmax(logits)
    probs = jnp.cumsum(probs, axis=-1)
    a = jax.random.uniform(rng, logits.shape[:-1] + (1,))
    out = jnp.argmin(a > probs, axis=-1)
    return out

def apply_repetition_penalty(
    sequences,
    logits,
    i,
    repetition_penalty,
    repetition_window,
    repetition_penalty_normalize,
):
    """Apply repetition penalty.

    Repetition penalty is introduced in Section 4.1. of the
    [CTRL paper](https://einstein.ai/presentations/ctrl.pdf).

    It involves reducing the logits corresponding to previously generated tokens.
    Given a list of generated tokens g, the next token probability is given by:

      p_{i + 1} propto exp(logit_i / (T * I_i))

    where I_i = repetition_penalty, if i in g,
                1, otherwise.

    Args:
      sequences: An array of shape (batch_size, max_seq_len) with the sequences.
      logits: An array of shape (batch_size, vocab_size) with next token logits.
      i: An array of shape () with the current sequence positions.
      repetition_penalty: A float indicating the repetition penalty.
      repetition_window: An int indicating the window size for repetition penalty.
        We reduce the probability of tokens that have been generated within
        repetition_window tokens prior to the token to be generated.
      repetition_penalty_normalize: A bool indicating whether to normalize
        the logits (log_softmax) before applying the repetition penalty. This
        ensures that all logits have the same sign.

    Returns:
      An array of shape (batch_size, vocab_size) with updated next token logits.
    """
    max_i = i  # We are currently generating a token for position i + 1.
    min_i = i - repetition_window + 1
    batch_size, vocab_size = logits.shape
    positions = jnp.arange(sequences.shape[1])
    positions = jnp.tile(
        positions[jnp.newaxis, :, jnp.newaxis], [batch_size, 1, vocab_size]
    )
    sequences_onehot = jnp.eye(vocab_size)[sequences]
    sequences_onehot = jnp.where(
        (positions >= min_i) * (positions <= max_i),
        sequences_onehot,
        jnp.zeros_like(sequences_onehot),
    )
    # Compute the indicator that a token appeared at least once in the
    # repetition window. Output shape: (batch_size, vocab_size).
    indicator = jnp.max(sequences_onehot, axis=1)
    # Compute a penalty tensor. The logits are divided by the penalty tensor.
    penalty_tensor = jnp.where(
        indicator, jnp.ones_like(logits) * repetition_penalty, jnp.ones_like(logits)
    )
    if repetition_penalty_normalize:
        logits = jax.nn.log_softmax(logits)
    # Dividing a negative logit by the penalty tensor actually increases the
    # resulting probability. Take the inverse for negative logits.
    penalty_tensor = jnp.where(logits > 0, penalty_tensor, 1 / penalty_tensor)

    logits = logits / penalty_tensor
    return logits

def temperature_sample(
    prompt,
    init_cache,
    tokens_to_logits,
    temperature=1.0,
    top_k=None,
    repetition_penalty=1,
    repetition_window=4,
    repetition_penalty_normalize=False,
    max_decode_len=512,
    rng=None,
    eos_token=None,
    pad_token=None,
    masked_tokens=None,
    use_lax_while_loop=True,
):
    """Temperature sampling for sequence models.

    Args:
      prompt: An array of shape (batch_size, prompt_length) containing the input
        prompt (the model consumes these tokens and starts generation after). For
        generic sampling, the prompt must be a single BOS token.
      init_cache: A flax.nn.attention.Cache object.
      tokens_to_logits: A fast autoregressive decoder function taking single token
        slices and the cache and returning next-token logits and updated cache.
      temperature: A float with the sampling temperature factor. As it approaches
        zero this sampling procedure becomes equivalent to greedy sampling.
      top_k: An int with the number of high probability tokens to keep for
        sampling. If None, keep all.
      repetition_penalty: A float with the repetition penalty. Values smaller
        (greater) than 1 encourage (discourage) repetition; 1 disables penalties.
      repetition_window: An int indicating the window size for repetition penalty.
        We reduce the probability of tokens that have been generated within
        repetition_window tokens prior to the token to be generated.
      repetition_penalty_normalize: A bool indicating whether to normalize the
        logits (log_softmax) before applying the repetition penalty. This ensures
        that all logits have the same sign.
      max_decode_len: An int indicating the maximum sequence length.
      rng: A jax.random.PRNGKey.
      eos_token: An int token id. If not None, we stop decoding a sequence after
        the first instance.
      pad_token: An int token used to pad sequences after the eos token. If none,
        we set pad_token to eos_token.
      masked_tokens: A list of int token id. If not None, we mask these token
        before sampling.
      use_lax_while_loop: A bool; whether to use a lax while loop or python loop.

    Returns:
      An array of shape (batch_size, max_decode_len) containing sequences. If
        variable-length, the sequences are right-padded with the EOS token.
    """
    batch_size = prompt.shape[0]
    eos_token = eos_token if eos_token is not None else -1
    if pad_token is None:
        logging.warn("Pad token is not provided. Using the EOS token.")
    pad_token = pad_token if pad_token is not None else eos_token
    end_marker = jnp.array(eos_token)
    temperature = jnp.array(temperature)

    # Initialize sampling loop state.
    rng0 = rng if rng is not None else random.PRNGKey(0)
    i0 = jnp.array(0)
    ended0 = jnp.zeros((batch_size, 1)).astype(jnp.bool_)

    # Initialize sequences with the prompt followed by the out_of_prompt_marker to
    # indicate when we can start generation.
    out_of_prompt_marker = jnp.array(-2)
    pad = (
        jnp.ones((batch_size, max_decode_len - prompt.shape[1]), dtype=jnp.int32)
        * out_of_prompt_marker
    )
    sequences0 = jnp.concatenate([prompt, pad], axis=1)
    token0 = sequences0[:, 0:1]

    # Sampling loop state is stored in a simple tuple.
    tokens_to_logits_state = None

    sampling_loop_init_state = (
        i0,
        sequences0,
        init_cache,
        token0,
        ended0,
        rng0,
        tokens_to_logits_state,
    )

    def sampling_loop_cond_fn(state):
        """Sampling loop termination condition."""
        (i, _, _, _, ended, _, _) = state
        # Have we reached max decoding length?
        not_at_end = i <= max_decode_len
        # Have all sampled sequences reached an end marker?
        all_sequences_ended = jnp.all(ended)
        return not_at_end & (~all_sequences_ended)

    def sampling_loop_body_fn(state):
        """Sampling loop state update."""
        i, sequences, cache, cur_token, ended, rng, tokens_to_logits_state = state

        # Split RNG for sampling.
        rng1, rng2 = random.split(rng)

        # Call fast-decoder model on current tokens to get raw next-position logits.
        logits, new_cache, new_tokens_to_logits_state = tokens_to_logits(
            cur_token, cache, internal_state=tokens_to_logits_state
        )
        logits = logits / temperature

        # Mask out the BOS token.
        if masked_tokens is not None:
            mask = common_utils.onehot(
                jnp.array(masked_tokens),
                num_classes=logits.shape[-1],
                on_value=LARGE_NEGATIVE,
            )
            mask = jnp.sum(mask, axis=0)[None, :]  # Combine multiple masks together
            logits = logits + mask

        # Apply the repetition penalty.
        if repetition_penalty != 1:
            logits = apply_repetition_penalty(
                sequences,
                logits,
                i,
                repetition_penalty=repetition_penalty,
                repetition_window=repetition_window,
                repetition_penalty_normalize=repetition_penalty_normalize,
            )

        # Mask out everything but the top-k entries.
        if top_k is not None:
            # Compute top_k_index and top_k_threshold with shapes (batch_size, 1).
            top_k_index = jnp.argsort(logits, axis=-1)[:, ::-1][:, top_k - 1 : top_k]
            top_k_threshold = jnp.take_along_axis(logits, top_k_index, axis=-1)
            logits = jnp.where(
                logits < top_k_threshold, jnp.full_like(logits, LARGE_NEGATIVE), logits
            )
        # Sample next token from logits.
        sample = multinomial(rng1, logits)
        next_token = sample.astype(jnp.int32)
        # Only use sampled tokens if we have past the out_of_prompt_marker.
        out_of_prompt = sequences[:, i + 1] == out_of_prompt_marker
        next_token = next_token * out_of_prompt + sequences[:, i + 1] * ~out_of_prompt
        # If end-marker reached for batch item, only emit padding tokens.
        next_token = next_token[:, None]
        next_token_or_endpad = jnp.where(
            ended, jnp.full_like(next_token, pad_token), next_token
        )
        ended |= next_token_or_endpad == end_marker
        # Add current sampled tokens to recorded sequences.
        new_sequences = lax.dynamic_update_slice(
            sequences, next_token_or_endpad, (0, i + 1)
        )
        return (
            i + 1,
            new_sequences,
            new_cache,
            next_token_or_endpad,
            ended,
            rng2,
            new_tokens_to_logits_state,
        )

    # Run sampling loop and collect final state.
    if use_lax_while_loop:
        final_state = lax.while_loop(
            sampling_loop_cond_fn, sampling_loop_body_fn, sampling_loop_init_state
        )
    else:
        final_state = sampling_loop_init_state
        while sampling_loop_cond_fn(final_state):
            final_state = sampling_loop_body_fn(final_state)

    final_sequences = final_state[1]
    # If generation ended early (all sampled sequences reached an end marker)
    # replace all remaining out_of_prompt_marker instances with the pad_token.
    final_sequences = jnp.where(
        final_sequences == out_of_prompt_marker,
        jnp.full_like(final_sequences, pad_token),
        final_sequences,
    )
    return final_sequences


#! /usr/bin/env python3
#
#  Copyright 2018 California Institute of Technology
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
# ISOFIT: Imaging Spectrometer Optimal FITting
# Author: David R Thompson, david.r.thompson@jpl.nasa.gov
#

from sys import platform
import json
import os
import re
import scipy as s
from common import json_load_ascii, combos, VectorInterpolator
from common import recursive_replace
from copy import deepcopy
from scipy.stats import norm as normal
from scipy.interpolate import interp1d
from rt_lut import TabularRT, FileExistsError

eps = 1e-5  # used for finite difference derivative calculations

class ModtranRT(TabularRT):
    """A model of photon transport including the atmosphere."""

    def __init__(self, config, instrument):

        TabularRT.__init__(self, config, instrument)

        self.modtran_dir = self.find_basedir(config)
        self.filtpath = os.path.join(self.lut_dir, 'wavelengths.flt')
        self.template = deepcopy(json_load_ascii(
            config['modtran_template_file'])['MODTRAN'])

        # Insert aerosol templates, if specified
        if 'aerosol_template_file' in config:
            self.template[0]['MODTRANINPUT']['AEROSOLS'] = \
                deepcopy(json_load_ascii(config['aerosol_template_file']))

        # Insert aerosol data, if specified
        if 'aerosol_model_file' in config:
            aer_data = s.loadtxt(config['aerosol_model_file'])
            self.aer_wl = aer_data[:, 0]
            aer_data = aer_data[:, 1:].T
            self.naer = int(len(aer_data)/3)
            aer_absc, aer_extc, aer_asym = [], [], []
            for i in range(self.naer):
                aer_extc.append(aer_data[i*3])
                aer_absc.append(aer_data[i*3+1])
                aer_asym.append(aer_data[i*3+2])
            self.aer_absc = s.array(aer_absc)
            self.aer_extc = s.array(aer_extc)
            self.aer_asym = s.array(aer_asym)

        # Build the lookup table
        self.build_lut(instrument)

    def find_basedir(self, config):
        '''Seek out a modtran base directory'''

        try:
            return config['modtran_directory']
        except KeyError:
            pass  # fall back to environment variable
        try:
            return os.getenv('MODTRAN_DIR')
        except KeyError:
            raise KeyError('I could not find the MODTRAN base directory')

    def load_tp6(self, infile):
        '''Load a .tp6 file.  This contains the solar geometry.  We
           Return cosine of mean solar zenith'''

        with open(infile, 'r') as f:
            ts, te = -1, -1  # start and end indices
            lines = []
            while len(lines) == 0 or len(lines[-1]) > 0:
                try:
                    lines.append(f.readline())
                except UnicodeDecodeError:
                    pass
            #lines = f.readlines()
            for i, line in enumerate(lines):
                if "SINGLE SCATTER SOLAR" in line:
                    ts = i+5
                if ts >= 0 and len(line) < 5:
                    te = i
                    break
            if ts < 0:
                raise ValueError('Could not find solar geometry in .tp6 file')
        szen = s.array([float(lines[i].split()[3])
                        for i in range(ts, te)]).mean()
        return szen

    def load_chn(self, infile, coszen):
        """Load a .chn output file and parse critical coefficient vectors.
           These are:
             wl      - wavelength vector
             sol_irr - solar irradiance
             sphalb  - spherical sky albedo at surface
             transm  - diffuse and direct irradiance along the
                          sun-ground-sensor path
             transup - transmission along the ground-sensor path only
           We parse them one wavelength at a time."""

        with open(infile) as f:
            sols, transms, sphalbs, wls, rhoatms, transups = [], [], [], [], [], []
            lines = f.readlines()
            for i, line in enumerate(lines):
                if i < 5:
                    continue
                toks = line.strip().split(' ')
                toks = re.findall(r"[\S]+", line.strip())
                wl, wid = float(toks[0]), float(toks[8])  # nm
                solar_irr = float(toks[18]) * 1e6 * \
                    s.pi / wid / coszen  # uW/nm/sr/cm2
                rdnatm = float(toks[4]) * 1e6  # uW/nm/sr/cm2
                rhoatm = rdnatm * s.pi / (solar_irr * coszen)
                sphalb = float(toks[23])
                transm = float(toks[22]) + float(toks[21])
                transup = float(toks[24])
                sols.append(solar_irr)
                transms.append(transm)
                sphalbs.append(sphalb)
                rhoatms.append(rhoatm)
                transups.append(rhoatm)
                wls.append(wl)
        params = [s.array(i) for i in
                  [wls, sols, rhoatms, transms, sphalbs, transups]]
        return tuple(params)

    def ext550_to_vis(self, ext550):
        return s.log(50.0) / (ext550 + 0.01159)

    def modtran_driver(self, overrides):
        """Write a MODTRAN 6.0 input file"""

        param = deepcopy(self.template)

        if hasattr(self, 'aer_absc'):
            fracs = s.zeros((self.naer))

        # Perform overrides
        for key, val in overrides.items():
            recursive_replace(param, key, val)

            if key.startswith('AER'):
                i = int(key.split('_')[-1])
                fracs[i] = val

            elif key == 'EXT550' or key == 'AOT550' or key == 'AOD550':
                # MODTRAN 6.0 convention treats negative visibility as AOT550
                recursive_replace(param, 'VIS', -val)

            elif key == 'FILTNM':
                param[0]['MODTRANINPUT']['SPECTRAL']['FILTNM'] = val

            elif key in ['ITYPE', 'H1ALT', 'IDAY', 'IPARM', 'PARM1',
                         'PARM2', 'GMTIME', 'TRUEAZ', 'OBSZEN']:
                param[0]['MODTRANINPUT']['GEOMETRY'][key] = val

        # For custom aerosols, specify final extinction and absorption
        # MODTRAN 6.0 convention treats negative visibility as AOT550
        if hasattr(self, 'aer_absc'):
            total_aot = fracs.sum()
            recursive_replace(param, 'VIS', -total_aot)
            total_extc = self.aer_extc.T.dot(fracs)
            total_absc = self.aer_absc.T.dot(fracs)
            norm_fracs = fracs/(fracs.sum())
            total_asym = self.aer_asym.T.dot(norm_fracs)

            # Normalize to 550 nm
            total_extc550 = interp1d(self.aer_wl, total_extc)(0.55)
            lvl0 = param[0]['MODTRANINPUT']['AEROSOLS']['IREGSPC'][0]
            lvl0['NARSPC'] = len(self.aer_wl)
            lvl0['VARSPC'] = [float(v) for v in self.aer_wl]
            lvl0['ASYM'] = [float(v) for v in total_asym]
            lvl0['EXTC'] = [float(v) / total_extc550 for v in total_extc]
            lvl0['ABSC'] = [float(v) / total_extc550 for v in total_absc]

        return json.dumps({"MODTRAN": param})

    def build_lut(self, instrument, rebuild=False):
        """ Each LUT is associated with a source directory.  We build a
            lookup table by:
              (1) defining the LUT dimensions, state vector names, and the grid
                  of values;
              (2) running modtran if needed, with each MODTRAN run defining a
                  different point in the LUT; and
              (3) loading the LUTs, one per key atmospheric coefficient vector,
                  into memory as VectorInterpolator objects."""

        # Regenerate MODTRAN input wavelength file
        if not os.path.exists(self.filtpath):
            self.wl2flt(instrument.wl, instrument.fwhm, self.filtpath)

        TabularRT.build_lut(self, instrument, rebuild)

    def rebuild_cmd(self, point, fn):

        vals = dict([(n, v) for n, v in zip(self.lut_names, point)])
        vals['DISALB'] = True
        vals['NAME'] = fn
        vals['FILTNM'] = os.path.normpath(self.filtpath)
        modtran_config_str = self.modtran_driver(dict(vals))

        # Check rebuild conditions: LUT is missing or from a different config
        infilename = 'LUT_'+fn+'.json'
        infilepath = os.path.join(self.lut_dir, infilename)
        outchnname = fn+'.chn'
        outchnpath = os.path.join(self.lut_dir, outchnname)
        if not os.path.exists(infilepath) or\
           not os.path.exists(outchnpath):
            rebuild = True
        else:
            with open(infilepath, 'r') as f:
                current = f.read()
                rebuild = (modtran_config_str.strip() != current.strip())

        if not rebuild:
            raise FileExistsError('File exists')

        # write_config_file
        with open(infilepath, 'w') as f:
            f.write(modtran_config_str)

        # Specify location of the proper MODTRAN 6.0 binary for this OS
        xdir = {'linux': 'linux', 'darwin': 'macos', 'windows': 'windows'}
        cmd = self.modtran_dir+'/bin/'+xdir[platform]+'/mod6c_cons '+infilename
        return cmd

    def load_rt(self, point, fn):
        tp6file = self.lut_dir+'/'+fn+'.tp6'
        solzen = self.load_tp6(tp6file)
        coszen = s.cos(solzen * s.pi / 180.0)
        chnfile = self.lut_dir+'/'+fn+'.chn'
        wl, sol, rhoatm, transm, sphalb, transup = self.load_chn(
            chnfile, coszen)
        return wl, sol, solzen, rhoatm, transm, sphalb, transup

    def wl2flt(self, wls, fwhms, outfile):
        """ helper function to generate Gaussian distributions around the center
            wavelengths """
        I = None
        sigmas = fwhms/2.355
        span = 2.0 * (wls[1]-wls[0])  # nm
        steps = 101

        with open(outfile, 'w') as fout:

            fout.write('Nanometer data for sensor\n')
            for wl, fwhm, sigma in zip(wls, fwhms, sigmas):

                ws = wl + s.linspace(-span, span, steps)
                vs = normal.pdf(ws, wl, sigma)
                vs = vs/vs[int(steps/2)]
                wns = 10000.0/(ws/1000.0)

                fout.write('CENTER:  %6.2f NM   FWHM:  %4.2f NM\n' %
                           (wl, fwhm))

                for w, v, wn in zip(ws, vs, wns):
                    fout.write(' %9.4f %9.7f %9.2f\n' % (w, v, wn))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SimEVLA.py: This module simulates the slewing motion of the EVLA for use in
determining optimized schedules. Based largely on the equivalent code created
by NRAO SSS for use in the Observation Planning Tool. Small modifications have
been implemented to make it more scriptable.
"""

__author__ = "Seth Bruzewski"
__email__ = "bruzewskis@unm.edu"
__license__ = "GPL"

# Built in Modules
import logging

# Third Party Modules
import numpy as np
from astropy.coordinates import SkyCoord, Angle, EarthLocation
import astropy.units as u

class MotionSimulator(object):
    '''
    A pythonic class simulator of the EVLA slewing system, for implementation
    in scheduling codes.
    '''

    # Generate a logger
    __logger = logging.getLogger(__name__)

    # Motion Ranges
    __MIN_AZ = Angle(-85.0, unit='deg')
    __MAX_AZ = Angle(445.0, unit='deg')
    __MIN_EL = Angle(8.0, unit='deg')
    __MAX_EL = Angle(125.0, unit='deg')

    # Velocity and acceleration rates
    __VEL_AZ = 40.0/60 * u.Unit('deg/s')
    __VEL_EL = 20.0/60 * u.Unit('deg/s')
    __ACC_AZ = 1000 * u.Unit('deg/s^2')
    __ACC_EL = 1000 * u.Unit('deg/s^2')

    # Setup and Settling Time
    __MIN_SETUP_TIME = 0 * u.s
    __SETTLING_TIME = 7.0 * u.s

    # Default currentAz/El to midpoint of their ranges.
    __currentAz = Angle(225.0, unit='deg')
    __currentEl = Angle(35.0, unit='deg')

    # These variables keep track of what the new Az/El will be set to during
    # the moveTime calculation
    __newAz = None
    __newEl = None

    # Keep track of error codes generated during moveTo
    __errors = []

    def __init__(self, startaz=None, startel=None, wrap=None):
        '''
        Creates a EvlaTelescopeMotionSimulator with an initial position of
        `startAz` and `startEl` that is on wrap `wrap`.

        Args
        =======
        startaz (Angle) - Azimuth to initialize the telescope with as an
            Astropy Angle object. If not provided, telescope is initialized
            to default Azimuth, which is 225 degrees.
        startel (Angle) - Elevation to initialize the telescope with as an
            Astropy Angle object. If not provided, telescope is initialized
            to default Elevation, which is 35 degrees.
        wrap (string) - Preferred wrap of the telescope, may be CLOCKWISE,
            COUNTERCLOCKWISE, or NO_PREFERENCE

        Returns
        =======
        None

        Raises
        =======
        TypeError - Raised by class functions setCurrentAntennaAzimuth or
            setCurrentAntennaElevation when inputs are not the correct type
        ValueError - Raised by class functions setCurrentAntennaAzimuth or
            setCurrentAntennaElevation when inputs are outside of range
        '''

        self.setCurrentAntennaAzimuth(startaz, wrap)
        self.setCurrentAntennaElevation(startel)

    def moveTo(self, azEnd, elEnd, wrap=None, goOverTheTop=False):
        '''
        Returns the estimated move time from our current saved position to
        `az`, `el`.  For the very first calculation, a starting position
        due south at 65 degrees elevation is assumed.  The method takes into
        account antenna wrap by calculating 4 different move times and
        returning the smallest valid move time.  The 4 cases are as follows:

        1) The position we're moving to is on the COUNTERCLOCKWISE (left)
        wrap. We simple rotate the antenna to that position and see how long
        it takes. This is a valid option only if `wrap` is not CLOCKWISE
        and `goOverTheTop` is false.

        2) The position is on the CLOCKWISE (right) wrap.  Again, we calculate
        how long it will take to rotate the antenna to that position on that
        wrap. This is a valid option only if {@code wrap} is not
        COUNTERCLOCKWISE and `goOverTheTop` is false.

        3) We rotate the antenna 180 degrees and then go over the top to get
        to our target position. (i.e. we're subtracting 180 degrees from the
        AZ in case 1 above.)  This is only valid if `goOverTheTop' is
        true and the elevation is greater than 55 degrees (180 - MAX_EL).

        4) The same as 3 above, but we add 180 to the az instead of
        subtracting.

        This method also clears and refills the list of errors encountered
        while attempting this move.

        Args
        =======
        azEnd (Angle) - Target azimuth to slew to as an Astropy Angle. Will
            be checked against telescope motion ranges.
        elEnd (Angle) - Target elevation to slew to as an Astropy Angle. Will
            be checked against telescope motion ranges.
        wrap (string) - Preferred wrap of the telescope, may be CLOCKWISE,
            COUNTERCLOCKWISE, or NO_PREFERENCE
        goOverTheTop (bool) - Whether or not to allow the telescope to
            traverse in elevation beyond the zenith. For standard observing
            this is not advised.

        Returns
        =======
        tmin (Quantity) - The total slew time expressed as an Astropy Quantity
            object with units of time.

        Raises
        =======
        None
        '''

        # Check type of inputs
        if not isinstance(azEnd, Angle) or not isinstance(elEnd, Angle):
            self.__errors.append('IMPROPER_INPUT_TYPE')
            self.__logger.error('azEnd or elEnd are not Astropy Angles')
            return np.nan

        # Start up stuff
        azStart = self.__currentAz.copy()
        elStart = self.__currentEl.copy()

        # Keeping track of things
        self.__newAz = None
        self.__newEl = None
        self.__errors = []

        # Initialize min to something large
        tmin = 10 * u.day

        # Initialize wrap to a default if it was None
        w = 'NO_PREFERENCE' if wrap is None else wrap

        toAz = azStart
        toAzEnd = azEnd

        # Put the toAz on the left or COUNTERCLOCKWISE wrap.
        if toAz > self.__MIN_AZ + 360*u.deg:
            toAz -= 360*u.deg

        # Make sure the end position is on the same wrap!  This is checked
        # separately from the above because we aren't guaranteed that just
        # because we need to subtract 360 from the start pos that we need to
        # do so to the end as well.  This is mainly a concern for scans that
        # cross the 0d/360d boundary.
        if toAzEnd > self.__MIN_AZ + 360*u.deg:
            toAzEnd -= 360*u.deg

        toEl = elStart
        toElEnd = elEnd

        # Debug stuff
        coords = [toAz,toEl,toAzEnd,toElEnd]
        logstr = "Moving : ({0:.2f},{1:.2f}) --> ({2:.2f},{3:.2f})"
        self.__logger.debug(logstr.format(*[a.deg for a in coords]))

        # Whereas Longitude is always gaurunteed to be in bounds of our
        # motion, Latitude (elevation) is not, so we have to check that here
        # before going any further.  We only check the min, because the max a
        # Latitude can hold is 90 and our MAX_EL is 125 or so.
        if toEl < self.__MIN_EL:
            self.__errors.append('ELEVATION_OUT_OF_RANGE')
            toEl = self.__MIN_EL.copy()

        # We need to do the same check for the elevation END position too!
        if toElEnd < self.__MIN_EL:
            self.__errors.append('SCAN_END_ELEVATION_OUT_OF_RANGE')
            toElEnd = self.__MIN_EL.copy()

        # This checks to make sure that the if the user selected a CLOCKWISE
        # preference, but the source is in the region where the wraps do NOT
        # overlap, we still get a valid calculation.  The reason this comes
        # up is because for much of the logic below we're considering a wrap
        # to be a full 360 degrees when it is only actually 265.  So, a
        # source could be in the CLOCKWISE wrap (180 - 445 degrees) but be
        # missed in the logic below if it's position is less than 275 degrees.
        # (which is within the range (-85, -85 + 360))
        if w=='CLOCKWISE' and toAz > 85*u.deg:
            w = 'NO_PREFERENCE'

        # tmp variables for calculations
        tmpAz = toAz.copy()
        tmpAzEnd = toAzEnd.copy()
        tmpEl = toEl.copy()
        tmpElEnd = toElEnd.copy()

        # Counterclockwise or No Pref
        if w!='CLOCKWISE':
            # Case 3:
            if goOverTheTop:
                print('case3')
                tmpAz = toAz - 180*u.deg
                tmpAzEnd = toAzEnd - 180*u.deg
                tmpEl = 180*u.deg - toEl
                tmpElEnd = 180*u.deg - toElEnd

                if (tmpAz >= self.__MIN_AZ and tmpEl <= self.__MAX_EL):
                    tmin, updateBool = self.__updateMin(tmin,tmpAzEnd,tmpElEnd)
                    if updateBool:
                        self.__newAz = tmpAzEnd
                        self.__newEl = tmpElEnd

                        time_str = tmin.to('s').value
                        case_str = "Case 3 min = {0:.2f} s"
                        self.__logger.debug(case_str.format(time_str))

            # Case 1:
            else:
                tmin, updateBool = self.__updateMin(tmin, tmpAzEnd, tmpElEnd)
                if updateBool:
                    self.__newAz = tmpAzEnd
                    self.__newEl = tmpElEnd

                    time_str = tmin.to('s').value
                    case_str = "Case 1 min = {0:.2f} s"
                    self.__logger.debug(case_str.format(time_str))

        # Clockwise or No Pref
        if w!='COUNTERCLOCKWISE':
            # Case 4
            if goOverTheTop:
                print('case4')
                tmpAz = toAz + 180*u.deg
                tmpEl = 180*u.deg - toEl
                tmpAzEnd = toAzEnd + 180.0*u.deg
                tmpElEnd = 180*u.deg - toElEnd

                if tmpAz <= self.__MAX_AZ and tmpEl <= self.__MAX_EL:
                    tmin, updateBool = self.__updateMin(tmin,tmpAzEnd,tmpElEnd)
                    if updateBool:
                        self.__newAz = tmpAzEnd
                        self.__newEl = tmpElEnd

                        time_str = tmin.to('s').value
                        case_str = "Case 4 min = {0:.2f} s"
                        self.__logger.debug(case_str.format(time_str))

            # Case 2
            else:
                tmpAz = toAz + 360*u.deg
                tmpAzEnd = toAzEnd + 360*u.deg
                tmpEl = toEl.copy()
                tmpElEnd = toElEnd.copy()
                if tmpAz < self.__MAX_AZ:
                    tmin, updateBool = self.__updateMin(tmin,tmpAzEnd,tmpElEnd)
                    if updateBool:
                        self.__newAz = tmpAzEnd
                        self.__newEl = tmpElEnd

                        time_str = tmin.to('s').value
                        case_str = "Case 2 min = {0:.2f} s"
                        self.__logger.debug(case_str.format(time_str))

        if self.__newAz is None or self.__newEl is None:
            self.__errors.append('MOVE_NOT_POSSIBLE')
        else:
            # Azimuth checks
            if self.__MIN_AZ > self.__newAz:
                self.__logger.error("newAz < min: " + self.__newAz.to_string())
                self.__errors.append('SCAN_END_AZIMUTH_OUT_OF_RANGE')
                self.__currentAz = self.__MIN_AZ.copy()
            elif self.__MAX_AZ < self.__newAz:
                self.__logger.error("newAz > max: " + self.__newAz.to_string())
                self.__errors.append('SCAN_END_AZIMUTH_OUT_OF_RANGE')
                self.__currentAz = self.__MAX_AZ.copy()
            else:
                self.__currentAz = self.__newAz

            # Elevation checks
            if self.__MIN_EL > self.__newEl:
                self.__logger.error("newEl < min: " + self.__newEl.to_string())
                self.__errors.append('SCAN_END_ELEVATION_OUT_OF_RANGE')
                self.__currentEl = self.__MIN_EL.copy()
            elif self.__MAX_EL < self.__newEl:
                self.__logger.error("newEl > max: " + self.__newEl.to_string())
                self.__errors.append('SCAN_END_ELEVATION_OUT_OF_RANGE')
                self.__currentEl = self.__MAX_EL.copy()
            else:
                self.__currentEl = self.__newEl

            # Add any settling time necessary.
            tmin += self.__SETTLING_TIME

            # If the slew took less time than the min. time it takes to
            # prepare the antenna for observing, use that min. instead.
            if tmin < self.__MIN_SETUP_TIME:
                tmin = self.__MIN_SETUP_TIME.copy()

        # Decompose units just to be safe
        tmin = tmin.decompose()

        return tmin

    def getErrors(self):
        '''
        Returns a list of any errrors that may have occured during the lastest
        call to the moveTo method. These errors are stored as simple strings
        for the user to parse or use as they see fit.

        Args
        =======
        None

        Returns
        =======
        self.__errors.copy() (list) - List of errors generated by moveTo

        Raises
        =======
        None
        '''

        return self.__errors.copy()

    def getCurrentAntennaWrap(self):
        '''
        Returns which AntennaWrap this simulator is currently on using the
        following rules.  If the currentAz is greater than 275 degrees (MIN_AZ
        + 360) then we're in the CLOCKWISE (right) wrap.  Otherwise we're in
        the COUNTERCLOCKWISE (left) wrap.  Note: This could be changed to split
        the range in half: greater than 180 degrees is CLOCKWISE, the rest is
        COUNTERCLOCKWISE.  As it is now, you are only in the CLOCKWISE wrap if
        you are in the overlapping region.

        Args
        =======
        None

        Returns
        =======
        wrap (string) - The current antenna wrap expressed as a string.

        Raises
        =======
        None
        '''

        if self.__currentAz > 180*u.deg:
            return 'CLOCKWISE' #AntennaWrap.CLOCKWISE
        else:
            return 'COUNTERCLOCKWISE' #AntennaWrap.COUNTERCLOCKWISE

    def getCurrentAntennaAzimuth(self):
        '''
        Returns a copy of the current antenna azimuth. This value should be
        in the range -85 to 445 degrees.

        Args
        =======
        None

        Returns
        =======
        self.__currentAz.copy() (Angle) - The current antenna azimuth
            expressed as an Astropy Angle object.

        Raises
        =======
        None
        '''

        return self.__currentAz.copy()

    def setCurrentAntennaAzimuth(self, a, w=None):
        '''
        Sets the current Antenna Az to a clone of Angle equivalent to
        `a` at wrap `w` if `a` is within range.

        Args
        =======
        None

        Returns
        =======
        self.__currentAz.copy() (Angle) - The current antenna azimuth
            expressed as an Astropy Angle object.

        Raises
        =======
        TypeError - Raised when inputs are not the correct type
        ValueError - Raised when inputs are outside of range
        '''

        # Check input type
        if not a is None:
            if not isinstance(a, Angle):
                raise TypeError('Initializing Azimuth is not Astropy Angle')

        # Check that input is in range
        if w is None:
            if a is not None and self.__MIN_AZ < a < self.__MAX_AZ:
                self.__currentAz = a.copy()
            elif a is None:
                pass
            else:
                raise ValueError('Invalid Antenna Azimuth: '+str(a.deg))
        else:
            self.setCurrentAntennaAzimuth(self.toAntennaAzimuth(a, w))

    def getCurrentAntennaElevation(self):
        '''
        Returns a copy of the current antenna elevation. This value should be
        in the range 8 to 125 degrees.

        Args
        =======
        None

        Returns
        =======
        self.__currentEl.copy() (Angle) - The current antenna elevation
            expressed as an Astropy Angle object.

        Raises
        =======
        None
        '''

        return self.__currentEl.copy()

    def setCurrentAntennaElevation(self, a):
        '''
        Sets the current Antenna El to a clone of `a` if a is within
        range.

        Args
        =======
        None

        Returns
        =======
        self.__currentEl.copy() (Angle) - The current antenna azimuth
            expressed as an Astropy Angle object.

        Raises
        =======
        TypeError - Raised when inputs are not the correct type
        ValueError - Raised when inputs are outside of range
        '''

        # Check input type
        if not a is None:
            if not isinstance(a, Angle):
                raise TypeError('Initializing Elevation is not Astropy Angle')

        # Check that input is in range
        if a is not None and self.__MIN_EL < a < self.__MAX_EL:
            self.__currentEl = a.copy()
        elif a is None:
            pass
        else:
            raise ValueError('Invalid Antenna Elevation: '+str(a.deg))

    def toAntennaAzimuth(self, az, w):
        '''
        Returns an Angle in degrees between -85 and 445 degrees that
        represents `az` at AntennaWrap `w`.

        Args
        =======
        az (Angle) - Initial azimuth to be converted
        w (string) - Current wrap of the antenna

        Returns
        =======
        antennaAz (Angle) - Proper azimuth inside the antenna wrap

        Raises
        =======
        None
        '''

        antennaAz = None

        if az is not None:
            # az is 0 - 360 degrees
            if w=='CLOCKWISE':
                # 180 to 445 degrees
                if az < self.__MAX_AZ-360*u.deg:
                    az += 360*u.deg
            elif w=='COUNTERCLOCKWISE':
                # -85 to 180 degrees
                if az > self.__MIN_AZ+360*u.deg:
                    az -= 360*u.deg

            antennaAz = az

        return antennaAz

    def getAzimuthMinimum(self):
        '''
        Returns the minimum azimuth value for EVLA antenna pointings.

        Args
        =======
        None

        Returns
        =======
        self.__MIN_AZ.copy (Angle) - Telescope minimum azmiuth

        Raises
        =======
        None
        '''

        return self.__MIN_AZ.copy()

    def getAzimuthMaximum(self):
        '''
        Returns the maximum azimuth value for EVLA antenna pointings.

        Args
        =======
        None

        Returns
        =======
        self.__MAX_AZ.copy (Angle) - Telescope maximum azmiuth

        Raises
        =======
        None
        '''
        return self.__MAX_AZ.copy()

    def getAzimuthDefault(self):
        '''
        Returns the default azimuth value for EVLA antenna pointings.

        Args
        =======
        None

        Returns
        =======
        DEF_AZ (Angle) - Telescope default azimuth

        Raises
        =======
        None
        '''
        return Angle(225, unit='deg')

    def getElevationMinimum(self):
        '''
        Returns the minimum elevation value for EVLA antenna pointings.

        Args
        =======
        None

        Returns
        =======
        self.__MIN_EL.copy (Angle) - Telescope minimum elevation

        Raises
        =======
        None
        '''
        return self.__MIN_EL.copy()

    def getElevationMaximum(self):
        '''
        Returns the maximum elevation value for EVLA antenna pointings.

        Args
        =======
        None

        Returns
        =======
        self.__MAX_AZ.copy (Angle) - Telescope maximum elevation

        Raises
        =======
        None
        '''
        return self.__MIN_EL.copy()

    def getElevationDefault(self):
        '''
        Returns the default elevation value for EVLA antenna pointings.

        Args
        =======
        None

        Returns
        =======
        DEF_EL (Angle) - Telescope default elevation

        Raises
        =======
        None
        '''
        return Angle(35, unit='deg')

    def __updateMin(self, currentMin, az, el):
        '''
        This method updates the passed in `currentMin` variable with
        the minimum value of 'tmin' in the moveTo function, the time it takes
        to move from currentAz to az, and the time it takes to move from
        currentEl to el. If the currentMin is changed, we return true.

        Args
        =======
        currentMin (Quantity) - The current slew time minimum expressed as an
            Astropy Quantity object with units of time.
        az (Angle) - Target azimuth to slew to.
        el (Angle) - Target elevation to slew to.

        Returns
        =======
        newCurrentMin (Quantity) - The new current slew time minimum expressed
            as an Astropy Quantity object with units of time.
        update (bool) - Whether or not `tmin` will have been changed

        Raises
        =======
        None
        '''
        taz = self.__calcAzMoveTime(self.__currentAz, az)
        tel = self.__calcElMoveTime(self.__currentEl, el)

        # t is the larger of taz and tel
        t = max([taz,tel])

        if currentMin > t:
            return t, True
        else:
            return currentMin, False

    def __calcAzMoveTime(self, f, t):
        '''
        Calculates the move time from angle `f` to angle `t` in Azimuth.

        Args
        =======
        f (Angle) - Initial azimuth angle
        t (Angle) - Final azimuth angle

        Returns
        =======
        telSlewAz (Quantity) - Slew time, accounting for factors like the
            telescope slew acceleration, expressed as an Astropy Quantity
            with units of time.

        Raises
        =======
        None
        '''
        azd = abs(f-t)

        # Time it takes to reach full speed
        # Accounts for both acceleration & decceleration.
        timeAccAz = 2 * self.__VEL_AZ / self.__ACC_AZ

        # Distance it takes to reach full speed
        distAccAz = (self.__VEL_AZ * self.__VEL_AZ) / self.__ACC_AZ

        # If the antenna never reaches full speed, use this equation.
        if azd < distAccAz:
            telSlewAz = 2 * np.sqrt(azd / self.__ACC_AZ)
        # otherwise, use this equation.
        else:
            telSlewAz = timeAccAz + (azd - distAccAz) / self.__VEL_AZ

        return telSlewAz

    def __calcElMoveTime(self, f, t):
        '''
        Calculates the move time from angle `f` to angle `t` in Elevation.

        Args
        =======
        f (Angle) - Initial elevation angle
        t (Angle) - Final elevation angle

        Returns
        =======
        telSlewEl (Quantity) - Slew time, accounting for factors like the
            telescope slew acceleration, expressed as an Astropy Quantity
            with units of time.

        Raises
        =======
        None
        '''
        eld = abs(f-t)

        # Time it takes to reach full speed
        # Accounts for both acceleration & decceleration.
        timeAccEl = 2 * self.__VEL_EL / self.__ACC_EL

        # Distance it takes to reach full speed
        distAccEl = (self.__VEL_EL * self.__VEL_EL) / self.__ACC_EL

        if eld < distAccEl:
            # If the antenna never reaches full speed, use this equation.
            telSlewEl = 2 * np.sqrt(eld / self.__ACC_EL)
        else:
            # otherwise, use this equation.
            telSlewEl = timeAccEl + (eld - distAccEl) / self.__VEL_EL

        return telSlewEl

def EquitorialToHorizontal(ra, dec, lst):
    '''
    This function converts Equitorial coordinates (RA/DEC) to a horizontal
    coordinate system local to the VLA (AZ/EL). Ideally utilizes astropy
    Angles, but will convert where possible. Also requires a specific time
    given in LST. Note that Azimuth here is measured clockwise from the North.

    Args
    =======
    ra (various) - Either an Astropy Angle class, or a string/number in
        hourangle units. The latter will be converted to an Astropy Angle.
    dec (various) - Either an Astropy Angle class, or a string/number in
        degree units. The latter will be converted to an Astropy Angle.
    lst (string) - The local sidereal time at the VLA to be converted for,
        ideally in format HH:MM:SS.SS. This will be converted into hourangle.

    Returns
    =======
    Az (Angle) - The resulting Azimuth generated from the inputs. Measured
        clockwise from due North.
    Alt (Angle) - The resulting Altitude/Elevation generated from the inputs.
        Measured from the horizon.

    Raises
    =======
    UnitTypeError - Raised if the inputs cannot be coerced into Astropy Angles
    '''

    # Check type of inputs for safety
    if not isinstance(ra, Angle):
        ra = Angle(ra, unit=u.hourangle)
    if not isinstance(dec, Angle):
        dec = Angle(dec, unit=u.deg)
    if not isinstance(lst, Angle):
        dec = Angle(lst, unit=u.hourangle)

    # Let's just assume we're at the VLA
    vla = EarthLocation.of_site('VLA')
    ha = lst - ra
    lat = vla.lat

    # Do Trigonometry
    cosAlt_sinAz = -np.cos(dec)*np.sin(ha)
    cosAlt_cosAz = np.sin(dec)*np.cos(lat) - np.cos(dec)*np.cos(ha)*np.sin(lat)
    sinAlt = np.sin(dec)*np.sin(lat) + np.cos(dec)*np.cos(ha)*np.cos(lat)

    Az = Angle(np.arctan2(cosAlt_sinAz,cosAlt_cosAz)).wrap_at('360d')
    Alt = Angle(np.arcsin(sinAlt))

    return Az, Alt

def PointToPointTime(start_sc, end_sc, lst):
    '''
    Calculate the time to slew between two coordinates on the sky. Note that
    this function will not check if the slew is possible, instead allowing
    MotionSimulator to throw errors where it likes. This is largely useful
    for diagnostic checking.

    Args
    =======
    start_sc (SkyCoord) - The coordinates of the first point on the sky, where
        the telescope will begin its slew. Must be an Astropy SkyCoord object
    end_sc (SkyCoord) - The coordinates of the second point on the sky, where
        the telescope will end its slew. Must be an Astropy SkyCoord object
    lst (string) - The local sidereal time at the VLA to be converted for,
        ideally in format HH:MM:SS.SS. This will be converted into hourangle.

    Returns
    =======
    slewtime (Quantity) - The time it takes to slew between the input points.
        This is returned as an Astropy Quantity with units of time, that way it
        can be easily worked with.

    Raises
    =======
    TypeError - Raised when start_sc or end_sc are not SkyCoords
    '''
    # Check type of inputs
    if not isinstance(start_sc, SkyCoord):
        raise TypeError('param \'start_sc\' is not a Astropy SkyCoord')
    if not isinstance(end_sc, SkyCoord):
        raise TypeError('param \'end_sc\' is not a Astropy SkyCoord')

    # Basic logging
    mlogger = logging.getLogger(__name__)

    # Translate Start Point
    start_az, start_el = EquitorialToHorizontal(start_sc.ra,start_sc.dec,lst)

    # Translate End Point
    end_az, end_el = EquitorialToHorizontal(end_sc.ra,end_sc.dec,lst)

    # Nice Debug
    coo_fmt = '(ra:{0:.2f},dec:{1:.2f}) --> (az:{2:.2f},el:{3:.2f})'
    coo1 = [start_sc.ra.deg, start_sc.dec.deg, start_az.deg, start_el.deg]
    coo2 = [end_sc.ra.deg, end_sc.dec.deg, end_az.deg, end_el.deg]
    mlogger.debug('Converting : '+coo_fmt.format(*coo1))
    mlogger.debug('Converting : '+coo_fmt.format(*coo2))

    # Generate Simulator
    sim = MotionSimulator(start_az, start_el)
    slewtime = sim.moveTo(end_az, end_el)
    mlogger.debug('Took {0:.2f} seconds to slew'.format(slewtime.value))

    # Errors
    err = sim.getErrors()
    mlogger.debug('Found {:d} errors'.format(len(err)))
    for e in err:
        mlogger.error(e)

    return slewtime

if __name__=='__main__':
    '''
    Psuedomain function for simple tests to give some idea how this works

    Args
    =======
    None

    Returns
    =======
    None

    Raises
    =======
    None
    '''

    # Test 1
    print('Test 1:')
    s1 = SkyCoord(21.915*u.hourangle, -30*u.deg)
    s2 = SkyCoord(21.915*u.hourangle, 30*u.deg)
    sc2str = lambda sc : sc.to_string('hmsdms')
    print('Slewing from {0:s} to {1:s}'.format(sc2str(s1), sc2str(s2)))

    lst0 = Angle('21:54:55', unit='hourangle')
    st = PointToPointTime(s1, s2, lst0)
    st_sec = st.to('s').value
    st_dst = s1.separation(s2).deg
    fmt_str = 'Took {0:.2f} seconds to slew {1:.2f} degrees'
    print(fmt_str.format(st_sec, st_dst))

    # Test 2
    print('\nTest 2:')
    sim = MotionSimulator()
    totaltime = 0*u.s
    for _ in range(10):
        new_az = Angle(np.random.uniform(0,360) * u.deg)
        new_el = Angle(np.random.uniform(10,90) * u.deg)
        slewt = sim.moveTo(new_az, new_el)
        totaltime += slewt

    t2_st = totaltime.to('s').value
    print('Slewed to 10 points over {0:.2f} seconds'.format(t2_st))

import numpy,import astropy,from astropy