import calendar
import logging
import math
import numba
import numpy as np
import pandas as pd
#import profile
import scipy.constants
import warnings

from indices_python import utils

#-----------------------------------------------------------------------------------------------------------------------
# list the objects that we'll make publicly visible from this module, as interpreted by 'import *'
# a good explanation of this: https://stackoverflow.com/questions/44834/can-someone-explain-all-in-python
__all__ = ['pdsi_from_climatology',
           'potential_evapotranspiration']

#-----------------------------------------------------------------------------------------------------------------------
# global variables
#TODO/FIXME eliminate all globals where reasonable (these are in place
# for convenience and to reflect some of what's happening in pdinew.f)

# the number of time steps (months) that'll be back filled when performing the assignment of index values via
# the backtracking process; doubles as an array index used with the various backtracking related arrays
_k8 = 0

# the current time step (month) being processed, should be read-only in all functions where
# it's used other than the main loop in _pdsi() where the value is updated at each time step
_i = 0

#-----------------------------------------------------------------------------------------------------------------------
# constants
_WET = 0
_DRY = 1

#-----------------------------------------------------------------------------------------------------------------------
# set up a basic, global _logger
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s %(levelname)s %(message)s',
                    datefmt='%Y-%m-%d  %H:%M:%S')
_logger = logging.getLogger(__name__)

# set numpy's print options so when array values are printed we can control the precision
np.set_printoptions(formatter={'float': lambda x: "{0:.2f}".format(x)})

#-----------------------------------------------------------------------------------------------------------------------
#@profile
def pdsi_from_climatology(precip_timeseries,
                          temp_timeseries,
                          awc,
                          latitude,
                          B,
                          H,
                          data_begin_year,
                          calibration_begin_year,
                          calibration_end_year):

    # calculate the negative tangent of the latitude which is used as an argument to the water balance function
    neg_tan_lat = -1 * math.tan(math.radians(latitude))

    # compute water balance values using the function translated from the Fortran pdinew.f
    pdat, pedat, pldat, prdat, rdat, tldat, etdat, rodat, prodat = \
        _water_balance(temp_timeseries, precip_timeseries, awc, neg_tan_lat, B, H)

    #TODO pdat should equal precipitation at this point, no need for duplicate array -- CONFIRM

    # compute CAFEC coefficients using translated functions from pdinew.f Fortran code
    alpha, beta, delta, gamma, t_ratio = _cafec_coefficients(precip_timeseries,
                                                             pedat,
                                                             etdat,
                                                             prdat,
                                                             rdat,
                                                             rodat,
                                                             prodat,
                                                             tldat,
                                                             pldat,
                                                             data_begin_year,
                                                             calibration_begin_year,
                                                             calibration_end_year)

    # compute the weighting factor (climatic characteristic) using the version translated from pdinew.f
    K = _climatic_characteristic(alpha,
                                 beta,
                                 gamma,
                                 delta,
                                 pdat,
                                 pedat,
                                 prdat,
                                 prodat,
                                 pldat,
                                 t_ratio,
                                 data_begin_year,
                                 calibration_begin_year,
                                 calibration_end_year)

    # compute the Z-index
    #TODO eliminate the P-hat once development/debugging is complete, since it's really an intermediate/internal value
    Z = _zindex(alpha,
                beta,
                gamma,
                delta,
                precip_timeseries,
                pedat,
                prdat,
                prodat,
                pldat,
                K)

    # compute the final Palmers
    PDSI, PHDI, PMDI = _pdsi(Z)

    return PDSI, PHDI, PMDI, Z, pedat

#-----------------------------------------------------------------------------------------------------------------------
#@profile
#@numba.jit   #TODO/FIXME numba compile not working as expected yet, AssertionError
def _cafec_coefficients(P,
                        PET,
                        ET,
                        PR,
                        R,
                        RO,
                        PRO,
                        L,
                        PL,
                        data_start_year,
                        calibration_start_year,
                        calibration_end_year):
    '''
    This function calculates CAFEC coefficients used for computing Palmer's Z index using inputs from
    the water balance function. Translated from Fortran pdinew.f

    :param P: 1-D numpy.ndarray of monthly precipitation observations, in inches, the number of array elements
              (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param PET: 1-D numpy.ndarray of monthly potential evapotranspiration values, in inches, the number of array elements
                (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param ET: 1-D numpy.ndarray of monthly evapotranspiration values, in inches, the number of array elements
               (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param PR: 1-D numpy.ndarray of monthly potential recharge values, in inches, the number of array elements
               (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param R: 1-D numpy.ndarray of monthly recharge values, in inches, the number of array elements
              (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param RO: 1-D numpy.ndarray of monthly runoff values, in inches, the number of array elements
               (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param PRO: 1-D numpy.ndarray of monthly potential runoff values, in inches, the number of array elements
                (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param L: 1-D numpy.ndarray of monthly loss values, in inches, the number of array elements
              (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param PL: 1-D numpy.ndarray of monthly potential loss values, in inches, the number of array elements
              (array size) should be a multiple of 12 (representing an ordinal number of full years)
    :param data_start_year: initial year of the input arrays, i.e. the first element of each of the input arrays
                            is assumed to correspond to January of this initial year
    :param calibration_start_year: initial year of the calibration period, should be greater than or equal to the data_start_year
    :param calibration_end_year: final year of the calibration period
    :return 1-D numpy.ndarray of CAFEC coefficient values, alpha, beta, delta, gamma, and the T ratio,
            with shape of these arrays == (12,), corresponding to calendar months (12 elements)
    :rtype: numpy.ndarray of floats
    '''

    # the potential (PET, ET, PR, PL) and actual (R, RO, S, L, P) water balance arrays are reshaped as 2-D arrays
    # (matrices) such that the rows of each matrix represent years and the columns represent calendar months
    PET = utils.reshape_to_years_months(PET)
    ET = utils.reshape_to_years_months(ET)
    PR = utils.reshape_to_years_months(PR)
    PL = utils.reshape_to_years_months(PL)
    R = utils.reshape_to_years_months(R)
    RO = utils.reshape_to_years_months(RO)
    PRO = utils.reshape_to_years_months(PRO)
    L = utils.reshape_to_years_months(L)
    P = utils.reshape_to_years_months(P)

    # ALPHA, BETA, GAMMA, DELTA CALCULATIONS
    # A calibration period is used to calculate alpha, beta, gamma, and
    # and delta, four coefficients dependent on the climate of the area being
    # examined. The NCDC and CPC use the calibration period January 1931
    # through December 1990 (cf. Karl, 1986; Journal of Climate and Applied
    # Meteorology, Vol. 25, No. 1, January 1986).

    #!!!!!!!!!!!!!
    # TODO make sure calibration years range is valid, i.e. within actual data years range

    # determine the array (year axis) indices for the calibration period
    total_data_years = int(P.shape[0] / 12)
    data_end_year = data_start_year + total_data_years - 1
    calibration_start_year_index = calibration_start_year - data_start_year
    calibration_end_year_index = calibration_end_year - data_start_year

    # get calibration period arrays
    if (calibration_start_year > data_start_year) or (calibration_end_year < data_end_year):
        P_calibration = P[calibration_start_year_index:calibration_end_year_index + 1]
        ET_calibration = ET[calibration_start_year_index:calibration_end_year_index + 1]
        PET_calibration = PET[calibration_start_year_index:calibration_end_year_index + 1]
        R_calibration = R[calibration_start_year_index:calibration_end_year_index + 1]
        PR_calibration = PR[calibration_start_year_index:calibration_end_year_index + 1]
        L_calibration = L[calibration_start_year_index:calibration_end_year_index + 1]
        PL_calibration = PL[calibration_start_year_index:calibration_end_year_index + 1]
        RO_calibration = RO[calibration_start_year_index:calibration_end_year_index + 1]
        PRO_calibration = PRO[calibration_start_year_index:calibration_end_year_index + 1]
    else:
        P_calibration = P
        ET_calibration = ET
        PET_calibration = PET
        R_calibration = R
        PR_calibration = PR
        L_calibration = L
        PL_calibration = PL
        RO_calibration = RO
        PRO_calibration = PRO

    # due to (innocuous/annoying) RuntimeWarnings raised by some math/numpy modules, we wrap the below code in a catcher
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)

        # get sums for each calendar month (compute sums over the year axis, giving a sum for each calendar month over all years)
        P_sum = np.nansum(P_calibration, axis=0)
        ET_sum = np.nansum(ET_calibration, axis=0)
        PET_sum = np.nansum(PET_calibration, axis=0)
        R_sum = np.nansum(R_calibration, axis=0)
        PR_sum = np.nansum(PR_calibration, axis=0)
        L_sum = np.nansum(L_calibration, axis=0)
        PL_sum = np.nansum(PL_calibration, axis=0)
        RO_sum = np.nansum(RO_calibration, axis=0)
        PRO_sum = np.nansum(PRO_calibration, axis=0)

        # (calendar) monthly CAFEC coefficients
        alpha = np.empty((12,))
        beta = np.empty((12,))
        gamma = np.empty((12,))
        delta = np.empty((12,))
        t_ratio = np.empty((12,))

        # compute the alpha, beta, gamma, and delta coefficients for each calendar month
        for i in range(12):

            #     ALPHA CALCULATION
            #     -----------------
            if PET_sum[i] != 0.0:
                alpha[i] = ET_sum[i] / PET_sum[i]
            else:
                if ET_sum[i] == 0.0:
                    alpha[i] = 1.0
                else:
                    alpha[i] = 0.0

            #
            #     BETA CALCULATION
            #     ----------------
            if PR_sum[i] != 0.0:
                beta[i] = R_sum[i] / PR_sum[i]
            else:
                if R_sum[i] == 0.0:
                    beta[i] = 1.0
                else:
                    beta[i] = 0.0

            #
            #     GAMMA CALCULATION
            #     -----------------
            if PRO_sum[i] != 0.0:
                gamma[i] = RO_sum[i] / PRO_sum[i]
            else:
                if RO_sum[i] == 0.0:
                    gamma[i] = 1.
                else:
                    gamma[i] = 0.0

            #
            #     DELTA CALCULATION
            #     -----------------
            if PL_sum[i] != 0.0:
                delta[i] = L_sum[i] / PL_sum[i]
            else:
                delta[i] = 0.0

            # 'T' ratio of average moisture demand to the average moisture supply in the month
            t_ratio[i] = (PET_sum[i] + R_sum[i] + RO_sum[i]) / (P_sum[i] + L_sum[i])

    return alpha, beta, delta, gamma, t_ratio

#-----------------------------------------------------------------------------------------------------------------------
#@profile
@numba.jit
def _climatic_characteristic(alpha,
                             beta,
                             gamma,
                             delta,
                             Pdat,
                             PEdat,
                             PRdat,
                             PROdat,
                             PLdat,
                             t_ratio,
                             begin_year,
                             calibration_begin_year,
                             calibration_end_year):

    SABSD = np.zeros((12,))

    number_calibration_years = calibration_end_year - calibration_begin_year + 1

    # loop over the calibration years, get the sum of the absolute values of the moisture departure SABSD for each month
    for j in range(calibration_begin_year - begin_year, calibration_end_year - begin_year + 1):
        for m in range(12):

            #-----------------------------------------------------------------------
            #     REREAD MONTHLY PARAMETERS FOR CALCULATION OF
            #     THE 'K' MONTHLY WEIGHTING FACTORS USED IN Z-INDEX CALCULATION
            #-----------------------------------------------------------------------
            PHAT = (alpha[m] * PEdat[j, m]) + (beta[m] * PRdat[j, m]) + (gamma[m] * PROdat[j, m]) - (delta[m] * PLdat[j, m])
            D = Pdat[j, m] - PHAT
            SABSD[m] = SABSD[m] + abs(D)

    # get the first approximation of K (weighting factor)
    SWTD = 0.0
    AKHAT = np.empty((12,))
    for m in range(12):
        DBAR = SABSD[m] / number_calibration_years
        AKHAT[m] = 1.5 * math.log10((t_ratio[m] + 2.8) / DBAR) + 0.5
        SWTD = SWTD + (DBAR * AKHAT[m])

    AK = np.empty((12,))
    for m in range(12):
        AK[m] = 17.67 * AKHAT[m] / SWTD

    return AK

