text:
from orbit.models.ktrlite import KTRLite
import pandas as pd
import numpy as np
import math
from scipy.stats import nct
from enum import Enum
import torch
import matplotlib.pyplot as plt
from copy import deepcopy

from ..constants.constants import (
    KTRTimePointPriorKeys,
    PredictMethod,
    TrainingMetaKeys,
    PredictionMetaKeys
)
from ..exceptions import IllegalArgument, ModelException, PredictionException
from ..utils.general import is_ordered_datetime
from ..utils.kernels import gauss_kernel, sandwich_kernel
from ..utils.features import make_seasonal_regressors
from .model_template import ModelTemplate
from ..estimators.pyro_estimator import PyroEstimatorSVI
from ..models import KTRLite
from orbit.constants.palette import OrbitPalette
from ..utils.knots import get_knot_idx, get_knot_dates
from ..utils.plot import orbit_style_decorator

class DataInputMapper(Enum):
    """
    mapping from object input to pyro input
    """
    # All of the following have default defined in DEFAULT_SLGT_FIT_ATTRIBUTES
    # ----------  Data Input ---------- #
    # observation related
    NUM_OF_VALID_RESPONSE = 'N_VALID_RES'
    WHICH_VALID_RESPONSE = 'WHICH_VALID_RES'
    RESPONSE_OFFSET = 'MEAN_Y'
    DEGREE_OF_FREEDOM = 'DOF'
    _RESIDUALS_SCALE_UPPER = 'RESID_SCALE_UB'
    # ----------  Level  ---------- #
    _NUM_KNOTS_LEVEL = 'N_KNOTS_LEV'
    LEVEL_KNOT_SCALE = 'LEV_KNOT_SCALE'
    _KERNEL_LEVEL = 'K_LEV'
    # ----------  Regression  ---------- #
    _NUM_KNOTS_COEFFICIENTS = 'N_KNOTS_COEF'
    _KERNEL_COEFFICIENTS = 'K_COEF'
    _NUM_OF_REGULAR_REGRESSORS = 'N_RR'
    _NUM_OF_POSITIVE_REGRESSORS = 'N_PR'
    _NUM_OF_NEGATIVE_REGRESSORS = 'N_NR'
    _REGULAR_REGRESSOR_MATRIX = 'RR'
    _POSITIVE_REGRESSOR_MATRIX = 'PR'
    _NEGATIVE_REGRESSOR_MATRIX = 'NR'
    _REGULAR_REGRESSOR_INIT_KNOT_LOC = 'RR_INIT_KNOT_LOC'
    _REGULAR_REGRESSOR_INIT_KNOT_SCALE = 'RR_INIT_KNOT_SCALE'
    _REGULAR_REGRESSOR_KNOT_SCALE = 'RR_KNOT_SCALE'
    _POSITIVE_REGRESSOR_INIT_KNOT_LOC = 'PR_INIT_KNOT_LOC'
    _POSITIVE_REGRESSOR_INIT_KNOT_SCALE = 'PR_INIT_KNOT_SCALE'
    _POSITIVE_REGRESSOR_KNOT_SCALE = 'PR_KNOT_SCALE'
    _NEGATIVE_REGRESSOR_INIT_KNOT_LOC = 'NR_INIT_KNOT_LOC'
    _NEGATIVE_REGRESSOR_INIT_KNOT_SCALE = 'NR_INIT_KNOT_SCALE'
    _NEGATIVE_REGRESSOR_KNOT_SCALE = 'NR_KNOT_SCALE'
    # ----------  Prior Specification  ---------- #
    _COEF_PRIOR_LIST = 'COEF_PRIOR_LIST'
    _LEVEL_KNOTS = 'LEV_KNOT_LOC'
    _SEAS_TERM = 'SEAS_TERM'

class BaseSamplingParameters(Enum):
    """
    The output sampling parameters related with the base model
    """
    LEVEL_KNOT = 'lev_knot'
    LEVEL = 'lev'
    YHAT = 'yhat'
    OBS_SCALE = 'obs_scale'

class RegressionSamplingParameters(Enum):
    """
    The output sampling parameters related with regression component.
    """
    COEFFICIENTS_KNOT = 'coef_knot'
    COEFFICIENTS_INIT_KNOT = 'coef_init_knot'
    COEFFICIENTS = 'coef'

# Defaults Values
DEFAULT_REGRESSOR_SIGN = '='
DEFAULT_COEFFICIENTS_INIT_KNOT_SCALE = 1.0
DEFAULT_COEFFICIENTS_INIT_KNOT_LOC = 0
DEFAULT_COEFFICIENTS_KNOT_SCALE = 0.1
DEFAULT_LOWER_BOUND_SCALE_MULTIPLIER = 0.01
DEFAULT_UPPER_BOUND_SCALE_MULTIPLIER = 1.0

