#!/usr/bin/env python

""" A set of tools for modelling the antenna reponse and for callibrating the amplitude vs frequency of the antennas
based on pyCRtools. see Schellart et al. Detecting cosmic rays with the LOFAR radio telescope,  and Nelles et al. Calibrating the absolute amplitude scale for air showers measured at LOFAR

Note: LBA_ant_calibrator still needs some work.

author: Brian hare
"""

##internal
import glob
from pickle import load
import datetime

##external
import numpy as np
from matplotlib import pyplot as plt
from scipy.interpolate import interp1d
from scipy.interpolate import RegularGridInterpolator, pchip_interpolate
from scipy.io import loadmat

##mine
from LoLIM.utilities import processed_data_dir, MetaData_directory, RTD, SId_to_Sname, v_air
import LoLIM.transmogrify as tmf
import LoLIM.IO.metadata as md

## psuedo-code example use of calibrator
# AC = ant_calibrator( "D20160712T173455.100Z" )
#
# AC.FFT_prep( antenna_name, even_antenna_data,  odd_antnna_data ) ## note that ant_calibrator can be re-used for multiple data sets
# AC.apply_GalaxyCal()  ## this applies the relative callibration, and corrects for defficencies in the antenna model
# AC.unravelAntennaResponce(azimuth_degrees, elivation_degrees) # apply antenna model. Is better to NOT do this, and to apply anntenna model to E-field model instead
# out_1, out_2 = AC.get_result() ## out_1 is zenith component, out_2 is azimuthal component (if unravelAntennaResponce was called,  else is still even and odd antennas)
# out_1_hilbertEnvelope = np.abs( out_1 ) ## apply_GalaxyCal automatically applies a hilbert transform
# out_1_reals = np.real( out_1 )

#### TODO: imporve this so that it has a workspace for N frequencies, so that it doesn't have to allocate memory every call
## make it automatic so that subsequent calls use same internal memeory if lenght is same (but output doesn't share memory....)
class LBA_antenna_model:
    """a class encapsulating the antenna model for the Low Band Antennas."""

    def __init__(self):
        voltage_theta = np.loadtxt(MetaData_directory+"/lofar/antenna_response_model/LBA_Vout_theta.txt", skiprows=1)
        voltage_phi   = np.loadtxt(MetaData_directory+"/lofar/antenna_response_model/LBA_Vout_phi.txt", skiprows=1)

        voltage_theta_responce = voltage_theta[:, 3] + 1j*voltage_theta[:, 4]
        voltage_phi_responce = voltage_phi[:, 3] + 1j*voltage_phi[:, 4]

        freq_start = 10.0 * 1.e6
        freq_step = 1.0 * 1.e6
        num_freq = 101

        theta_start = 0.0
        theta_step = 5.0
        num_theta = 19

        phi_start = 0.0
        phi_step = 10.0
        num_phi = 37

        frequency_samples = np.arange(num_freq)*freq_step + freq_start
        theta_samples = np.arange(num_theta)*theta_step + theta_start
        phi_samples = np.arange(num_phi)*phi_step + phi_start

        voltage_theta_responce = voltage_theta_responce.reshape( (num_freq, num_theta, num_phi) )
        voltage_phi_responce = voltage_phi_responce.reshape( (num_freq, num_theta, num_phi) )

        self.theta_responce_interpolant = RegularGridInterpolator((frequency_samples, theta_samples, phi_samples),   voltage_theta_responce)
        self.phi_responce_interpolant =   RegularGridInterpolator((frequency_samples, theta_samples, phi_samples),   voltage_phi_responce)

    def JonesMatrix(self, frequency, zenith, azimuth):
        """return the Jones Matrix for a single frequency (in Hz), for a wave with a zenith and azimuth angle in degrees. Dot the jones matrix with the electric field vector, first component of vector is Zenith component of
        electric field and second component is azimuthal electric field, then the first component of the resulting vector will be voltage on odd  (X) antenna and second component will be voltage on even (Y) antenna.
        Returns identity matrix where frequency is outside of 10 to 100 MHz"""

        jones_matrix = np.zeros( (2,2), dtype=complex )

        if frequency < 10.0E6 or frequency>100.0E6: ##if frequency is outside of range, then return some invertable nonsense
            jones_matrix[0,0] = 1.0
            jones_matrix[1,1] = 1.0
            return jones_matrix

        ## calculate for X dipole
        azimuth += 135 # put azimuth in coordinates of the X antenna
        while azimuth > 360: ## normalize the azimuthal angle
            azimuth -= 360
        while azimuth < 0:
            azimuth += 360

        jones_matrix[0, 0] = self.theta_responce_interpolant(  [frequency, zenith, azimuth] )
        jones_matrix[0, 1] = -1*self.phi_responce_interpolant( [frequency, zenith, azimuth] ) ## I don't really know why this -1 must be here

        ## calculate for Y dipole
        azimuth += 90.0 # put azimuth in coordinates of the Y antenna
        while azimuth > 360: ## normalize the azimuthal angle
            azimuth -= 360
        while azimuth < 0:
            azimuth += 360
        jones_matrix[1, 0] = -1*self.theta_responce_interpolant(  [frequency, zenith, azimuth] ) ## I don't really know why this -1 must be here
        jones_matrix[1, 1] = self.phi_responce_interpolant( [frequency, zenith, azimuth] )

        return jones_matrix

    def JonesMatrix_MultiFreq(self, frequencies, zenith, azimuth, out=None):
        """same as JonesMatrix, except that frequencies is expected to be an array. Returns an array of jones matrices"""

        if out is None:
            out_JM = np.zeros( (len(frequencies), 2,2), dtype=complex )
        else:
            out_JM = out

        good_frequencies = np.logical_and( frequencies>10.0E6, frequencies<100E6)
        num_freqs = np.sum( good_frequencies )

        points = np.zeros( (num_freqs, 3) ) ## figure out how to not need this
        points[:, 0] = frequencies[ good_frequencies ]
        points[:, 1] = zenith

        ## calculate for X dipole
        points[:, 2] = azimuth + 135 # put azimuth in coordinates of the X antenna
        while np.any( points[:, 2] > 360 ): ## normalize the azimuthal angle
            points[:, 2] [ points[:, 2]>360 ] -= 360
        while np.any( points[:, 2] <0 ):
            points[:, 2] [ points[:, 2]<0 ] += 360

        out_JM[good_frequencies, 0, 0] = self.theta_responce_interpolant( points )
        out_JM[good_frequencies, 0, 1] = -1*self.phi_responce_interpolant( points )

        ## calculate for Y dipole
        points[:, 2] += 90.0 # put azimuth in coordinates of the Y antenna
        while np.any( points[:, 2] > 360 ): ## normalize the azimuthal angle
            points[:, 2] [ points[:, 2]>360 ] -= 360
        while np.any( points[:, 2] <0 ):
            points[:, 2] [ points[:, 2]<0 ] += 360

        out_JM[good_frequencies, 1, 0] = -1*self.theta_responce_interpolant( points )
        out_JM[good_frequencies, 1, 1] = self.phi_responce_interpolant( points )

        ## set the frequencies outide 10 to 100 MHz to just identity matix
        out_JM[ np.logical_not(good_frequencies), 0, 0] = 1.0
        out_JM[ np.logical_not(good_frequencies), 1, 1] = 1.0

#        fi = np.argmin( np.abs(frequencies-60.0E6) )

        return out_JM

def get_LBA_frequency_calibration(frequencies, fill=0.0):
    """ given a set of frequencies, in units of Hz, return the antenna callibration"""

    Calibration_curve = np.zeros(101)

#    Calibration_curve[29:82] = np.array([0,  1.09124663e-06,   1.11049910e-06,   1.11101995e-06,
#                 1.14234774e-06,   1.15149299e-06,   1.17121699e-06,
#                 1.18578121e-06,   1.19696124e-06,   1.20458122e-06,
#                 1.24675978e-06,   1.27966600e-06,   1.32418333e-06,
#                 1.32115453e-06,   1.33871075e-06,   1.34295545e-06,
#                 1.34157430e-06,   1.37660390e-06,   1.39226359e-06,
#                 1.39827006e-06,   1.51409426e-06,   1.61610247e-06,
#                 1.74643510e-06,   1.74588169e-06,   1.73061463e-06,
#                 1.69229172e-06,   1.64633321e-06,   1.60982965e-06,
#                 1.59572009e-06,   1.64618678e-06,   1.81628916e-06,
#                 2.09520281e-06,   2.17610590e-06,   2.20907337e-06,
#                 2.12050148e-06,   2.04923844e-06,   2.06549879e-06,
#                 2.24906987e-06,   2.40356459e-06,   2.52199062e-06,
#                 2.48380048e-06,   2.40835417e-06,   2.38248922e-06,
#                 2.48599834e-06,   2.60617662e-06,   2.66466169e-06,
#                 2.78010597e-06,   2.90548503e-06,   3.08686745e-06,
#                 3.26101312e-06,   3.50261561e-06,   3.74739666e-06, 0])