#-----------------------------------------------------------------------------------------------------------------------
#@profile
@numba.jit
def _water_balance(T,
                   P,
                   AWC,
                   TLA,
                   B,
                   H,
                   begin_year=1895):

    '''
    Computes a water balance accounting for monthly time series. Translated from the Fortran code pdinew.f

    :param T: monthly average temperature values, starting in January of the initial year
    :param P: monthly total precipitation values, starting in January of the initial year
    :param AWC: available water capacity, below (not including) the top inch
    :param B: read from soil constants file
    :param H: read from soil constants file
    :param begin_year: initial year of the dataset
    :param TLA: negative tangent of the latitude
    :return: P, SP, PE, PL, PR, R, L, ET, RO, T, SSs, SSu
    :rtype: numpy arrays with shape (total_years, 12)
    '''
    #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    #TODO document the differences between PRO and SP, as the SP return value (spdat) is being used as PRO (potential
    # runoff) in later parts of the code, let's verify that this is intended, or perhaps spdat and prodat are
    # synonymous/duplicates and one of these can be eliminated, suspicion is that the PRO variable in pdinew.f
    # refers to a probability value, so the variable for potential runoff is named SP instead (?)
    #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    #TODO determine the total number of years from the original shape of the input array, assume shape == (total months)

    #EXPERIMENTAL/DEBUG ONLY -- REMOVE
    if AWC >= 1.0:
        AWC = AWC - 1.0

    # reshape the precipitation array from 1-D (assumed to be total months) to (years, 12) with the second
    # dimension being calendar months, and the final/missing monthly values of the final year padded with NaNs
    T = utils.reshape_to_years_months(T)
    P = utils.reshape_to_years_months(P)
    total_years = P.shape[0]

    WCTOP = 1.0
    SS  = WCTOP
    SU = AWC
    WCTOT = AWC + WCTOP

    #TODO document this, where do these come from?
    PHI = np.array([-0.3865982, -0.2316132, -0.0378180, 0.1715539, 0.3458803, 0.4308320, \
                     0.3916645, 0.2452467, 0.0535511, -0.15583436, -0.3340551, -0.4310691])

    # initialize the data arrays with NaNs
    pdat = np.full((total_years, 12), np.NaN)
    pedat = np.full((total_years, 12), np.NaN)
    pldat = np.full((total_years, 12), np.NaN)
    prdat = np.full((total_years, 12), np.NaN)
    rdat = np.full((total_years, 12), np.NaN)
    tldat = np.full((total_years, 12), np.NaN)
    etdat = np.full((total_years, 12), np.NaN)
    rodat = np.full((total_years, 12), np.NaN)
    prodat = np.full((total_years, 12), np.NaN)

    # loop on years and months
    end_year = begin_year + total_years
    years_range = range(begin_year, end_year)
    for year_index, year in enumerate(years_range):

        for month_index in range(12):

            temperature = T[year_index, month_index]
            precipitation = P[year_index, month_index]

            #-----------------------------------------------------------------------
            #     HERE START THE WATER BALANCE CALCULATIONS
            #-----------------------------------------------------------------------
            SP = SS + SU
            PR = AWC + WCTOP - SP

            # PRO is the potential runoff. According to Alley (1984),
            # PRO = AWC - PR = Ss + Su; with Ss and Su referring to those values at
            # the beginning of the month: Ss0 and Su0.
            PRO = SP

            #-----------------------------------------------------------------------
            #     1 - CALCULATE PE (POTENTIAL EVAPOTRANSPIRATION)
            #-----------------------------------------------------------------------
            if temperature <= 32.0:
                PE = 0.0
            else:
                DUM = PHI[month_index] * TLA
                DK = math.atan(math.sqrt(1.0 - (DUM * DUM)) / DUM)
                if DK < 0.0:
                    DK = 3.141593 + DK
                DK = (DK + 0.0157) / 1.57
                if temperature >= 80.0:
                    PE = (math.sin((temperature / 57.3) - 0.166) - 0.76) * DK
                else:
                    DUM = math.log(temperature - 32.0)
                    PE = math.exp(-3.863233 + (B * 1.715598) - (B * math.log(H)) + (B * DUM)) * DK

            #-----------------------------------------------------------------------
            #     CONVERT DAILY TO MONTHLY
            #-----------------------------------------------------------------------
            PE = PE * calendar.monthrange(year, month_index + 1)[1]

            #-----------------------------------------------------------------------
            #     2 - PL  POTENTIAL LOSS
            #-----------------------------------------------------------------------
            if SS >= PE:
                PL  = PE
            else:
                PL = ((PE - SS) * SU) / (AWC + WCTOP) + SS
                PL = min(PL, SP)

            #-----------------------------------------------------------------------
            #     3 - CALCULATE RECHARGE, RUNOFF, RESIDUAL MOISTURE, LOSS TO BOTH
            #         SURFACE AND UNDER LAYERS, DEPENDING ON STARTING MOISTURE
            #         CONTENT AND VALUES OF PRECIPITATION AND EVAPORATION.
            #-----------------------------------------------------------------------
            if precipitation >= PE:
                #     ----------------- PRECIP EXCEEDS POTENTIAL EVAPORATION
                ET = PE
                TL = 0.0
                if (precipitation - PE) > (WCTOP - SS):
                    #         ------------------------------ EXCESS PRECIP RECHARGES
                    #                                        UNDER LAYER AS WELL AS UPPER
                    RS = WCTOP - SS
                    SSS = WCTOP
                    if (precipitation - PE - RS) < (AWC - SU):
                    #             ---------------------------------- BOTH LAYERS CAN TAKE
                    #                                                THE ENTIRE EXCESS
                        RU = precipitation - PE - RS
                        RO = 0.0
                    else:
                        #             ---------------------------------- SOME RUNOFF OCCURS
                        RU = AWC - SU
                        RO = precipitation - PE - RS - RU

                    SSU = SU + RU
                    R   = RS + RU
                else:
                    #         ------------------------------ ONLY TOP LAYER RECHARGED
                    R  = precipitation - PE
                    SSS = SS + precipitation - PE
                    SSU = SU
                    RO  = 0.0

            else:
                #     ----------------- EVAPORATION EXCEEDS PRECIPITATION
                R  = 0.0
                if SS >= (PE - precipitation):
                #         ----------------------- EVAP FROM SURFACE LAYER ONLY
                    SL  = PE - precipitation
                    SSS = SS - SL
                    UL  = 0.0
                    SSU = SU
                else:
                    #         ----------------------- EVAP FROM BOTH LAYERS
                    SL  = SS
                    SSS = 0.0
                    UL  = (PE - precipitation - SL) * SU / (WCTOT)  # this includes the fix suggested by Jacobi 2013
                    UL  = min(UL, SU)
                    SSU = SU - UL

                TL  = SL + UL
                RO  = 0.0
                ET  = precipitation  + SL + UL

            # set the climatology and water balance data array values for this year/month time step
            pdat[year_index, month_index] = precipitation
            pedat[year_index, month_index] = PE
            pldat[year_index, month_index] = PL
            prdat[year_index, month_index] = PR
            rdat[year_index, month_index] = R
            tldat[year_index, month_index] = TL
            etdat[year_index, month_index] = ET
            rodat[year_index, month_index] = RO
            prodat[year_index, month_index] = PRO

            # reset the upper and lower soil moisture values
            SS = SSS
            SU = SSU

    return pdat, pedat, pldat, prdat, rdat, tldat, etdat, rodat, prodat

#-----------------------------------------------------------------------------------------------------------------------
#@profile
@numba.jit
def _zindex(alpha,
            beta,
            gamma,
            delta,
            P,
            PE,
            PR,
            PRO,
            PL,
            K):
    '''
    Compute the Z-Index for a time series.

    :param alpha: array of the monthly "alpha" CAFEC coefficients, 1-D with 12 elements, one per calendar month
    :param beta: array of the monthly "beta" CAFEC coefficients, 1-D with 12 elements, one per calendar month
    :param gamma: array of the monthly "delta" CAFEC coefficients, 1-D with 12 elements, one per calendar month
    :param delta: array of the monthly "gamma" CAFEC coefficients, 1-D with 12 elements, one per calendar month
    :param P: array of monthly precipitation values, 1-D with total size matching the corresponding 2-D arrays
    :param PE: array of monthly potential evapotranspiration values, 2-D with shape (# of years, 12)
    :param PR: array of monthly potential recharge values, 2-D with shape (# of years, 12)
    :param PRO: array of monthly potential runoff values, 2-D with shape (# of years, 12)
    :param PL: array of monthly potential loss values, 2-D with shape (# of years, 12)
    :param K: array of the monthly climatic characteristic, 1-D with 12 elements, one per calendar month
    :return Z-Index and CAFEC precipitation (P-hat)
    :rtype: two floats
    '''

    # reshape the precipitation array, expected to have 1-D shape, to 2-D shape (years, 12)
    P = utils.reshape_to_years_months(P)

    # allocate the Z-Index and P-hat arrays we'll build and return
    Z = np.full(P.shape, np.NaN)
    P_hat = np.full(P.shape, np.NaN)

    # loop over all years and months of the time series
    for j in range(P.shape[0]):

        for m in range(12):

            # compute the CAFEC precipitation (P-hat)
            ET_hat = alpha[m] * PE[j, m]             # eq. 10, Palmer 1965
            R_hat = beta[m] * PR[j, m]               # eq. 11, Palmer 1965
            RO_hat = gamma[m] * PRO[j, m]            # eq. 12, Palmer 1965
            L_hat = delta[m] * PL[j, m]              # eq. 13, Palmer 1965
            P_hat = ET_hat + R_hat + RO_hat - L_hat  # eq. 14, Palmer 1965

            # moisture departure
            d = P[j, m] - P_hat      # eq. 15, Palmer 1965

            # Z-Index
            Z[j, m] = K[m] * d    # eq. 19, Palmer 1965

    return Z

#-----------------------------------------------------------------------------------------------------------------------
#@profile
@numba.jit
def _zindex_from_climatology(temp_timeseries,
                             precip_timeseries,
                             awc,
                             neg_tan_lat,
                             B,
                             H,
                             data_begin_year,
                             calibration_begin_year,
                             calibration_end_year):
    '''
    Compute the Z-Index for a time series.

    :return Z-Index
    :rtype: array of floats
    '''

#     # calculate the negative tangent of the latitude which is used as an argument to the water balance function
#     neg_tan_lat = -1 * math.tan(math.radians(latitude))

    # compute water balance values using the function translated from the Fortran pdinew.f
    pdat, pedat, pldat, prdat, rdat, tldat, etdat, rodat, prodat = \
        _water_balance(temp_timeseries, precip_timeseries, awc, neg_tan_lat, B, H)

    #NOTE we need to compute CAFEC coefficients for use later/below
    # compute PDSI etc. using translated functions from pdinew.f Fortran code
    alpha, beta, delta, gamma, t_ratio = _cafec_coefficients(precip_timeseries,
                                                             pedat,
                                                             etdat,
                                                             prdat,
                                                             rdat,
                                                             rodat,
                                                             prodat,
                                                             tldat,
                                                             pldat,
                                                             data_begin_year,
                                                             calibration_begin_year,
                                                             calibration_end_year)

    # compute the weighting factor (climatic characteristic) using the version translated from pdinew.f
    K = _climatic_characteristic(alpha,
                                 beta,
                                 gamma,
                                 delta,
                                 pdat,
                                 pedat,
                                 prdat,
                                 prodat,
                                 pldat,
                                 t_ratio,
                                 data_begin_year,
                                 calibration_begin_year,
                                 calibration_end_year)

    # compute the Z-index
    #TODO eliminate the P-hat once development/debugging is complete, since it's really an intermediate/internal value
    return _zindex(alpha,
                   beta,
                   gamma,
                   delta,
                   precip_timeseries,
                   pedat,
                   prdat,
                   prodat,
                   pldat,
                   K)