class KTRModel(ModelTemplate):
    """Base KTR model object with shared functionality for PyroVI method
    Parameters
    ----------
    level_knot_scale : float
        sigma for level; default to be .1
    level_segments : int
        the number of segments partitioned by the knots of level (trend)
    level_knot_distance : int
        the distance between every two knots of level (trend)
    level_knot_dates : array like
        list of pre-specified dates for the level knots
    seasonality : int, or list of int
        multiple seasonality
    seasonality_fs_order : int, or list of int
        fourier series order for seasonality
    seasonality_segments : int
        the number of segments partitioned by the knots of seasonality
    seasonal_initial_knot_scale : float
        scale parameter for seasonal regressors initial coefficient knots; default to be 1
    seasonal_knot_scale : float
        scale parameter for seasonal regressors drift of coefficient knots; default to be 0.1.
    regressor_col : array-like strings
        regressor columns
    regressor_sign : list
        list of signs with '=' for regular regressor, '+' for positive regressor, and '-' for negative regressor.
    regressor_init_knot_loc : list
        list of regressor knot pooling mean priors, default to be 0's
    regressor_init_knot_scale : list
        list of regressor knot pooling sigma's to control the pooling strength towards the grand mean of regressors;
        default to be 1.
    regressor_knot_scale : list
        list of regressor knot sigma priors; default to be 0.1.
    regression_segments : int
        the number of segments partitioned by the knots of regression
    regression_knot_distance : int
        the distance between every two knots of regression
    regression_knot_dates : array-like
        list of pre-specified dates for regression knots
    regression_rho : float
        sigma in the Gaussian kernel for the regression term
    degree of freedom : int
        degree of freedom for error t-distribution
    date_freq : str
        date frequency; if not supplied, the minimum timestamp difference in the date would be used.
    coef_prior_list : list of dicts
        each dict in the list should have keys as
        'name', prior_start_tp_idx' (inclusive), KTRTimePointPriorKeys.PRIOR_END_TP_IDX.value (not inclusive),
        KTRTimePointPriorKeys.PRIOR_MEAN.value, KTRTimePointPriorKeys.PRIOR_SD.value, and KTRTimePointPriorKeys.PRIOR_REGRESSOR_COL.value
    residuals_scale_upper : float
    flat_multiplier : bool
        Default set as True. If False, we will adjust knot scale with a multiplier based on regressor volume
        around each knot; When True, set all multiplier as 1
    ktrlite_optim_args : dict
        the optimizing config for the ktrlite model (to fit level/seasonality). Default to be dict().
    """
    _data_input_mapper = DataInputMapper
    # stan or pyro model name (e.g. name of `*.stan` file in package)
    _model_name = 'ktr'
    _supported_estimator_types = [PyroEstimatorSVI]

    def __init__(self,
                 # level
                 level_knot_scale=0.1,
                 level_segments=10,
                 level_knot_distance=None,
                 level_knot_dates=None,
                 # seasonality
                 seasonality=None,
                 seasonality_fs_order=None,
                 seasonality_segments=2,
                 seasonal_initial_knot_scale=1.0,
                 seasonal_knot_scale=0.1,
                 # regression
                 regressor_col=None,
                 regressor_sign=None,
                 regressor_init_knot_loc=None,
                 regressor_init_knot_scale=None,
                 regressor_knot_scale=None,
                 regression_segments=5,
                 regression_knot_distance=None,
                 regression_knot_dates=None,
                 regression_rho=0.15,
                 # shared
                 degree_of_freedom=30,
                 date_freq=None,
                 # time-based coefficient priors
                 coef_prior_list=None,
                 flat_multiplier=True,
                 residuals_scale_upper=None,
                 ktrlite_optim_args=dict(),
                 **kwargs):
        super().__init__(**kwargs)  # create estimator in base class

        # level configurations
        self.level_knot_scale = level_knot_scale
        self.level_segments = level_segments
        self.level_knot_distance = level_knot_distance
        self.level_knot_dates = level_knot_dates
        self._level_knot_dates = self.level_knot_dates
        self.level_knots = None
        self._level_knots = None

        self._kernel_level = None
        self._num_knots_level = None
        self.knots_tp_level = None

        # seasonality configurations
        self.seasonality = seasonality
        self.seasonality_fs_order = seasonality_fs_order
        self._seasonality = self.seasonality
        # used to name different seasonal components in prediction
        self._seasonality_labels = list()
        self._seasonality_fs_order = self.seasonality_fs_order
        self.seasonal_initial_knot_scale = seasonal_initial_knot_scale
        self.seasonal_knot_scale = seasonal_knot_scale
        self.seasonality_segments = seasonality_segments
        self._seas_term = 0

        self._seasonality_coef_knot_dates = None
        self._seasonality_coef_knots = None

        # regression configurations
        self.regressor_col = regressor_col
        self.regressor_sign = regressor_sign
        self.regressor_init_knot_loc = regressor_init_knot_loc
        self.regressor_init_knot_scale = regressor_init_knot_scale
        self.regressor_knot_scale = regressor_knot_scale
        self.regression_knot_distance = regression_knot_distance
        self.regression_segments = regression_segments
        self._regression_knot_dates = regression_knot_dates

        self.regression_rho = regression_rho
        self.flat_multiplier = flat_multiplier

        # set private var to arg value
        # if None set default in _set_default_args()
        self._regressor_sign = self.regressor_sign
        self._regressor_init_knot_loc = self.regressor_init_knot_loc
        self._regressor_init_knot_scale = self.regressor_init_knot_scale
        self._regressor_knot_scale = self.regressor_knot_scale

        self.coef_prior_list = coef_prior_list
        self._coef_prior_list = []
        self._regression_knots_idx = None
        self._num_of_regressors = 0

        # positive regressors
        self._num_of_positive_regressors = 0
        self._positive_regressor_col = list()
        self._positive_regressor_init_knot_loc = list()
        self._positive_regressor_init_knot_scale = list()
        self._positive_regressor_knot_scale_1d = list()
        self._positive_regressor_knot_scale = list()
        # negative regressors
        self._num_of_negative_regressors = 0
        self._negative_regressor_col = list()
        self._negative_regressor_init_knot_loc = list()
        self._negative_regressor_init_knot_scale = list()
        self._negative_regressor_knot_scale_1d = list()
        self._negative_regressor_knot_scale = list()
        # regular regressors
        self._num_of_regular_regressors = 0
        self._regular_regressor_col = list()
        self._regular_regressor_init_knot_loc = list()
        self._regular_regressor_init_knot_scale = list()
        self._regular_regressor_knot_scale_1d = list()
        self._regular_regressor_knot_scale = list()
        self._regressor_col = list()

        # init dynamic data attributes
        # the following are set by `_set_dynamic_attributes()` and generally set during fit()
        # from input df
        # response data
        self._is_valid_response = None
        self._which_valid_response = None
        self._num_of_valid_response = 0

        # regression data
        self._knots_tp_coefficients = None
        self._positive_regressor_matrix = None
        self._negative_regressor_matrix = None
        self._regular_regressor_matrix = None

        # other configurations
        self.date_freq = date_freq
        self.degree_of_freedom = degree_of_freedom
        self.residuals_scale_upper = residuals_scale_upper
        self._residuals_scale_upper = residuals_scale_upper
        self.ktrlite_optim_args = ktrlite_optim_args

        self._set_static_attributes()
        self._set_model_param_names()

    def _set_model_param_names(self):
        """Overriding base template functions. Model parameters to extract"""
        self._model_param_names += [param.value for param in BaseSamplingParameters]
        if self._num_of_regressors > 0:
            self._model_param_names += [param.value for param in RegressionSamplingParameters]

    def _set_default_args(self):
        """Set default attributes for None"""
        # default checks for seasonality and seasonality_fs_order will be conducted
        # in ktrlite model and we will extract them from ktrlite model directly later
        if self.coef_prior_list is not None:
            self._coef_prior_list = deepcopy(self.coef_prior_list)

        # if no regressors, end here #
        if self.regressor_col is None:
            # regardless of what args are set for these, if regressor_col is None
            # these should all be empty lists
            self._regressor_sign = list()
            self._regressor_init_knot_loc = list()
            self._regressor_init_knot_scale = list()
            self._regressor_knot_scale = list()

            return

        def _validate_params_len(params, valid_length):
            for p in params:
                if p is not None and len(p) != valid_length:
                    raise IllegalArgument('Wrong dimension length in Regression Param Input')

        # regressor defaults
        num_of_regressors = len(self.regressor_col)

        _validate_params_len([
            self.regressor_sign, self.regressor_init_knot_loc,
            self.regressor_init_knot_scale, self.regressor_knot_scale],
            num_of_regressors
        )

        if self.regressor_sign is None:
            self._regressor_sign = [DEFAULT_REGRESSOR_SIGN] * num_of_regressors

        if self.regressor_init_knot_loc is None:
            self._regressor_init_knot_loc = [DEFAULT_COEFFICIENTS_INIT_KNOT_LOC] * num_of_regressors

        if self.regressor_init_knot_scale is None:
            self._regressor_init_knot_scale = [DEFAULT_COEFFICIENTS_INIT_KNOT_SCALE] * num_of_regressors

        if self.regressor_knot_scale is None:
            self._regressor_knot_scale = [DEFAULT_COEFFICIENTS_KNOT_SCALE] * num_of_regressors

        self._num_of_regressors = num_of_regressors

    def _set_static_regression_attributes(self):
        # if no regressors, end here
        if self._num_of_regressors == 0:
            return

        for index, reg_sign in enumerate(self._regressor_sign):
            if reg_sign == '+':
                self._num_of_positive_regressors += 1
                self._positive_regressor_col.append(self.regressor_col[index])
                # used for 'pr_knot_loc' sampling in pyro
                self._positive_regressor_init_knot_loc.append(self._regressor_init_knot_loc[index])
                self._positive_regressor_init_knot_scale.append(self._regressor_init_knot_scale[index])
                # used for 'pr_knot' sampling in pyro
                self._positive_regressor_knot_scale_1d.append(self._regressor_knot_scale[index])
            elif reg_sign == '-':
                self._num_of_negative_regressors += 1
                self._negative_regressor_col.append(self.regressor_col[index])
                # used for 'nr_knot_loc' sampling in pyro
                self._negative_regressor_init_knot_loc.append(self._regressor_init_knot_loc[index])
                self._negative_regressor_init_knot_scale.append(self._regressor_init_knot_scale[index])
                # used for 'nr_knot' sampling in pyro
                self._negative_regressor_knot_scale_1d.append(self._regressor_knot_scale[index])
            else:
                self._num_of_regular_regressors += 1
                self._regular_regressor_col.append(self.regressor_col[index])
                # used for 'rr_knot_loc' sampling in pyro
                self._regular_regressor_init_knot_loc.append(self._regressor_init_knot_loc[index])
                self._regular_regressor_init_knot_scale.append(self._regressor_init_knot_scale[index])
                # used for 'rr_knot' sampling in pyro
                self._regular_regressor_knot_scale_1d.append(self._regressor_knot_scale[index])
        # regular first, then positive, then negative
        self._regressor_col = self._regular_regressor_col + self._positive_regressor_col + self._negative_regressor_col
        # numpy conversion
        self._positive_regressor_init_knot_loc = np.array(self._positive_regressor_init_knot_loc)
        self._positive_regressor_init_knot_scale = np.array(self._positive_regressor_init_knot_scale)
        self._positive_regressor_knot_scale_1d = np.array(self._positive_regressor_knot_scale_1d)
        self._negative_regressor_init_knot_loc = np.array(self._negative_regressor_init_knot_loc)
        self._negative_regressor_init_knot_scale = np.array(self._negative_regressor_init_knot_scale)
        self._negative_regressor_knot_scale_1d = np.array(self._negative_regressor_knot_scale_1d)
        self._regular_regressor_init_knot_loc = np.array(self._regular_regressor_init_knot_loc)
        self._regular_regressor_init_knot_scale = np.array(self._regular_regressor_init_knot_scale)
        self._regular_regressor_knot_scale_1d = np.array(self._regular_regressor_knot_scale_1d)

    @staticmethod
    def _validate_coef_prior(coef_prior_list):
        for test_dict in coef_prior_list:
            if set(test_dict.keys()) != set([
                KTRTimePointPriorKeys.NAME.value,
                KTRTimePointPriorKeys.PRIOR_START_TP_IDX.value,
                KTRTimePointPriorKeys.PRIOR_END_TP_IDX.value,
                KTRTimePointPriorKeys.PRIOR_MEAN.value,
                KTRTimePointPriorKeys.PRIOR_SD.value,
                KTRTimePointPriorKeys.PRIOR_REGRESSOR_COL.value
            ]):
                raise IllegalArgument('wrong key name in inserted prior dict')
            len_insert_prior = list()
            for key, val in test_dict.items():
                if key in [
                    KTRTimePointPriorKeys.PRIOR_MEAN.value,
                    KTRTimePointPriorKeys.PRIOR_SD.value,
                    KTRTimePointPriorKeys.PRIOR_REGRESSOR_COL.value,
                ]:
                    len_insert_prior.append(len(val))
            if not all(len_insert == len_insert_prior[0] for len_insert in len_insert_prior):
                raise IllegalArgument('wrong dimension length in inserted prior dict')

    # @staticmethod
    # def _validate_level_knot_inputs(level_knot_dates, level_knots):
    #     if len(level_knots) != len(level_knot_dates):
    #         raise IllegalArgument('level_knots and level_knot_dates should have the same length')

    def _set_coef_prior_idx(self):
        if self._coef_prior_list and len(self._regressor_col) > 0:
            for x in self._coef_prior_list:
                prior_regressor_col_idx = [
                    np.where(np.array(self._regressor_col) == col)[0][0]
                    for col in x[KTRTimePointPriorKeys.PRIOR_REGRESSOR_COL.value]
                ]
                x.update({'prior_regressor_col_idx': prior_regressor_col_idx})

    def _set_static_attributes(self):
        """model data input based on args at instantiation or computed from args at instantiation"""
        self._set_default_args()
        self._set_static_regression_attributes()

        # self._validate_level_knot_inputs(self.level_knot_dates, self.level_knots)

        if self._coef_prior_list:
            self._validate_coef_prior(self._coef_prior_list)
            self._set_coef_prior_idx()

    def _set_valid_response_attributes(self, training_meta):
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        response = training_meta[TrainingMetaKeys.RESPONSE.value]

        if self._seasonality:
            max_seasonality = np.round(np.max(self._seasonality)).astype(int)
            if num_of_observations < max_seasonality:
                raise ModelException(
                    "Number of observations {} is less than max seasonality {}".format(
                        num_of_observations, max_seasonality))
        # get some reasonable offset to regularize response to make default priors scale-insensitive
        if self._seasonality:
            max_seasonality = np.round(np.max(self._seasonality)).astype(int)
            self.response_offset = np.nanmean(response[:max_seasonality])
        else:
            self.response_offset = np.nanmean(response)

        self.is_valid_response = ~np.isnan(response)
        # [0] to convert tuple back to array
        self.which_valid_response = np.where(self.is_valid_response)[0]
        self.num_of_valid_response = len(self.which_valid_response)

    def _set_regressor_matrix(self, df, training_meta):
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        # validate regression columns
        if self.regressor_col is not None and \
                not set(self.regressor_col).issubset(df.columns):
            raise ModelException(
                "DataFrame does not contain specified regressor column(s)."
            )

        # init of regression matrix depends on length of response vector
        self._positive_regressor_matrix = np.zeros((num_of_observations, 0), dtype=np.double)
        self._negative_regressor_matrix = np.zeros((num_of_observations, 0), dtype=np.double)
        self._regular_regressor_matrix = np.zeros((num_of_observations, 0), dtype=np.double)

        # update regression matrices
        if self._num_of_positive_regressors > 0:
            self._positive_regressor_matrix = df.filter(
                items=self._positive_regressor_col, ).values

        if self._num_of_negative_regressors > 0:
            self._negative_regressor_matrix = df.filter(
                items=self._negative_regressor_col, ).values

        if self._num_of_regular_regressors > 0:
            self._regular_regressor_matrix = df.filter(
                items=self._regular_regressor_col, ).values

    def _set_coefficients_kernel_matrix(self, df, training_meta):
        """Derive knots position and kernel matrix and other related meta data"""
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]
        # date_col = training_meta[TrainingMetaKeys.DATE_COL.value]

        # placeholder
        self._kernel_coefficients = np.zeros((num_of_observations, 0), dtype=np.double)
        self._num_knots_coefficients = 0

        if self._num_of_regressors > 0:
            self._regression_knots_idx = get_knot_idx(
                date_array=date_array,
                num_of_obs=num_of_observations,
                knot_dates=self._regression_knot_dates,
                knot_distance=self.regression_knot_distance,
                num_of_segments=self.regression_segments,
                date_freq=self.date_freq,
            )

            tp = np.arange(1, num_of_observations + 1) / num_of_observations
            self._knots_tp_coefficients = (1 + self._regression_knots_idx) / num_of_observations
            self._kernel_coefficients = gauss_kernel(tp, self._knots_tp_coefficients, rho=self.regression_rho)
            self._num_knots_coefficients = len(self._knots_tp_coefficients)
            if self.date_freq is None:
                self.date_freq = date_array.diff().min()
            self._regression_knot_dates = get_knot_dates(date_array[0], self._regression_knots_idx, self.date_freq)

    def _set_knots_scale_matrix(self, df, training_meta):
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        if self._num_of_positive_regressors > 0:
            # calculate average local absolute volume for each segment
            local_val = np.ones((self._num_of_positive_regressors, self._num_knots_coefficients))
            if self.flat_multiplier:
                multiplier = np.ones(local_val.shape)
            else:
                multiplier = np.ones(local_val.shape)
                # store local value for the range on the left side since last knot
                for idx in range(len(self._regression_knots_idx)):
                    if idx < len(self._regression_knots_idx) - 1:
                        str_idx = self._regression_knots_idx[idx]
                        end_idx = self._regression_knots_idx[idx + 1]
                    else:
                        str_idx = self._regression_knots_idx[idx]
                        end_idx = num_of_observations

                    local_val[:, idx] = np.mean(np.fabs(self._positive_regressor_matrix[str_idx:end_idx]), axis=0)

                global_mean = np.expand_dims(np.mean(np.fabs(self._positive_regressor_matrix), axis=0), -1)
                test_flag = local_val < 0.01 * global_mean

                # adjust knot scale with the multiplier derive by the average value and shift by 0.001 to avoid zeros in
                # scale parameters
                multiplier[test_flag] = DEFAULT_LOWER_BOUND_SCALE_MULTIPLIER
                # replace entire row of nan (when 0.1 * global_mean is equal to global_min) with upper bound
                multiplier[np.isnan(multiplier).all(axis=-1)] = 1.0

            # geometric drift i.e. 0.1 = 10% up-down in 1 s.d. prob.
            # self._positive_regressor_knot_scale has shape num_of_pr x num_of_knot
            self._positive_regressor_knot_scale = (
                    multiplier * np.expand_dims(self._positive_regressor_knot_scale_1d, -1)
            )
            # keep a lower bound of scale parameters
            self._positive_regressor_knot_scale[self._positive_regressor_knot_scale < 1e-4] = 1e-4
            # TODO: we change the type here, maybe we should change it earlier?
            self._positive_regressor_init_knot_scale = np.array(self._positive_regressor_init_knot_scale)
            self._positive_regressor_init_knot_scale[self._positive_regressor_init_knot_scale < 1e-4] = 1e-4

        if self._num_of_negative_regressors > 0:
            # calculate average local absolute volume for each segment
            local_val = np.ones((self._num_of_negative_regressors, self._num_knots_coefficients))
            if self.flat_multiplier:
                multiplier = np.ones(local_val.shape)
            else:
                multiplier = np.ones(local_val.shape)
                # store local value for the range on the left side since last knot
                for idx in range(len(self._regression_knots_idx)):
                    if idx < len(self._regression_knots_idx) - 1:
                        str_idx = self._regression_knots_idx[idx]
                        end_idx = self._regression_knots_idx[idx + 1]
                    else:
                        str_idx = self._regression_knots_idx[idx]
                        end_idx = num_of_observations

                    local_val[:, idx] = np.mean(np.fabs(self._negative_regressor_matrix[str_idx:end_idx]), axis=0)

                global_mean = np.expand_dims(np.mean(np.fabs(self._negative_regressor_matrix), axis=0), -1)
                test_flag = local_val < 0.01 * global_mean

                # adjust knot scale with the multiplier derive by the average value and shift by 0.001 to avoid zeros in
                # scale parameters
                multiplier[test_flag] = DEFAULT_LOWER_BOUND_SCALE_MULTIPLIER
                # replace entire row of nan (when 0.1 * global_mean is equal to global_min) with upper bound
                multiplier[np.isnan(multiplier).all(axis=-1)] = 1.0

            # geometric drift i.e. 0.1 = 10% up-down in 1 s.d. prob.
            self._negative_regressor_knot_scale = (
                    multiplier * np.expand_dims(self._negative_regressor_knot_scale_1d, -1)
            )
            # keep a lower bound of scale parameters
            self._negative_regressor_knot_scale[self._negative_regressor_knot_scale < 1e-4] = 1e-4
            # TODO: we change the type here, maybe we should change it earlier?
            self._negative_regressor_init_knot_scale = np.array(self._negative_regressor_init_knot_scale)
            self._negative_regressor_init_knot_scale[self._negative_regressor_init_knot_scale < 1e-4] = 1e-4

        if self._num_of_regular_regressors > 0:
            # do the same for regular regressor
            # calculate average local absolute volume for each segment
            local_val = np.ones((self._num_of_regular_regressors, self._num_knots_coefficients))
            if self.flat_multiplier:
                multiplier = np.ones(local_val.shape)
            else:
                multiplier = np.ones(local_val.shape)
                # store local value for the range on the left side since last knot
                for idx in range(len(self._regression_knots_idx)):
                    if idx < len(self._regression_knots_idx) - 1:
                        str_idx = self._regression_knots_idx[idx]
                        end_idx = self._regression_knots_idx[idx + 1]
                    else:
                        str_idx = self._regression_knots_idx[idx]
                        end_idx = num_of_observations

                    local_val[:, idx] = np.mean(np.fabs(self._regular_regressor_matrix[str_idx:end_idx]), axis=0)

            # adjust knot scale with the multiplier derive by the average value and shift by 0.001 to avoid zeros in
            # scale parameters
            global_mean = np.expand_dims(np.mean(np.fabs(self._regular_regressor_matrix), axis=0), -1)
            test_flag = local_val < 0.01 * global_mean
            multiplier[test_flag] = DEFAULT_LOWER_BOUND_SCALE_MULTIPLIER
            # replace entire row of nan (when 0.1 * global_mean is equal to global_min) with upper bound
            multiplier[np.isnan(multiplier).all(axis=-1)] = 1.0

            # geometric drift i.e. 0.1 = 10% up-down in 1 s.d. prob.
            # self._regular_regressor_knot_scale has shape num_of_pr x num_of_knot
            self._regular_regressor_knot_scale = (
                    multiplier * np.expand_dims(self._regular_regressor_knot_scale_1d, -1)
            )
            # keep a lower bound of scale parameters
            self._regular_regressor_knot_scale[self._regular_regressor_knot_scale < 1e-4] = 1e-4
            # TODO: we change the type here, maybe we should change it earlier?
            self._regular_regressor_init_knot_scale = np.array(self._regular_regressor_init_knot_scale)
            self._regular_regressor_init_knot_scale[self._regular_regressor_init_knot_scale < 1e-4] = 1e-4

    def _generate_tp(self, training_meta, prediction_date_array):
        """Used in _generate_seas"""
        training_end = training_meta[TrainingMetaKeys.END.value]
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]
        prediction_start = prediction_date_array[0]
        output_len = len(prediction_date_array)
        if prediction_start > training_end:
            start = num_of_observations
        else:
            start = pd.Index(date_array).get_loc(prediction_start)

        new_tp = np.arange(start + 1, start + output_len + 1) / num_of_observations
        return new_tp

    def _generate_insample_tp(self, training_meta, date_array):
        """Used in _generate_seas"""
        train_date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        idx = np.nonzero(np.in1d(train_date_array, date_array))[0]
        tp = (idx + 1) / num_of_observations
        return tp

    # def _generate_coefs(self, training_meta, prediction_date_array, coef_knot_dates, coef_knot):
    #     """Used in _generate_seas"""
    #     new_tp = self._generate_tp(training_meta, prediction_date_array)
    #     knots_tp_coef = self._generate_insample_tp(training_meta, coef_knot_dates)
    #     kernel_coef = sandwich_kernel(new_tp, knots_tp_coef)
    #     coefs = np.squeeze(np.matmul(coef_knot, kernel_coef.transpose(1, 0)), axis=0).transpose(1, 0)
    #     return coefs

    def _generate_seas(self, df, training_meta, coef_knot_dates, coef_knots,
                        seasonality, seasonality_fs_order, seasonality_labels):
        """To calculate the seasonality term based on the _seasonal_knots_input.
        Parameters
        ----------
        df : pd.DataFrame
            input df
        training_meta: dict
            meta dictionary for the training input
        coef_knot_dates : 1-D array like
            dates for seasonality coefficient knots
        coef_knots : dict
            dict of seasonal coefficient knots from each seasonality
        seasonality : list
            seasonality input; list of float
        seasonality_fs_order : list
            seasonality_fs_order input list of int

        Returns
        -----------
        dict :
            a dictionary contains seasonal regression components mapped by each seasonality
        """
        df = df.copy()
        # store each component as a dictionary
        seas_decomp = dict()

        if seasonality is not None and len(seasonality) > 0:

            date_col = training_meta[TrainingMetaKeys.DATE_COL.value]
            date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]
            training_end = training_meta[TrainingMetaKeys.END.value]
            num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]

            prediction_date_array = df[date_col].values
            prediction_start = prediction_date_array[0]

            if prediction_start > training_end:
                # time index for prediction start
                start = num_of_observations
            else:
                # time index for prediction start
                start = pd.Index(date_array).get_loc(prediction_start)

            # dictionary
            seas_regressors = make_seasonal_regressors(
                n=df.shape[0],
                periods=seasonality,
                orders=seasonality_fs_order,
                labels=seasonality_labels,
                shift=start,
            )

            new_tp = self._generate_tp(training_meta, prediction_date_array)
            knots_tp_coef = self._generate_insample_tp(training_meta, coef_knot_dates)
            coef_kernel = sandwich_kernel(new_tp, knots_tp_coef)

            # init of regression matrix depends on length of response vector
            total_seas_regression = np.zeros((1, df.shape[0]), dtype=np.double)

            for k in seasonality_labels:
                seas_regresor_matrix = seas_regressors[k]
                coef_knot = coef_knots[k]
                # time-step x coefficients
                seas_coef = np.squeeze(np.matmul(coef_knot, coef_kernel.transpose(1, 0)), axis=0).transpose(1, 0)
                seas_regression = np.sum(seas_coef * seas_regresor_matrix, axis=-1)
                seas_decomp[k] = np.expand_dims(seas_regression, 0)
                total_seas_regression += seas_regression
        else:
            total_seas_regression = np.zeros((1, df.shape[0]), dtype=np.double)

        return total_seas_regression, seas_decomp

    def _set_levs_and_seas(self, df, training_meta):
        response_col = training_meta['response_col']
        date_col = training_meta[TrainingMetaKeys.DATE_COL.value]
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]

        # use ktrlite to derive levs and seas
        ktrlite = KTRLite(
            response_col=response_col,
            date_col=date_col,
            level_knot_scale=self.level_knot_scale,
            level_segments=self.level_segments,
            level_knot_dates=self.level_knot_dates,
            level_knot_distance=self.level_knot_distance,
            seasonality=self.seasonality,
            seasonality_fs_order=self.seasonality_fs_order,
            seasonal_initial_knot_scale=self.seasonal_initial_knot_scale,
            seasonal_knot_scale=self.seasonal_knot_scale,
            seasonality_segments=self.seasonality_segments,
            degree_of_freedom=self.degree_of_freedom,
            date_freq=self.date_freq,
            estimator='stan-map',
            **self.ktrlite_optim_args
        )
        ktrlite.fit(df=df)
        # self._ktrlite_model = ktrlite
        ktrlite_pt_posteriors = ktrlite.get_point_posteriors()
        ktrlite_obs_scale = ktrlite_pt_posteriors['map']['obs_scale']

        # load _seasonality and _seasonality_fs_order
        self._seasonality = ktrlite._model._seasonality
        self._seasonality_fs_order = ktrlite._model._seasonality_fs_order
        for seas in self._seasonality:
            self._seasonality_labels.append('seasonality_{}'.format(seas))

        # if input None for upper bound of residuals scale, use data-driven input
        if self.residuals_scale_upper is None:
            # make it 5 times to have some buffer in case we over-fit in KTRLite
            self._residuals_scale_upper = min(ktrlite_obs_scale * 5, training_meta['response_sd'])

        # this part is to extract level and seasonality result from KTRLite
        self._level_knots = np.squeeze(ktrlite_pt_posteriors['map']['lev_knot'])
        self._level_knot_dates = ktrlite._model._level_knot_dates
        tp = np.arange(1, num_of_observations + 1) / num_of_observations
        # # trim level knots dates when they are beyond training dates
        # lev_knot_dates = list()
        # lev_knots = list()
        # for i, x in enumerate(self.level_knot_dates):
        #     if (x <= df[date_col].max()) and (x >= df[date_col].min()):
        #         lev_knot_dates.append(x)
        #         lev_knots.append(self._level_knots[i])
        # self._level_knot_dates = pd.to_datetime(lev_knot_dates)
        # self._level_knots = np.array(lev_knots)

        self._level_knots_idx = get_knot_idx(
            date_array=date_array,
            num_of_obs=None,
            knot_dates=self._level_knot_dates,
            knot_distance=None,
            num_of_segments=None,
            date_freq=self.date_freq,
        )
        self.knots_tp_level = (1 + self._level_knots_idx) / num_of_observations
        self._kernel_level = sandwich_kernel(tp, self.knots_tp_level)
        self._num_knots_level = len(self._level_knot_dates)

        if self._seasonality:
            self._seasonality_coef_knot_dates = ktrlite._model._coef_knot_dates
            coef_knots_flatten = ktrlite_pt_posteriors['map']['coef_knot']
            coef_knots = dict()
            pos = 0
            for idx, label in enumerate(self._seasonality_labels):
                order = self._seasonality_fs_order[idx]
                coef_knots[label] = coef_knots_flatten[..., pos:(pos + 2 * order), :]
                pos += 2 * order
            self._seasonality_coef_knots = coef_knots

            # we just need total here and because of
            self._seas_term, _ = self._generate_seas(
                df,
                training_meta,
                self._seasonality_coef_knot_dates,
                self._seasonality_coef_knots,
                self._seasonality,
                self._seasonality_fs_order,
                self._seasonality_labels)
            # remove batch size as an input for models
            self._seas_term = np.squeeze(self._seas_term, 0)

    def _filter_coef_prior(self, df):
        if self._coef_prior_list and len(self._regressor_col) > 0:
            # iterate over a copy due to the removal operation
            for test_dict in self._coef_prior_list[:]:
                prior_regressor_col = test_dict[KTRTimePointPriorKeys.PRIOR_REGRESSOR_COL.value]
                m = test_dict[KTRTimePointPriorKeys.PRIOR_MEAN.value]
                sd = test_dict[KTRTimePointPriorKeys.PRIOR_SD.value]
                end_tp_idx = min(test_dict[KTRTimePointPriorKeys.PRIOR_END_TP_IDX.value], df.shape[0])
                start_tp_idx = min(test_dict[KTRTimePointPriorKeys.PRIOR_START_TP_IDX.value], df.shape[0])
                if start_tp_idx < end_tp_idx:
                    expected_shape = (end_tp_idx - start_tp_idx, len(prior_regressor_col))
                    test_dict.update({KTRTimePointPriorKeys.PRIOR_END_TP_IDX.value: end_tp_idx})
                    test_dict.update({KTRTimePointPriorKeys.PRIOR_START_TP_IDX.value: start_tp_idx})
                    # mean/sd expanding
                    test_dict.update({KTRTimePointPriorKeys.PRIOR_MEAN.value: np.full(expected_shape, m)})
                    test_dict.update({KTRTimePointPriorKeys.PRIOR_SD.value: np.full(expected_shape, sd)})
                else:
                    # removing invalid prior
                    self._coef_prior_list.remove(test_dict)

    def set_dynamic_attributes(self, df, training_meta):
        """Overriding: func: `~orbit.models.BaseETS._set_dynamic_attributes"""
        self._set_regressor_matrix(df, training_meta)
        self._set_coefficients_kernel_matrix(df, training_meta)
        self._set_knots_scale_matrix(df, training_meta)
        self._set_levs_and_seas(df, training_meta)
        self._filter_coef_prior(df)
        self._set_valid_response_attributes(training_meta)

    @staticmethod
    def _concat_regression_coefs(pr_beta=None, rr_beta=None):
        """Concatenates regression posterior matrix
        In the case that `pr_beta` or `rr_beta` is a 1d tensor, transform to 2d tensor and
        concatenate.
        Args
        ----
        pr_beta : array like
            postive-value constrainted regression betas
        rr_beta : array like
            regular regression betas
        Returns
        -------
        array like
            concatenated 2d array of shape (1, len(rr_beta) + len(pr_beta))
        """
        regressor_beta = None
        if pr_beta is not None and rr_beta is not None:
            pr_beta = pr_beta if len(pr_beta.shape) == 2 else pr_beta.reshape(1, -1)
            rr_beta = rr_beta if len(rr_beta.shape) == 2 else rr_beta.reshape(1, -1)
            regressor_beta = torch.cat((rr_beta, pr_beta), dim=1)
        elif pr_beta is not None:
            regressor_beta = pr_beta
        elif rr_beta is not None:
            regressor_beta = rr_beta

        return regressor_beta

    def predict(self, posterior_estimates, df, training_meta, prediction_meta,
                coefficient_method="smooth",
                include_error=False, store_prediction_array=False, **kwargs):
        """Vectorized version of prediction math
        Parameters
        ----
        coefficient_method : str
            either "smooth" or "empirical". when "empirical" is used, curves are sampled/aggregated directly
            from beta posteriors; when "smooth" is used, first extract sampled/aggregated posteriors of knots
            then beta.
            this mainly impacts the aggregated estimation method; full bayesian should not be impacted
        include_error : bool
            if generating the noise samples
        store_prediction_array : bool
            if storing the prediction array
        """
        ################################################################
        # Model Attributes
        ################################################################
        # FIXME: do we still need this?
        model = deepcopy(posterior_estimates)
        arbitrary_posterior_value = list(model.values())[0]
        num_sample = arbitrary_posterior_value.shape[0]

        ################################################################
        # Prediction Attributes
        ################################################################
        output_len = prediction_meta[PredictionMetaKeys.PREDICTION_DF_LEN.value]
        prediction_start = prediction_meta[PredictionMetaKeys.START.value]
        date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        training_end = training_meta[TrainingMetaKeys.END.value]

        # Here assume dates are ordered and consecutive
        # if prediction_meta[PredictionMetaKeys.START.value] > self.training_end,
        # assume prediction starts right after train end
        if prediction_start > training_end:
            # time index for prediction start
            start = num_of_observations
        else:
            start = pd.Index(date_array).get_loc(prediction_start)

        new_tp = np.arange(start + 1, start + output_len + 1) / num_of_observations
        if include_error:
            # in-sample knots
            lev_knot_in = model.get(BaseSamplingParameters.LEVEL_KNOT.value)
            # TODO: hacky way; let's just assume last two knot distance is knots distance for all knots
            lev_knot_width = self.knots_tp_level[-1] - self.knots_tp_level[-2]
            # check whether we need to put new knots for simulation
            if new_tp[-1] >= self.knots_tp_level[-1] + lev_knot_width:
                # derive knots tp
                knots_tp_level_out = np.arange(self.knots_tp_level[-1] + lev_knot_width, new_tp[-1], lev_knot_width)
                new_knots_tp_level = np.concatenate([self.knots_tp_level, knots_tp_level_out])
                lev_knot_out = np.random.laplace(0, self.level_knot_scale,
                                                 size=(lev_knot_in.shape[0], len(knots_tp_level_out)))
                lev_knot_out = np.cumsum(np.concatenate([lev_knot_in[:, -1].reshape(-1, 1), lev_knot_out],
                                                        axis=1), axis=1)[:, 1:]
                lev_knot = np.concatenate([lev_knot_in, lev_knot_out], axis=1)
            else:
                new_knots_tp_level = self.knots_tp_level
                lev_knot = lev_knot_in
            kernel_level = sandwich_kernel(new_tp, new_knots_tp_level)
        else:
            lev_knot = model.get(BaseSamplingParameters.LEVEL_KNOT.value)
            kernel_level = sandwich_kernel(new_tp, self.knots_tp_level)
        obs_scale = model.get(BaseSamplingParameters.OBS_SCALE.value)
        obs_scale = obs_scale.reshape(-1, 1)

        # if self._seasonality is not None:
        # condition of seasonality is checked inside
        total_seas, seas_decomp = self._generate_seas(df, training_meta,
                                                      self._seasonality_coef_knot_dates,
                                                      self._seasonality_coef_knots,
                                                      self._seasonality,
                                                      self._seasonality_fs_order,
                                                      self._seasonality_labels)

        #     # seas is 1-d array, add the batch size back
        #     seas = np.expand_dims(seas, 0)
        # else:
        #     # follow component shapes
        #     seas = np.zeros((1, output_len))

        trend = np.matmul(lev_knot, kernel_level.transpose((1, 0)))
        regression = np.zeros(trend.shape)
        if self._num_of_regressors > 0:
            regressor_matrix = df.filter(items=self._regressor_col, ).values
            regressor_betas = self._get_regression_coefs_matrix(
                training_meta,
                posterior_estimates,
                coefficient_method,
                date_array=prediction_meta[TrainingMetaKeys.DATE_ARRAY.value]
            )
            regression = np.sum(regressor_betas * regressor_matrix, axis=-1)

        if include_error:
            epsilon = nct.rvs(self.degree_of_freedom, nc=0, loc=0,
                              scale=obs_scale, size=(num_sample, len(new_tp)))
            trend += epsilon

        pred_array = trend + total_seas + regression

        # if decompose output dictionary of components
        decomp_dict = {
            'prediction': pred_array,
            'trend': trend,
            'regression': regression
        }

        # this is an input from ktrlite
        decomp_dict.update(seas_decomp)

        if store_prediction_array:
            self.pred_array = pred_array
        else:
            self.pred_array = None

        return decomp_dict

    def _get_regression_coefs_matrix(self, training_meta, posteriors, coefficient_method='smooth', date_array=None):
        """internal function to provide coefficient matrix given a date array
        Args
        ----
        posteriors : dict
            posterior samples
        date_array : array like
            array of date stamp
        coefficient_method : str
            either "smooth" or "empirical". when "empirical" is used, curves are sampled/aggregated directly
            from beta posteriors; when "smooth" is used, first extract sampled/aggregated posteriors of knots
            then beta.
            this mainly impacts the aggregated estimation method; full bayesian should not be impacted.
        """
        num_of_observations = training_meta[TrainingMetaKeys.NUM_OF_OBS.value]
        training_start = training_meta[TrainingMetaKeys.START.value]
        training_end = training_meta[TrainingMetaKeys.END.value]
        train_date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]

        if self._num_of_regular_regressors + self._num_of_positive_regressors + self._num_of_negative_regressors == 0:
            return None

        # if date_array not specified, coefficients in the training period will be retrieved
        if date_array is None:
            if coefficient_method == 'smooth':

                coef_knots = posteriors.get(RegressionSamplingParameters.COEFFICIENTS_KNOT.value)
                # only 1 knot for 0 segments
                if self.regression_segments == 0:
                    coef_knots = np.expand_dims(coef_knots, -1)
                if len(self._regressor_col) == 1:
                    coef_knots = np.expand_dims(coef_knots, 1)

                # result in batch x time step x regressor size shape
                regressor_betas = np.matmul(coef_knots, self._kernel_coefficients.transpose((1, 0)))

                # if len(self._regressor_col) == 1:
                #     regressor_betas = np.expand_dims(regressor_betas, 0)
                regressor_betas = regressor_betas.transpose((0, 2, 1))
            elif coefficient_method == 'empirical':
                regressor_betas = posteriors.get(RegressionSamplingParameters.COEFFICIENTS.value)
            else:
                raise IllegalArgument('Wrong coefficient_method:{}'.format(coefficient_method))
        else:
            date_array = pd.to_datetime(date_array).values
            output_len = len(date_array)
            train_len = num_of_observations
            # some validation of date array
            if not is_ordered_datetime(date_array):
                raise IllegalArgument('Datetime index must be ordered and not repeat')
            prediction_start = date_array[0]

            if prediction_start < training_start:
                raise PredictionException('Prediction start must be after training start.')

            # If we cannot find a match of prediction range, assume prediction starts right after train end
            if prediction_start > training_end:
                # time index for prediction start
                start = train_len
                coef_repeats = [0] * (start - 1) + [output_len]
            else:
                # time index for prediction start
                start = pd.Index(train_date_array).get_loc(prediction_start)
                if output_len <= train_len - start:
                    coef_repeats = [0] * start + [1] * output_len + [0] * (train_len - start - output_len)
                else:
                    coef_repeats = [0] * start + [1] * (train_len - start - 1) + [output_len - train_len + start + 1]
            new_tp = np.arange(start + 1, start + output_len + 1) / num_of_observations

            if coefficient_method == 'smooth':
                kernel_coefficients = gauss_kernel(new_tp, self._knots_tp_coefficients, rho=self.regression_rho)

                coef_knots = posteriors.get(RegressionSamplingParameters.COEFFICIENTS_KNOT.value)
                if len(self._regressor_col) == 1:
                    coef_knots = np.expand_dims(coef_knots, -1)
                # only 1 knot for 0 segments
                if self.regression_segments == 0:
                    coef_knots = np.expand_dims(coef_knots, -1)

                regressor_betas = np.matmul(coef_knots, kernel_coefficients.transpose((1, 0)))
                if len(regressor_betas.shape) == 2:
                    regressor_betas = np.expand_dims(regressor_betas, 0)
                regressor_betas = regressor_betas.transpose((0, 2, 1))
            elif coefficient_method == 'empirical':
                regressor_betas = posteriors.get(RegressionSamplingParameters.COEFFICIENTS.value)
                regressor_betas = np.repeat(regressor_betas, repeats=coef_repeats, axis=1)
            else:
                raise IllegalArgument('Wrong coefficient_method:{}'.format(coefficient_method))

        return regressor_betas

    def get_regression_coefs(self, training_meta, point_method, point_posteriors, posterior_samples,
                             coefficient_method='smooth', date_array=None,
                             include_ci=False, lower=0.05, upper=0.95
                             ):
        """Return DataFrame regression coefficients.

        Parameters
        ----------
        coefficient_method : str
            either "smooth" or "empirical". when "empirical" is used, curves are sampled/aggregated directly
            from beta posteriors; when "smooth" is used, first extract sampled/aggregated posteriors of knots
            then beta.
        date_array : array-like
            the list of dates for which the regressio coefficients will be reported.
            Default to be None. When it's None, all the dates in the training data will be used.
        include_ci : bool
            if including the confidence intervals for the regression coefficients
        lower : float between (0, 1). default to be 0.05
            lower bound for the CI
        upper : float between (0, 1). default to be 0.95.
            upper bound for the CI

        Returns
        -------
        Pandas data frame holding the dynamic regression coefficients
        """
        date_col = training_meta[TrainingMetaKeys.DATE_COL.value]
        reg_df = pd.DataFrame()
        if self._num_of_regressors == 0:
            return reg_df

        _point_method = point_method
        if point_method is None:
            _point_method = PredictMethod.MEDIAN.value

        posteriors = point_posteriors.get(_point_method)
        coefs = np.squeeze(self._get_regression_coefs_matrix(training_meta,
                                                             posteriors,
                                                             coefficient_method=coefficient_method,
                                                             date_array=date_array))
        if len(coefs.shape) == 1:
            coefs = np.expand_dims(coefs, -1)

        reg_df = pd.DataFrame(data=coefs, columns=self._regressor_col)
        if date_array is not None:
            reg_df[date_col] = date_array
        else:
            reg_df[date_col] = training_meta[TrainingMetaKeys.DATE_ARRAY.value]
        # re-arrange columns
        reg_df = reg_df[[date_col] + self._regressor_col]

        if include_ci:
            posteriors = posterior_samples
            coefs = self._get_regression_coefs_matrix(training_meta,
                                                      posteriors,
                                                      coefficient_method=coefficient_method,
                                                      date_array=date_array)
            coefficients_lower = np.quantile(coefs, lower, axis=0)
            coefficients_upper = np.quantile(coefs, upper, axis=0)
            reg_df_lower = reg_df.copy()
            reg_df_upper = reg_df.copy()
            for idx, col in enumerate(self._regressor_col):
                reg_df_lower[col] = coefficients_lower[:, idx]
                reg_df_upper[col] = coefficients_upper[:, idx]
            return reg_df, reg_df_lower, reg_df_upper

        return reg_df

    def get_regression_coef_knots(self, training_meta, point_method, point_posteriors, posterior_samples):
        """Return DataFrame regression coefficient knots
        """
        date_col = training_meta[TrainingMetaKeys.DATE_COL.value]
        _point_method = point_method
        if point_method is None:
            _point_method = PredictMethod.MEDIAN.value

        # init dataframe
        knots_df = pd.DataFrame()
        # end if no regressors
        if self._num_of_regular_regressors + self._num_of_positive_regressors + self._num_of_negative_regressors == 0:
            return knots_df

        knots_df[date_col] = self._regression_knot_dates
        # TODO: make the label as a constant
        knots_df['step'] = self._regression_knots_idx
        # batch size x regressor size x knot size
        coef_knots = point_posteriors \
            .get(_point_method) \
            .get(RegressionSamplingParameters.COEFFICIENTS_KNOT.value)

        # only 1 knot for 0 segments
        if self.regression_segments == 0:
            coef_knots = np.expand_dims(coef_knots, -1)
        if len(self._regressor_col) == 1:
            coef_knots = np.expand_dims(coef_knots, 1)

        for idx, col in enumerate(self._regressor_col):
            knots_df[col] = np.transpose(coef_knots[:, idx])

        return knots_df

    @orbit_style_decorator
    def plot_regression_coefs(self, training_meta, point_method, point_posteriors, posterior_samples,
                              coefficient_method='smooth', date_array=None,
                              include_ci=False, lower=0.05, upper=0.95,
                              with_knot=False, is_visible=True,
                              ncol=2, ylim=None, markersize=200, figsize=(16, 8)):
        """Plot regression coefficients.

        Parameters
        ----------
        coefficient_method : str
            either "smooth" or "empirical". when "empirical" is used, curves are sampled/aggregated directly
            from beta posteriors; when "smooth" is used, first extract sampled/aggregated posteriors of knots
            then beta.
        date_array : array-like
            the list of dates for which the regressio coefficients will be reported.
            Default to be None. When it's None, all the dates in the training data will be used.
        include_ci : bool
            if including the confidence intervals for the regression coefficients
        lower : float between (0, 1). default to be 0.05
            lower bound for the CI
        upper : float between (0, 1). default to be 0.95.
            upper bound for the CI
        with_knot : bool
            if plotting the regression knots in the graph
        ncol : int
            number of columns of the panel grid
                is_visible : boolean
            whether we want to show the plot. If called from unittest, is_visible might = False.
        is_visible : bool
            whether we want to show the plot. If called from unittest, is_visible might = False.
        markersize : int; optional
            knot marker size
        figsize : tuple; optional
            figsize passed to `matplotlib.pyplot.figure()`

        """
        # assume your first column is the date; this way can use a static method
        if include_ci:
            coef_df, coef_df_lower, coef_df_upper = self.get_regression_coefs(
                training_meta, point_method, point_posteriors, posterior_samples,
                coefficient_method=coefficient_method, date_array=date_array,
                include_ci=include_ci, lower=lower, upper=upper
            )
        else:
            coef_df = self.get_regression_coefs(
                training_meta, point_method, point_posteriors, posterior_samples,
                coefficient_method=coefficient_method, date_array=date_array,
                include_ci=include_ci, lower=lower, upper=upper
            )
            coef_df_lower, coef_df_upper = None, None
        if with_knot:
            knot_df = self.get_regression_coef_knots(training_meta, point_method, point_posteriors, posterior_samples)
        else:
            knot_df = None

        regressor_col = coef_df.columns.tolist()[1:]
        nrow = math.ceil(len(regressor_col) / ncol)
        fig, axes = plt.subplots(nrow, ncol, figsize=figsize, squeeze=False)

        for idx, col in enumerate(regressor_col):
            row_idx = idx // ncol
            col_idx = idx % ncol
            coef = coef_df[col]
            axes[row_idx, col_idx].plot(coef, alpha=.8, label='coefficients', color=OrbitPalette.BLUE.value)
            if coef_df_lower is not None and coef_df_upper is not None:
                coef_lower = coef_df_lower[col]
                coef_upper = coef_df_upper[col]
                axes[row_idx, col_idx].fill_between(np.arange(0, coef_df.shape[0]), coef_lower, coef_upper,
                                                    alpha=.3, color=OrbitPalette.BLUE.value)
            if knot_df is not None:
                step = knot_df['step']
                knots = knot_df[col].values
                axes[row_idx, col_idx].scatter(x=step, y=knots, marker='^', s=markersize,
                                               color=OrbitPalette.GREEN.value, alpha=0.5)
            if ylim is not None:
                axes[row_idx, col_idx].set_ylim(ylim)
            axes[row_idx, col_idx].set_title('{}'.format(col))
            axes[row_idx, col_idx].ticklabel_format(useOffset=False)

        plt.tight_layout()

        if is_visible:
            plt.show()
        else:
            plt.close()

        return axes

    # TODO: need a unit test of this function
    def get_level_knots(self, training_meta, point_method, point_posteriors, posterior_samples):
        """Given posteriors, return knots and correspondent date"""
        date_col = training_meta[TrainingMetaKeys.DATE_COL.value]
        _point_method = point_method
        if point_method is None:
            _point_method = PredictMethod.MEDIAN.value
        lev_knots = point_posteriors \
            .get(_point_method) \
            .get(BaseSamplingParameters.LEVEL_KNOT.value)
        lev_knots = np.squeeze(lev_knots, 0)
        out = {
            date_col: self._level_knot_dates,
            BaseSamplingParameters.LEVEL_KNOT.value: lev_knots,
        }

        return pd.DataFrame(out)

    def get_levels(self, training_meta, point_method, point_posteriors, posterior_samples):
        date_col = training_meta[TrainingMetaKeys.DATE_COL.value]
        date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]
        _point_method = point_method
        if point_method is None:
            _point_method = PredictMethod.MEDIAN.value
        levs = point_posteriors \
            .get(_point_method) \
            .get(BaseSamplingParameters.LEVEL.value)
        levs = np.squeeze(levs, 0)
        out = {
            date_col: date_array,
            BaseSamplingParameters.LEVEL.value: levs,
        }

        return pd.DataFrame(out)

    @orbit_style_decorator
    def plot_lev_knots(self, training_meta, point_method, point_posteriors, posterior_samples,
                       path=None, is_visible=True, title="", fontsize=16,
                       markersize=250, figsize=(16, 8)):
        """ Plot the fitted level knots along with the actual time series.
        Parameters
        ----------
        path : str; optional
            path to save the figure
        is_visible : boolean
            whether we want to show the plot. If called from unittest, is_visible might = False.
        title : str; optional
            title of the plot
        fontsize : int; optional
            fontsize of the title
        markersize : int; optional
            knot marker size
        figsize : tuple; optional
            figsize passed to `matplotlib.pyplot.figure()`
       Returns
        -------
            matplotlib axes object
        """
        date_col = training_meta[TrainingMetaKeys.DATE_COL.value]
        date_array = training_meta[TrainingMetaKeys.DATE_ARRAY.value]
        response = training_meta[TrainingMetaKeys.RESPONSE.value]

        levels_df = self.get_levels(training_meta, point_method, point_posteriors, posterior_samples)
        knots_df = self.get_level_knots(training_meta, point_method, point_posteriors, posterior_samples)

        fig, ax = plt.subplots(1, 1, figsize=figsize)
        ax.plot(date_array, response, color=OrbitPalette.BLUE.value, lw=1, alpha=0.7, label='actual')
        ax.plot(levels_df[date_col], levels_df[BaseSamplingParameters.LEVEL.value],
                color=OrbitPalette.BLACK.value, lw=1, alpha=0.8,
                label=BaseSamplingParameters.LEVEL.value)
        ax.scatter(knots_df[date_col], knots_df[BaseSamplingParameters.LEVEL_KNOT.value],
                   color=OrbitPalette.GREEN.value, lw=1, s=markersize, marker='^', alpha=0.8,
                   label=BaseSamplingParameters.LEVEL_KNOT.value)
        ax.legend()
        ax.grid(True, which='major', c='grey', ls='-', lw=1, alpha=0.5)
        ax.set_title(title, fontsize=fontsize)
        if path:
            fig.savefig(path)
        if is_visible:
            plt.show()
        else:
            plt.close()
        return ax