#    Calibration_curve[29:82] = np.array([0,
#         5.17441583458e-07,  5.26112551988e-07,   5.54435685779e-07,
#         5.85016839251e-07,  6.10870583894e-07,   6.47309035485e-07,
#         6.5091456311e-07,   6.92446769829e-07,   7.24758841756e-07,
#         7.68263947385e-07,  7.91148165708e-07,   8.40734630391e-07,
#         8.43296509798e-07,  8.9313176062e-07,    9.2381731899e-07,
#         9.49238145305e-07,  9.75552986024e-07,   9.90712789075e-07,
#         1.05134677402e-06,  1.07555421962e-06,   1.09540733982e-06,
#         1.11101225721e-06,  1.13901998747e-06,   1.20840289547e-06,
#         1.25099501879e-06,  1.28820016316e-06,   1.42132795525e-06,
#         1.58730515709e-06,  1.69274022328e-06,   1.79771448679e-06,
#         1.69009894816e-06,  1.59916753151e-06,   1.52598688512e-06,
#         1.35488689628e-06,  1.25281147938e-06,   1.25843707081e-06,
#         1.28115364631e-06,  1.29327138907e-06,   1.30804348155e-06,
#         1.30108509354e-06,  1.30377069039e-06,   1.30654052835e-06,
#         1.30565358077e-06,  1.31712312807e-06,   1.31343342371e-06,
#         1.32267337872e-06,  1.33943924503e-06,   1.36389183621e-06,
#         1.4202852471e-06,   1.45672364953e-06,   1.5184892913e-06, 0])

    Calibration_curve[29:82] = np.array([fill, 1.37321451961e-05,
                                             1.39846332239e-05,
                                             1.48748993821e-05,
                                             1.54402170354e-05,
                                             1.60684568225e-05,
                                             1.66241942741e-05,
                                             1.67039066047e-05,
                                             1.74480931848e-05,
                                             1.80525736486e-05,
                                             1.87066855054e-05,
                                             1.88519099831e-05,
                                             1.99625051386e-05,
                                             2.01878566584e-05,
                                             2.11573680797e-05,
                                             2.15829455528e-05,
                                             2.20133824866e-05,
                                             2.23736319125e-05,
                                             2.24484419697e-05,
                                             2.37802483891e-05,
                                             2.40581543111e-05,
                                             2.42020383477e-05,
                                             2.45305869187e-05,
                                             2.49399905965e-05,
                                             2.63774023804e-05,
                                             2.70334253414e-05,
                                             2.78034857678e-05,
                                             3.07147991391e-05,
                                             3.40755705892e-05,
                                             3.67311849851e-05,
                                             3.89987440028e-05,
                                             3.72257913465e-05,
                                             3.54293510934e-05,
                                             3.35552370942e-05,
                                             2.96529815929e-05,
                                             2.79271252352e-05,
                                             2.8818544973e-05,
                                             2.92478843809e-05,
                                             2.98454768706e-05,
                                             3.07045462103e-05,
                                             3.07210553534e-05,
                                             3.16442871206e-05,
                                             3.2304638838e-05,
                                             3.33203882046e-05,
                                             3.46651060935e-05,
                                             3.55193137077e-05,
                                             3.73919275937e-05,
                                             3.97397037914e-05,
                                             4.30625048727e-05,
                                             4.74612081994e-05,
                                             5.02345866124e-05,
                                             5.53621848304e-05, fill])

    Calibration_curve_interp = interp1d(np.linspace(0.e6,100e6,101), Calibration_curve, kind='linear', bounds_error=False, fill_value=fill)
    return Calibration_curve_interp( frequencies )

def invert_2X2_matrix_list( matrices ):
    """ if matrices is an array of 2x2 matrices, then return the array of inverse matrices """
    num = len(matrices)
    out = np.zeros( (num, 2,2), dtype=matrices.dtype)

    out[:, 0,0] = matrices[:, 1,1]
    out[:, 0,1] = -matrices[:, 0,1]
    out[:, 1,0] = -matrices[:, 1,0]
    out[:, 1,1] = matrices[:, 0,0]

    determinants = matrices[:, 0,0]*matrices[:, 1,1] - matrices[:, 0,1]*matrices[:, 1,0]

    out /= determinants[:, np.newaxis, np.newaxis]

    return out

def fourier_series( x, p):
    """Evaluates a partial Fourier series

        F(x) \\approx \\frac{a_{0}}{2} + \\sum_{n=1}^{\\mathrm{order}} a_{n} \\sin(nx) + b_{n} \\cos(nx)
    """

    r = p[0] / 2

    order = int( (len(p) - 1) / 2 )

    for i in range(order):

        n = i + 1

        r += p[2*i + 1] * np.sin(n * x) + p[2*i + 2] * np.cos(n * x)

    return r

def getGalaxyCalibrationData(antenna_noise_power, timestamp, channel_width, antenna_type="outer"):
    """return factor to correct for amplitude shifts. Essenturally returns sqrt( P_{expected} / P_{measured} ). Where P is noise power.
    for antenna_type outer it returns factor for Y/X dipoles, for "inner" returns "X/Y". antenna_noise_power is an array of measured powers for each antenna.
    Even/odd indecies should be Y/X dipole for outer and oppisite for inner.
    timestamp should be posix timestamp"""

    longitude = 6.869837540/RTD

    ## this is in outer order:  Y,X
    coefficients_lba = [ np.array([ 0.01489468, -0.00129305,  0.00089477, -0.00020722, -0.00046507]),   ## for Y antennas
                         np.array([ 0.01347391, -0.00088765,  0.00059822,  0.00011678, -0.00039787])  ] ## for X antennas

    # Convert timestamp to datetime object
    t = datetime.datetime.utcfromtimestamp(timestamp)
    # Calculate JD(UT1)
    ut = tmf.gregorian2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + float(t.second) / 3600.) / 24.))
    # Calculate JD(TT)
    dtt = tmf.delta_tt_utc(tmf.date2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + float(t.second) / 3600.) / 24.)))
    tt = tmf.gregorian2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + (float(t.second) + dtt / 3600.)) / 24.))
    # Calculate Local Apparant Sidereal Time
    last = tmf.rad2circle(tmf.last(ut, tt, longitude))

    galactic_noise_power =[  fourier_series(last, coefficients_lba[0]),  fourier_series(last, coefficients_lba[1])  ]

    antenna_noise_power[ antenna_noise_power==0 ] = np.nan

    if antenna_type == 'outer':
        Y_measured_powers = antenna_noise_power[0::2]
        X_measured_powers = antenna_noise_power[1::2]

        Y_expected_power = galactic_noise_power[0]
        X_expected_power = galactic_noise_power[1]

    else:
        X_measured_powers = antenna_noise_power[0::2]
        Y_measured_powers = antenna_noise_power[1::2]

        X_expected_power = galactic_noise_power[0]
        Y_expected_power = galactic_noise_power[1]

    ## note this should make new arrays
    Y_factors = Y_expected_power * channel_width / Y_measured_powers
    X_factors = X_expected_power * channel_width / X_measured_powers

    np.sqrt(Y_factors, out=Y_factors)
    np.sqrt(X_factors, out=X_factors)

    if antenna_type == 'outer':
        return Y_factors, X_factors
    else:
        return X_factors, Y_factors

class LBA_ant_calibrator:
    """ This is a class for callibrating the LBA antennas and removing the antenna responce function. Only valid between 30 to 80 MHz. NOTE: removing the antenna responce function is ill-conditioned. A better approach is to callibrate the data,
    then filter a model using the antenna responce. Using this class will inherently do a hilbert transform (negative frequencies will be set to zero)"""

    def __init__(self, load_gal_cal=True, timeID=None, findRFI=None, cal_curve=get_LBA_frequency_calibration, jones_matrices=None):
        """if load_gal_cal is true, then timeID and  findRFI should indicate RFI info. If load_gal_cal is false, than data is not normalized relative to noise.
        cal_curve should take numpy arrays of frequency and return frequency dependant adjustment of the data. can be None.
        Jones matrices should be a function that has three parameters: 1) numpy array of frequencies in Hz, 2) zenith in degrees, and 3) azimuth in degrees, then returns array of jones matrices. Default is pycrtools model """

        self.cal_curve = cal_curve

        if load_gal_cal:

            if findRFI is None:
                findRFI = "/findRFI/findRFI_results"

            if isinstance(findRFI, str): ## load findRFI data from file
                galcal_data_loc = processed_data_dir(timeID) + findRFI
                with open( galcal_data_loc, 'rb' ) as fin:
                    findRFI = load(fin)

            self.calibration_factors = {}
            for findRFI_info in findRFI.values():
                antenna_names = findRFI_info["antenna_names"]
                num_antennas = len( antenna_names )
                cleaned_power = findRFI_info["cleaned_power"]
                timestamp = findRFI_info["timestamp"]
                analyzed_blocksize = findRFI_info["blocksize"]

                even_cal_factors, odd_cal_factors = getGalaxyCalibrationData(cleaned_power,  timestamp, 5.0E-9/analyzed_blocksize )

                for ant_i in range(0, int(num_antennas/2)):
                    ant_name = antenna_names[ ant_i*2 ]
                    self.calibration_factors[ ant_name ] = (even_cal_factors[ant_i], odd_cal_factors[ant_i])
        else:
            self.calibration_factors = None

        ### Galaxy callibration data ###
#        self.calibration_factors = {}
#        galcal_data_loc = processed_data_dir(timeID) + '/cal_tables/galaxy_cal'
#        galcal_fpaths = glob.glob(galcal_data_loc + '/*.gcal')
#        for fpath in galcal_fpaths:
#            with open(fpath, 'rb') as fin:
#                data = np.load(fin)
#                ant_names = data["arr_0"].astype(str, copy=False)
#                factors = data["arr_1"]
#            ant_i = 0
#            while ant_i<len(ant_names):
#                self.calibration_factors[ ant_names[ant_i] ] = [factors[ant_i], factors[ant_i+1]]
#                ant_i += 2

        if jones_matrices is None:
            self.antenna_model = LBA_antenna_model().JonesMatrix_MultiFreq

    def FFT_prep(self, even_ant_name, even_pol_data, odd_pol_data):
        """ prepare to apply callibrations to a pair of dipoles. Essentually just takes FFT. Assume a han window is not needed."""

        if len(even_pol_data) != len(odd_pol_data):
            raise ValueError('even and odd polarization data need to be same length')

        self.even_ant_name = even_ant_name
        self.N_points = len(even_pol_data)

        ##FFT
        self.even_pol_FFT = np.fft.fft(even_pol_data)
        self.odd_pol_FFT =  np.fft.fft(odd_pol_data)

        ##Frequencies
        self.frequencies = np.fft.fftfreq(self.N_points, 5.0E-9)