#-----------------------------------------------------------------------------------------------------------------------
#@profile
#@numba.jit  #FIXME not yet working
def _pdsi(Z):
    '''
    :param Z: 2-D array of Z-Index values, corresponding in total size to Z
    '''

    # indicate that we'll use the globals for the backtracking months count and the current time step
    global _k8
    global _i

    # initialize some effective wetness values (previous and current)
    PV = 0.0
    V = 0.0

    # provisional severity index values computed for each month/timestep, typically carried forward for use
    # in the next timestep, since current month X values are computed using the previous month's X values
    X1 = 0.0   # index appropriate to a wet spell that is becoming established, as well as the percent chance that a wet spell has begun
    X2 = 0.0   # index appropriate to a drought that is becoming established, as well as the percent chance that a drought has begun
    X3 = 0.0   # index appropriate to a wet spell or drought that has already been established

    # percentage probability that an established weather spell (wet or dry) has ended (Pe in Palmer 1965, eq. 30)
    prob_ended = 0.0   ##NOTE this was PRO in pdinew.f, now the variable name PRO is used for potential runoff

    # create a DataFrame to use as a container for the arrays and values we'll use throughout the computation loop below
    Z = np.reshape(Z, (Z.size, 1))
    df = pd.DataFrame(data=Z,
                      index=range(0, Z.size),
                      columns=['Z'])

    # create a list of coluumn names that match to the intermediate work arrays
    column_names = ['PPR', 'PDSI', 'PHDI', 'WPLM', 'SX', 'SX1', 'SX2', 'SX3', 'X', 'PX1', 'PX2', 'PX3']
    for column_name in column_names:

        # get the array corresponding to the current column
        column_array = np.full(Z.shape, np.NaN).flatten()

        # add the column to the DataFrame (as a Series)
        df[column_name] = pd.Series(column_array)

    # add a month index column used to keep track of which month index to use for the start of backtracking
    column_array = np.full(Z.shape, 0, dtype=int).flatten()
    df['index_i'] = pd.Series(column_array)

    # the total number of backtracking months, i.e. when performing backtracking we'll back fill this many months
    _k8 = 0

    # loop over all years and months of the time series
    for _i in range(Z.size):

#         # DEBUGGING ONLY -- REMOVE
#         print('_i: {0}'.format(_i))

        # keep track of final backtracking index, meaningful once _k8 > 0, where _k8 > 0 indicates that backtracking is required
        df.index_i[_k8] = _i

        # original (all-caps) comments from pdinew.f, left in place to facilitate development

        # if the percentage probability (PPR in Palmer 1965) is 100% or 0%
        if prob_ended == 100.0 or prob_ended == 0.0:
        #     ------------------------------------ NO ABATEMENT UNDERWAY
        #                                          WET OR DROUGHT WILL END IF
        #                                             -0.5 =< X3 =< 0.5

            # "near normal" is defined as the range [-0.5 ... 0.5], Palmer 1965 p. 29
            if abs(X3) <= 0.5:
            #         ---------------------------------- END OF DROUGHT OR WET
                PV = 0.0
                df.PPR[_i] = 0.0
                df.PX3[_i] = 0.0
                #             ------------ BUT CHECK FOR NEW WET OR DROUGHT START FIRST
                # GOTO 200 in pdinew.f
                df, X1, X2, X3, V, prob_ended = _compute_X(df, X1, X2, PV)

            elif X3 > 0.5:
                #         ----------------------- WE ARE IN A WET SPELL
                if df.Z[_i] >= 0.15:
                    #              ------------------ THE WET SPELL INTENSIFIES
                    #GO TO 210 in pdinew.f
                    df, X1, X2, X3, V, PV, prob_ended = _between_0s(df, X3)

                else:
                    #             ------------------ THE WET STARTS TO ABATE (AND MAY END)
                    #GO TO 170 in pdinew.f
                    df, X1, X2, X3, V, prob_ended = _wet_spell_abatement(df, V, prob_ended, X1, X2, X3)

            elif X3 < -0.5:
                #         ------------------------- WE ARE IN A DROUGHT

                # in order to start to pull out of a drought the Z-Index for the month needs to be >= -0.15 (Palmer 1965, eq. 29)
                if df.Z[_i] <= -0.15:  #NOTE pdinew.f uses <= here, rather than <: "IF (Z(j,m).LE.-.15) THEN..."
                    #              -------------------- THE DROUGHT INTENSIFIES
                    #GO TO 210
                    df, X1, X2, X3, V, PV, prob_ended = _between_0s(df, X3)

                else:
                    # Palmer 1965, p. 29: "any value of Z >= -0.15 will tend to end a drought"
                    #             ------------------ THE DROUGHT STARTS TO ABATE (AND MAY END)
                    #GO TO 180
                    df, X1, X2, X3, V, prob_ended = _dry_spell_abatement(df, V, X1, X2, X3, prob_ended)

        else:
            #     ------------------------------------------ABATEMENT IS UNDERWAY
            if X3 > 0.0:

                #         ----------------------- WE ARE IN A WET SPELL
                #GO TO 170 in pdinew.f
                df, X1, X2, X3, V, prob_ended = _wet_spell_abatement(df, V, prob_ended, X1, X2, X3)

            else:  # if X3 <= 0.0:

                #         ----------------------- WE ARE IN A DROUGHT
                #GO TO 180
                df, X1, X2, X3, V, prob_ended = _dry_spell_abatement(df, V, X1, X2, X3, prob_ended)

    # clear out the remaining values from the backtracking array, if any are left, assigning into the indices arrays
    for x in range(_k8):

        # get the next backtrack month index
        y = df.index_i[x]

        df.PDSI[y] = df.X[y]
        df.PHDI[y] = df.PX3[y]

        if df.PX3[y] == 0.0:

            df.PHDI[y] = df.X[y]

        df.WPLM[y] = _case(df.PPR[y], df.PX1[y], df.PX2[y], df.PX3[y])

    # return the Palmer severity index values as 1-D arrays
    return df.PDSI.values, df.PHDI.values, df.WPLM.values

#-----------------------------------------------------------------------------------------------------------------------
# from label 190 in pdinew.f
#@profile
@numba.jit
def _get_PPR_PX3(df,
                 prob_ended,
                 Ze,
                 V,
                 PV,
                 X3):

    # indicate that we'll use the globals for backtrack months count and current time step
    global _k8
    global _i

    #-----------------------------------------------------------------------
    #     PROB(END) = 100 * (V/Q)  WHERE:
    #             V = SUM OF MOISTURE EXCESS OR DEFICIT (UD OR UW)
    #                 DURING CURRENT ABATEMENT PERIOD
    #             Q = TOTAL MOISTURE ANOMALY REQUIRED TO END THE
    #                 CURRENT DROUGHT OR WET SPELL
    #-----------------------------------------------------------------------
    if prob_ended == 100.0:
        #     --------------------- DROUGHT OR WET CONTINUES, CALCULATE
        #                           PROB(END) - VARIABLE Ze
        Q = Ze
    else:
        Q = Ze + V

    # percentage probability that an established drought or wet spell has ended
    df.PPR[_i] = (PV / Q) * 100.0   # eq. 30 Palmer 1965
    if df.PPR[_i] >= 100.0:

        df.PPR[_i] = 100.0
        df.PX3[_i] = 0.0
    else:
        # eq. 25, Palmer (1965); eq. 4, Wells et al (2003)
        df.PX3[_i] = (0.897 * X3) + (df.Z[_i] / 3.0)

    return df

#-----------------------------------------------------------------------------------------------------------------------
# from label 170 in pdinew.f
#@profile
@numba.jit
def _wet_spell_abatement(df,
                         V,
                         prob_ended,
                         X1,
                         X2,
                         X3):

    # indicate that we'll use the globals for backtrack months count and current time step
    global _k8
    global _i

    #-----------------------------------------------------------------------
    #      WET SPELL ABATEMENT IS POSSIBLE
    #-----------------------------------------------------------------------
    Ud = df.Z[_i] - 0.15
    PV = Ud + min(V, 0.0)
    if PV >= 0.0:

        # GO TO label 210 in pdinew.f
        df, X1, X2, X3, V, PV, prob_ended = _between_0s(df, X3)

    else:
        #     ---------------------- DURING A WET SPELL, PV => 0 IMPLIES
        #                            PROB(END) HAS RETURNED TO 0
        Ze = -2.691 * X3 + 1.5

        # compute the PPR and PX3 values
        df = _get_PPR_PX3(df, prob_ended, Ze, V, PV, X3)  # label 190 in pdinew.f

        # continue at label 200 in pdinew.f
        # recompute the X values and other intermediates
        df, X1, X2, X3, V, prob_ended = _compute_X(df, X1, X2, PV)

    return df, X1, X2, X3, V, prob_ended

#-----------------------------------------------------------------------------------------------------------------------
# from label 180 in pdinew.f
#@profile
@numba.jit
def _dry_spell_abatement(df,
                         V,   # accumulated effective wetness
                         X1,
                         X2,
                         X3,
                         prob_ended):
    '''

    '''
    # indicate that we'll use the globals for backtrack months count and current time step
    global _k8
    global _i

    #-----------------------------------------------------------------------
    #      DROUGHT ABATEMENT IS POSSIBLE
    #-----------------------------------------------------------------------

    # the "effective wetness" of the current month
    Uw = df.Z[_i] + 0.15   # Palmer 1965, eq. 29

    # add this month's effective wetness to the previously accumulated effective wetness (if it was positive),
    # this value will eventually be used as the numerator in the equation for calculating percentage probability
    # that an established weather spell has ended (eq. 30, Palmer 1965)
    PV = Uw + max(V, 0.0)

    # if this month's cumulative effective wetness value isn't positive then there is
    # insufficient moisture this month to end or even abate the established dry spell
    if PV <= 0:
        #     ---------------------- DURING A DROUGHT, PV =< 0 IMPLIES
        #                            PROB(END) HAS RETURNED TO 0
        # GOTO 210 in pdinew.f
        df, X1, X2, X3, V, PV, prob_ended = _between_0s(df, X3)

    else:
        # abatement is underway

        # Calculate the Z value which corresponds to an amount of moisture that is sufficient to end
        # the currently established drought in a single month. Once this is known then we can compare
        # against the actual Z value, to see if we've pulled out of the established drought or not
        Ze = (-2.691 * X3) - 1.5   # eq. 28, Palmer 1965

        # compute the percentage probability that the established spell has ended (PPR),
        # and the severity index for the established spell (PX3)
        df = _get_PPR_PX3(df, prob_ended, Ze, V, PV, X3)  # label 190 in pdinew.f

        # continue at label 200 in pdinew.f
        # recompute the X values and other intermediates
        df, X1, X2, X3, V, prob_ended = _compute_X(df, X1, X2, PV)

    return df, X1, X2, X3, V, prob_ended