max_stars_repo_name: uber/orb

max_issues_repo_name: uber/orb

max_forks_repo_name: uber/orb

reason: import numpy,from scipy


text:
'''Backtest Moving Average (MA) crossover strategies
'''

import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from MA import MABacktester

class MAShortOnlyBacktester(MABacktester):
    '''Backtest a Moving Average (MA) crossover strategy that only goes short

    Parameters:
    series: (Panda Series) a list of CLOSE prices by date
    ms: (int) short moving average
    ml: (int) long moving average
    ema: (boolean) True if you want exponential MA's
    '''

    def __init__(self, series, ms=1, ml=10, ema=False):
        self._ms = ms
        self._ml = ml
        self._ema = ema
        super(MAShortOnlyBacktester,self).__init__(series,ms=ms,ml=ml,long_only=False,ema=ema)

    def __str__(self):
        return "MA Short Only Backtest Strategy (ms=%d, ml=%d, ema=%s, start=%s, end=%s)" % (
            self._ms, self._ml, str(self._ema), str(self._start_date), str(self._end_date))

    def _trade_logic(self):
        '''Implements the trade logic in order to come up with
        a set of stances
        '''

        self._indicators()

        self._df['stance'] = np.where(self._df['mdiff'] < 0, -1, 0)

ext: py