#
    def apply_time_shift(self, even_time_shift, odd_time_shift):
        """apply phase shifts to the data, corresponding to an amount of time (in seconds)"""
        self.even_pol_FFT *= np.exp( self.frequencies*(-1j*2*np.pi*even_time_shift) )
        self.odd_pol_FFT *= np.exp( self.frequencies*(-1j*2*np.pi*odd_time_shift) )

    def apply_GalaxyCal(self):
        """applies the galaxy calibration to this data. Can be applied independantly of unravelAntennaResponce. Negative frequencies are set to zero."""

        ### first we apply the correction factors ###
        if self.calibration_factors is not None:
            PolE_factor, PolO_factor = self.calibration_factors[self.even_ant_name]

            self.even_pol_FFT *= PolE_factor
            self.odd_pol_FFT  *= PolO_factor
        else:
             PolE_factor = 1.0
             PolO_factor = 1.0

        if self.cal_curve is not None:
            ## callibration curve to correct for frequency dependent model defeciencies
            Calibration_curve_interp = get_LBA_frequency_calibration( self.frequencies )

            self.even_pol_FFT *= Calibration_curve_interp
            self.odd_pol_FFT  *= Calibration_curve_interp

        return [PolE_factor, PolO_factor]

    def unravelAntennaResponce(self, zenith, azimuth):
        """given a direction to source (azimuth off X and zenith from Z, in degrees ), if call this function, then apply_GalaxyCal MUST also be applied to the data
        Note that this function assumes the data is LBA_outer, which has flipped polarizations compared to LBA inner"""

        jones_matrices = self.antenna_model(self.frequencies, zenith, azimuth)

        inverse_jones_matrix = invert_2X2_matrix_list( jones_matrices )

        ### apply the Jones matrix.  Note that the polarities (even and odd) are flipped)
        zenith_component = self.odd_pol_FFT*inverse_jones_matrix[:, 0,0] +  self.even_pol_FFT*inverse_jones_matrix[:, 0,1]
        azimuth_component = self.odd_pol_FFT*inverse_jones_matrix[:, 1,0] +  self.even_pol_FFT*inverse_jones_matrix[:, 1,1]

        self.even_pol_FFT = zenith_component
        self.odd_pol_FFT = azimuth_component

    def getResult(self):
        """get the results of analysis. Essentially preforms inverse FFT.
        If unravelAntennaResponce was called, first return is zenith component, second is azimuthal. Else first return is even polarization, second is odd"""

        return np.fft.ifft(self.even_pol_FFT),    np.fft.ifft(self.odd_pol_FFT)

def plot_responce():
    N_azimuth = 100
    N_zenith = 500
    frequency = 58.0E6

    zeniths = np.linspace( 0, 90, N_zenith)
    azimuths = np.linspace( 0, 360, N_azimuth )

    resulting_grid = np.zeros((N_zenith, N_azimuth))
    AM = LBA_antenna_model()
#    AM = pycrtools_antenna_model()

    for ze_i in range(N_zenith):
        print(ze_i, '/', N_zenith)
        for az_i in range(N_azimuth):

            JM = AM.JonesMatrix(frequency, zeniths[ze_i], azimuths[az_i])
            resulting_grid[ze_i, az_i] = np.abs( JM[0,0] )

    plt.imshow(resulting_grid)
    plt.show()

if __name__ == "__main__":
    # frequencies_MHz = np.linspace(0.0, 100.0, 10000)

    AM = LBA_antenna_model()
    # matrix = AM.JonesMatrix_MultiFreq(frequencies_MHz*1E6, 0.0, 0.0)

    # plt.semilogy(frequencies_MHz, np.abs(matrix[:, 0,0]), linewidth=7 )
    # plt.show()

    elivations = np.linspace(0.5,15,100)

    responses = []
    for e in elivations:
        M = AM.JonesMatrix(60e6, 90-e, 0)
        responses.append(np.abs(M[0,0]))

    plt.plot( elivations, responses )
    plt.show()


# Copyright 2018 The Cirq Developers
#
# 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
#
#     https://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.

"""Wavefunction simulator specialized to Google's xmon gate set.

This class should not be used directly, see instead XmonSimulator in the
xmon_simulator class.
"""

import math
import multiprocessing
import multiprocessing.dummy as dummy

from typing import Any, Dict, List, Union, Tuple

import numpy as np

from cirq.google.sim import mem_manager

I_PI_OVER_2 = 0.5j * np.pi

def ensure_pool(func):
    """Decorator that ensures a pool is available for a stepper."""
    def func_wrapper(*args, **kwargs):
        if len(args) == 0 or not isinstance(args[0], Stepper):
            raise Exception('@ensure_pool can only be used on Stepper methods.')
        if args[0]._pool is None:
            with args[0]:
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)
    return func_wrapper