#-----------------------------------------------------------------------------------------------------------------------
# label 200 in pdinew.f
#@profile
#@numba.jit  #FIXME not working yet
def _compute_X(df,
               X1,
               X2,
               PV):

    # indicate that we'll use the globals for backtrack months count and current time step
    global _k8
    global _i

    #-----------------------------------------------------------------------
    #     CONTINUE X1 AND X2 CALCULATIONS.
    #     IF EITHER INDICATES THE START OF A NEW WET OR DROUGHT,
    #     AND IF THE LAST WET OR DROUGHT HAS ENDED, USE X1 OR X2
    #     AS THE NEW X3.
    #-----------------------------------------------------------------------

    # compute PX1 using the general formula for X, and make sure we come up with a positive value
    df.PX1[_i] = (0.897 * X1) + (df.Z[_i] / 3.0)
    df.PX1[_i] = max(df.PX1[_i], 0.0)

    # let's see if pre-computing the PX2 value here prevents us from getting into the error condition where we end up with a NaN for X2
    df.PX2[_i] = (0.897 * X2) + (df.Z[_i] / 3.0)
    df.PX2[_i] = min(df.PX2[_i], 0.0)

    #TODO explain these conditionals, refer to literature
    if df.PX1[_i] >= 1.0:

        if df.PX3[_i] == 0.0:
            #         ------------------- IF NO EXISTING WET SPELL OR DROUGHT
            #                             X1 BECOMES THE NEW X3
            df.X[_i] = df.PX1[_i]
            df.PX3[_i] = df.PX1[_i]
            df.PX1[_i] = 0.0
            iass = 1
            df = _assign(df, iass)

            #GOTO 220 in pdinew.f
            V = PV
            prob_ended = df.PPR[_i]
            X1 = df.PX1[_i]
            X2 = df.PX2[_i]
            X3 = df.PX3[_i]

            return df, X1, X2, X3, V, prob_ended

    #TODO explain these conditionals, refer to literature
    if df.PX2[_i] <= -1.0:

        if df.PX3[_i] == 0.0:
            #         ------------------- IF NO EXISTING WET SPELL OR DROUGHT
            #                             X2 BECOMES THE NEW X3
            df.X[_i]   = df.PX2[_i]
            df.PX3[_i] = df.PX2[_i]
            df.PX2[_i] = 0.0
            iass = 2
            df = _assign(df, iass)

            #GOTO 220 in pdinew.f
            V = PV
            prob_ended = df.PPR[_i]
            X1 = df.PX1[_i]
            X2 = df.PX2[_i]
            X3 = df.PX3[_i]

            return df, X1, X2, X3, V, prob_ended

    #TODO explain these conditionals, refer to literature
    if df.PX3[_i] == 0.0:
        #    -------------------- NO ESTABLISHED DROUGHT (WET SPELL), BUT X3=0
        #                         SO EITHER (NONZERO) X1 OR X2 MUST BE USED AS X3
        if df.PX1[_i] == 0.0:

            df.X[_i] = df.PX2[_i]
            iass = 2
            df = _assign(df, iass)

            #GOTO 220 in pdinew.f
            V = PV
            prob_ended = df.PPR[_i]
            X1 = df.PX1[_i]
            X2 = df.PX2[_i]
            X3 = df.PX3[_i]

            return df, X1, X2, X3, V, prob_ended

        elif df.PX2[_i] == 0:

            df.X[_i] = df.PX1[_i]
            iass = 1
            df = _assign(df, iass)

            #GOTO 220 in pdinew.f
            V = PV
            prob_ended = df.PPR[_i]
            X1 = df.PX1[_i]
            X2 = df.PX2[_i]
            X3 = df.PX3[_i]

            return df, X1, X2, X3, V, prob_ended

    #-----------------------------------------------------------------------
    #     AT THIS POINT THERE IS NO DETERMINED VALUE TO ASSIGN TO X,
    #     ALL VALUES OF X1, X2, AND X3 ARE SAVED IN SX* arrays. AT A LATER
    #     TIME X3 WILL REACH A VALUE WHERE IT IS THE VALUE OF X (PDSI).
    #     AT THAT TIME, THE ASSIGN SUBROUTINE BACKTRACKS THROUGH SX* arrays
    #     CHOOSING THE APPROPRIATE X1 OR X2 TO BE THAT MONTH'S X.
    #-----------------------------------------------------------------------

    df.SX1[_k8] = df.PX1[_i]
    df.SX2[_k8] = df.PX2[_i]
    df.SX3[_k8] = df.PX3[_i]
    df.X[_i] = df.PX3[_i]
    _k8 = _k8 + 1

    #-----------------------------------------------------------------------
    #     SAVE THIS MONTHS CALCULATED VARIABLES (V,PRO,X1,X2,X3) FOR
    #     USE WITH NEXT MONTHS DATA
    #-----------------------------------------------------------------------
    V = PV
    prob_ended = df.PPR[_i]
    X1 = df.PX1[_i]
    X2 = df.PX2[_i]
    X3 = df.PX3[_i]

    return df, X1, X2, X3, V, prob_ended

#-----------------------------------------------------------------------------------------------------------------------
# from label 210 in pdinew.f
#@profile
@numba.jit
def _between_0s(df,
                X3):
    '''
    Compute index values when a weather spell has been established and no abatement is underway. In this case we
    calculate the X3 for use as the current month's severity index, if no backtracking is called for, and if
    backtracking is called for then for all the backtracked months we'll use the X3 values computed for those months.

    :param df: a pandas DataFrame containing the various work arrays used within
    :param K8: the number of backtracking steps currently called for, i.e. the number of previous months
               for which a conclusive severity index has not yet been determined and which require back fill
    :param X3: the previous month's severity index for any weather spell that has become established
    :param i: month index, assuming the timeseries work arrays have a single dimension with shape: (total months)
    :return: seven values: 1) the same pandas DataFrame used as the first argument, now with an updated state
                           2) the X1 for this month (always 0.0), to be used within the next month's index calculations
                           3) the X2 for this month (always 0.0), to be used within the next month's index calculations
                           4) the X3 for this month, to be used within the next month's index calculations
                           5) the accumulated "effective wetness" after this month's accounting (always 0.0)
                           6) the percentage probability that an established weather spell has ended (always 0.0)
                           7) the number of months requiring back fill via the backtracking process (always 0)
    '''

    # indicate that we'll use the globals for backtrack months count and current time step
    global _k8
    global _i

    #-----------------------------------------------------------------------
    #     PROB(END) RETURNS TO 0.  A POSSIBLE ABATEMENT HAS FIZZLED OUT,
    #     SO WE ACCEPT ALL STORED VALUES OF X3
    #-----------------------------------------------------------------------

    # reset the previous cumulative effective wetness value
    PV = 0.0

    # the probability that we've reached the end of a weather spell (PPR, or Pe in Palmer 1965) is zero
    df.PPR[_i] = 0.0

    # X3 is the index value in play here since our Pe is zero, so only calculate this month's X3 and use this as the X
    df.PX1[_i] = 0.0
    df.PX2[_i] = 0.0
    df.PX3[_i] = (0.897 * X3) + (df.Z[_i] / 3.0)
    df.X[_i] = df.PX3[_i]

    if _k8 == 0:  # no backtracking required

        df.PDSI[_i] = df.X[_i]
        df.PHDI[_i] = df.PX3[_i]

        if df.PX3[_i] == 0.0:

            df.PHDI[_i] = df.X[_i]

        df.WPLM[_i] = _case(df.PPR[_i], df.PX1[_i], df.PX2[_i], df.PX3[_i])

    else:  # perform backtracking, assigning the stored X3 values as our new backtrack array values

        iass = 3
        df = _assign(df, iass)

    #-----------------------------------------------------------------------------------------------
    #     SAVE THIS MONTHS CALCULATED VARIABLES (V, PRO, X1, X2, X3) FOR USE WITH NEXT MONTHS DATA
    #-----------------------------------------------------------------------------------------------

    # accumulated effective wetness reset to zero
    V = 0.0
    X1 = df.PX1[_i]  # zero
    X2 = df.PX2[_i]  # zero
    X3 = df.PX3[_i]  # the X3 calculated above
    prob_ended = df.PPR[_i]

    return df, X1, X2, X3, V, PV, prob_ended

#-----------------------------------------------------------------------------------------------------------------------
#@profile
#@numba.jit  #FIXME not working yet
def _assign(df,
            which_X):

    '''
    :df pandas DataFrame containing the timeseries arrays we'll use to perform the appropriate assignment of values
    :param which_X: flag to determine which of the df.SX* values to save into the main backtracking array, df.SX,
                    valid values are 1, 2, or 3 (this corresponds to iass/ISAVE in pdinew.f)
    :param K8: the number of months we need to back fill when performing backtracking assignment of values
                    (this is K8 in pdinew.f)
    :param i: month index, assuming the timeseries work arrays have a single dimension with shape: (total months)
    :return: pandas DataFrame, the same one as was passed as the first argument, but now with updated values
     '''

    # indicate that we'll use the globals for the backtracking months count and the current time step
    global _k8
    global _i

    #-----------------------------------------------------------------------
    #     FIRST FINISH OFF FILE 8 WITH LATEST VALUES OF PX3, Z,X
    #     X=PX1 FOR I=1, PX2 FOR I=2, PX3,  FOR I=3
    #-----------------------------------------------------------------------
    df.SX[_k8] = df.X[_i]

    if _k8 == 0:  # no backtracking required

        # set the PDSI to X
        df.PDSI[_i] = df.X[_i]

        # the PHDI is X3 if not zero, otherwise use X
        if df.PX3[_i] == 0.0:
            df.PHDI[_i] = df.X[_i]
        else:
            df.PHDI[_i] = df.PX3[_i]

        # select the best fit for WPLM
        df.WPLM[_i] = _case(df.PPR[_i], df.PX1[_i], df.PX2[_i], df.PX3[_i])

    else:  # perform backtracking

        #-------------------------------------------------------------------------------------------------------
        # Here we're filling an array (SX) with the final PDSI values for each
        # of the months that are being back filled as part of the backtracking process.
        # Each month has had three X values computed and stored in the SX? arrays,
        # and based on which of these is specified we'll fill the SX array with the
        # stored X values appropriate for the month. In the case of either X1 or X2
        # then we'll first check to make sure the X value has not gone to zero,
        # which indicates a state change and which flips the X that'll be saved
        # for the month and subsequent months until the value goes to zero (at which
        # point it would flip again).
        #-------------------------------------------------------------------------------------------------------

        if which_X == 3:
            #     ---------------- USE ALL X3 VALUES

            # fill the stored backtrack values array (SX) with the stored X3 values for all the backtracked months
            for Mm in range(_k8):

                df.SX[Mm] = df.SX3[Mm]  # label 10 in assign() within pdinew.f

        else:
            #     -------------- BACKTRACK THRU ARRAYS, STORING ASSIGNED X1 (OR X2)
            #                    IN SX UNTIL IT IS ZERO, THEN SWITCHING TO THE OTHER
            #                    UNTIL IT IS ZERO, ETC.

            # loop from the final month of the backtracking stack through to the first
            for Mm in range(_k8 - 1, -1, -1):  # corresponds with line 1100 in pdinew.f

                #-------------------------------------------------------------------------------------------------------
                # Based on which of the X values we're told to store we first check to make sure that it's non-zero,
                # and if so then it's stored in SX, otherwise we start using the other X until it too becomes zero,
                # at which time the X being used (which_X) switches again, and so on.
                #-------------------------------------------------------------------------------------------------------

                if which_X == 1: # then GO TO 20 in pdinew.f

                    # we should assign the stored X1 values into the main backtracking array until
                    # we find a stored X1 that is zero, which indicates a state change (wet to dry)

                    if df.SX1[Mm] == 0:  # then GO TO 50 in pdinew.f

                        # the X1, or severity index for a wet spell that's being established, is zero
                        # for this backtrack month, indicating that we need to switch to using the X2 values
                        which_X = 2
                        df.SX[Mm] = df.SX2[Mm]

                    else:
                        # assign the stored X1 into this backtrack month's position within the main backtracking array
                        df.SX[Mm] = df.SX1[Mm]

                elif which_X == 2: # then GO TO 40 in pdinew.f

                    # we should assign the stored X2 values into the main backtracking array until
                    # we find a stored X2 that is zero, which indicates a state change (dry to wet)

                    if df.SX2[Mm] == 0: # then GO TO 30 in pdinew.f

                        # the X2, or severity index for a dry spell that's being established, is zero
                        # for this backtrack month, indicating that we need to switch to using the X1 values
                        which_X = 1
                        df.SX[Mm] = df.SX1[Mm]

                    else:

                        # assign the stored X2 into this backtrack month's position within the main backtracking array
                        df.SX[Mm] = df.SX2[Mm]

        # label 70 from pdinew.f
        #-----------------------------------------------------------------------
        #     PROPER ASSIGNMENTS TO ARRAY SX HAVE BEEN MADE,
        #     OUTPUT THE MESS
        #-----------------------------------------------------------------------

        # assignment of stored X values as the final Palmer index values for the unassigned/backtracking months
        for n in range(_k8 + 1):  # this matches with pdinew.f line 1116 label 70 in assign() subroutine, DO 80 N = 1,K8

            # get the 1-D index equivalent to the j/m index for the final arrays
            # (PDSI, PHDI, etc.) that we'll assign to in the backtracking process
            ix = df.index_i[n]

            # assign the backtracking array's value for the current backtrack month as that month's final PDSI value
            df.PDSI[ix] = df.SX[n]

            # the PHDI is X3 if not zero, otherwise use X
            #TODO literature reference for this?
            df.PHDI[ix] = df.PX3[ix]
            if df.PX3[ix] == 0.0:
                df.PHDI[ix] = df.SX[n]

            # select the best fit for WPLM
            df.WPLM[ix] = _case(df.PPR[ix],
                                df.PX1[ix],
                                df.PX2[ix],
                                df.PX3[ix])

        # reset the backtracking month count / array index
        _k8 = 0

    return df