max_issues_repo_name: dcimring/crypto-price-analysis

max_forks_repo_name: dcimring/crypto-price-analysis

reason: import numpy


text:
"""Utils for monoDepth.
"""
import sys
import re
import numpy as np
import cv2
import torch

def read_pfm(path):
    """Read pfm file.

    Args:
        path (str): path to file

    Returns:
        tuple: (data, scale)
    """
    with open(path, "rb") as file:

        color = None
        width = None
        height = None
        scale = None
        endian = None

        header = file.readline().rstrip()
        if header.decode("ascii") == "PF":
            color = True
        elif header.decode("ascii") == "Pf":
            color = False
        else:
            raise Exception("Not a PFM file: " + path)

        dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
        if dim_match:
            width, height = list(map(int, dim_match.groups()))
        else:
            raise Exception("Malformed PFM header.")

        scale = float(file.readline().decode("ascii").rstrip())
        if scale < 0:
            # little-endian
            endian = "<"
            scale = -scale
        else:
            # big-endian
            endian = ">"

        data = np.fromfile(file, endian + "f")
        shape = (height, width, 3) if color else (height, width)

        data = np.reshape(data, shape)
        data = np.flipud(data)

        return data, scale

def write_pfm(path, image, scale=1):
    """Write pfm file.

    Args:
        path (str): pathto file
        image (array): data
        scale (int, optional): Scale. Defaults to 1.
    """

    with open(path, "wb") as file:
        color = None

        if image.dtype.name != "float32":
            raise Exception("Image dtype must be float32.")

        image = np.flipud(image)

        if len(image.shape) == 3 and image.shape[2] == 3:  # color image
            color = True
        elif (
            len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
        ):  # greyscale
            color = False
        else:
            raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")

        file.write("PF\n" if color else "Pf\n".encode())
        file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))

        endian = image.dtype.byteorder

        if endian == "<" or endian == "=" and sys.byteorder == "little":
            scale = -scale

        file.write("%f\n".encode() % scale)

        image.tofile(file)