class Stepper(object):
    """A wave function simulator for quantum circuits with the xmon gate set.

    Xmons have a natural gate set made up of

    * Single qubit phase gates, exp(-i t Z)
    * Single qubit gates about a operation in the Pauli X/Y plane,
      exp(-i t (cos(theta) X + sin(theta) Y)
    * Two qubit phase gates exp(i t |11><11|)

    This stepper will do sharded simulation of the wave function using
    python's multiprocessing module.

    This stepper can be used like a context manager:
        with Stepper(num_qubits=3) as s:
            s.simulate_phases((1, 0.25))
            s.simulate_w(2, 0.25, 0.25)
            ...
    In this case the stepper will shut down the multiprocessing pool upon
    exiting the with context.

    If the  stepper is not used as a context manager, then it is required that
    __exit__ be called in order to ensure that the multiprocessing pool is
    properly closed (__enter__ does not need to be called).
    """

    def __init__(self,
                 num_qubits: int,
                 num_prefix_qubits: int = None,
                 initial_state: Union[int, np.ndarray] = 0,
                 min_qubits_before_shard: int = 18,
                 use_processes=False) -> None:
        """Construct a new XmonSimulator.

        Args:
          num_qubits: The number of qubits to simulate.
          num_prefix_qubits: The wavefunction of the qubits is sharded into
              (2 ** num_prefix_qubits) parts. If this is None, then this will
              shard over the nearest power of two below the cpu count. If less
              than 10 qubits are being simulated then no sharding is done,
              depending on whether the shard_for_small_num_qubits is set or
              not.
          initial_state: If this is an int, then this is the state to
              initialize the stepper to, expressed as an integer of the
              computational basis. The 1s bit of the integer determines the
              values of the last qubit, the 2s bit determines the value of the
              second-to-last-qubit, and so forth. This sounds odd, but it
              matches how people write numbers: the smallest value digit goes
              last.
              Otherwise, if this is a np.ndarray it is the full initial state
              and this must be the correct size, normalized (an L2 norm of 1),
              and have dtype of np.complex64. An array with zeroes everywhere,
              except for a 1 at index k, is equivalent to state prepared when
              the initial state is set to the integer k.
          min_qubits_before_shard: Sharding will be done only for this number
              of qubits or more. The default is 18.
          use_processes: Whether or not to use processes instead of threads.
              Processes can improve the performance slightly (varies by machine
              but on the order of 10 percent faster).  However this varies
              significantly by architecture, and processes should not be used
              for interactive python use on Windows.
        """
        self._num_qubits = num_qubits
        if num_prefix_qubits is None:
            num_prefix_qubits = int(math.log(multiprocessing.cpu_count(), 2))
        if num_prefix_qubits > num_qubits:
            num_prefix_qubits = num_qubits
        if num_qubits < min_qubits_before_shard:
            num_prefix_qubits = 0
        self._num_prefix_qubits = num_prefix_qubits
        # Each shard is of a dimension equal to 2 ** num_shard_qubits.
        self._num_shard_qubits = self._num_qubits - self._num_prefix_qubits

        self._num_shards = 2 ** self._num_prefix_qubits
        self._shard_size = 2 ** self._num_shard_qubits

        # TODO(dabacon): This could be parallelized.
        self._init_shared_mem(initial_state)
        self._pool = None  # type: Union[ThreadlessPool, Any]
        self._pool_fn = multiprocessing.Pool if use_processes else dummy.Pool

    def _init_shared_mem(self, initial_state: int):
        self._shared_mem_dict = {}  # type: Dict[str, int]
        self.init_z_vects()
        self._init_scratch()
        self._init_state(initial_state)

    def init_z_vects(self):
        """Initializes bitwise vectors which is precomputed in shared memory.

        There are two types of vectors here, a zero one vectors and a pm
        (plus/minus) vectors. The pm vectors have rows that are Pauli Z
        operators acting on the all ones vector. The column-th row corresponds
        to the Pauli Z acting on the column'th-qubit.  Example for three shard
        qubits:
             [[1, -1, 1, -1, 1, -1, 1, -1],
              [1, 1, -1, -1, 1, 1, -1, -1],
              [1, 1, 1, 1, -1, -1, -1, -1]]
        The zero one vectors are the pm vectors with 1 replacing -1 and 0
        replacing 1.

        There are number of shard qubit zero one vectors and each of these is
        of size equal to the shard size. For the zero one vectors, the ith one
        of these vectors has a  kth index value that is equal to 1 if the i'th
        bit of k is set and zero otherwise. The vector directly encode the
        little-endian binary digits of its index in the list:
        v[j][i] = (i >> j) & 1. For the pm vectors, the ith one of these
        vectors has a  k'th index value that is equal to -1 if the i'th bit of
        k is set and 1 otherwise.
        """
        shard_size = 2 ** self._num_shard_qubits

        a, b = np.indices((shard_size, self._num_shard_qubits))
        a >>= b
        a &= 1
        zero_one_vects = np.ascontiguousarray(a.transpose())
        zero_one_vects_handle = mem_manager.SharedMemManager.create_array(
            zero_one_vects)
        self._shared_mem_dict['zero_one_vects_handle'] = zero_one_vects_handle

        pm_vects = 1 - 2 * zero_one_vects
        pm_vects_handle = mem_manager.SharedMemManager.create_array(pm_vects)
        self._shared_mem_dict['pm_vects_handle'] = pm_vects_handle

    def _init_scratch(self):
        """Initializes a scratch pad equal in size to the wavefunction."""
        scratch = np.zeros((self._num_shards, self._shard_size),
                           dtype=np.complex64)
        scratch_handle = mem_manager.SharedMemManager.create_array(
            scratch.view(dtype=np.float32))
        self._shared_mem_dict['scratch_handle'] = scratch_handle

    def _init_state(self, initial_state: Union[int, np.ndarray]):
        """Initializes a the shard wavefunction and sets the initial state."""
        state = np.reshape(
            decode_initial_state(initial_state, self._num_qubits),
            (self._num_shards, self._shard_size))
        state_handle = mem_manager.SharedMemManager.create_array(
            state.view(dtype=np.float32))
        self._shared_mem_dict['state_handle'] = state_handle

    def __del__(self):
        for handle in self._shared_mem_dict.values():
            mem_manager.SharedMemManager.free_array(handle)

    def __enter__(self):
        if self._pool is None:
            self._pool = (self._pool_fn(processes=self._num_shards)
                          if self._num_prefix_qubits > 0 else ThreadlessPool())
        return self

    def __exit__(self, *args):
        # Terminate is safe here since all work should have been completed.
        if self._pool is not None:
            self._pool.terminate()
            self._pool.join()
            self._pool = None

    def _shard_num_args(self,
                        constant_dict: Dict[str, Any] = None
                        ) -> List[Dict[str, Any]]:
        """Helper that returns a list of dicts including a num_shard entry.

        The dict for each entry also includes shared_mem_dict, the number of
        shards, the number of shard qubits, and the supplied constant dict.

        Args:
            constant_dict: Dictionary that will be updated to every element of
                the returned list of dictionaries.

        Returns:
            A list of dictionaries. Each dictionary is constant except for the
            'shard_num' key which ranges from 0 to number of shards - 1.
            Included keys are 'num_shards' and 'num_shard_qubits' along with
            all the keys in constant_dict.
        """
        args = []
        for shard_num in range(self._num_shards):
            append_dict = dict(constant_dict) if constant_dict else {}
            append_dict['shard_num'] = shard_num
            append_dict['num_shards'] = self._num_shards
            append_dict['num_shard_qubits'] = self._num_shard_qubits
            append_dict.update(self._shared_mem_dict)
            args.append(append_dict)
        return args

    @property
    def current_state(self):
        """Returns the current wavefunction."""
        return self._current_state()

    @ensure_pool
    def _current_state(self):
        return np.array(
            self._pool.map(_state_shard, self._shard_num_args())).flatten()

    @ensure_pool
    def reset_state(self, reset_state):
        """Reset the state to the given initial state.

        Args:
            reset_state: If this is an int, then this is the state to reset
                the stepper to, expressed as an integer of the computational
                basis. Integer to bitwise indices is little endian. Otherwise
                if this is a np.ndarray this must be the correct size, be
                normalized (L2 norm of 1), and have dtype of np.complex64.

        Raises:
            ValueError if the state is incorrectly sized or not of the correct
            dtype.
        """
        # If the pool has been closed, recreate to calculate state.'
        if isinstance(reset_state, int):
            self._pool.map(_reset_state,
                           self._shard_num_args({'reset_state': reset_state}))
        elif isinstance(reset_state, np.ndarray):
            check_state(reset_state, self._num_qubits)
            args = []
            for kwargs in self._shard_num_args():
                shard_num = kwargs['shard_num']
                shard_size = 1 << kwargs['num_shard_qubits']
                start = shard_num * shard_size
                end = start + shard_size
                kwargs['reset_state'] = reset_state[start:end]
                args.append(kwargs)
            self._pool.map(_reset_state, args)

    @ensure_pool
    def simulate_phases(self, phase_map: Dict[Tuple[int, ...], float]):
        """Simulate a set of phase gates on the xmon architecture.

        Args:
            phase_map: A map from a tuple of indices to a value, one for each
                phase gate being simulated. If the tuple key has one index, then
                this is a Z phase gate on the index-th qubit with a rotation
                angle of pi times the value of the map. If the tuple key has two
                indices, then this is a |11> phasing gate, acting on the qubits
                at the two indices, and a rotation angle of pi times the value
                of the map.
        """
        self._pool.map(_clear_scratch, self._shard_num_args())
        # Iterate over the map of phase data.
        for indices, half_turns in phase_map.items():
            args = self._shard_num_args(
                {'indices': indices, 'half_turns': half_turns})
            if len(indices) == 1:
                self._pool.map(_single_qubit_accumulate_into_scratch, args)
            elif len(indices) == 2:
                self._pool.map(_two_qubit_accumulate_into_scratch, args)
        # Exponentiate the phases and add them into the state.
        self._pool.map(_apply_scratch_as_phase, self._shard_num_args())

    @ensure_pool
    def simulate_w(self,
                   index: int,
                   half_turns: float,
                   axis_half_turns: float):
        """Simulate a single qubit rotation gate about a X + b Y.

        The gate simulated is U = exp(-i pi/2 W half_turns)
            where W = cos(pi axis_half_turns) X + sin(pi axis_half_turns) Y

        Args:
          index: The qubit to act on.
          half_turns: The amount of the overall rotation, see the formula
              above.
          axis_half_turns: The angle between the pauli X and Y operators,
              see the formula above.
        """
        args = self._shard_num_args({
            'index': index,
            'half_turns': half_turns,
            'axis_half_turns': axis_half_turns
        })
        if index >= self._num_shard_qubits:
            # W gate spans shards.
            self._pool.map(_clear_scratch, args)
            self._pool.map(_w_between_shards, args)
            self._pool.map(_copy_scratch_to_state, args)
        else:
            # W gate is within a shard.
            self._pool.map(_w_within_shard, args)

        # Normalize after every w.
        norm_squared = np.sum(self._pool.map(_norm_squared, args))
        args = self._shard_num_args({
            'norm_squared': norm_squared
        })
        self._pool.map(_renorm, args)

    @ensure_pool
    def simulate_measurement(self, index: int) -> bool:
        """Simulates a single qubit measurement in the computational basis.

        Args:
            index: Which qubit is measured.

        Returns:
            True iff the measurement result corresponds to the |1> state.
        """
        args = self._shard_num_args({'index': index})
        prob_one = np.sum(self._pool.map(_one_prob_per_shard, args))
        result = bool(np.random.random() <= prob_one)

        args = self._shard_num_args({
            'index': index,
            'result': result,
            'prob_one': prob_one
        })
        self._pool.map(_collapse_state, args)
        return result

    def sample_measurements(
            self,
            indices: List[int],
            repetitions: int=1) -> List[List[bool]]:
        """Samples from measurements in the computational basis.

        Note that this does not collapse the wave function.

        Args:
            indices: Which qubits are measured.

        Returns:
            Measurement results with True corresponding to the |1> state.
            The outer list is for repetitions, and the inner corresponds to
            measurements ordered by the input indices.

        Raises:
            ValueError if repetitions is less than one.
        """
        if repetitions < 1:
            raise ValueError(
                'Number of repetitions cannot be negative. Was {}'.format(

                    repetitions))
        if len(indices) == 0:
            return [[]]

        # Calculate probabilities and reshape to tensor of qubits.
        tensor = np.reshape(np.abs(self.current_state) ** 2,
                            self._num_qubits * [2])

        # Tensor axis order is reverse of index order, so we transpose here.
        tensor = np.transpose(tensor)

        # Indices that should be summed over.
        sum_indices = tuple(
            x for x in range(self._num_qubits) if x not in indices)

        # Sum over those indices, and reshape into a a tensor of len(indices)
        # qubits.
        probs = np.reshape(np.sum(tensor, axis=sum_indices), [2] * len(indices))

        # Calculate how the indices not summed over should be reordered.
        index_map = {v: k for k,v in enumerate(sorted(indices))}
        perm = [index_map[x] for x in indices]
        # Apply this permutation to the probabilities and flatten.,
        probs = np.reshape(np.transpose(probs, perm), -1)

        # We now have the probability vector, correctly ordered, so sample over
        # it. Note that we us ints here, since numpy's choice does not allow for
        # choosing from a list of tuples or list of lists.
        result = np.random.choice(2 ** len(indices), size=repetitions, p=probs)
        # Convert to bools and note also one final reverse of list to get
        # ordering correct.
        return np.transpose(
            [(1 & (result >> i)).astype(np.bool) for i in range(len(indices))][
            ::-1]).tolist()

def decode_initial_state(initial_state: Union[int, np.ndarray],
                         num_qubits: int) -> np.ndarray:
    """Verifies the initial_state is valid and converts it to ndarray form."""
    if isinstance(initial_state, np.ndarray):
        if len(initial_state) != 2 ** num_qubits:
            raise ValueError(
                'initial state was of size {} '
                'but expected state for {} qubits'.format(
                    len(initial_state), num_qubits))
        state = initial_state
    elif isinstance(initial_state, int):
        if initial_state < 0:
            raise ValueError('initial_state must be positive')
        elif initial_state >= 2 ** num_qubits:
            raise ValueError(
                'initial state was {} but expected state for {} qubits'.format(
                    initial_state, num_qubits))
        else:
            state = np.zeros(2 ** num_qubits, dtype=np.complex64)
            state[initial_state] = 1.0
    else:
        raise TypeError('initial_state was not of type int or ndarray')
    check_state(state, num_qubits)
    return state

def check_state(state: np.ndarray, num_qubits: int):
    """Validates that the given state is a valid wave function."""
    if state.size != 1 << num_qubits:
        raise ValueError(
            'State has incorrect size. Expected {} but was {}.'.format(
                1 << num_qubits, state.size))
    if state.dtype != np.complex64:
        raise ValueError(
            'State has invalid dtype. Expected {} but was {}'.format(
                np.complex64, state.dtype))
    norm = np.sum(np.abs(state) ** 2)
    if not np.isclose(norm, 1):
        raise ValueError('State is not normalized instead had norm %s' % norm)

def _state_shard(args: Dict[str, Any]) -> np.ndarray:
    state_handle = args['state_handle']
    return mem_manager.SharedMemManager.get_array(state_handle).view(
        dtype=np.complex64)[args['shard_num']]