#-----------------------------------------------------------------------------------------------------------------------
#@profile
@numba.jit
def _case(PROB,
          X1,
          X2,
          X3):
    #
    #     THIS SUBROUTINE SELECTS THE PRELIMINARY (OR NEAR-REAL TIME)
    #     PALMER DROUGHT SEVERITY INDEX (PDSI) FROM THE GIVEN X VALUES
    #     DEFINED BELOW AND THE PROBABILITY (PROB) OF ENDING EITHER A
    #     DROUGHT OR WET SPELL.
    #
    #     X1   - INDEX FOR INCIPIENT WET SPELLS (ALWAYS POSITIVE)
    #     X2   - INDEX FOR INCIPIENT DRY SPELLS (ALWAYS NEGATIVE)
    #     X3   - SEVERITY INDEX FOR AN ESTABLISHED WET SPELL (POSITIVE) OR DROUGHT (NEGATIVE)
    #     PDSI - THE SELECTED PDSI (EITHER PRELIMINARY OR FINAL)
    #
    #   This subroutine written and provided by CPC (Tom Heddinghaus & Paul Sabol).
    #

    if X3 == 0.0: #) GO TO 10 in pdinew.f
        #     IF X3=0 THE INDEX IS NEAR NORMAL AND EITHER A DRY OR WET SPELL
        #     EXISTS.  CHOOSE THE LARGEST ABSOLUTE VALUE OF X1 OR X2.
        PDSI = X1
        if abs(X2) > abs(X1):
            PDSI = X2

    elif 0.0 < PROB < 100.0: # GO TO 20 in pdinew.f

        # put the probability value into 0..1 range
        probability = PROB / 100.0
        if X3 > 0.0: #) GO TO 30
            #     TAKE THE WEIGHTED SUM OF X3 AND X2
            PDSI = (1.0 - probability) * X3 + probability * X2
        else:
            #     TAKE THE WEIGHTED SUM OF X3 AND X1
            PDSI = (1.0 - probability) * X3 + probability * X1

    else:
        #     A WEATHER SPELL IS ESTABLISHED AND PDSI=X3 AND IS FINAL
        PDSI = X3

    return PDSI

#-----------------------------------------------------------------------------------------------------------------------
# from pdinew.f
def _pe(temperature,
        month_index,
        latitude,
        data_start_year,
        B,
        H):
    """
    Computes potential evapotranspiration based on the method used in original PDSI code pdi.f

    :param temperature: temperature value in degrees Fahrenheit
    :param month_index: index into monthly timeseries, with 0 corresponding to January of the initial year
    :param latitude: latitude value in degrees north
    :param data_start_year: initial year of the data being computed
    :param B: ?
    :param H: ?
    """
    #TODO document this, where do these come from?
    PHI = np.array([-0.3865982, -0.2316132, -0.0378180, 0.1715539, 0.3458803, 0.4308320, \
                     0.3916645, 0.2452467, 0.0535511, -0.15583436, -0.3340551, -0.4310691])

    TLA = -1 * math.tan(math.radians(latitude))
#     TLA = math.tan(math.radians(latitude))

    #-----------------------------------------------------------------------
    #     1 - CALCULATE PE (POTENTIAL EVAPOTRANSPIRATION)
    #-----------------------------------------------------------------------
    if temperature <= 32.0:
        PE = 0.0
    else:
        DUM = PHI[month_index % 12] * TLA

        try:
            DK = math.atan(math.sqrt(1.0 - (DUM * DUM)) / DUM)
        except ValueError:
            _logger.exception('Failed to complete', exc_info=True)
            raise

        if DK < 0.0:
            DK = 3.141593 + DK
        DK = (DK + 0.0157) / 1.57
        if temperature >= 80.0:
            PE = (math.sin((temperature / 57.3) - 0.166) - 0.76) * DK
        else:
            DUM = math.log(temperature - 32.0)
            PE = math.exp(-3.863233 + (B * 1.715598) - (B * math.log(H)) + (B * DUM)) * DK

        #-----------------------------------------------------------------------
        #     CONVERT DAILY TO MONTHLY
        #-----------------------------------------------------------------------
        year = data_start_year + int((month_index + 1) / 12)
        month = month_index % 12
        PE = PE * calendar.monthrange(year, month + 1)[1]

    return PE

#-----------------------------------------------------------------------------------------------------------------------
def potential_evapotranspiration(monthly_temps_celsius,
                                 latitude,
                                 data_start_year,
                                 B,
                                 H):

    # assumes monthly_temps_celsius, B, and H have same dimensions, etc.

    pet = np.full(monthly_temps_celsius.shape, np.NaN)
    monthly_temps_fahrenheit = scipy.constants.convert_temperature(monthly_temps_celsius, 'C', 'F')
    for i in range(monthly_temps_celsius.size):
        pet[i] = _pe(monthly_temps_fahrenheit[i], i, latitude, data_start_year, B, H)
    return pet


import pygame,sys,random

pygame.init()

size = width,height = 600,600
screen = pygame.display.set_mode(size)
font = pygame.font.SysFont("monospace",20)

#mitu ruutu on üles-alla
boxesX = int(width/25) - 2
boxesY = int(height/25) - 2

#mitu miini on väljakul
numOfMines = 35

#2-dimensiooniline array mängu staatuse hoidmiseks
mineField = [0] * (boxesX)
for b in range(boxesY):
    mineField[b] = [0] * (boxesX)

#peamine img,kust võtan subsurfacete kaudu paar pilti
mainImg = pygame.image.load("tile.png")
imgs =  []

for a in range(4):
    imgs.append(mainImg.subsurface(a * 25,0,25,25))

#paneme random asukohtadesse miiniväljal miinid.
def populate(mineField, numOfMines):
    minePos = []
    possibleMines = [k for k in range(11,19)]

    posX,posY = random.randint(0,len(mineField[0]) - 1),random.randint(0,len(mineField[1]) - 1)

    for a in range(numOfMines):
        #while loopiga vaatan, et positsioonid oleksid kõik erinevad
        while [posX,posY] in minePos:
            posX = random.randint(0,len(mineField[0]) - 1)
            posY = random.randint(0,len(mineField[1]) - 1)
        minePos.append([posX,posY])

    for pos in minePos:
        mineField[pos[0]][pos[1]] = 5

        #siin toimub numbrite leidmise maagia
        for a in range(-1,2):
            if pos[1] + a < len(mineField[0]) - 1 and pos[0] + 1 < len(mineField) - 1:
                mPos = mineField[pos[0] + 1][pos[1] + a]
                if mPos != 5 and mPos not in possibleMines:
                    mineField[pos[0] + 1][pos[1] + a] = 11
                elif mPos in possibleMines:
                    mineField[pos[0] + 1][pos[1] + a] +=  1

        for a in range(-1,2):
            if pos[1] + a < len(mineField) - 1 and pos[0] - 1 < len(mineField[0]) - 1:
                mPos = mineField[pos[0] - 1][pos[1] + a]

                if  mPos != 5 and mPos not in possibleMines:
                    mineField[pos[0] - 1][pos[1] + a] = 11
                elif mPos in possibleMines:
                    mineField[pos[0] - 1][pos[1] + a] +=  1

        for a in range(-1,2):
            for b in range(-1,2):
                pass

        if pos[1] - 1 > 0:
            if mineField[pos[0]][pos[1] - 1] != 5 and mineField[pos[0]][pos[1] - 1] not in possibleMines:
                mineField[pos[0]][pos[1] -1] = 11
            elif mineField[pos[0]][pos[1] - 1] in possibleMines:
                mineField[pos[0]][pos[1] - 1] += 1

        if pos[1] + 1 < len(mineField[0]) - 1:
            if mineField[pos[0]][pos[1] +1] != 5 and mineField[pos[0]][pos[1]+ 1] not in possibleMines:
                mineField[pos[0]][pos[1] +1] = 11
            elif mineField[pos[0]][pos[1] + 1] in possibleMines:
                mineField[pos[0]][pos[1] + 1] += 1

def reveal(screen,mineField, x,y):

    def flood_fill(screen, mineField, x, y):

        mPos = mineField[y - 1][x - 1]

        if mPos == 0:
            mineField[y - 1][x - 1] = 1
            screen.blit(imgs[1], (x*25,y * 25))
            if x < 22:
                if mineField[y - 1][x] == 0:
                    flood_fill(screen,mineField,x + 1,y)
            if y < 22:
                if mineField[y][x - 1] == 0:
                    flood_fill(screen,mineField,x,y + 1)
            if y >= 0:
                if mineField[y - 2][x - 1] == 0:
                    flood_fill(screen,mineField,x,y - 1)
            if x >= 0:
                if mineField[y - 1][x - 2] == 0:
                    flood_fill(screen,mineField,x - 1,y)

    flood_fill(screen, mineField, x, y)

#5 - miinid
#0 - avamata tile
#1 - avatud tile
#3 - lipp
#11:18 - miinide arv ruudu läheduses

populate(mineField,numOfMines)
print(mineField)

#paneme algse miinivälja ekraanile
for a in range(boxesY):
    for b in range(boxesX):
        #if mineField[a][b] == 0:
        screen.blit(imgs[0],(25*(a+1),25*(b+1)))

#subsurface järelejäänud flagide uuendamiseks korralikult ekraanil
#muidu peaks terve ekraani koguaeg üle joonistama
screenSub = screen.subsurface((0,0,25,25))

mineFieldRect = pygame.rect.Rect(25,25, 25*boxesX, 25*boxesY)

#main loop
while True:
    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            pygame.quit()
            sys.exit(0)
        elif e.type == pygame.MOUSEBUTTONDOWN:
            pos = pygame.mouse.get_pos()
            #kui click on minefieldi peal
            if mineFieldRect.collidepoint(pos):
                nearestX = pos[0] // 25
                nearestY = pos[1] // 25

                #y ja x telg on vahetunud praegu,mingi jama
                tileAtPos = mineField[nearestY - 1][nearestX - 1]

                if e.button == 1:
                    if tileAtPos == 5:
                        screen.blit(imgs[2], (25 * nearestX, 25 * nearestY))
                        print("Game over!")
                    elif tileAtPos > 10:
                        numOfNearbyMines = str(tileAtPos)[1:]
                        screen.blit(imgs[1], (25*nearestX,25*nearestY))
                        screen.blit(font.render(numOfNearbyMines,False, (255,0,0)), (25 * nearestX, 25 * nearestY))
                    elif tileAtPos != 3:
                        reveal(screen,mineField,nearestX,nearestY)
                        screen.blit(imgs[1],(25*nearestX,25*nearestY))
                        mineField[nearestY- 1][nearestX - 1] = 1

                elif e.button == 3:
                    if tileAtPos == 3:
                        #mineField[nearestY - 1][nearestX - 1] = 1
                        screen.blit(imgs[0],(25*nearestX,25*nearestY))
                        numOfMines += 1
                    elif tileAtPos != 1:
                        if numOfMines > 0:
                            screen.blit(imgs[3],(25*nearestX,25*nearestY))
                            mineField[nearestY - 1][nearestX - 1] = 3
                            numOfMines -= 1

    screenSub.fill((0,0,0))
    text = font.render(str(numOfMines),False, (255,255,255))
    screen.blit(text, (0,0))

    pygame.display.flip()