def read_image(path):
    """Read image and output RGB image (0-1).

    Args:
        path (str): path to file

    Returns:
        array: RGB image (0-1)
    """
    img = cv2.imread(path)

    if img.ndim == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0

    return img

def process_camera_img(img):

    if img.ndim == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0

    return img

def resize_image(img):
    """Resize image and make it fit for network.

    Args:
        img (array): image

    Returns:
        tensor: data ready for network
    """
    height_orig = img.shape[0]
    width_orig = img.shape[1]

    if width_orig > height_orig:
        scale = width_orig / 384
    else:
        scale = height_orig / 384

    height = (np.ceil(height_orig / scale / 32) * 32).astype(int)
    width = (np.ceil(width_orig / scale / 32) * 32).astype(int)

    img_resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)

    img_resized = (
        torch.from_numpy(np.transpose(img_resized, (2, 0, 1))).contiguous().float()
    )
    img_resized = img_resized.unsqueeze(0)

    return img_resized

def resize_depth(depth, width, height):
    """Resize depth map and bring to CPU (numpy).

    Args:
        depth (tensor): depth
        width (int): image width
        height (int): image height

    Returns:
        array: processed depth
    """
    depth = torch.squeeze(depth[0, :, :, :]).to("cpu")

    depth_resized = cv2.resize(
        depth.numpy(), (width, height), interpolation=cv2.INTER_CUBIC
    )

    return depth_resized