def _scratch_shard(args: Dict[str, Any]) -> np.ndarray:
    scratch_handle = args['scratch_handle']
    return mem_manager.SharedMemManager.get_array(scratch_handle).view(
        dtype=np.complex64)[args['shard_num']]

def _pm_vects(args: Dict[str, Any]) -> np.ndarray:
    return mem_manager.SharedMemManager.get_array(args['pm_vects_handle'])

def _zero_one_vects(args: Dict[str, Any]) -> np.ndarray:
    return mem_manager.SharedMemManager.get_array(
        args['zero_one_vects_handle'])

def _kth_bit(x: int, k: int) -> int:
    """Returns 1 if the kth bit of x is set, 0 otherwise."""
    return (x >> k) & 1

def _reset_state(args: Dict[str, Any]):
    shard_num = args['shard_num']
    shard_size = 2 ** args['num_shard_qubits']
    reset_state = args['reset_state']

    if isinstance(reset_state, int):
        _state_shard(args).fill(0)
        if shard_num == reset_state // shard_size:
            _state_shard(args)[reset_state % shard_size] = 1.0
    else:
        np.copyto(_state_shard(args), reset_state)

def _clear_scratch(args: Dict[str, Any]):
    """Sets all of the scratch shard to zero."""
    _scratch_shard(args).fill(0)

def _single_qubit_accumulate_into_scratch(args: Dict[str, Any]):
    """Accumulates single qubit phase gates into the scratch shards."""
    index = args['indices'][0]
    shard_num = args['shard_num']
    half_turns = args['half_turns']
    num_shard_qubits = args['num_shard_qubits']
    scratch = _scratch_shard(args)

    # ExpZ = exp(-i pi Z half_turns / 2).
    if index >= num_shard_qubits:
        # Acts on prefix qubits.
        sign = 1 - 2 * _kth_bit(shard_num, index - num_shard_qubits)
        scratch -= half_turns * sign
    else:
        # Acts on shard qubits.
        scratch -= half_turns * _pm_vects(args)[index]

def _one_projector(args: Dict[str, Any], index: int) -> Union[int, np.ndarray]:
    """Returns a projector onto the |1> subspace of the index-th qubit."""
    num_shard_qubits = args['num_shard_qubits']
    shard_num = args['shard_num']
    if index >= num_shard_qubits:
        return _kth_bit(shard_num, index - num_shard_qubits)
    return _zero_one_vects(args)[index]

def _two_qubit_accumulate_into_scratch(args: Dict[str, Any]):
    """Accumulates two qubit phase gates into the scratch shards."""
    index0, index1 = args['indices']
    half_turns = args['half_turns']
    scratch = _scratch_shard(args)

    projector = _one_projector(args, index0) * _one_projector(args, index1)
    # Exp11 = exp(-i pi |11><11| half_turns), but we accumulate phases as
    # pi / 2.
    scratch += 2 * half_turns * projector

def _apply_scratch_as_phase(args: Dict[str, Any]):
    """Takes scratch shards and applies them as exponentiated phase to state.
    """
    state = _state_shard(args)
    state *= np.exp(I_PI_OVER_2 * _scratch_shard(args))

def _w_within_shard(args: Dict[str, Any]):
    """Applies a W gate when the gate acts only within a shard."""
    index = args['index']
    half_turns = args['half_turns']
    axis_half_turns = args['axis_half_turns']
    state = _state_shard(args)
    pm_vect = _pm_vects(args)[index]
    num_shard_qubits = args['num_shard_qubits']
    shard_size = 2 ** num_shard_qubits

    reshape_tuple = (2 ** (num_shard_qubits - 1 - index), 2, 2 ** index)
    perm_state = np.reshape(
        np.reshape(state, reshape_tuple)[:, ::-1, :], shard_size)
    cos = np.cos(-0.5 * np.pi * half_turns)
    sin = np.sin(-0.5 * np.pi * half_turns)

    cos_axis = np.cos(np.pi * axis_half_turns)
    sin_axis = np.sin(np.pi * axis_half_turns)

    new_state = cos * state + 1j * sin * perm_state * (
        cos_axis - 1j * sin_axis * pm_vect)
    np.copyto(state, new_state)

def _w_between_shards(args: Dict[str, Any]):
    """Applies a W gate when the gate acts between shards."""
    shard_num = args['shard_num']
    state = _state_shard(args)
    num_shard_qubits = args['num_shard_qubits']
    index = args['index']
    half_turns = args['half_turns']

    axis_half_turns = args['axis_half_turns']

    perm_index = shard_num ^ (1 << (index - num_shard_qubits))
    perm_state = mem_manager.SharedMemManager.get_array(
        args['state_handle']).view(np.complex64)[perm_index]

    cos = np.cos(-0.5 * np.pi * half_turns)
    sin = np.sin(-0.5 * np.pi * half_turns)

    cos_axis = np.cos(np.pi * axis_half_turns)
    sin_axis = np.sin(np.pi * axis_half_turns)

    scratch = _scratch_shard(args)
    z_op = (1 - 2 * _kth_bit(shard_num, index - num_shard_qubits))
    np.copyto(scratch, state * cos + 1j * sin * perm_state *
              (cos_axis - 1j * sin_axis * z_op))

def _copy_scratch_to_state(args: Dict[str, Any]):
    """Copes scratch shards to state shards."""
    np.copyto(_state_shard(args), _scratch_shard(args))

def _one_prob_per_shard(args: Dict[str, Any]) -> float:
    """Returns the probability of getting a one measurement on a state shard.
    """
    index = args['index']

    state = _state_shard(args) * _one_projector(args, index)
    norm = np.linalg.norm(state)
    return norm * norm

def _norm_squared(args: Dict[str, Any]) -> float:
    """Returns the norm for each state shard."""
    state = _state_shard(args)
    return np.sum(np.abs(state) ** 2)

def _renorm(args: Dict[str, Any]):
    """Renormalizes the state using the norm arg."""
    state = _state_shard(args)
    # If our gate is so bad that we have norm of zero, we have bigger problems.
    state /= np.sqrt(args['norm_squared'])

def _collapse_state(args: Dict[str, Any]):
    """Projects state shards onto the appropriate post measurement state.

    This function makes no assumptions about the interpretation of quantum
    theory.

    Args:
        args: The args from shard_num_args.
    """
    index = args['index']
    result = args['result']
    prob_one = args['prob_one']

    state = _state_shard(args)
    normalization = np.sqrt(prob_one if result else 1 - prob_one)
    state *= (_one_projector(args, index) * result +
              (1 - _one_projector(args, index)) * (1 - result))
    state /= normalization

class ThreadlessPool(object):
    """A Pool that does not use any processes or threads.

    Only supports map, close, and join, the later two being trivial.
    No enforcement of closing or joining is done, so map can be called
    repeatedly.
    """

    # noinspection PyMethodMayBeStatic
    def map(self, func, iterable, chunksize=None):
        assert chunksize is None, 'Chunking not supported by SimplePool'
        return [func(x) for x in iterable]

    def terminate(self):
        pass

    def join(self):
        pass


#!/usr/bin/env python3
"""
The data types for an Ocean Salinity DTBXY
"""

import numpy as np

# data type list of lists. The inner lists are there for each file sub-block.
datatype = [
            [('MaxValid', 'float32')],
            [('MinValid', 'float32')],

            # REGION
            [('region_count', 'uint32')],
            # Region ID, start and stop snap time, start and stop snap id
            [('Region_ID', 'uint32'), ('Days', 'int32'), ('Seconds', 'uint32'), ('Microseconds', 'uint32'),
             ('stop_Days', 'int32'), ('stop_Seconds', 'uint32'), ('stop_Microseconds', 'uint32'),
             ('Start_Snapshot_ID', 'uint32'), ('Stop_Snapshot_ID', 'uint32')],
            # stats (just one number) repeated for the 3 models, 8 pols, 12 fov zones
            [('mean', 'float32'), ('median', 'float32'), ('min', 'float32'), ('max', 'float32'), ('std', 'float32')],
            # counts, dTb, std_dTb, flags repeated 129 x 129
            [('count_deltaTB', 'uint32'), ('deltaTB', 'float32'), ('std_deltaTB', 'float32'), ('flags', 'ushort')],

            # SNAPSHOTS
            [('snap_count', 'uint32')],
            # snapshot general info
            [('Snapshot_ID', 'uint32'), ('Snapshot_OBET', 'uint64'), ('Snapshot_Latitude', 'float32'),
             ('Snapshot_Longitude', 'float32'), ('Snapshot_Altitude', 'float32'), ('Snapshot_Flags', 'ushort'),
             ('L1c_TEC', 'int16')],
            [('measurement_count', 'ushort')],
            # measured Tb mean and std
            [('L1cTB', 'ushort'), ('std_L1cTB', 'ushort')],
            # BOA fwd model components
            [('atmosTB', 'int16'), ('std_atmosTB', 'ushort'), ('flatSeaTB', 'int16'), ('std_flatSeaTB', 'ushort'),
             ('roughTB', 'int16'), ('std_roughTB', 'ushort'), ('galTB', 'int16'), ('std_galTB', 'ushort'),
             ('sunTB', 'int16'), ('std_sunTB', 'ushort'), ('sumTB', 'int16'), ('std_sumTB', 'ushort')],
            # TOA fwd model components with L1c TEC
            [('atmosTB', 'int16'), ('std_atmosTB', 'ushort'), ('flatSeaTB', 'int16'), ('std_flatSeaTB', 'ushort'),
             ('roughTB', 'int16'), ('std_roughTB', 'ushort'), ('galTB', 'int16'), ('std_galTB', 'ushort'),
             ('sunTB', 'int16'), ('std_sunTB', 'ushort'), ('sumTB', 'int16'), ('std_sumTB', 'ushort')],
            # TOA fwd model components with A3 TEC
            [('atmosTB', 'int16'), ('std_atmosTB', 'ushort'), ('flatSeaTB', 'int16'), ('std_flatSeaTB', 'ushort'),
             ('roughTB', 'int16'), ('std_roughTB', 'ushort'), ('galTB', 'int16'), ('std_galTB', 'ushort'),
             ('sunTB', 'int16'), ('std_sunTB', 'ushort'), ('sumTB', 'int16'), ('std_sumTB', 'ushort')],
            # geophysics
            [('SSS', 'int16'), ('std_SSS', 'ushort'), ('SST', 'int16'), ('std_SST', 'ushort'), ('WS', 'int16'),
             ('std_WS', 'ushort'), ('A3TEC', 'int16'), ('std_A3TEC', 'ushort'), ('Tair', 'int16'),
             ('std_Tair', 'ushort'), ('SP', 'int16'), ('std_SP', 'ushort'), ('TCWV', 'int16'), ('std_TCWV', 'ushort'),
             ('HS', 'int16'), ('std_HS', 'ushort')],
            # flags
            [('coast', 'ushort'), ('sun_point', 'ushort'), ('sun_tails', 'ushort'), ('rfi', 'ushort'),
             ('rain', 'ushort'), ('ice', 'ushort')],

            [('gp_count', 'uint32')],
            # grid points
            [('Grid_Point_ID', 'uint32'), ('Grid_Point_Latitude', 'float32'), ('Grid_Point_Longitude', 'float32')],
            [('measurement_count', 'ushort')],
            [('Snapshot_Index', 'ushort'), ('Zone_Bits', 'ushort')]
            ]