from tools.scraper import Scraper
from tools.state import State

def reply(activity, bot, data):
    print(data.get_entities())
    credentials = activity["message"]["text"].split(" ")
    rollno = credentials[0]
    wak = credentials[1]
    if len(wak) > 5:
        bot.send_text_activity(activity, "Please check your roll no. and web access key again.")
        bot.send_text_activity(activity, "Enter roll no. and web access key seperated by a single space.")
        return
    scraper = Scraper()
    scraper.authenticate(rollno,wak)
    if scraper.authenticated:
        state = State()
        studentdata = scraper.get_studentdata()
        studentdata["wak"] = wak
        studentdata["RowKey"] = activity["sender"]["id"]
        state.insertOrUpdateStudent(studentdata)
        bot.send_text_activity(activity, "Authentication Successful!")
    else:
        bot.send_text_activity(activity, "Please check your roll no. and web access key again.")
        bot.send_text_activity(activity, "Enter roll no. and web access key seperated by a single space.")


#======================================================================
#
#     This module contains routines to postprocess the VFI
#     solutions.
#
#     Simon Scheidegger, 11/16 ; 07/17
#======================================================================

import TasmanianSG
import numpy as np

from parameters import *

#======================================================================
# Routine to plot sparse grid solutions
def plot_routine(n_agents, grid, dim, num_points):
    f=open("forplot.txt", 'w')
    k_points=0.5*(k_bar + k_up)*np.ones((num_points, n_agents))

    k_dim=np.linspace(k_bar, k_up, num_points)

    k_points[:,[dim]]=k_dim.reshape(num_points,1)

    vals=grid.evaluateBatch(k_points)

    to_print=np.hstack((k_points, vals))
    np.savetxt(f, to_print, fmt= '% 2.5f')

    f.close()
    return

#======================================================================
# Routine compute the errors
def ls_error(n_agents, t1, t2, num_points):
    file=open('errors.txt', 'w')

    np.random.seed(0)
    unif=np.random.rand(num_points, n_agents)
    k_sample=k_bar+(unif)*(k_up-k_bar)
    to_print=np.empty((1,3))

    for i in range(t1, t2):
        sum_diffs=0
        diff = 0

        v_prev=TasmanianSG.TasmanianSparseGrid()
        v_next=TasmanianSG.TasmanianSparseGrid()

        v_prev.read("valnew_1." + str(i) + ".txt")
        v_next.read("valnew_1." + str(i+1) + ".txt")

        diff=v_next.evaluateBatch(k_sample) - v_prev.evaluateBatch(k_sample)
        max_abs_diff=np.amax(np.fabs(diff))
        average = np.average(np.fabs(diff))

        to_print[0,0]=i+1
        to_print[0,1]=max_abs_diff
        to_print[0,2]= average

        np.savetxt(file, to_print, fmt='%2.16f')

    file.close()

    return

#======================================================================


from .base import *

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = True

ANALYTICS_ID = 'UA-xxxxxx-1'

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = 'ml3&38def!wfp=eyu(=71gd(_tlmpujwpb958^7wq)5^=lw5*#'

# SECURITY WARNING: define the correct hosts in production!
ALLOWED_HOSTS = ['*']

# Base URL to use when referring to full URLs within the Wagtail admin backend -
# e.g. in notification emails. Don't include '/admin' or a trailing slash
BASE_URL = 'localhost:8000'

EMAIL_BACKEND = 'ictcg.email.NotifyEmailBackend'
EMAIL_NOTIFY_API_KEY = os.getenv('EMAIL_NOTIFY_API_KEY', None)
EMAIL_NOTIFY_BASIC_TEMPLATE = "f4ff560b-eec5-4b1c-b6b4-3dc1ed4a1664"

STATICFILES_STORAGE = 'django.contrib.staticfiles.storage.StaticFilesStorage'

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.postgresql_psycopg2',
        'NAME': 'gdmp-ictcg-dev',
    }
}

INSTALLED_APPS += [
    'debug_toolbar'
]

MIDDLEWARE += [
    'debug_toolbar.middleware.DebugToolbarMiddleware'
]

INTERNAL_IPS = [
    '127.0.0.1'
]

try:
    from .local import *
except ImportError:
    pass


# vim: fdm=indent
'''
author:     Fabio Zanini
date:       02/03/20
content:    Scan parameters for GBM analysis.
'''
import os
import sys
import requests
import pandas as pd
import numpy as np
import loompy

from matplotlib import gridspec
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns

# Make sure you can import leidenalg with fixed nodes
sys.path.insert(0, os.path.abspath('../..')+'/leidenalg/build/lib.linux-x86_64-3.8') # Fabio's laptop

# Make sure you can import northstar by adding its parent folder to your PYTHONPATH aka sys.path
sys.path.insert(0, os.path.abspath('../..')+'/northstar') # Fabio's laptop
import northstar

def ingest_gbm_data():
    # GBM data folder
    savedir_init = os.path.abspath('..')+'/data/'
    gbmdir = savedir_init+'GBM_data_and_metadata/'
    gbm_loomfn = gbmdir+'GBM_data.loom'

    # only download if this is the first time running this notebook:
    if not os.path.exists(gbmdir):
        print('Downloading GBM raw data...')

        #darmanis glioblastoma (GBM) dataset (to annotate based on brain atlas)
        url = 'http://storage.googleapis.com/gbmseqrawdata/rawData.zip'

        from io import BytesIO
        from zipfile import ZipFile
        from urllib.request import urlopen

        resp = urlopen(url)
        zipfile = ZipFile(BytesIO(resp.read()))
        files = zipfile.namelist()
        print('Avalable glioblastoma files:')
        for f in files:
            print(f)

        print('Extracting GBM data from ZIP archive...')
        for item in [1,2,3,4]:
            zipfile.extract(files[item],path=savedir_init)

        print('Reading GBM data from CSV files')
        GBM_count_path = savedir_init+files[4]
        GBM_counttable = pd.read_csv(GBM_count_path,sep=' ',index_col=0)
        GBM_meta = pd.read_csv(savedir_init+files[2],sep=' ',index_col=0)
        GBMtsne = pd.read_csv(savedir_init+files[1],sep=' ',index_col=0)
        # add annotation from separate file
        GBMmeta2= pd.read_excel(savedir_init+files[3],index_col=0,header=19,sep=' ')[:3589] # read from xls file

        print('Standardize annotations, add original t-SNE coordinates, and other minor manipulations...')
        GBM_meta = GBM_meta.join(GBMtsne)
        GBM_meta = GBM_meta.join(GBMmeta2['characteristics: cell type'])
        GBM_meta.rename(
                index=str,
                columns={'characteristics: cell type': 'Cell_type2'},
                inplace=True)

        # correct spelling and systematize cell types
        name_lut = {
            'Astocyte': 'Astrocyte',
            'microglia': 'Microglia',
            'Vascular': 'Endothelial',
            'Astrocytes': 'Astrocyte',
            'Oligodendrocyte': 'Oligodendrocyte',
            'Neurons': 'Neuron',
            'Neoplastic': 'Neoplastic',
            'Immune cell': 'Immune cell',
            'OPC': 'OPC',
            'Neuron': 'Neuron',
        }
        GBM_meta['Cell_type'] = GBM_meta['Cell_type2'].map(name_lut)

    elif os.path.isfile(gbm_loomfn):
        print('Reading GBM data from loom file')
        with loompy.connect(gbm_loomfn) as dsl:
            cells = dsl.ca['CellID']
            genes = dsl.ra['GeneName']
            GBM_meta = pd.DataFrame([], index=cells)
            for col in dsl.ca:
                GBM_meta[col] = dsl.ca[col]
            GBM_counttable = pd.DataFrame(
                data=dsl[:, :],
                index=genes,
                columns=cells,
                )

    else:
        print('Reading GBM data from CSV files')
        GBMtsne = pd.read_csv(gbmdir+'GBM_TSNE.csv',sep=' ',index_col=0)
        GBM_meta = pd.read_csv(gbmdir+'GBM_metadata.csv',sep=' ',index_col=0)
        GBMmeta2= pd.read_excel(gbmdir+'GEO_upload/spyros.darmanis_metadata_GBM.xls',index_col=0,header=19,sep=' ')[:3589]
        GBM_counttable = pd.read_csv(gbmdir+'GBM_raw_gene_counts.csv',sep=' ',index_col=0)

        print('Standardize annotations, add original t-SNE coordinates, and other minor manipulations...')
        GBM_meta = GBM_meta.join(GBMtsne)
        GBM_meta = GBM_meta.join(GBMmeta2['characteristics: cell type'])
        GBM_meta.rename(index=str,columns={'characteristics: cell type':'Cell_type2'},inplace=True)

        # correct spelling and systematize cell types
        name_lut = {
            'Astocyte': 'Astrocyte',
            'microglia': 'Microglia',
            'Vascular': 'Endothelial',
            'Astrocytes': 'Astrocyte',
            'Oligodendrocyte': 'Oligodendrocyte',
            'Neurons': 'Neuron',
            'Neoplastic': 'Neoplastic',
            'Immune cell': 'Immune cell',
            'OPC': 'OPC',
            'Neuron':'Neuron',
        }
        GBM_meta['Cell_type'] = GBM_meta['Cell_type2'].map(name_lut)

        matrix = GBM_counttable.values.astype(np.float32)
        col_attrs = {'CellID': GBM_counttable.columns.values}
        for col in GBM_meta:
            col_attrs[col] = GBM_meta[col].values
        row_attrs = {'GeneName': GBM_counttable.index.values}
        loompy.create(
            gbm_loomfn,
            layers={'': matrix},
            row_attrs=row_attrs,
            col_attrs=col_attrs,
            )

    return {
        'counts': GBM_counttable,
        'meta': GBM_meta,
        }

def define_accuracy(true_types, cell_types):
    ind_known = ~pd.Index(true_types).isin(['Immune cell', 'Neoplastic'])

    # TODO: implement better classification
    match = (true_types == cell_types)[ind_known].mean()

    return match

def test_conditions(all_names, conditions, atlas_sub, gbm, repeats=1):
    import time
    results = []
    ncomb = len(conditions)
    for ic, comb in enumerate(conditions):
        print('{:} / {:}: {:}'.format(ic + 1, ncomb, comb))
        for ir in range(repeats):
            kwargs = dict(zip(all_names, comb))
            t0 = time.time()
            no = northstar.Subsample(
                atlas=atlas_sub,
                **kwargs,
                )
            cell_types = no.fit_transform(gbm['counts'])
            t1 = time.time()
            acc = define_accuracy(gbm['meta']['Cell_type'].values, cell_types)
            kwargs['accuracy'] = acc
            kwargs['runtime'] = t1 - t0
            kwargs['repeat'] = ir + 1
            results.append(kwargs)

    return pd.DataFrame(results)

def get_combinations(params_dict):
    import itertools as it
    all_names = sorted(params_dict)
    combinations = list(it.product(*(params_dict[x] for x in all_names)))
    return (all_names, combinations)

def chunker(seq, n_chunks, constants=None):
    groups = [{'conditions': []} for i in range(n_chunks)]
    i = 0
    for x in seq:
        groups[i]['conditions'].append(x)
        i = (i + 1) % n_chunks
    for g in groups:
        g['constants'] = constants
    return groups

def worker_scan_function(group):
    conditions = group['conditions']
    all_names, atlas_sub, gbm, repeats = group['constants']
    return test_conditions(
            all_names, conditions, atlas_sub, gbm, repeats=repeats)