def write_depth(path, depth, bits=1):
    """Write depth map to pfm and png file.

    Args:
        path (str): filepath without extension
        depth (array): depth
    """
    write_pfm(path + ".pfm", depth.astype(np.float32))

    depth_min = depth.min()
    depth_max = depth.max()

    max_val = (2**(8*bits))-1

    if depth_max - depth_min > np.finfo("float").eps:
        out = max_val * (depth - depth_min) / (depth_max - depth_min)
    else:
        out = 0

    if bits == 1:
        cv2.imwrite(path + ".png", out.astype("uint8"))
    elif bits == 2:
        cv2.imwrite(path + ".png", out.astype("uint16"))

    return

max_stars_repo_name: sharypovandrey/MiDaS

max_issues_repo_name: sharypovandrey/MiDaS

max_forks_repo_name: sharypovandrey/MiDaS


text:
import pickle
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
from tqdm import tqdm

class Classifier(object):

    def __init__(self):
        super().__init__()

    def _fit(self, train_data, train_labels):
        raise NotImplementedError()

    def _score(self, test_data, test_labels):
        raise NotImplementedError()

    def fit(self, train_data, train_labels):
        print(f"    fit")
        start = time.time()
        self._fit(train_data, train_labels)
        end = time.time()
        runtime = int(end - start)
        print(f'    completed in {runtime // 3600}h {(runtime % 3600) // 60}m {(runtime % 60)}s')
        return self

    def score(self, test_data, test_labels):
        print(f"    score")
        start = time.time()
        score = self._score(test_data, test_labels)
        end = time.time()
        runtime = int(end - start)
        print(f'    -> {score}')
        print(f'    completed in {runtime // 3600}h {(runtime % 3600) // 60}m {(runtime % 60)}s')
        return score

    def save(self, output_path):
        output_dir = Path(output_path).parent.mkdir(exist_ok=True, parents=True)
        output_path.write_bytes(pickle.dumps(self))

    @staticmethod
    def load(in_path):
        instance = pickle.loads(Path(in_path).read_bytes())
        return instance

def read_dataset(datasets_dir, dataset_name, subset_to_size=None, flatten=True):
    print(f"[+] Read from {dataset_name}")
    dataset_dir = datasets_dir / dataset_name

    labels = np.load(dataset_dir.joinpath("labels.npy"))

    if not subset_to_size:
        # read full dataset
        imgs = []
        for idx in tqdm(range(labels.size), bar_format='    {l_bar}{bar:30}{r_bar}'):
            img_path = dataset_dir.joinpath(f'{idx:06}.npy')
            img = np.load(img_path)
            imgs.append(img)
        imgs = np.stack(imgs, 0)
        if flatten:
            imgs = imgs.reshape(labels.size, -1)
        return imgs, labels

    else:
        # subset dataset
        size_per_label = subset_to_size // np.unique(labels).size

        subset_data = []
        subset_labels = []

        sizes_per_label = defaultdict(int)
        p_bar = tqdm(total=subset_to_size, bar_format='    {l_bar}{bar:30}{r_bar}')
        for idx, label in enumerate(labels):

            if sizes_per_label[label] < size_per_label:
                img_path = dataset_dir.joinpath(f'{idx:06}.npy')
                subset_data.append(np.load(img_path))
                subset_labels.append(label)
                p_bar.update(1)
                sizes_per_label[label] += 1

            if len(subset_data) == subset_to_size:
                p_bar.close()
                break

        else:
            raise Exception('[!] ran out of images')

        subset_data = np.stack(subset_data, 0)
        if flatten:
            subset_data = subset_data.reshape(subset_to_size, -1)
        return subset_data, subset_labels

max_stars_repo_name: vwesselkamp/deepfake-fingerprint-atacks

max_issues_repo_name: vwesselkamp/deepfake-fingerprint-atacks

max_forks_repo_name: vwesselkamp/deepfake-fingerprint-atacks


text:
# I believe this code is based on a tutorial at:
# https://www.pyimagesearch.com/2015/06/01/home-surveillance-and-motion-detection-with-the-raspberry-pi-python-and-opencv/

from picamera.array import PiRGBArray
import datetime
import imutils
import os
import socket

import cv2
import numpy as np

class TempClient():

    def __init__(self, basePath="/tmp/"):
        self.basePath = basePath

    def upload(self, frame, timestamp):
        # write the image to temporary file
        t = os.path.join(self.basePath, "motion.png")
        cv2.imwrite(t, frame)
        print( f"Wrote image to {t}" )

def getImage():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        host ="127.0.0.1"
        port = 12346
        s.connect((host,port))
        s.sendall("image".encode())
# need to shutdown our write end of the socket so the server knows it's time to send the image
        s.shutdown(socket.SHUT_WR)
        sfile = s.makefile('rb')
        jpeg_byte_string = sfile.read()
        sfile.close()
        s.close()
        return jpeg_byte_string
    except Exception as inst:
        raise

if __name__ == "__main__":

    conf = {
        "min_upload_seconds": 3.0,
        "min_motion_frames": 4,
        "delta_thresh": 5,
        "resolution": [640, 480],
        "min_area": 5000
        }

    client = TempClient()

    avg = None
    lastUploaded = datetime.datetime.now()
    motionCounter = 0