# Filename: pair.py
import sys
import numpy
import time
from operator import add
import pandas as pd
import datetime
import json
import urllib2
import pytz
from bs4 import BeautifulSoup
from datetime import datetime
from pandas.io.data import DataReader
import riak
import numpy as np
import statsmodels.api as stat
import statsmodels.tsa.stattools as ts
from riak import RiakClient, RiakObject
import boto, urllib2
from   boto.ec2 import connect_to_region
import os
import sys

#start is furthest day back and end is most recent day.  Grabs all data in between and stores in riak as json
#runs through a file of ticker values
def getData(tickerFile, dataSource, start, end, riakIP):

    rc = RiakClient(protocol='pbc',host = riakIP, pb_port=8087)#set up riak connection
    added = []#list of successful adds
    notAdded = []#list of unsuccessful adds
    stock = pd.read_csv(tickerFile,sep='\t',header=None)#read in stock tickers

    #loop over all stock tickers
    for i in range(0,len(stock.head(100))):

        ticker = stock.ix[i,0]
        if getDataByTicker(ticker,dataSource,start,end,riakIP) == 0:
            notAdded.append(ticker)
        else:
            added.append(ticker)
    return added, notAdded

#start is furthest day back and end is closest to today, Store single stock data in riak
#only grabs one stock
def getDataByTicker(ticker, dataSource, start, end, riakIP):

    rc = RiakClient(protocol='pbc',host = riakIP, pb_port=8087)
    #get daily data for each ticker
    gtemp = pd.DataFrame()
    bucket = rc.bucket('stocks')
    try:
        gtemp = DataReader(ticker,  dataSource, start, end)
        print ticker
    except:
        pass

        #didnt get any data
    if len(gtemp) == 0:
        return 0
    #got data
    else:

        for j in range(0,len(gtemp.index)):

            #upload json to Riak Bucket
            date = gtemp.index[j].date()
            riakKey = str(ticker + '_' + str(date))
            riakVal = {'OPEN': gtemp.values[j,0],\
                        'HIGH': gtemp.values[j,1],\
                        'LOW': gtemp.values[j,2], \
                        'CLOSE': gtemp.values[j,3], \
                        'VOLUME': gtemp.values[j,4],\
                        'DATE': str(date),\
                        'TICKER': str(ticker)}

            obj = RiakObject(rc, bucket, riakKey)

            obj.add_index("ticker_bin", str(ticker))
            obj.add_index("year_int", int(date.year))
            obj.add_index("month_int", int(date.month))
            obj.add_index("day_int", int(date.day))

            obj.content_type = 'text/json'
            #obj.data = riakVal
            obj.data = json.dumps(riakVal)
            obj.store()

    return len(gtemp.index)

def downloadStock(ticker,dataSource,start,end):
    gtemp = pd.DataFrame()
    try:
        gtemp = DataReader(ticker,  dataSource, start, end)
        print ticker
    except:
        pass
    return gtemp

def writeHistory(ticker, data, riakIP):
    rc = RiakClient(protocol='pbc',host = riakIP, pb_port=8087)
    bucket = rc.bucket('stocks')
    gtemp = data
    if len(gtemp) == 0:
        return 0
    else:

        for j in range(0,len(gtemp.index)):

                #upload json to Riak Bucket
                date = gtemp.index[j].date()
                riakKey = str(ticker + '_' + str(date))
                riakVal = {'OPEN': gtemp.values[j,0],\
                            'HIGH': gtemp.values[j,1],\
                            'LOW': gtemp.values[j,2], \
                            'CLOSE': gtemp.values[j,3], \
                            'VOLUME': gtemp.values[j,4],\
                            'DATE': str(date),\
                            'TICKER': str(ticker)}

                obj = RiakObject(rc, bucket, riakKey)

                obj.add_index("ticker_bin", str(ticker))
                obj.add_index("year_int", int(date.year))
                obj.add_index("month_int", int(date.month))
                obj.add_index("day_int", int(date.day))

                obj.content_type = 'text/json'
                #obj.data = riakVal
                obj.data = json.dumps(riakVal)
                obj.store()

    return len(gtemp.index)

#searches riak bucket via 2i query and returns a dict of the data
def riakSearchData(searchBucket, searchTerm, searchVal1, searchVal2,riakIP):
    myData = {}#empty dict
    myBucket = RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(searchBucket)
    #check wether 1 or 2 search terms
    if searchVal2 != None:
        for key in myBucket.get_index(searchTerm, searchVal1, searchVal2): #get all keys with 2i match
            myData[key] = json.loads(myBucket.get(key).data)#store data for each key
    else:
        for key in myBucket.get_index(searchTerm, searchVal1):#get all keys with 2i match
            myData[key] = json.loads(myBucket.get(key).data)#store data for each key
    return myData

#store an individual key value pair in a bucket
def storeKV(myBucket, myKey, myVal, riakIP):
    riak.RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(myBucket).new(myKey, data = myVal).store()
    return

#delete a key from a bucket, provide feedback to ensure deletion
def deleteKey(delBucket, delKey,riakIP):
    riak.RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(delBucket).delete(delKey)
    if riak.RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(delBucket).get(delKey).data == None:
        print 'Successful delete: %s' % delKey
    else:
        print 'Failed delete: %s' % delKey
    return

#delete key from bucket, no feedback
def quickDeleteKey(delBucket,delKey, riakIP):
    riak.RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(delBucket).delete(delKey)
    return

#delete all keys in a bucket, no feedback
def quickDeleteAllKeys(delBucket,riakIP):
    for keys in  riak.RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(delBucket).stream_keys():
        for delKey in keys:
            quickDeleteKey(delBucket, delKey,riakIP)
    print 'Done'
    return

#delete all keys in a bucket, with feedback
def deleteAllKeys(delBucket,riakIP):
    delList = []
    try:
        for keys in  riak.RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(delBucket).stream_keys():
            for delKey in keys:
                deleteKey(delBucket, delKey,riakIP)
                delList.append(delKey)
    except:
        print 'delete error'
        pass
    return delList

#get all key value pairs from a bucket
def getAllKV(myBucket,riakIP):
    myData = {}
    riak_bucket = riak.RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(myBucket)
    for keys in riak_bucket.stream_keys():
        for key in keys:
            tempData = riak_bucket.get(key).data
            print('Key: %s Value: %s' % (key, tempData))
            myData[key] = tempData
    return myData

#get single value for a key in a bucket
def getValue(myBucket, myKey,riakIP):
    myVal = json.loads(riak.RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket(myBucket).get(myKey).data)
    return myVal

#Take a tuple of tuples in and return something
def pairAnalysis(pairTuple, ndays, beginDay = 0, zThresh = 2, critLevel = '5%'):

    #pair tuple looks like ([tickerA, [data]],[tickerB,[data]])
    #input is assumed to be same length and sorted by date with most recent date first

    #unwrap first stock ticker and data
    stockA = pairTuple[0]
    stockAData = list(stockA[1])

    #unwrap the data for stockA
    stockADates = [x[2] for x in stockAData]
    stockAClose = [x[0] for x in stockAData]
    stockAVolume = [x[1] for x in stockAData]

    #unwrap second stock ticker and data
    stockB = pairTuple[1]
    stockBData = list(stockB[1])

    #unwrap stockB data
    stockBDates = [x[2] for x in stockBData]
    stockBClose = [x[0] for x in stockBData]
    stockBVolume = [x[1] for x in stockBData]

    pair = pairCalc(stockAClose,stockBClose,beginDay,ndays, zThresh, critLevel)
    #if pair tradeable, add some more info
    if type(pair) is list:
            pair.insert(0,stockADates[beginDay])
            pair.insert(0,stockB[0])
            pair.insert(0,stockA[0])
            return pair
    else:
        return pair

def pairCalc(tsA,tsB,beginDay,ndays,zThresh = 2, critLevel = '5%'):

    #perform engle granger cointegration test
    if beginDay < 0 or beginDay >= ndays or zThresh < 0 or not (critLevel in ['1%','5%','10%']):
        print 'input error'
        return 0

    coint = egct(tsA[beginDay:ndays],tsB[beginDay:ndays], critLevel)

    #if coint return 0, then the two timeseries are not cointegrated
    if (coint[0] != 1):
        return 0
    #else calculate stuff
    else:
        #signal = tsA[0] - beta*tsB[0] - CONSTANT = normal gaussian with mean 0
        signal = [a - coint[1][1]*b - coint[1][0] for a in tsA[beginDay:ndays] for b in tsB[beginDay:ndays]]
        sigMean = numpy.mean(signal)
        sigStd = numpy.std(signal)
        #zscore is (signal - signalMean) / signalStd
        zscore = (signal[beginDay] - sigMean)/sigStd
        #if current zscore is larger than zThresh, possible pair to trade
        if abs(zscore) > zThresh:
            return [tsA[0],tsB[0], zscore, coint[1][1], sigMean, sigStd]
    return 1