def test_conditions_parallel(
        all_names, conditions, atlas_sub, gbm, repeats=1, n_procs=None):
    import multiprocessing
    from multiprocessing import Pool

    if n_procs is None:
        n_cpus = multiprocessing.cpu_count()
        n_procs = min(n_cpus - 1, 40)

    groups = chunker(
            conditions, n_procs,
            constants=[all_names, atlas_sub, gbm, repeats])

    print('Starting pool of {:} CPUs'.format(n_procs))
    p = Pool(n_procs)
    results = p.map(worker_scan_function, groups)
    results = pd.concat(results, axis=0)

    return results

def test_combinations(params_dict, atlas_sub, gbm, repeats=1, n_procs=None):
    (all_names, combinations) = get_combinations(params_dict)

    if n_procs != 1:
        df = test_conditions_parallel(
            all_names, combinations, atlas_sub, gbm, repeats=repeats, n_procs=n_procs)
    else:
        df = test_conditions(
            all_names, combinations, atlas_sub, gbm, repeats=repeats)
    return df

if __name__ == '__main__':

    af = northstar.fetch_atlas.AtlasFetcher()
    atlas_sub = af.fetch_atlas('Darmanis_2015_nofetal', kind='subsample')
    print('Rename a few cell types in the atlas')
    atlas_sub['cell_types'] = atlas_sub['cell_types'].map({
        'Oligodendrocyte': 'Oligodendrocyte',
        'Vascular': 'Endothelial',
        'Astrocyte': 'Astrocyte',
        'Neuron': 'Neuron',
        'OPC': 'OPC',
        'microglia': 'Immune cell'},
        )

    gbm = ingest_gbm_data()

    print('Classify and cluster cells with default parameters')
    no = northstar.Subsample(
        atlas=atlas_sub,
        )
    cell_types = no.fit_transform(gbm['counts'])
    acc = define_accuracy(gbm['meta']['Cell_type'].values, cell_types)
    print(acc)

    if False:
        print('Scan broad parameter space')
        params_dict = dict(
            n_features_per_cell_type=[10, 30, 50],
            n_features_overdispersed=[50, 300, 1000],
            n_pcs=[10, 20, 40],
            n_neighbors=[10, 20, 50],
            distance_metric=['correlation'],
            threshold_neighborhood=[0.8],
            resolution_parameter=[0.0001, 0.001, 0.01],
            )

        print('Perform parameter scan with parallelism')
        df = test_combinations(params_dict, atlas_sub, gbm, repeat=1, n_procs=None)

        #print('Classify and cluster cells with parameter scan')
        #df_np = test_combinations(params_dict, atlas_sub, gbm, repeats=1, n_procs=1)

        savedir_init = os.path.abspath('..')+'/data/'
        gbmdir = savedir_init+'GBM_data_and_metadata/'
        df.to_csv(gbmdir+'results_grid.tsv', sep='\t', index=True)

        print('Plot results')
        fig, ax = plt.subplots(figsize=(3, 3))
        ax.hist(df['accuracy'], bins=np.linspace(0, 1, 21))
        ax.set_xlabel('Accuracy')
        ax.set_xticks([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 1.0])
        ax.grid(False)
        ax.set_title('All parameter combinations')
        fig.tight_layout()

        print('Identify enrichments')
        threshold = 0.8
        for col in params_dict.keys():
            vc1 = df.loc[df['accuracy'] < threshold, col].value_counts()
            vc2 = df.loc[df['accuracy'] >= threshold, col].value_counts()

            ratio = 1.0 * vc2 / (vc1 + vc2)
            print(col)
            print(ratio)

        from scipy.stats import spearmanr, pearsonr
        print('Spearman rho btw runtime and accuracy: {:}'.format(spearmanr(df['accuracy'].values, df['runtime'].values)))
        print('Pearson r btw runtime and accuracy: {:}'.format(pearsonr(df['accuracy'].values, df['runtime'].values)))

        print('Print the top parameters')
        print(df.nlargest(10, 'accuracy'))

    if False:
        print('Scan parameter subspace of good rough sketches')
        params_dict2 = dict(
            n_features_per_cell_type=[5, 10, 20, 30, 50, 80],
            n_features_overdispersed=[300],
            n_pcs=[10, 20, 30, 40],
            n_neighbors=[5, 10, 20, 35, 50],
            distance_metric=['correlation'],
            threshold_neighborhood=[0.8],
            resolution_parameter=[0.0001],
            )
        df2 = test_combinations(params_dict2, atlas_sub, gbm)

        print('Print the top parameters')
        print(df2.nlargest(10, 'accuracy'))

        from scipy.stats import spearmanr, pearsonr
        print('Spearman rho btw runtime and accuracy: {:}'.format(spearmanr(df2['accuracy'].values, df2['runtime'].values)))
        print('Pearson r btw runtime and accuracy: {:}'.format(pearsonr(df2['accuracy'].values, df2['runtime'].values)))

        print('Plot results for subscan')
        fig, ax = plt.subplots(figsize=(3, 3))
        ax.hist(df2['accuracy'], bins=np.linspace(0, 1, 51))
        ax.set_xlabel('Accuracy')
        ax.set_xticks([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 1.0])
        ax.grid(False)
        for tk in ax.get_xticklabels():
            tk.set_rotation(90)
        ax.set_title('Subscan')
        fig.tight_layout()

    print('Vice versa, modify only the two key parameters and see')
    params_dict3 = dict(
        n_features_per_cell_type=[15],
        n_pcs=[15],
        n_neighbors=[20],
        distance_metric=['correlation'],
        threshold_neighborhood=[0.8],
        resolution_parameter=[1e-5, 2e-5, 3e-5, 6e-5, 8e-5, 0.0001, 0.0002, 0.0003, 0.0006, 0.001, 0.002, 0.003, 0.005, 0.008, 0.01, 0.02, 0.05],
        n_features_overdispersed=[50, 75, 100, 150, 200, 250, 300, 400, 600, 800, 1000],
        )
    df3 = test_combinations(params_dict3, atlas_sub, gbm, repeats=2)
    accs = df3[['n_features_overdispersed', 'resolution_parameter', 'accuracy']].groupby(['n_features_overdispersed', 'resolution_parameter']).mean()['accuracy'].unstack('resolution_parameter')

    print('Plot heatmap for those two params only')
    fig, ax = plt.subplots(figsize=(5, 5))
    sns.heatmap(
        np.maximum(accs, 0.9),
        ax=ax,
        vmin=0.9,
        vmax=1,
        )
    for tk in ax.get_yticklabels():
        tk.set_rotation(0)
    fig.tight_layout()
    fig.savefig('../figures/SupplFig4.svg')
    fig.savefig('../figures/SupplFig4.png')

    plt.ion()
    plt.show()


#list=[0,1,2,3,4,5,7,8,9]
class obj:
    def solve(self,nums):
        arr=[0]*(len(nums)+1)
        for i in nums:
            arr[i]=arr[i]+1
        missing = []
        for i in range (len(arr)):
            if arr[i] == 0 and i!=0:
                missing.append(i)
        return missing
a = obj()
#b= obj()
#print(a.solve([0,1,3,5,7,9]))
print(a.solve([0,1,2,3,4,5,6,8,9]))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Feb 17 23:11:31 2020