# Read images from the camera daemon
    while True:
        jpeg_byte_string = getImage()
        cv2_img_flag=cv2.IMREAD_COLOR
        img_array = np.asarray(bytearray(jpeg_byte_string), dtype=np.uint8)
        frame = cv2.imdecode(img_array, cv2_img_flag)

        # the timestamp and occupied/unoccupied text
        timestamp = datetime.datetime.now()
        text = "Unoccupied"

        # resize the frame, convert it to grayscale, and blur it
        frame = imutils.resize(frame, width=500)
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        gray = cv2.GaussianBlur(gray, (21, 21), 0)

        # if the average frame is None, initialize it
        if avg is None:
            print("[INFO] starting background model...", flush=True)
            avg = gray.copy().astype("float")
            continue

        # accumulate the weighted average between the current frame and
        # previous frames, then compute the difference between the current
        # frame and running average
        cv2.accumulateWeighted(gray, avg, 0.5)
        frameDelta = cv2.absdiff(gray, cv2.convertScaleAbs(avg))

        # threshold the delta image, dilate the thresholded image to fill
        # in holes, then find contours on thresholded image
        thresh = cv2.threshold(frameDelta, conf["delta_thresh"], 255, cv2.THRESH_BINARY)[1]
        thresh = cv2.dilate(thresh, None, iterations=2)
        cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cnts = cnts[0] if imutils.is_cv2() else cnts[1]

        # loop over the contours
        for c in cnts:
            # if the contour is too small, ignore it
            if cv2.contourArea(c) < conf["min_area"]:
                continue

            # compute the bounding box for the contour, draw it on the frame,
            # and update the text
            (x, y, w, h) = cv2.boundingRect(c)
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            text = "Occupied"

        # draw the text and timestamp on the frame
        ts = timestamp.strftime("%A %d %B %Y %I:%M:%S%p")
        cv2.putText(frame, "Room Status: {}".format(text), (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
        cv2.putText(frame, ts, (10, frame.shape[0] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.35, (0, 0, 255), 1)

        # check to see if the room is occupied
        if text == "Occupied":
            # check to see if enough time has passed between uploads
            if (timestamp - lastUploaded).seconds >= conf["min_upload_seconds"]:
                # increment the motion counter
                motionCounter += 1

                # check to see if the number of frames with consistent motion is
                # high enough
                if motionCounter >= conf["min_motion_frames"]:
                    client.upload(frame,ts)

                    # update the last uploaded timestamp and reset the motion
                    # counter
                    lastUploaded = timestamp
                    motionCounter = 0

        # otherwise, the room is not occupied
        else:
            motionCounter = 0

max_stars_repo_name: Bleyddyn/RpiCameraDaemon

max_issues_repo_name: Bleyddyn/RpiCameraDaemon

max_forks_repo_name: Bleyddyn/RpiCameraDaemon


text:
#!/usr/bin/env python

import numpy as np
import openbabel as ob
import pybel as pb
import os, sys, copy
from aqml.cheminfo.rw.ctab import write_ctab
import scipy.spatial.distance as ssd
import tempfile as tpf
from aqml.cheminfo.core import *
from aqml.cheminfo.molecule.elements import Elements

T,F = True,False

elements_mmff94 = [ 'H','C','N','O','F','Si','P','S','Cl', 'Br','I' ]
zs_mmff94 = [ chemical_symbols.index(el) for el in elements_mmff94 ]

elements_c_uff = ['B','Al','Ga','In','Tl', \
                      'Si','Ge','Sn', \
                      'P', 'As','Sb','Bi', \
                           'Se','Te', ]
zsc_uff = [ chemical_symbols.index(el) for el in elements_mmff94 ]

element_cans = copy.copy(elements_c_uff)

nves = {}
nves = dict(list(zip(['B','Al','Ga','In','Tl'],[3]*5)))
nves.update( dict(list(zip(['Si','Ge','Sn'], [4]*3))) )
nves.update( dict(list(zip(['P','As','Sb','Bi'], [3]*4))) )
nves.update( dict(list(zip(['Se','Te'], [2]*2))) )

def get_coords(m):
    assert type(m) is ob.OBMol
    na = m.NumAtoms()
    coords = []
    for i in range(1,na+1): # atom idx starts from 1 in openbabel
        ai = m.GetAtom(i)
        coords.append( [ ai.GetX(), ai.GetY(), ai.GetZ() ] )
    return np.array(coords)

def get_bom(m):
    na = m.NumAtoms()
    bom = np.zeros((na,na),dtype=int)
    for i in range(na):
        ai = m.GetAtomById(i)
        for bond in ob.OBAtomBondIter(ai):
            ab, ae = bond.GetBeginAtom(), bond.GetEndAtom()
            ib1 = ab.GetIdx()-1; ie1 = ae.GetIdx()-1
            #ib2 = ab.GetId(); ie2 = ae.GetId()
            #print i, ' -- ', ib1,ie1, ib2,ie2
            bo = bond.GetBO()
            bom[ib1,ie1] = bo; bom[ie1,ib1] = bo
    return bom

def to_oechem_can(smi):#,hack_uff=F):
    """ Note that some SMILES such as C=N=O cannot be recognized
    correctly by openbabel. With OEChem, it is interpretted as
    C=[NH]=O, as was desired """

    from openeye import oechem

    if smi in element_cans:
        return '['+smi+'H%d]'%( nves[smi] )
    m = oechem.OEGraphMol()
    assert oechem.OESmilesToMol(m,smi)
    assert oechem.OEAddExplicitHydrogens(m)
    #atyps = []
    #for ai in m.GetAtoms():
    #    zi = ai.GetAtomicNum()
    #    vi = ai.GetValence()
    #    nhi = ai.GetTotalHCount()
    #    if zi in [7,15] and vi==5 and nhi>0:
    #        for aj in ai.GetAtomIter():
    #            zj = aj.GetAtomicNum()
    #            if zj==1: aj.SetIsotope(T)
    #    elif zi in [5] and vi==3 and nhi>0:
    #        for aj in ai.GetAtomIter():
    #            zj = aj.GetAtomicNum()
    #            if zj==1: aj.SetIsotope(T)
    #
    # "OESMILESFlag_ImpHCount" is indispensible!!
    # Otherwise, B=C won't be processed correctly by openbabel, i.e., somehow obabel
    # tries to add two H's (instead of one) to B. While things are ok with [BH]=C
    flavor = oechem.OESMILESFlag_Isotopes | oechem.OESMILESFlag_Canonical | oechem.OESMILESFlag_ImpHCount
    smi = oechem.OECreateSmiString(m, flavor)
    # OECreateIsoSmiString() # ISOMERIC .eq. Isotopes | AtomStereo | BondStereo | Canonical | AtomMaps | RGroups
    #fout = tpf.NamedTemporaryFile(dir='/tmp/').name + '.sdf'
    #ofs = oemolostream(fout)
    #ofs.SetFormat( OEFormat_SDF )
    #OEWriteMolecule(ofs, m)
    #_m = pb.readstring( 'sdf', open(fout,'r').read() )
    #m = _m.OBMol
    #os.system('rm %s'%fout)
    return smi # m

class Mol(object):

    def __init__(self, obj, hack_smi=F, hack_env=F, fmt=None, addh=T,  \
                 make3d=F, ichg=F, irad=F, ff='mmff94', isotope=F):
        """
        openbabel has problem intepreting smiles like O=N=C (in reality, it
        should be O=[NH]=C). We can hack this by calling `to_oechem_can() first.
        (if hack_smi=T)

        Another issue is the lack of uff parameters for envs like As=, As#,
        Sb=, et al. We can hack this by replacing these elements by the element
        of the same group in periodic table. (hack_env=T)
        """
        conv = ob.OBConversion()
        m = ob.OBMol()
        typ = type(obj)
        self.istat = T
        self.bom = None
        repl = []
        if typ is str:
            ismi = F
            if obj[-3:] in ['sdf','mol','pdb']:
                fmt = obj[-3:]
                #conversion from file to OBMol object
                conv.SetInFormat(fmt)
                conv.ReadFile(m, obj)
            else:
                ismi = T
                smi = obj
                if hack_smi:
                    smi = to_oechem_can(obj)
                if fmt is None: fmt = 'smi'
                conv.SetInFormat(fmt)
                conv.ReadString(m, smi)
                m.AddHydrogens()
        else:
            print('#ERROR: input not supported yet!')
            raise

        #print '#1'

        na = m.NumAtoms()
        self.na = na
        if self.bom is None:
            self.bom = get_bom(m)

        # set isotope to 0
        # otherwise, we'll encounter SMILES like 'C[2H]',
        # and error correspondently.
        # In deciding which atoms should be have spin multiplicity
        # assigned, hydrogen atoms which have an isotope specification
        # (D,T or even 1H) do not count. So SMILES N[2H] is NH2D (spin
        # multiplicity left at 0, so with a full content of implicit
        # hydrogens), whereas N[H] is NH (spin multiplicity=3). A
        # deuterated radical like NHD is represented by [NH][2H].
        if not isotope:
            for i in range(na):
                ai = m.GetAtomById(i); ai.SetIsotope(0)

        zs = []; chgs = []
        for i in range(1,na+1):
            ai = m.GetAtom(i)
            zi = ai.GetAtomicNum()
            zs.append( zi )
            chgs.append( ai.GetFormalCharge() )
        self.zs = np.array(zs, np.int)
        self.chgs = np.array(chgs,dtype=int)
        if not ismi:
            coords = get_coords(m)

        istat_rad = T
        if not irad:
            if sum(zs)%2 > 0:
                istat_rad = F

        istat_chg = T
        if not ichg:
            if sum(chgs) != 0:
                istat_chg = F

        if (not istat_rad) or (not istat_chg):
            self.istat = F
            print( '    [ERROR] Dected radical/charged mol' )
            return

        if ff in ['mmff94','mmff94s']:
            compl = np.setdiff1d(zs, zs_mmff94)
            if len(compl) > 0:
                print(( '    [ERROR] Detected symbols not supported by MMFF94: ', chemical_symbols[compl] ))
                self.istat = F
                return
        #print '#2'

        # get coarse geometry
        # Note: openbabel fails to process SMILES like "Br[pH]1n[pH2]n[p]n1", i.e.,
        #       after calling `make3d(), the program adds one more H atom to the
        #       last P atom. To circumvent this,
        repls = []
        nr = 0
        if ismi and make3d:
            # find atom whose type is not covered by UFF
            if hack_env:
                m2, repls = self.get_repl(m)
                nr = len(repls)
                if nr > 0: print(' *info: detected envs not supported by UFF/MMFF94')
            else:
                m2 = m
            m3 = pb.Molecule(m2)
            m3.make3D(forcefield=ff, steps=50)
            assert len(m3.atoms) == na, '#ERROR: make3d() added extra H atoms!'
            if hack_env and nr > 0:
                m = self.restore(m3.OBMol)
                print(' *info: now envs are restored')
            else:
                m = m3.OBMol

        M = pb.Molecule(m)
        coords = np.array([ ai.coords for ai in M.atoms ])
        #print '#3'
        # check coords
        if np.any(coords.astype(np.str) == 'nan'):
            print('    [ERROR] %s cannot process some envs, e.g., -As='%ff)
            self.istat = F
            return
        self.m = m
        self.M = M
        self.coords = coords
        self.repls = repls

    def get_repl(self,m):
        """ get envs to be replaced
        This has to be done in case the env is not supported in
        UFF/MMFF94. E.g., [As]# in [As]#C"""
        m2 = copy.copy(m)
        assert np.all(self.bom==self.bom.astype(np.int)), '#ERROR: BO != integer?'
        repls = []
        for ia in range(self.na):
            ai = m2.GetAtomById(ia)
            zi = self.zs[ia]
            _bosi = self.bom[ia]
            bosi = _bosi[_bosi>0]; bosi.sort()
            atyp = ''.join(['%d'%boi for boi in bosi])
            if zi in zsc_uff:
                if atyp != '1'*nves[ chemical_symbols[zi] ]:
                    zj = subs[zi]
                    ai.SetAtomicNum(zj)
                    repls.append( (ia,zi,zj) )
        return m2, repls

    def restore(self,m):
        m2 = copy.copy(m)
        return m2

    def is_overcrowded(self):
        ds = ssd.squareform( ssd.pdist(self.coords) )
        non_bonds = np.where( np.logical_and(self.bom==0, ds!=0.) )
        rcs = Elements().rcs[ self.zs ]
        dsmin = rcs[..., np.newaxis] + [rcs]
        return np.any(ds[non_bonds]<dsmin[non_bonds])

    def clone(self):
        m2 = pb.Molecule(self.m).clone
        return m2.OBMol

    def Atoms(self):
        """
        convert to ase object
        """
        import ase
        mu = pb.Molecule( self.m )
        m2 = ase.Atoms([])
        for ai in mu.atoms:
            m2.append(ase.Atom(ai.atomicnum, ai.coords))
        return m2

    def to_RDKit(self):
        """
        convert OBMol to RDKit_Mol
        """
        from rdkit import Chem

        c = ob.OBConversion()
        c.SetOutFormat('sdf')
        ctab = c.WriteString(self.m)
        mu = Chem.MolFromMolBlock( ctab, removeHs=False ) # plz keep H's
        return mu

    def write(self, fout):
        """ this will keep stereochem """
        c = ob.OBConversion()
        c.SetOutFormat(fout[-3:])
        iok = c.WriteFile(self.m, fout)

    def write_nostereo(self, fout):
        write_ctab(self.zs, self.chgs, self.bom, self.coords, sdf=fout)

    def optg_c2(self, smartss, iass, ff='MMFF94', step=500):
        """
        optg with constraints specified by `smartss

        vars
        ================
        smartss -- ['[#1]', 'C(=O)[OH]']
        iass    -- [[0,], [1,2,3] ]
        """
        # Define constraints
        c = ob.OBFFConstraints()

        iast = list(range(self.na))
        iasf = []
        for i, smarts_i in enumerate(smartss):
            #if 'H' in smarts_i: assert self.addh
            q = pb.Smarts( smarts_i )
            assert q.obsmarts.Match(self.m), '#ERROR: no match?'
            for match in q.obsmarts.GetMapList():
                idxs = [ self.m.GetAtom(ia).GetIdx() for ia in match ]
                for j in iass[i]:
                    iasf.append( idxs[j] )

        for ia in list( set(iast)^set(iasf) ):
            c.AddAtomConstraint(idxs[j])

        # Setup the force field with the constraints
        f = ob.OBForceField.FindForceField(ff)
        assert f.Setup(self.m, c), '#ERROR: ForceFiled setup failure [contains non-mmff94 elements]?'
        f.SetConstraints(c)
        #if optimizer in ['cg','CG','ConjugateGradients']:
        f.ConjugateGradients(step)
        #elif optimizer in ['sd','SD','SteepestDescent']:
        #    f.SteepestDescent(step)
        f.GetCoordinates(self.m)
        self.M = pb.Molecule(m)
        self.coords = get_coords(m)

    def optg_c(self, steps=60, ff="mmff94", optimizer='cg'):
        """
        Opt geometry by ff with constrained torsions
        ==============================================
        defaut ff: mmff94, mmff94s seems to be
                   problematic, sometimes it
                   leads to weird geometries
        steps: ff steps for constraint opt
        """
        m = self.m #copy.copy( self.m )
        # Define constraints
        c = ob.OBFFConstraints()
        #c.AddDistanceConstraint(1, 8, 10)        # Angstroms
        #c.AddAngleConstraint(1, 2, 3, 120.0)     # Degrees
        # constrain torsion only
        torsions = self.get_all_torsions()
        for torsion in torsions:
            i,j,k,l = torsion
            ang = torsions[torsion]
            c.AddTorsionConstraint(i,j,k,l, ang) # Degrees
        # Setup the force field with the constraints
        obff = ob.OBForceField.FindForceField(ff)
        obff.Setup(m, c)
        obff.SetConstraints(c)
        obff.ConjugateGradients(steps)
        obff.GetCoordinates(m)
        self.m = m
        self.M = pb.Molecule(m)
        self.coords = get_coords(m)

    def optg(self, cparam=[False,'slow'], ff="MMFF94", steps=2500):
        """ geoemtry optimization
        ================
        cparam: conformer parameters to be used
                cparam = [False] -- no conformer generation
                cparam = [True, algo, num_conf, N] -- gen confs and `N ff steps for each conf
        """
        # ff opt
        m = self.m
        obff = ob.OBForceField.FindForceField(ff)
        obff.Setup(m)
        #obff.SteepestDescent(1200, 1.0e-4) # a quick cleanup before searching
        obff.ConjugateGradients(600, 1.0e-3) # fast & coarse
        if cparam[0]:
            if cparam[1] in ['slow']:
                obff.FastRotorSearch(True)
            elif algo in ['slowest']:
                nc, n = cparam[1:]
                obff.WeightedRotorSearch(nc, n)
            else:
                print('unknow algo')
                raise
        obff.ConjugateGradients(steps, 1.0e-6) # fast & coarse
        #obff.SteepestDescent(500, 1.0e-6) # slower &
        obff.GetCoordinates(m)
        self.m = m
        self.M = pb.Molecule(m)
        self.coords = get_coords(m)

    def get_all_torsions(self):
        """
        enumerate Torsions in a molecule
        """

        m = self.m
        #m.DeleteHydrogens()

        torsionSmarts = '[!$(*#*)&!D1]~[!$(*#*)&!D1]'
        q = pb.Smarts( torsionSmarts )
        iok = q.obsmarts.Match(m)

        torsions = {}; tts = []
        for match in q.obsmarts.GetMapList():
            ia2,ia3 = match
            a2, a3 = m.GetAtom(ia2), m.GetAtom(ia3)
            hyb2, hyb3 = a2.GetHyb(), a3.GetHyb()

            if (hyb2 not in [2,3]) or (hyb3 not in [2,3]):
                # no formal torsion here !
                continue

            for a1 in ob.OBAtomAtomIter( a2 ):
                ia1 = a1.GetIdx()
                if ia1 == ia3:
                    continue

                for a4 in ob.OBAtomAtomIter( a3 ):
                    ia4 = a4.GetIdx()
                    if ia4 in [ia2, ia1]: # when ia4 == ia1, it means 3-membered ring
                        continue
                    else:
                        torsion_1 = (ia1,ia2,ia3,ia4)
                        torsion_2 = (ia4,ia3,ia2,ia1)
                        if torsion_1 not in tts:
                            torsions[ torsion_1 ] = m.GetTorsion(ia1,ia2,ia3,ia4)
                            tts.append(torsion_1); tts.append(torsion_2)
        return torsions

    def get_bom(self):
        """
        get connectivity table
        """
        bom = np.zeros((self.na,self.na), np.int)
        for i in range(1,self.na+1):
            ai = self.m.GetAtom(i)
            for bond in ob.OBAtomBondIter(ai):
                ia1 = bond.GetBeginAtomIdx()-1; ia2 = bond.GetEndAtomIdx()-1
                bo = bond.GetBO()
                bom[ia1,ia2] = bo; bom[ia2,ia1] = bo
        return bom

if __name__ == '__main__':

    import os, sys

    # Note that file in `fs represent
    fs = sys.argv[1:]
    for f in fs:
        #print ' -- f = ', f
        assert f[-3:] == 'sdf'
        si = f.split('/'); nsi = len(si)
        if nsi == 2:
            f2 = '%s/raw/%s_raw.sdf'%(si[0], si[-1][:-4])
        elif nsi == 1:
            f2 = 'raw/%s_raw.sdf'%( si[-1][:-4] )
        else:
            print('#ERROR:')
            raise

        iok = True
        try:
            obj = Mol( f )
            g = perceive_g(obj.zs,obj.coords) # graph from geometry
            bom = obj.get_bom() # bond orders in SDF file
            g0 = (bom > 0).astype(np.int)
            if not np.all(g == g0):
                iok = False
        except:
            print(' -- Parsing failed for %s'%f)
            iok = False # fails to retrieve coordinates!

        if not iok:
            obj2 = Mol( f2 )
            obj2.optg_c(ff="MMFF94", steps=120)
            obj2.write( f )

reason: import numpy,import scipy


text:
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Sun Oct 21 20:16:53 2018

@author: guille
"""
import tensorflow as tf
import numpy as np
import gym

# Enviroment parameters
env = gym.make('LunarLander-v2')
STATE_SIZE=env.observation_space.shape[0]
ACTION_SPACE=env.action_space.n
ACTION_SIZE=1

# Evaluator parameters:
EPISODES_TO_EVALUATE=20
EPOCHS_PER_EPISODE_TO_EVALUATE=500
VISUALIZE_EVALUATOR=True
SAVE_PATH="tmp/lander_weights"

# Program start
tf.reset_default_graph()
sess=tf.Session()

# Recreate network
state_input_tensor=tf.placeholder(tf.float32, shape=(None,STATE_SIZE),name="state_input_tensor")

with tf.variable_scope("main_network"):
    h1=tf.layers.dense(state_input_tensor,256,activation=tf.nn.relu,name="layer1")
    h2=tf.layers.dense(h1,256,activation=tf.nn.relu,name="layer2")
    h3=tf.layers.dense(h2,256,activation=tf.nn.relu,name="layer3")
    output=tf.layers.dense(h3,ACTION_SPACE,activation=None,name="output_layer")

saver = tf.train.Saver()

tf.get_default_graph().finalize()

# Restore variables from disk.
saver.restore(sess,SAVE_PATH)

# Evaluate model
for e in range(EPISODES_TO_EVALUATE):
    done=False
    state=env.reset()
    while not done:
        if VISUALIZE_EVALUATOR:
			env.render()
        Qs=sess.run(output,feed_dict={state_input_tensor:[state]})
        action=np.argmax(Qs)
        state,reward,done, _=env.step(action)
env.render(close=True)


text:
import pandas as pd
import numpy as np

from pandas.core.nanops import nanmean

def get_seasonality(row, window_size):
    """
    Returns seasonality component of a time series
    :param row: pandas series containing the time series to extract seasonality from
    :param window_size: length of moving window
    :return: pandas series of seasonality component
    """
    averages = np.array([nanmean(row[i::window_size], axis=0) for i in range (window_size)])
    averages -= np.mean(averages, axis=0)
    season = np.tile(averages.T, len(row)//window_size+1).T[:len(row)]

    return pd.Series(season)

def get_trend(row, window_size, center=True):
    """
    Returns trend component of a time series
    :param row: pandas series containing the time series to extract seasonality from
    :param window_size: length of moving window
    :param center:
    :return:
    """
    trend = row.rolling(window_size, center=center).mean()

    return trend

max_stars_repo_name: DavidMayett/timeseries

max_issues_repo_name: DavidMayett/timeseries

max_forks_repo_name: DavidMayett/timeseries