#write tradeable pair back into riak
def writePairs(pairList,bucketName,riakIP):

    #tradeable pairs are lists
    tradeable = [x for x in pairList if type(x) is list]

    for pair in tradeable:
        writeSinglePair(pair,bucketName,riakIP)

    #return a list of written pairs
    return tradeable

#write a signle pair to riak
#assumes pair is in a list of values
def writeSinglePair(pair,bucketName,riakIP):

    rc = RiakClient(protocol='pbc',host = riakIP, pb_port=8087)
    bucket = rc.bucket(bucketName)

    #create key value pairs to stock in riak
    key = str(str(pair[0])+ '_' + str(pair[1]))
    val = {'StockA': pair[0], \
                'StockB': pair[1], \
                'Date': pair[2],\
                'CloseA': pair[3], \
                'CloseB': pair[4], \
                'ZScore': pair[5],\
                'Beta': pair[6],\
                'SignalMean': pair[7],\
                'SignalSD': pair[8]}
    myDate = pair[2].split('-')
    obj = RiakObject(rc, bucket, key)

    #add 2i tags
    obj.add_index("stocka_bin", str(pair[0]))
    obj.add_index("stockb_bin", str(pair[3]))
    obj.add_index("year_int", int(myDate[0]))
    obj.add_index("month_int", int(myDate[1]))
    obj.add_index("day_int", int(myDate[2]))
    obj.content_type = 'text/json'
    obj.data = val
    obj.data = json.dumps(val)
    #store
    obj.store()

    #return a list of written pairs
    return pair

#return 1 if the two series are cointegrated and 0 otherwise, return regression parameters either way
#assumes y,x are aligned and of equal length
#critLevel can be '1%', '5%' or '10%'
def egct(y, x,critLevel):

    #must add a constant row of 1s to dependent variable, its a multidimensional regression thing
    x = stat.add_constant(x)
    #get residuals
    result = stat.OLS(y, x).fit()
    #regression parameters, slope and intercept
    regPar = result.params
    #run augmented dickey fuller test of stationarity of residuals
    #null hypothesis is stationaity of timeseries
    adfResults = ts.adfuller(result.resid, maxlag=0, regression='c', autolag=None, store=False, regresults=True)
    #test statistic
    tstat = adfResults[0]
    #critical value
    critVal = adfResults[2][critLevel]
    #if test stat is less than critical value, accept null hyptohesis of stationarity
    if tstat < critVal:
        return [1,regPar]
    else:
        return [0,regPar]

#get all values for a stock from riak
#return close,volume,date values in a list of list
def riakGetStock(searchVal,riakIP):
    myData = []
    myBucket = RiakClient(protocol='pbc',host = riakIP, pb_port=8087).bucket('stocks')
    for key in myBucket.get_index('ticker_bin', searchVal): # get all from 2002 to 2012
        value = json.loads(myBucket.get(key).data)
        myData.append([(value['CLOSE']), (value['VOLUME']), str(value['DATE'])])
    return myData

#quick function to sort a list of list on the inner list 3 value(date)
def mySort(s,n):

    try:
        sortList = list(s)
        sortList.sort(key=lambda x: x[n], reverse=True)
    except:
        print "error using mySort"
        return 0

    return sortList

#cut length of time series to n
def myFilter(s,n):
    if type(s) is list:
        try:
            return s[0:n]
        except:
            print 'error using myFilter'
            return 0
    else:
        print 'not a list'
        return 0

def bootCluster(accessKey,secretKey,region,instanceType):

    conn = boto.ec2.connect_to_region("us-east-1", aws_access_key_id=accessKey,aws_secret_access_key=secretKey)
    instances = [i for r in conn.get_all_instances() for i in r.instances]

    #start all non running instances
    myInst = []
    awsHosts = []
    awsIPs = []

    for i in instances:
        if i.state == 'stopped' and i.instance_type == instanceType:
            conn.start_instances(i.id)
            myInst.append(str(i.id))

    for i in myInst:
        i.update()

        while i.state != 'running':
            print 'waiting for: ' + str(i.id)
            time.sleep(2)
            i.update()

        awsHosts.append(str(i.dns_name))
        awsIPs.append(str(i.private_ip_address))

        print str(i.id)+ ' is running'

    return myInst, awsHosts, awsIPs

def stopCluster(accessKey,secretKey,region,instanceType):

    conn = boto.ec2.connect_to_region("us-east-1", aws_access_key_id=accessKey,aws_secret_access_key=secretKey)
    instances = [i for r in conn.get_all_instances() for i in r.instances]

    myInst = []

    for i in instances:
        if i.state == 'running' and i.instance_type == instanceType:
            conn.stop_instances(i.id)
            myInst.append(str(i.id))

    for i in myInst:
        i.update()

        while i.state != 'stopped':
            print 'waiting for: ' + str(i.id)
            time.sleep(2)
            i.update()
        print str(i.id)+ ' is stopped'

    return myInst

def submitSparkJob(sparkJob):

    os.system('fab -R worker '+ sparkJob)

    return 'Submitted: ' + sparkJob

def getDNSIP(accessKey,secretKey,region,instanceType):

    conn = boto.ec2.connect_to_region("us-east-1", aws_access_key_id=accessKey,aws_secret_access_key=secretKey)
    awsHosts = []
    awsIPs = []
    instances = [i for r in conn.get_all_instances() for i in r.instances]

    for i in instances:
        if i.state == 'running' and i.instance_type == 't2.medium':
            awsHosts.append(str(i.dns_name))
            awsIPs.append(str(i.private_ip_address))

    return awsHosts, awsIPs

def updateDate(riakIP):

    print '::::Writing Date of New Update::::'
    newUpdate = {'Year': datetime.now().year,\
                   'Month': datetime.now().month,\
                   'Day': datetime.now().day,\
                   'Hour': datetime.now().hour,\
                   'Minute': datetime.now().minute}
    print newUpdate
    storeKV("meta", "update", json.dumps(newUpdate), riakIP)

    return

# End of pair.py


from __future__ import division
import numba as nb
import numpy as np

from .utils import get_func, isstr, aggregate_common_doc, funcs_no_separate_nan
from .utils_numpy import aliasing, input_validation, check_dtype, check_fill_value

class AggregateOp(object):
    """
    Every subclass of AggregateOp handles a different aggregation operation. There are
    several private class methods that need to be overwritten by the subclasses
    in order to implement different functionality.

    On object instantiation, all necessary static methods are compiled together into
    two jitted callables, one for scalar arguments, and one for arrays. Calling the
    instantiated object picks the right cached callable, does some further preprocessing
    and then executes the actual aggregation operation.
    """

    forced_fill_value = None
    counter_fill_value = 1
    counter_dtype = bool
    mean_fill_value = None
    mean_dtype = np.float64
    outer = False
    reverse = False
    nans = False

    def __init__(self, func=None, **kwargs):
        if func is None:
            func = type(self).__name__.lower()
        self.func = func
        self.__dict__.update(kwargs)
        # Cache the compiled functions, so they don't have to be recompiled on every call
        self._jit_scalar = self.callable(self.nans, self.reverse, scalar=True)
        self._jit_non_scalar = self.callable(self.nans, self.reverse, scalar=False)

    def __call__(self, group_idx, a, size=None, fill_value=0, order='C',
                 dtype=None, axis=None, ddof=0):
        iv = input_validation(group_idx, a, size=size, order=order, axis=axis, check_bounds=False)
        group_idx, a, flat_size, ndim_idx, size = iv

        # TODO: The typecheck should be done by the class itself, not by check_dtype
        dtype = check_dtype(dtype, self.func, a, len(group_idx))
        check_fill_value(fill_value, dtype, func=self.func)
        input_dtype = type(a) if np.isscalar(a) else a.dtype
        ret, counter, mean, outer = self._initialize(flat_size, fill_value, dtype, input_dtype, group_idx.size)
        group_idx = np.ascontiguousarray(group_idx)

        if not np.isscalar(a):
            a = np.ascontiguousarray(a)
            jitfunc = self._jit_non_scalar
        else:
            jitfunc = self._jit_scalar
        jitfunc(group_idx, a, ret, counter, mean, outer, fill_value, ddof)
        self._finalize(ret, counter, fill_value)

        if self.outer:
            return outer

        # Deal with ndimensional indexing
        if ndim_idx > 1:
            ret = ret.reshape(size, order=order)
        return ret

    @classmethod
    def _initialize(cls, flat_size, fill_value, dtype, input_dtype, input_size):
        if cls.forced_fill_value is None:
            ret = np.full(flat_size, fill_value, dtype=dtype)
        else:
            ret = np.full(flat_size, cls.forced_fill_value, dtype=dtype)

        counter = mean = outer = None
        if cls.counter_fill_value is not None:
            counter = np.full_like(ret, cls.counter_fill_value, dtype=cls.counter_dtype)
        if cls.mean_fill_value is not None:
            dtype = cls.mean_dtype if cls.mean_dtype else input_dtype
            mean = np.full_like(ret, cls.mean_fill_value, dtype=dtype)
        if cls.outer:
            outer = np.full(input_size, fill_value, dtype=dtype)

        return ret, counter, mean, outer

    @classmethod
    def _finalize(cls, ret, counter, fill_value):
        if cls.forced_fill_value is not None and fill_value != cls.forced_fill_value:
            if cls.counter_dtype == bool:
                ret[counter] = fill_value
            else:
                ret[~counter.astype(bool)] = fill_value

    @classmethod
    def callable(cls, nans=False, reverse=False, scalar=False):
        """ Compile a jitted function doing the hard part of the job """
        _valgetter = cls._valgetter_scalar if scalar else cls._valgetter
        valgetter = nb.njit(_valgetter)
        outersetter = nb.njit(cls._outersetter)

        _cls_inner = nb.njit(cls._inner)
        if nans:
            def _inner(ri, val, ret, counter, mean):
                if not np.isnan(val):
                    _cls_inner(ri, val, ret, counter, mean)
            inner = nb.njit(_inner)
        else:
            inner = _cls_inner

        def _loop(group_idx, a, ret, counter, mean, outer, fill_value, ddof):
            # fill_value and ddof need to be present for being exchangeable with loop_2pass
            size = len(ret)
            rng = range(len(group_idx) - 1, -1 , -1) if reverse else range(len(group_idx))
            for i in rng:
                ri = group_idx[i]
                if ri < 0:
                    raise ValueError("negative indices not supported")
                if ri >= size:
                    raise ValueError("one or more indices in group_idx are too large")
                val = valgetter(a, i)
                inner(ri, val, ret, counter, mean)
                outersetter(outer, i, ret[ri])
        return nb.njit(_loop, nogil=True)

    @staticmethod
    def _valgetter(a, i):
        return a[i]

    @staticmethod
    def _valgetter_scalar(a, i):
        return a

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        raise NotImplementedError("subclasses need to overwrite _inner")

    @staticmethod
    def _outersetter(outer, i, val):
        pass