@author: iqrah
"""
from __future__ import division

import sys
import lex
#import ply.lex as lex

tokens = [

    'T_Identifier',
	'T_IntConstant',
	'T_DoubleConstant',
	'T_BoolConstant',
	'T_StringConstant',
	'Plus',
	'Minus',
	'Multiply',
	'Divide',
	'Modulus',
	'T_GreaterEqual',
	'Greater',
	'T_LessEqual',
	'Less',
	'T_Equal',
	'T_NotEqual',
	'Eq',
	'T_And',
	'T_Or',
	'Not',
	'Semicolon',
	'Comma',
	'Dot',
	'LeftParen',
	'RightParan',
    'LeftCurly',
	'RightCurly',
	'Comment'
]
reserved = {
	'void' : 'T_Void',
	'int' : 'T_Int',
	'double' : 'T_Double',
	'bool' : 'T_Bool',
	'string' : 'T_String',
	'null' : 'T_Null',
	'for' : 'T_For',
	'while' : 'T_While',
	'if' : 'T_If',
	'else' : 'T_Else',
	'return' : 'T_Return',
	'break' : 'T_Break',
	'Print' : 'T_Print',
	'ReadInteger' : 'T_ReadInteger',
	'ReadLine' : 'T_ReadLine'
}
tokens += list(reserved.values())

operators = {
	'Plus' : '+',
	'Minus' : '-',
	'Multiply' : '*',
	'Divide' : '/',
	'Modulus' : '%',
	'Greater' : '>',
	'Less' : '<',
	'Eq' : '=',
	'Not' : '!',
	'Semicolon' : ';',
	'Comma' : ',',
	'Dot' : '.',
	'LeftParen' : '(',
	'RightParan' : ')',
    'LeftCurly' : '{',
	'RightCurly' : '}'
}

t_ignore = ' \t'
t_Plus = r'\+'
t_Minus = r'\-'
t_Multiply = r'\*'
t_Divide = r'\/'
t_Modulus = r'\%'
t_Greater = r'>'
t_Less = r'<'
t_Eq = r'='
t_Not = r'\!'
t_Semicolon = r';'
t_Comma = r','
t_Dot = r'\.'
t_LeftParen = r'\('
t_RightParan = r'\)'
t_LeftCurly = r'{'
t_RightCurly = r'}'

def t_GreaterEqual(t):
	r'>='
	t.type = 'T_GreaterEqual'
	return t

def t_LessEqual(t):
	r'<='
	t.type = 'T_LessEqual'
	return t

def t_Equal(t):
	r'=='
	t.type = 'T_Equal'
	return t

def t_NotEqual(t):
	r'\!='
	t.type = 'T_NotEqual'
	return t

def t_And(t):
	r'\&\&'
	t.type = 'T_And'
	return t

def t_Or(t):
	r'\|\|'
	t.type = 'T_Or'
	return t

def t_BoolConstant(t):
	r'(true|false)'
	t.type = 'T_BoolConstant'
	return t

def t_Identifier(t):
	r'[a-zA-Z_][a-zA-Z_0-9]*'
	if reserved.get(t.value) != None:
		t.type = reserved.get(t.value)
		return t
	t.type = 'T_Identifier'
	if len(t.value) > 31:
		t_error(t)
		return
	return t

def t_StringConstant(t):
	r'"[^\n|"]*(")?'
	t.type = 'T_StringConstant'
	if ((t.value[-1] != '\"') or (t.value[0] == '\"' and len(t.value) == 1)):
		t_error(t)
		t.lexer.lineno += 1
		return
	return t

def t_DoubleConstant(t): #decimal points done 2.10 and 1200.0 vs 1200
	r'\d+\.\d*([eE]?)[+-]?\d+'
	#t.value = float(t.value)
	t.type = 'T_DoubleConstant'
	return t

def t_IntConstant(t): # solved hex vs int
	#r'\d+'
	r'(0[xX][\da-fA-F]+)|\d+'
	t.type = 'T_IntConstant'
	return t

def t_Comment(t):
	r'(/\*([^*]|[\r\n]|(\*+([^*/]|[\r\n])))*\*+/)|(//.*)'
	lines = t.value.count('\n')
	t.lexer.lineno += lines
	pass

def t_newline(t):
     r'\n+'
     t.lexer.lineno += len(t.value)

# Compute column.
#     input is the input text string
#     token is a token instance
def find_column(input, token):
	line_start = input.rfind('\n', 0, token.lexpos) + 1
	return (token.lexpos - line_start) + 1

def t_error(t):
	#print("\n*** Error line", t.lineno, ".")
	write_to_file("\n*** Error line " + str(t.lineno) + ".")
	if (t.type == 'T_Identifier'):
		#print("*** Identifier too long: \"", t.value, "\"\n")
		write_to_file("*** Identifier too long: \"" + t.value + "\"\n")
		corrected_value = "(truncated to " + t.value[:31] + ")"
		#print(t.value, "line ", t.lineno, "cols ", find_column(input_str, t), "-", find_column(input_str, t)+len(str(t.value))-1, " is ", t.type, corrected_value)
		write_to_file(t.value + " line " + str(t.lineno) + " cols " + str(find_column(input_str, t)) + "-" + str(find_column(input_str, t)+len(str(t.value))-1) + " is " + t.type + " " + corrected_value)
	elif (t.type == 'T_StringConstant'):
		#print("*** Unterminated string constant: %s" % t.value, "\n")
		write_to_file("*** Unterminated string constant: " + t.value + "\n")
	else:
		#print("*** Unrecognized char: '%s'" % t.value[0], "\n")
		write_to_file("*** Unrecognized char: '%s'"  % str(t.value[0]) + "\n")
	t.lexer.skip(1)

lexer = lex.lex()

input_str = ''
'''
input_str = "120 45 012 0xa23 001 10 0 12.00 12.E+2 12.e+0"
lexer.input(input_str)
'''
#input_file = open("/home/iqrah/Desktop/Spring_02_2020/Compilers/Compiler_Project_Iqrah/samples/number.frag", "r")
input_file = open(str(sys.argv[1]), "r")
if input_file.mode == 'r':
	input_str = input_file.read()

lines = input_str.splitlines()
#print(lines)

lexer.input(input_str)

def get_next_token():
	return lexer.token()

#output_file = open(str(sys.argv[2]), "a+")
#output_file.truncate(0)
def write_to_file(output_str):
	#terminal = sys.stdout
	#sys.stdout = output_file
	print(output_str)
	#sys.stdout = terminal
"""
while True:
	t = lexer.token()
	if not t:
		break

	op_type = t.type
	if( operators.get(t.type) != None ):
		op_type = "'" + operators.get(t.type) + "'"

	value = ""
	if ( t.type == "T_BoolConstant" or t.type == "T_StringConstant" ):
		value = "(value = " + str(t.value) + ")"

	elif ( t.type == "T_IntConstant" ):
		if(t.value[0:2] == '0X' or t.value[0:2] == '0x'):
			int_val = str(int(t.value, 16))
		else:
			int_val = str(int(t.value))
		value = "(value = " + int_val + ")"

	elif ( t.type == "T_DoubleConstant" ):
		double_val = str(float(t.value))
		if(double_val[-2:] == '.0'):
			double_val = double_val[:-2]
		value = "(value = " + double_val + ")"
	#print(t)
	#print(t.value, "\t\tline ", t.lineno, "cols ", t.lexpos+1, "-", t.lexpos+len(str(t.value)), " ", t.type)
	#print(t.value, "\t\tline ", t.lineno, "cols ", find_column(input_str, t), "-",  find_column(input_str, t)+len(str(t.value))-1, " is ", op_type, value)
	write_to_file(t.value.ljust(12) + " line " + str(t.lineno) + " cols " + str(find_column(input_str, t)) + "-" + str(find_column(input_str, t)+len(str(t.value))-1) + " is " + op_type + " " + value)

	"""


import scipy
import numpy as np

def stft(x, fs, framesz, hop, noise=0):
    # framesamp = int(framesz * fs)
    # hopsamp = int(hop * fs)
    framesamp = framesz
    hopsamp = hop
    w = scipy.hanning(framesamp)
    X = scipy.array([scipy.fft(w * x[i:i + framesamp])[:(framesamp + 1) // 2]
                     / framesamp + np.ones((framesamp + 1) // 2) * noise
                     for i in range(0, len(x) - framesamp, hopsamp)])
    return X


import paho.mqtt.client as mqtt  # import the client1
import json
import pymysql.cursors
from credentials import server_ip

# Connect to the database
connection = pymysql.connect(host='localhost',
                             user='root',
                             password='',
                             db='bugbox_db',
                             charset='utf8mb4',
                             cursorclass=pymysql.cursors.DictCursor)

broker_address = server_ip
queue_name = "queue_agriculture"
count = 0

def on_message(client, userdata, message):
    global count
    msg = json.loads(str(message.payload.decode("utf-8")))
    # print(str(msg))
    with connection.cursor() as cursor:
        sql = "INSERT INTO `agriculture_data` (`time`, `greenhouse_id`, `humidity`, `temperature`, `moisture`) " \
              "VALUES ('"+str(msg['time'])+"',"+str(msg['greenhouse_id'])+","+ str(msg['humidity'])+","+str(msg['temperature'])+","+ str(msg['moisture'])+")"
        # print(sql)
        cursor.execute(sql)
        count = count + 1

    if count % 500 == 0:
        connection.commit()
        print("Received 500 more requests")
        count = 0

def on_connect(mqtt_client, obj, flags, rc):
    print("Subscribing to topic ", queue_name)
    mqtt_client.subscribe(queue_name)

print("creating new instance")
client = mqtt.Client("P2")  # create new instance
client.on_message = on_message  # attach function to callback
client.on_connect = on_connect
print("connecting to broker")
client.connect(broker_address, port=1883)  # connect to broker

client.loop_forever()


import glob
import os
import shutil
import sys

import cffi

from . import coreclr, hostfxr, mono, netfx

__all__ = ["ffi", "load_coreclr", "load_hostfxr", "load_mono", "load_netfx"]

ffi = cffi.FFI()

for cdef in coreclr.cdef + hostfxr.cdef + mono.cdef + netfx.cdef:
    ffi.cdef(cdef)

def load_coreclr(runtime):
    dll_name = _get_dll_name("coreclr")
    dll_path = os.path.join(runtime.path, dll_name)
    return ffi.dlopen(dll_path)

def load_hostfxr(dotnet_root):
    hostfxr_name = _get_dll_name("hostfxr")
    hostfxr_path = os.path.join(dotnet_root, "host", "fxr", "3.*", hostfxr_name)

    for hostfxr_path in reversed(sorted(glob.glob(hostfxr_path))):
        try:
            return ffi.dlopen(hostfxr_path)
        except Exception:
            pass

    raise RuntimeError(f"Could not find a suitable hostfxr library in {dotnet_root}")

def load_mono(path=None, gc=None):
    if path is None:
        from ctypes.util import find_library

        path = find_library(f"mono{gc or ''}-2.0")
        if path is None:
            raise RuntimeError("Could not find libmono")

    return ffi.dlopen(path, ffi.RTLD_GLOBAL)

def load_netfx():
    if sys.platform != "win32":
        raise RuntimeError(".NET Framework is only supported on Windows")

    dirname = os.path.join(os.path.dirname(__file__), "dlls")
    if sys.maxsize > 2 ** 32:
        arch = "amd64"
    else:
        arch = "x86"

    path = os.path.join(dirname, arch, "ClrLoader.dll")

    return ffi.dlopen(path)

def _get_dll_name(name):
    if sys.platform == "win32":
        return f"{name}.dll"
    elif sys.platform == "darwin":
        return f"lib{name}.dylib"
    else:
        return f"lib{name}.so"


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

'''
Created on 2019-3-21

@author: suning
'''
import sys
import os
basepath = os.path.dirname(os.path.abspath(sys.argv[0]))+"/../../"
sys.path.append(basepath)

import suning.api as api

a=api.OrderdiscountsideCreateRequest()

a.setDomainInfo("openpre.cnsuning.com","80")
a.setAppInfo("a13b8bd0efb06a770c57d1c370ce8ee7", "f08ce9836c4bcfc708194594081f6690")
# 如果使用oauth认证方式，那么调用下面方法来添加accessToken
#a.setAccessToken("4caf7fa30dd8")
a.activityName='中文zimu@#￥'
a.allShopType='0'
a.appStoreCode='S03017113'
a.baseQuantifierType='1'
a.channelInfo='55,56'
a.endTime='2019-03-04 20:25:59'
a.productList=[{"productCode":"761134378"}]
a.rewardQuantifierType='1'
a.ruleList=[{"baseAmount":"9999","rewardAmount":"999"}]
a.startTime='2019-03-04 20:23:00'
a.storeCode='S03017113'

try:
    f = a.getResponse()
    print(f)
except Exception as e:
    print(e)


from odoo import models, api, fields
from odoo.exceptions import ValidationError
from datetime import date

class InventoryValuation(models.TransientModel):
    _name = 'wiz.inventory.valuation'
    date = fields.Date(string='Date', default=date.today())

    def print_inventory_valuation(self):
        items = self.env['inventory.item'].search([])
        if len(items) < 1:
            raise ValidationError("There are no items to report on")
        if fields.Date.from_string(self.date) > date.today():
            raise ValidationError("You cannot select a date in the future")

        report = {
            'type': 'ir.actions.report.xml',
            'report_type': 'xlsx',
            'report_name': 'inventory.valuation.xlsx',
            'context': dict(self._context, xlsx_export=True),
            'datas': {'ids': [self.id]},
        }
        return report


#  optimizing warehouse flows with q-learning

# importing the necessary libraries
import numpy as np

# Setting the hyperparameters
discount_rate = 0.75
learning_rate = 0.9

# Step 1 : Defining the enviornment

# a ) Defining the states
location_to_states = { 'A':0, 'B':1, 'C':2, 'D':3,
                    'E':4, 'F':5, 'G':6, 'H':7,
                    'I':8, 'J':9, 'K':10, 'L':11 }

# b ) Defining the actions
actions = list(range(12))

# c ) Defining the rewards
#              A,B,C,D,E,F,G,H,I,J,K,L
R = np.array([[0,1,0,0,0,0,0,0,0,0,0,0],
              [1,0,1,0,0,1,0,0,0,0,0,0],
              [0,1,0,0,0,0,1,0,0,0,0,0],
              [0,0,0,0,0,0,0,1,0,0,0,0],
              [0,0,0,0,0,0,0,0,1,0,0,0],
              [0,1,0,0,0,0,0,0,0,1,0,0],
              [0,0,1,0,0,0,1,1,0,0,0,0],
              [0,0,0,1,0,0,1,0,0,0,0,1],
              [0,0,0,0,1,0,0,0,0,1,0,0],
              [0,0,0,0,0,1,0,0,1,0,1,0],
              [0,0,0,0,0,0,0,0,0,1,0,1],
              [0,0,0,0,0,0,0,1,0,0,1,0],
              ])

# Step 2 : Building the solution with Q-Learning

# Making the mapping from state_to_location
state_to_location = {state:location for location, state in location_to_states.items()}

# Making a final fucntion which will return route
# Note : This starting_location and ending_location will actually be letters representing the location
def route(starting_location, ending_location):
    route = [starting_location]

    # Now it will dynamically update the rewards
    R_new = np.copy(R)
    ending_state = location_to_states[ending_location]
    R_new[ending_state,ending_state] = 1000

    # a) Initialization the Q-Values
    Q = np.array(np.zeros([12,12])) # since there are 12 possible states and actions

    # b) Implementing the Q-learning
    for i in range(1000):
        current_state = np.random.randint(0,12)
        possible_actions = []
        for j in range(12):
            if R_new[current_state, j] > 0:
                possible_actions.append(j)

        next_state = np.random.choice(possible_actions)
        temporal_diff = R_new[current_state, next_state] + discount_rate * Q[next_state , np.argmax(Q[next_state,])] - Q[current_state, next_state]
        Q[current_state,next_state] += learning_rate * temporal_diff

    next_location = starting_location # Initially
    while next_location != ending_location:
        starting_state = location_to_states[starting_location]
        next_state = np.argmax(Q[starting_state,])
        next_location = state_to_location[next_state]
        route.append(next_location)
        starting_location = next_location

    return route

# There exists two paths from E-G, one through F and other through K
# If we want the path with K as intermediary_location then we can do following things:
# 1 ) Increase the reward from J to K (reward > 1 and reward < 1000)
# 2 ) Add a penalty from J to F (reward < 1), so it will never take that route
# 3 ) Use the route function two times, first from starting_location to intermediary_location
#     and then from intermediary location to ending_location.

# Step 3 : Going into production

def best_route(starting_location, intermediary_location, ending_location):
    return route(starting_location, intermediary_location) + route(intermediary_location, ending_location)[1:]
    # First route will return from E to K and then second return from K to G.
    # Since the first route returns K so we don't want K from second route, therefore [1:]

# Printing the results
print('Route : {}'.format(best_route('A','K','H')))