class Aggregate2pass(AggregateOp):
    """Base class for everything that needs to process the data twice like mean, var and std."""
    @classmethod
    def callable(cls, nans=False, reverse=False, scalar=False):
        # Careful, cls needs to be passed, so that the overwritten methods remain available in
        # AggregateOp.callable
        loop = super(Aggregate2pass, cls).callable(nans=nans, reverse=reverse, scalar=scalar)

        _2pass_inner = nb.njit(cls._2pass_inner)
        def _loop2(ret, counter, mean, fill_value, ddof):
            for ri in range(len(ret)):
                if counter[ri]:
                    ret[ri] = _2pass_inner(ri, ret, counter, mean, ddof)
                else:
                    ret[ri] = fill_value
        loop2 = nb.njit(_loop2)

        def _loop_2pass(group_idx, a, ret, counter, mean, outer, fill_value, ddof):
            loop(group_idx, a, ret, counter, mean, outer, fill_value, ddof)
            loop2(ret, counter, mean, fill_value, ddof)
        return nb.njit(_loop_2pass)

    @staticmethod
    def _2pass_inner(ri, ret, counter, mean, ddof):
        raise NotImplementedError("subclasses need to overwrite _2pass_inner")

    @classmethod
    def _finalize(cls, ret, counter, fill_value):
        """Copying the fill value is already done in the 2nd pass"""
        pass

class AggregateNtoN(AggregateOp):
    """Base class for cumulative functions, where the output size matches the input size."""
    outer = True

    @staticmethod
    def _outersetter(outer, i, val):
        outer[i] = val

class AggregateGeneric(AggregateOp):
    """Base class for jitting arbitrary functions."""
    counter_fill_value = None

    def __init__(self, func, **kwargs):
        self.func = func
        self.__dict__.update(kwargs)
        self._jitfunc = self.callable(self.nans)

    def __call__(self, group_idx, a, size=None, fill_value=0, order='C',
                 dtype=None, axis=None, ddof=0):
        iv = input_validation(group_idx, a, size=size, order=order, axis=axis, check_bounds=False)
        group_idx, a, flat_size, ndim_idx, size = iv

        # TODO: The typecheck should be done by the class itself, not by check_dtype
        dtype = check_dtype(dtype, self.func, a, len(group_idx))
        check_fill_value(fill_value, dtype, func=self.func)
        input_dtype = type(a) if np.isscalar(a) else a.dtype
        ret, _, _, _= self._initialize(flat_size, fill_value, dtype, input_dtype, group_idx.size)
        group_idx = np.ascontiguousarray(group_idx)

        sortidx = np.argsort(group_idx, kind='mergesort')
        self._jitfunc(sortidx, group_idx, a, ret)

        # Deal with ndimensional indexing
        if ndim_idx > 1:
            ret = ret.reshape(size, order=order)
        return ret

    def callable(self, nans=False):
        """Compile a jitted function and loop it over the sorted data."""
        jitfunc = nb.njit(self.func, nogil=True)

        def _loop(sortidx, group_idx, a, ret):
            size = len(ret)
            group_idx_srt = group_idx[sortidx]
            a_srt = a[sortidx]

            indices = step_indices(group_idx_srt)
            for i in range(len(indices) - 1):
                start_idx, stop_idx =  indices[i], indices[i + 1]
                ri = group_idx_srt[start_idx]
                if ri < 0:
                    raise ValueError("negative indices not supported")
                if ri >= size:
                    raise ValueError("one or more indices in group_idx are too large")
                ret[ri] = jitfunc(a_srt[start_idx:stop_idx])
        return nb.njit(_loop, nogil=True)

class Sum(AggregateOp):
    forced_fill_value = 0

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] = 0
        ret[ri] += val

class Prod(AggregateOp):
    forced_fill_value = 1

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] = 0
        ret[ri] *= val

class Len(AggregateOp):
    forced_fill_value = 0

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] = 0
        ret[ri] += 1

class All(AggregateOp):
    forced_fill_value = 1

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] = 0
        ret[ri] &= bool(val)

class Any(AggregateOp):
    forced_fill_value = 0

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] = 0
        ret[ri] |= bool(val)

class Last(AggregateOp):
    counter_fill_value = None

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        ret[ri] = val

class First(Last):
    reverse = True

class AllNan(AggregateOp):
    forced_fill_value = 1

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] = 0
        ret[ri] &= val == val

class AnyNan(AggregateOp):
    forced_fill_value = 0

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] = 0
        ret[ri] |= val != val

class Max(AggregateOp):
    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        if counter[ri]:
            ret[ri] = val
            counter[ri] = 0
        elif ret[ri] < val:
            ret[ri] = val

class Min(AggregateOp):
    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        if counter[ri]:
            ret[ri] = val
            counter[ri] = 0
        elif ret[ri] > val:
            ret[ri] = val

class ArgMax(AggregateOp):
    mean_fill_value = np.nan

    @staticmethod
    def _valgetter(a, i):
        return a[i], i

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        cmp_val, arg = val
        if counter[ri]:
            mean[ri] = cmp_val
            ret[ri] = arg
            counter[ri] = 0
        elif mean[ri] < cmp_val:
            mean[ri] = cmp_val
            ret[ri] = arg

class ArgMin(ArgMax):
    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        cmp_val, arg = val
        if counter[ri]:
            mean[ri] = cmp_val
            ret[ri] = arg
            counter[ri] = 0
        elif mean[ri] > cmp_val:
            mean[ri] = cmp_val
            ret[ri] = arg

class Mean(Aggregate2pass):
    forced_fill_value = 0
    counter_fill_value = 0
    counter_dtype = int

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] += 1
        ret[ri] += val

    @staticmethod
    def _2pass_inner(ri, ret, counter, mean, ddof):
        return ret[ri] / counter[ri]

class Std(Mean):
    mean_fill_value = 0

    @staticmethod
    def _inner(ri, val, ret, counter, mean):
        counter[ri] += 1
        mean[ri] += val
        ret[ri] += val * val

    @staticmethod
    def _2pass_inner(ri, ret, counter, mean, ddof):
        mean2 = mean[ri] * mean[ri]
        return np.sqrt((ret[ri] - mean2 / counter[ri]) / (counter[ri] - ddof))

class Var(Std):
    @staticmethod
    def _2pass_inner(ri, ret, counter, mean, ddof):
        mean2 = mean[ri] * mean[ri]
        return (ret[ri] - mean2 / counter[ri]) / (counter[ri] - ddof)

class CumSum(AggregateNtoN, Sum):
    pass

class CumProd(AggregateNtoN, Prod):
    pass

class CumMax(AggregateNtoN, Max):
    pass

class CumMin(AggregateNtoN, Min):
    pass

def get_funcs():
    funcs = dict()
    for op in (Sum, Prod, Len, All, Any, Last, First, AllNan, AnyNan, Min, Max,
               ArgMin, ArgMax, Mean, Std, Var,
               CumSum, CumProd, CumMax, CumMin):
        funcname = op.__name__.lower()
        funcs[funcname] = op(funcname)
        if funcname not in funcs_no_separate_nan:
            funcname = 'nan' + funcname
            funcs[funcname] = op(funcname, nans=True)
    return funcs

_impl_dict = get_funcs()
_default_cache = {}

def aggregate(group_idx, a, func='sum', size=None, fill_value=0, order='C',
              dtype=None, axis=None, cache=None, **kwargs):
    func = get_func(func, aliasing, _impl_dict)
    if not isstr(func):
        if cache in (None, False):
            aggregate_op = AggregateGeneric(func)
        else:
            if cache is True:
                cache = _default_cache
            aggregate_op = cache.setdefault(func, AggregateGeneric(func))
        return aggregate_op(group_idx, a, size, fill_value, order, dtype, axis, **kwargs)
    else:
        func = _impl_dict[func]
        return func(group_idx, a, size, fill_value, order, dtype, axis, **kwargs)

aggregate.__doc__ = """
    This is the numba implementation of aggregate.
    """ + aggregate_common_doc

@nb.njit(nogil=True, cache=True)
def step_count(group_idx):
    """Return the amount of index changes within group_idx."""
    cmp_pos = 0
    steps = 1
    if len(group_idx) < 1:
        return 0
    for i in range(len(group_idx)):
        if group_idx[cmp_pos] != group_idx[i]:
            cmp_pos = i
            steps += 1
    return steps

@nb.njit(nogil=True, cache=True)
def step_indices(group_idx):
    """Return the edges of areas within group_idx, which are filled with the same value."""
    ilen = step_count(group_idx) + 1
    indices = np.empty(ilen, np.int64)
    indices[0] = 0
    indices[-1] = group_idx.size
    cmp_pos = 0
    ri = 1
    for i in range(len(group_idx)):
        if group_idx[cmp_pos] != group_idx[i]:
            cmp_pos = i
            indices[ri] = i
            ri += 1
    return indices