# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import abc
import numpy as np
from warnings import warn
from sklearn.base import clone
from sklearn.model_selection import GroupKFold
from scipy.stats import norm
from sklearn.linear_model import (ElasticNetCV, LassoCV, LogisticRegressionCV)
from ...sklearn_extensions.linear_model import (StatsModelsLinearRegression, WeightedLassoCVWrapper)
from ...sklearn_extensions.model_selection import WeightedStratifiedKFold
from ...dml.dml import _FirstStageWrapper, _FinalWrapper
from ..._cate_estimator import TreatmentExpansionMixin, LinearModelFinalCateEstimatorMixin
from ..._ortho_learner import _OrthoLearner
from ...utilities import (_deprecate_positional, add_intercept,
                          broadcast_unit_treatments, check_high_dimensional,
                          cross_product, deprecated, fit_with_groups,
                          hstack, inverse_onehot, ndim, reshape,
                          reshape_treatmentwise_effects, shape, transpose,
                          get_feature_names_or_default, check_input_arrays,
                          filter_none_kwargs)

def _get_groups_period_filter(groups, n_periods):
    group_counts = {}
    group_period_filter = {i: [] for i in range(n_periods)}
    for i, g in enumerate(groups):
        if g not in group_counts:
            group_counts[g] = 0
        group_period_filter[group_counts[g]].append(i)
        group_counts[g] += 1
    return group_period_filter

class _DynamicModelNuisance:
    """
    Nuisance model fits the model_y and model_t at fit time and at predict time
    calculates the residual Y and residual T based on the fitted models and returns
    the residuals as two nuisance parameters.
    """

    def __init__(self, model_y, model_t, n_periods):
        self._model_y = model_y
        self._model_t = model_t
        self.n_periods = n_periods

    def fit(self, Y, T, X=None, W=None, sample_weight=None, groups=None):
        """Fit a series of nuisance models for each period or period pairs."""
        assert Y.shape[0] % self.n_periods == 0, \
            "Length of training data should be an integer multiple of time periods."
        period_filters = _get_groups_period_filter(groups, self.n_periods)
        self._model_y_trained = {}
        self._model_t_trained = {j: {} for j in np.arange(self.n_periods)}
        for t in np.arange(self.n_periods):
            self._model_y_trained[t] = clone(self._model_y, safe=False).fit(
                self._index_or_None(X, period_filters[t]),
                self._index_or_None(
                    W, period_filters[t]),
                Y[period_filters[self.n_periods - 1]])
            for j in np.arange(t, self.n_periods):
                self._model_t_trained[j][t] = clone(self._model_t, safe=False).fit(
                    self._index_or_None(X, period_filters[t]),
                    self._index_or_None(W, period_filters[t]),
                    T[period_filters[j]])
        return self

    def predict(self, Y, T, X=None, W=None, sample_weight=None, groups=None):
        """Calculate nuisances for each period or period pairs.

        Returns
        -------
        Y_res : (n, d_y) matrix or vector of length n
            Y residuals for each period in panel format.
            This shape is required for _OrthoLearner's crossfitting.
        T_res : (n, d_t, n_periods) matrix
            T residuals for pairs of periods (t, j), where the data is in panel format for t
            and in index form for j. For example, the residuals for (t, j) can be retrieved via
            T_res[np.arange(n) % n_periods == t, ..., j]. For t < j, the entries of this
            matrix are np.nan.
            This shape is required for _OrthoLearner's crossfitting.
        """
        assert Y.shape[0] % self.n_periods == 0, \
            "Length of training data should be an integer multiple of time periods."
        period_filters = _get_groups_period_filter(groups, self.n_periods)
        Y_res = np.full(Y.shape, np.nan)
        T_res = np.full(T.shape + (self.n_periods, ), np.nan)
        shape_formatter = self._get_shape_formatter(X, W)
        for t in np.arange(self.n_periods):
            Y_slice = Y[period_filters[self.n_periods - 1]]
            Y_pred = self._model_y_trained[t].predict(
                self._index_or_None(X, period_filters[t]),
                self._index_or_None(W, period_filters[t]))
            Y_res[period_filters[t]] = Y_slice\
                - shape_formatter(Y_slice, Y_pred)
            for j in np.arange(t, self.n_periods):
                T_slice = T[period_filters[j]]
                T_pred = self._model_t_trained[j][t].predict(
                    self._index_or_None(X, period_filters[t]),
                    self._index_or_None(W, period_filters[t]))
                T_res[period_filters[j], ..., t] = T_slice\
                    - shape_formatter(T_slice, T_pred)
        return Y_res, T_res

    def score(self, Y, T, X=None, W=None, sample_weight=None, groups=None):
        assert Y.shape[0] % self.n_periods == 0, \
            "Length of training data should be an integer multiple of time periods."
        period_filters = _get_groups_period_filter(groups, self.n_periods)
        if hasattr(self._model_y, 'score'):
            Y_score = np.full((self.n_periods, ), np.nan)
            for t in np.arange(self.n_periods):
                Y_score[t] = self._model_y_trained[t].score(
                    self._index_or_None(X, period_filters[t]),
                    self._index_or_None(W, period_filters[t]),
                    Y[period_filters[self.n_periods - 1]])
        else:
            Y_score = None
        if hasattr(self._model_t, 'score'):
            T_score = np.full((self.n_periods, self.n_periods), np.nan)
            for t in np.arange(self.n_periods):
                for j in np.arange(t, self.n_periods):
                    T_score[j][t] = self._model_t_trained[j][t].score(
                        self._index_or_None(X, period_filters[t]),
                        self._index_or_None(W, period_filters[t]),
                        T[period_filters[j]])
        else:
            T_score = None
        return Y_score, T_score

    def _get_shape_formatter(self, X, W):
        if (X is None) and (W is None):
            return lambda x, x_pred: np.tile(x_pred.reshape(1, -1), (x.shape[0], 1)).reshape(x.shape)
        return lambda x, x_pred: x_pred.reshape(x.shape)

    def _index_or_None(self, X, filter_idx):
        return None if X is None else X[filter_idx]

class _DynamicModelFinal:
    """
    Final model at fit time, fits a residual on residual regression with a heterogeneous coefficient
    that depends on X, i.e.

        .. math ::
            Y - E[Y | X, W] = \\theta(X) \\cdot (T - E[T | X, W]) + \\epsilon

    and at predict time returns :math:`\\theta(X)`. The score method returns the MSE of this final
    residual on residual regression.
    Assumes model final is parametric with no intercept.
    """
    # TODO: update docs

    def __init__(self, model_final, n_periods):
        self._model_final = model_final
        self.n_periods = n_periods
        self._model_final_trained = {k: clone(self._model_final, safe=False) for k in np.arange(n_periods)}

    def fit(self, Y, T, X=None, W=None, Z=None, nuisances=None, sample_weight=None, sample_var=None, groups=None):
        # NOTE: sample weight, sample var are not passed in
        period_filters = _get_groups_period_filter(groups, self.n_periods)
        Y_res, T_res = nuisances
        self._d_y = Y.shape[1:]
        for t in np.arange(self.n_periods - 1, -1, -1):
            Y_adj = Y_res[period_filters[t]].copy()
            if t < self.n_periods - 1:
                Y_adj -= np.sum(
                    [self._model_final_trained[j].predict_with_res(
                        X[period_filters[0]] if X is not None else None,
                        T_res[period_filters[j], ..., t]
                    ) for j in np.arange(t + 1, self.n_periods)], axis=0)
            self._model_final_trained[t].fit(
                X[period_filters[0]] if X is not None else None, T[period_filters[t]],
                T_res[period_filters[t], ..., t], Y_adj)

        return self

    def predict(self, X=None):
        """
        Return shape: m x dy x (p*dt)
        """
        d_t_tuple = self._model_final_trained[0]._d_t
        d_t = d_t_tuple[0] if d_t_tuple else 1
        x_dy_shape = (X.shape[0] if X is not None else 1, ) + \
            self._model_final_trained[0]._d_y
        preds = np.zeros(
            x_dy_shape +
            (self.n_periods * d_t, )
        )
        for t in range(self.n_periods):
            preds[..., t * d_t: (t + 1) * d_t] = \
                self._model_final_trained[t].predict(X).reshape(
                x_dy_shape + (d_t, )
            )
        return preds

    def score(self, Y, T, X=None, W=None, Z=None, nuisances=None, sample_weight=None, sample_var=None, groups=None):
        assert Y.shape[0] % self.n_periods == 0, \
            "Length of training data should be an integer multiple of time periods."
        Y_res, T_res = nuisances
        scores = np.full((self.n_periods, ), np.nan)
        period_filters = _get_groups_period_filter(groups, self.n_periods)
        for t in np.arange(self.n_periods - 1, -1, -1):
            Y_adj = Y_res[period_filters[t]].copy()
            if t < self.n_periods - 1:
                Y_adj -= np.sum(
                    [self._model_final_trained[j].predict_with_res(
                        X[period_filters[0]] if X is not None else None,
                        T_res[period_filters[j], ..., t]
                    ) for j in np.arange(t + 1, self.n_periods)], axis=0)
            Y_adj_pred = self._model_final_trained[t].predict_with_res(
                X[period_filters[0]] if X is not None else None,
                T_res[period_filters[t], ..., t])
            if sample_weight is not None:
                scores[t] = np.mean(np.average((Y_adj - Y_adj_pred)**2, weights=sample_weight, axis=0))
            else:
                scores[t] = np.mean((Y_adj - Y_adj_pred) ** 2)
        return scores

class _LinearDynamicModelFinal(_DynamicModelFinal):
    """Wrapper for the DynamicModelFinal with StatsModelsLinearRegression final model.

    The final model is a linear model with (d_t*n_periods) coefficients.
    This model is defined after the coefficients and covariance are calculated.
    """

    def __init__(self, model_final, n_periods):
        super().__init__(model_final, n_periods)
        self.model_final_ = StatsModelsLinearRegression(fit_intercept=False)

    def fit(self, Y, T, X=None, W=None, Z=None, nuisances=None, sample_weight=None, sample_var=None, groups=None):
        super().fit(Y, T, X=X, W=W, Z=Z, nuisances=nuisances,
                    sample_weight=sample_weight, sample_var=sample_var, groups=groups)
        # Compose final model
        cov = self._get_cov(nuisances, X, groups)
        coef = self._get_coef_()
        self.model_final_._n_out = self._d_y[0] if self._d_y else 0
        self.model_final_._param_var = cov / (Y.shape[0] / self.n_periods)
        self.model_final_._param = coef.T if self.model_final_._n_out else coef

    def _get_coef_(self):
        period_coefs = np.array([self._model_final_trained[t]._model.coef_ for t in range(self.n_periods)])
        if self._d_y:
            return np.array([
                np.array([period_coefs[k, i, :] for k in range(self.n_periods)]).flatten()
                for i in range(self._d_y[0])
            ])
        return period_coefs.flatten()

    def _get_cov(self, nuisances, X, groups):
        if self._d_y:
            return np.array(
                [self._fit_single_output_cov((nuisances[0][:, i], nuisances[1]), X, i, groups)
                 for i in range(self._d_y[0])]
            )
        return self._fit_single_output_cov(nuisances, X, -1, groups)

    def _fit_single_output_cov(self, nuisances, X, y_index, groups):
        """ Calculates the covariance (n_periods*n_treatments)
            x (n_periods*n_treatments) matrix for a single outcome.
        """
        Y_res, T_res = nuisances
        # Calculate auxiliary quantities
        period_filters = _get_groups_period_filter(groups, self.n_periods)
        # X ⨂ T_res
        XT_res = np.array([
            [
                self._model_final_trained[0]._combine(
                    X[period_filters[0]] if X is not None else None,
                    T_res[period_filters[t], ..., j],
                    fitting=False
                )
                for j in range(self.n_periods)
            ]
            for t in range(self.n_periods)
        ])
        d_xt = XT_res.shape[-1]
        # sum(model_final.predict(X, T_res))
        Y_diff = np.array([
            np.sum([
                self._model_final_trained[j].predict_with_res(
                    X[period_filters[0]] if X is not None else None,
                    T_res[period_filters[j], ..., t]
                ) for j in np.arange(t, self.n_periods)],
                axis=0
            )
            for t in np.arange(self.n_periods)
        ])
        J = np.zeros((self.n_periods * d_xt,
                      self.n_periods * d_xt))
        Sigma = np.zeros((self.n_periods * d_xt,
                          self.n_periods * d_xt))
        for t in np.arange(self.n_periods):
            res_epsilon_t = (Y_res[period_filters[t]] -
                             (Y_diff[t][:, y_index] if y_index >= 0 else Y_diff[t])
                             ).reshape(-1, 1, 1)
            resT_t = XT_res[t][t]
            for j in np.arange(self.n_periods):
                # Calculating the (t, j) block entry (of size n_treatments x n_treatments) of matrix Sigma
                res_epsilon_j = (Y_res[period_filters[j]] -
                                 (Y_diff[j][:, y_index] if y_index >= 0 else Y_diff[j])
                                 ).reshape(-1, 1, 1)
                resT_j = XT_res[j][j]
                cov_resT_tj = resT_t.reshape(-1, d_xt, 1) @ resT_j.reshape(-1, 1, d_xt)
                sigma_tj = np.mean((res_epsilon_t * res_epsilon_j) * cov_resT_tj, axis=0)
                Sigma[t * d_xt:(t + 1) * d_xt,
                      j * d_xt:(j + 1) * d_xt] = sigma_tj
                if j >= t:
                    # Calculating the (t, j) block entry (of size n_treatments x n_treatments) of matrix J
                    m_tj = np.mean(
                        XT_res[j][t].reshape(-1, d_xt, 1) @ resT_t.reshape(-1, 1, d_xt),
                        axis=0)
                    J[t * d_xt:(t + 1) * d_xt,
                      j * d_xt:(j + 1) * d_xt] = m_tj
        return np.linalg.inv(J) @ Sigma @ np.linalg.inv(J).T

class _DynamicFinalWrapper(_FinalWrapper):

    def predict_with_res(self, X, T_res):
        fts = self._combine(X, T_res, fitting=False)
        prediction = self._model.predict(fts)
        if self._intercept is not None:
            prediction -= self._intercept
        return reshape(prediction, (prediction.shape[0],) + self._d_y)

class DynamicDML(LinearModelFinalCateEstimatorMixin, _OrthoLearner):
    """CATE estimator for dynamic treatment effect estimation.

    This estimator is an extension of the Double ML approach for treatments assigned sequentially
    over time periods.

    The estimator is a special case of an :class:`_OrthoLearner` estimator, so it follows the two
    stage process, where a set of nuisance functions are estimated in the first stage in a crossfitting
    manner and a final stage estimates the CATE model. See the documentation of
    :class:`._OrthoLearner` for a description of this two stage process.

    Parameters
    ----------
    model_y: estimator or 'auto', optional (default is 'auto')
        The estimator for fitting the response to the features. Must implement
        `fit` and `predict` methods.
        If 'auto' :class:`.WeightedLassoCV`/:class:`.WeightedMultiTaskLassoCV` will be chosen.

    model_t: estimator or 'auto', optional (default is 'auto')
        The estimator for fitting the treatment to the features.
        If estimator, it must implement `fit` and `predict` methods;
        If 'auto', :class:`~sklearn.linear_model.LogisticRegressionCV` will be applied for discrete treatment,
        and :class:`.WeightedLassoCV`/:class:`.WeightedMultiTaskLassoCV`
        will be applied for continuous treatment.

    featurizer : :term:`transformer`, optional, default None
        Must support fit_transform and transform. Used to create composite features in the final CATE regression.
        It is ignored if X is None. The final CATE will be trained on the outcome of featurizer.fit_transform(X).
        If featurizer=None, then CATE is trained on X.

    fit_cate_intercept : bool, optional, default True
        Whether the linear CATE model should have a constant term.

    linear_first_stages: bool
        Whether the first stage models are linear (in which case we will expand the features passed to
        `model_y` accordingly)

    discrete_treatment: bool, optional (default is ``False``)
        Whether the treatment values should be treated as categorical, rather than continuous, quantities

    categories: 'auto' or list, default 'auto'
        The categories to use when encoding discrete treatments (or 'auto' to use the unique sorted values).
        The first category will be treated as the control treatment.

    cv: int, cross-validation generator or an iterable, optional (Default=2)
        Determines the cross-validation splitting strategy.
        Possible inputs for cv are:

        - None, to use the default 3-fold cross-validation,
        - integer, to specify the number of folds.
        - :term:`CV splitter`
        - An iterable yielding (train, test) splits as arrays of indices.
          Iterables should make sure a group belongs to a single split.

        For integer/None inputs, :class:`~sklearn.model_selection.GroupKFold` is used

        Unless an iterable is used, we call `split(X, T, groups)` to generate the splits.

    mc_iters: int, optional (default=None)
        The number of times to rerun the first stage models to reduce the variance of the nuisances.

    mc_agg: {'mean', 'median'}, optional (default='mean')
        How to aggregate the nuisance value for each sample across the `mc_iters` monte carlo iterations of
        cross-fitting.

    random_state: int, :class:`~numpy.random.mtrand.RandomState` instance or None, optional (default=None)
        If int, random_state is the seed used by the random number generator;
        If :class:`~numpy.random.mtrand.RandomState` instance, random_state is the random number generator;
        If None, the random number generator is the :class:`~numpy.random.mtrand.RandomState` instance used
        by :mod:`np.random<numpy.random>`.

    Examples
    --------
    A simple example with default models:

    .. testcode::
        :hide:

        import numpy as np
        np.set_printoptions(suppress=True)

    .. testcode::

        from econml.dynamic.dml import DynamicDML

        np.random.seed(123)

        n_panels = 100  # number of panels
        n_periods = 3  # number of time periods per panel
        n = n_panels * n_periods
        groups = np.repeat(a=np.arange(n_panels), repeats=n_periods, axis=0)
        X = np.random.normal(size=(n, 1))
        T = np.random.normal(size=(n, 2))
        y = np.random.normal(size=(n, ))
        est = DynamicDML()
        est.fit(y, T, X=X, W=None, groups=groups, inference="auto")

    >>> est.const_marginal_effect(X[:2])
    array([[-0.336..., -0.048..., -0.061...,  0.042..., -0.204...,
         0.00667271],
        [-0.101...,  0.433...,  0.054..., -0.217..., -0.101...,
         -0.159...]])
    >>> est.effect(X[:2], T0=0, T1=1)
    array([-0.601..., -0.091...])
    >>> est.effect(X[:2], T0=np.zeros((2, n_periods*T.shape[1])), T1=np.ones((2, n_periods*T.shape[1])))
    array([-0.601..., -0.091...])
    >>> est.coef_
    array([[ 0.112...],
       [ 0.231...],
       [ 0.055...],
       [-0.125...],
       [ 0.049...],
       [-0.079...]])
    >>> est.coef__interval()
    (array([[-0.063...],
           [-0.009...],
           [-0.114...],
           [-0.413...],
           [-0.117...],
           [-0.262...]]), array([[0.289...],
           [0.471...],
           [0.225...],
           [0.163...],
           [0.216...],
           [0.103...]]))
    """

    def __init__(self, *,
                 model_y='auto', model_t='auto',
                 featurizer=None,
                 fit_cate_intercept=True,
                 linear_first_stages=False,
                 discrete_treatment=False,
                 categories='auto',
                 cv=2,
                 mc_iters=None,
                 mc_agg='mean',
                 random_state=None):
        self.fit_cate_intercept = fit_cate_intercept
        self.linear_first_stages = linear_first_stages
        self.featurizer = clone(featurizer, safe=False)
        self.model_y = clone(model_y, safe=False)
        self.model_t = clone(model_t, safe=False)
        super().__init__(discrete_treatment=discrete_treatment,
                         discrete_instrument=False,
                         categories=categories,
                         cv=GroupKFold(cv) if isinstance(cv, int) else cv,
                         mc_iters=mc_iters,
                         mc_agg=mc_agg,
                         random_state=random_state)

    def _gen_featurizer(self):
        return clone(self.featurizer, safe=False)

    def _gen_model_y(self):
        if self.model_y == 'auto':
            model_y = WeightedLassoCVWrapper(random_state=self.random_state)
        else:
            model_y = clone(self.model_y, safe=False)
        return _FirstStageWrapper(model_y, True, self._gen_featurizer(),
                                  self.linear_first_stages, self.discrete_treatment)

    def _gen_model_t(self):
        if self.model_t == 'auto':
            if self.discrete_treatment:
                model_t = LogisticRegressionCV(cv=WeightedStratifiedKFold(random_state=self.random_state),
                                               random_state=self.random_state)
            else:
                model_t = WeightedLassoCVWrapper(random_state=self.random_state)
        else:
            model_t = clone(self.model_t, safe=False)
        return _FirstStageWrapper(model_t, False, self._gen_featurizer(),
                                  self.linear_first_stages, self.discrete_treatment)

    def _gen_model_final(self):
        return StatsModelsLinearRegression(fit_intercept=False)

    def _gen_ortho_learner_model_nuisance(self, n_periods):
        return _DynamicModelNuisance(
            model_t=self._gen_model_t(),
            model_y=self._gen_model_y(),
            n_periods=n_periods)

    def _gen_ortho_learner_model_final(self, n_periods):
        wrapped_final_model = _DynamicFinalWrapper(
            StatsModelsLinearRegression(fit_intercept=False),
            fit_cate_intercept=self.fit_cate_intercept,
            featurizer=self.featurizer,
            use_weight_trick=False)
        return _LinearDynamicModelFinal(wrapped_final_model, n_periods=n_periods)

    def _prefit(self, Y, T, *args, groups=None, only_final=False, **kwargs):
        u_periods = np.unique(np.unique(groups, return_counts=True)[1])
        if len(u_periods) > 1:
            raise AttributeError(
                "Imbalanced panel. Method currently expects only panels with equal number of periods. Pad your data")
        self._n_periods = u_periods[0]
        # generate an instance of the final model
        self._ortho_learner_model_final = self._gen_ortho_learner_model_final(self._n_periods)
        if not only_final:
            # generate an instance of the nuisance model
            self._ortho_learner_model_nuisance = self._gen_ortho_learner_model_nuisance(self._n_periods)
        TreatmentExpansionMixin._prefit(self, Y, T, *args, **kwargs)

    def _postfit(self, Y, T, *args, **kwargs):
        super()._postfit(Y, T, *args, **kwargs)
        # Set _d_t to effective number of treatments
        self._d_t = (self._n_periods * self._d_t[0], ) if self._d_t else (self._n_periods, )

    def _strata(self, Y, T, X=None, W=None, Z=None,
                sample_weight=None, sample_var=None, groups=None,
                cache_values=False, only_final=False, check_input=True):
        # Required for bootstrap inference
        return groups

    def fit(self, Y, T, *, X=None, W=None, sample_weight=None, sample_var=None, groups,
            cache_values=False, inference='auto'):
        """Estimate the counterfactual model from data, i.e. estimates function :math:`\\theta(\\cdot)`.

        The input data must contain groups with the same size corresponding to the number
        of time periods the treatments were assigned over.

        The data should be preferably in panel format, with groups clustered together.
        If group members do not appear together, the following is assumed:

        * the first instance of a group in the dataset is assumed to correspond to the first period of that group
        * the second instance of a group in the dataset is assumed to correspond to the
          second period of that group

        ...etc.

        Only the value of the features X at the first period of each unit are used for
        heterogeneity. The value of X in subseuqnet periods is used as a time-varying control
        but not for heterogeneity.

        Parameters
        ----------
        Y: (n, d_y) matrix or vector of length n
            Outcomes for each sample (required: n = n_groups * n_periods)
        T: (n, d_t) matrix or vector of length n
            Treatments for each sample (required: n = n_groups * n_periods)
        X: optional(n, d_x) matrix or None (Default=None)
            Features for each sample (Required: n = n_groups * n_periods). Only first
            period features from each unit are used for heterogeneity, the rest are
            used as time-varying controls together with W
        W: optional(n, d_w) matrix or None (Default=None)
            Controls for each sample (Required: n = n_groups * n_periods)
        sample_weight: optional(n,) vector or None (Default=None)
            Weights for each samples
        sample_var: optional(n,) vector or None (Default=None)
            Sample variance for each sample
        groups: (n,) vector, required
            All rows corresponding to the same group will be kept together during splitting.
            If groups is not None, the `cv` argument passed to this class's initializer
            must support a 'groups' argument to its split method.
        cache_values: bool, default False
            Whether to cache inputs and first stage results, which will allow refitting a different final model
        inference: string,:class:`.Inference` instance, or None
            Method for performing inference.  This estimator supports 'bootstrap'
            (or an instance of :class:`.BootstrapInference`) and 'auto'
            (or an instance of :class:`.LinearModelFinalInference`).

        Returns
        -------
        self: DynamicDML instance
        """
        if sample_weight is not None or sample_var is not None:
            warn("This CATE estimator does not yet support sample weights and sample variance. "
                 "These inputs will be ignored during fitting.",
                 UserWarning)
        return super().fit(Y, T, X=X, W=W,
                           sample_weight=None, sample_var=None, groups=groups,
                           cache_values=cache_values,
                           inference=inference)

    def score(self, Y, T, X=None, W=None, sample_weight=None, *, groups):
        """
        Score the fitted CATE model on a new data set. Generates nuisance parameters
        for the new data set based on the fitted residual nuisance models created at fit time.
        It uses the mean prediction of the models fitted by the different crossfit folds.
        Then calculates the MSE of the final residual Y on residual T regression.

        If model_final does not have a score method, then it raises an :exc:`.AttributeError`

        Parameters
        ----------
        Y: (n, d_y) matrix or vector of length n
            Outcomes for each sample (required: n = n_groups * n_periods)
        T: (n, d_t) matrix or vector of length n
            Treatments for each sample (required: n = n_groups * n_periods)
        X: optional(n, d_x) matrix or None (Default=None)
            Features for each sample (Required: n = n_groups * n_periods)
        W: optional(n, d_w) matrix or None (Default=None)
            Controls for each sample (Required: n = n_groups * n_periods)
        groups: (n,) vector, required
            All rows corresponding to the same group will be kept together during splitting.

        Returns
        -------
        score: float
            The MSE of the final CATE model on the new data.
        """
        if not hasattr(self._ortho_learner_model_final, 'score'):
            raise AttributeError("Final model does not have a score method!")
        Y, T, X, W, groups = check_input_arrays(Y, T, X, W, groups)
        self._check_fitted_dims(X)
        X, T = super()._expand_treatments(X, T)
        n_iters = len(self._models_nuisance)
        n_splits = len(self._models_nuisance[0])

        # for each mc iteration
        for i, models_nuisances in enumerate(self._models_nuisance):
            # for each model under cross fit setting
            for j, mdl in enumerate(models_nuisances):
                nuisance_temp = mdl.predict(Y, T, **filter_none_kwargs(X=X, W=W, groups=groups))
                if not isinstance(nuisance_temp, tuple):
                    nuisance_temp = (nuisance_temp,)

                if i == 0 and j == 0:
                    nuisances = [np.zeros((n_iters * n_splits,) + nuis.shape) for nuis in nuisance_temp]

                for it, nuis in enumerate(nuisance_temp):
                    nuisances[it][i * n_iters + j] = nuis

        for it in range(len(nuisances)):
            nuisances[it] = np.mean(nuisances[it], axis=0)
        return self._ortho_learner_model_final.score(Y, T, nuisances=nuisances,
                                                     **filter_none_kwargs(X=X, W=W,
                                                                          sample_weight=sample_weight, groups=groups))

    def cate_treatment_names(self, treatment_names=None):
        """
        Get treatment names for each time period.

        If the treatment is discrete, it will return expanded treatment names.

        Parameters
        ----------
        treatment_names: list of strings of length T.shape[1] or None
            The names of the treatments. If None and the T passed to fit was a dataframe,
            it defaults to the column names from the dataframe.

        Returns
        -------
        out_treatment_names: list of strings
            Returns (possibly expanded) treatment names.
        """
        slice_treatment_names = super().cate_treatment_names(treatment_names)
        treatment_names_out = []
        for k in range(self._n_periods):
            treatment_names_out += [f"({t})$_{k}$" for t in slice_treatment_names]
        return treatment_names_out

    def cate_feature_names(self, feature_names=None):
        """
        Get the output feature names.

        Parameters
        ----------
        feature_names: list of strings of length X.shape[1] or None
            The names of the input features. If None and X is a dataframe, it defaults to the column names
            from the dataframe.

        Returns
        -------
        out_feature_names: list of strings or None
            The names of the output features :math:`\\phi(X)`, i.e. the features with respect to which the
            final constant marginal CATE model is linear. It is the names of the features that are associated
            with each entry of the :meth:`coef_` parameter. Not available when the featurizer is not None and
            does not have a method: `get_feature_names(feature_names)`. Otherwise None is returned.
        """
        if self._d_x is None:
            # Handles the corner case when X=None but featurizer might be not None
            return None
        if feature_names is None:
            feature_names = self._input_names["feature_names"]
        if self.original_featurizer is None:
            return feature_names
        return get_feature_names_or_default(self.original_featurizer, feature_names)

    def _expand_treatments(self, X, *Ts):
        # Expand treatments for each time period
        outTs = []
        base_expand_treatments = super()._expand_treatments
        for T in Ts:
            if ndim(T) == 0:
                one_T = base_expand_treatments(X, T)[1]
                one_T = one_T.reshape(-1, 1) if ndim(one_T) == 1 else one_T
                T = np.tile(one_T, (1, self._n_periods, ))
            else:
                assert (T.shape[1] == self._n_periods if self.transformer else T.shape[1] == self._d_t[0]), \
                    f"Expected a list of time period * d_t, instead got a treatment array of shape {T.shape}."
                if self.transformer:
                    T = np.hstack([
                        base_expand_treatments(
                            X, T[:, [t]])[1] for t in range(self._n_periods)
                    ])
            outTs.append(T)
        return (X,) + tuple(outTs)

    @property
    def bias_part_of_coef(self):
        return self.ortho_learner_model_final_._model_final._fit_cate_intercept

    @property
    def fit_cate_intercept_(self):
        return self.ortho_learner_model_final_._model_final._fit_cate_intercept

    @property
    def original_featurizer(self):
        # NOTE: important to use the _ortho_learner_model_final_ attribute instead of the
        #       attribute so that the trained featurizer will be passed through
        return self.ortho_learner_model_final_._model_final_trained[0]._original_featurizer

    @property
    def featurizer_(self):
        # NOTE This is used by the inference methods and has to be the overall featurizer. intended
        # for internal use by the library
        return self.ortho_learner_model_final_._model_final_trained[0]._featurizer

    @property
    def model_final_(self):
        # NOTE This is used by the inference methods and is more for internal use to the library
        #      We need to use the _ortho_learner's copy to retain the information from fitting
        return self.ortho_learner_model_final_.model_final_

    @property
    def model_final(self):
        return self._gen_model_final()

    @model_final.setter
    def model_final(self, model):
        if model is not None:
            raise ValueError("Parameter `model_final` cannot be altered for this estimator!")

    @property
    def models_y(self):
        return [[mdl._model_y for mdl in mdls] for mdls in super().models_nuisance_]

    @property
    def models_t(self):
        return [[mdl._model_t for mdl in mdls] for mdls in super().models_nuisance_]

    @property
    def nuisance_scores_y(self):
        return self.nuisance_scores_[0]

    @property
    def nuisance_scores_t(self):
        return self.nuisance_scores_[1]

    @property
    def residuals_(self):
        """
        A tuple (y_res, T_res, X, W), of the residuals from the first stage estimation
        along with the associated X and W. Samples are not guaranteed to be in the same
        order as the input order.
        """
        if not hasattr(self, '_cached_values'):
            raise AttributeError("Estimator is not fitted yet!")
        if self._cached_values is None:
            raise AttributeError("`fit` was called with `cache_values=False`. "
                                 "Set to `True` to enable residual storage.")
        Y_res, T_res = self._cached_values.nuisances
        return Y_res, T_res, self._cached_values.X, self._cached_values.W


#Standard python libraries
import os
import warnings
import copy
import time
import itertools
import functools

#Dependencies - numpy, scipy, matplotlib, pyfftw
import numpy as np
import matplotlib.pyplot as plt
import pyfftw
from pyfftw.interfaces.numpy_fft import fft, fftshift, ifft, ifftshift, fftfreq
from scipy.interpolate import interp1d as sinterp1d
import scipy
from scipy.sparse import save_npz, load_npz, eye, csr_matrix
from scipy.sparse.linalg import eigs

from ufss import DiagramGenerator
from scipy.integrate import RK45

class RK_rho_container:
    def __init__(self,t,rho,pulse_number,manifold_key,*,interp_kind='linear',
                 optical_gap = 0):
        self.pulse_number = pulse_number
        self.n, self.M = rho.shape
        self.manifold_key = manifold_key
        self.optical_gap = optical_gap
        if t.size == 1:
            self.M = self.M+2
            self.t = np.array([-1,0,1],dtype='float') * np.spacing(t[0]) + t[0]
            rho_new = np.zeros((self.n,3),dtype='complex')
            rho_new[:,0] = 0
            rho_new[:,1] = 0.5 * rho[:,0]
            rho_new[:,2] = rho[:,0]
            self.rho = rho_new

            self.interp = self.make_interpolant(kind='zero')

        else:
            self.t = t
            self.rho = rho

            self.interp = self.make_interpolant(kind=interp_kind)

        self.t_checkpoint = t
        self.rho_checkpoint = rho

    def make_interpolant(self,*, kind='cubic'):
        """Interpolates density matrix
"""
        return sinterp1d(self.t,self.rho,fill_value = (0,np.nan),bounds_error = False,
                         assume_sorted=True,kind=kind)

    def one_time_step(self,rho0,t0,tf,*,find_best_starting_time = True):
        if find_best_starting_time and tf < self.t_checkpoint[-1]:
            diff1 = tf - t0

            diff2 = tf - self.t[-1]

            closest_t_checkpoint_ind = np.argmin(np.abs(self.t_checkpoint - tf))
            closest_t_checkpoint = self.t_checkpoint[closest_t_checkpoint_ind]
            diff3 = tf - closest_t_checkpoint

            rho0s = [rho0,self.rho[:,-1],self.rho_checkpoint[:,closest_t_checkpoint_ind]]

            neighbor_ind = closest_t_checkpoint_ind - 1
            if neighbor_ind >= 0:
                neighbor = self.t_checkpoint[closest_t_checkpoint_ind-1]
                diff4 = tf - neighbor
                rho0s.append(self.rho_checkpoint[:,neighbor_ind])
            else:
                neighbor = np.nan
                diff4 = np.inf

            t0s = np.array([t0,self.t[-1],closest_t_checkpoint,neighbor])
            diffs = np.array([diff1,diff2,diff3,diff4])

            for i in range(diffs.size):
                if diffs[i] < 0:
                    diffs[i] = np.inf

            if np.allclose(diffs,np.inf):
                raise ValueError('Method extend is only valid for times after the pulse has ended')

            t0 = t0s[np.argmin(diffs)]
            rho0 = rho0s[np.argmin(diffs)]

        elif find_best_starting_time and tf > self.t_checkpoint[-1]:
            if self.t_checkpoint[-1] > t0:
                t0 = self.t_checkpoint[-1]
                rho0 = self.rho_checkpoint[:,-1]
            else:
                pass

        else:
            pass
        # RWA_gap = self.manifold.dot(np.array([1,-1])) * self.optical_gap
        return self.one_time_step_function(rho0,t0,tf,manifold_key=self.manifold_key)#,RWA_gap = RWA_gap)

    def extend(self,t):
        ans = np.zeros((self.n,t.size),dtype='complex')

        if t[0] >= self.t_checkpoint[0]:

            t_intersect, t_inds, t_checkpoint_inds = np.intersect1d(t,self.t_checkpoint,return_indices=True)

            ans[:,t_inds] = self.rho_checkpoint[:,t_checkpoint_inds]

            if t_inds.size == t.size:
                return ans
            else:
                all_t_inds = np.arange(t.size)
                other_t_inds = np.setdiff1d(all_t_inds,t_inds)
                t0 = self.t_checkpoint[-1]
                rho0 = self.rho_checkpoint[:,-1]
                if t[other_t_inds[0]] >= t0:
                    find_best_starting_time = False
                else:
                    find_best_starting_time = True
                for t_ind in other_t_inds:
                    tf = t[t_ind]
                    ans[:,t_ind] = self.one_time_step(rho0,t0,tf,find_best_starting_time = find_best_starting_time)
                    t0 = tf
                    rho0 = ans[:,t_ind]

        elif t[0] >= self.t[-1]:
            t0 = self.t[-1]
            rho0 = self.rho[:,-1]
            for i in range(len(t)):
                ans[:,i] = self.one_time_step(rho0,t0,t[i],find_best_starting_time = True)
                t0 = t[i]
                rho0 = ans[:,i]
        else:
            raise ValueError('Method extend is only valid for times after the pulse has ended')

        self.rho_checkpoint = ans
        self.t_checkpoint = t
        return ans

    def __call__(self,t):
        """Assumes t is sorted """
        if type(t) is np.ndarray:
            pass
        elif type(t) is list:
            t = np.array(t)
        else:
            t = np.array([t])
        extend_inds = np.where(t>self.t[-1])
        interp_inds = np.where(t<=self.t[-1])
        ta = t[interp_inds]
        tb = t[extend_inds]
        if ta.size > 0:
            ans_a_flag = True
            if ta.size == self.M and np.allclose(ta,self.t):
                ans_a = self.rho
            else:
                ans_a = self.interp(ta)
        else:
            ans_a_flag = False
        if tb.size > 0:
            ans_b = self.extend(tb)
            ans_b_flag = True
        else:
            ans_b_flag = False

        if ans_a_flag and ans_b_flag:
            ans = np.hstack((ans_a,ans_b))
        elif ans_a_flag:
            ans = ans_a
        elif ans_b_flag:
            ans = ans_b
        else:
            ans = None
        return ans

    def __getitem__(self,inds):
        return self.rho[:,inds]

class RKE_DensityMatrices(DiagramGenerator):
    """This class is designed to calculate perturbative wavepackets in the
        light-matter interaction given the eigenvalues of the unperturbed
        hamiltonian and the material dipole operator evaluated in the
        eigenbasis of the unperturbed hamiltonian.

    Args:
        file_path (string): path to folder containing eigenvalues and the
            dipole operator for the system Hamiltonian
        detection_type (string): options are 'polarization' (default) or 'fluorescence'

"""
    def __init__(self,file_path,*,detection_type = 'polarization',
                 conserve_memory=False):
        self.slicing_time = 0
        self.interpolation_time = 0
        self.expectation_time = 0
        self.RK45_step_time = 0
        self.dipole_dot_rho_time = 0
        self.dipole_time = 0
        self.automation_time = 0
        self.diagram_to_signal_time = 0

        self.base_path = file_path

        self.undersample_factor = 1

        self.gamma_res = 6.91

        self.sparsity_threshold = 0.1

        self.conserve_memory = conserve_memory

        self.load_L()

        self.set_rho_shapes()

        if not self.conserve_memory:
            self.load_mu()

        try:
            self.load_H_mu()
            # more efficient if H_mu is available
            self.dipole_down = self.dipole_down_H_mu

        except:
            # generally less efficient - mostly here for backwards compatibility
            self.dipole_down = self.dipole_down_L_mu

        self.optical_gap = 0

        self.atol = 1E-6
        self.rtol = 1E-5

        if detection_type == 'polarization':
            self.rho_to_signal = self.polarization_detection_rho_to_signal
            self.return_complex_signal = False

        elif detection_type == 'complex_polarization':
            self.rho_to_signal = self.polarization_detection_rho_to_signal
            self.return_complex_signal = True
            detection_type = 'polarization'

        elif detection_type == 'integrated_polarization':
            raise Exception('detection_type: Integrated polarization is not implemented for Open RKE')
            self.rho_to_signal = self.integrated_polarization_detection_rho_to_signal

        elif detection_type == 'fluorescence':
            self.rho_to_signal = self.fluorescence_detection_rho_to_signal

        DiagramGenerator.__init__(self,detection_type=detection_type)
        self.KB_dict = {'Bu':self.bra_up,'Ku':self.ket_up,'Kd':self.ket_down,'Bd':self.bra_down}

        # Code will not actually function until the following three empty lists are set by the user
        self.efields = [] #initialize empty list of electric field shapes
        self.efield_times = [] #initialize empty list of times assoicated with each electric field shape
        self.dts = [] #initialize empty list of time spacings associated with each electric field shape
        self.polarization_sequence = [] #initialize empty polarization sequence
        self.pulse_times = [] #initialize empty list of pulse arrival times
        self.centers = [] #initialize empty list of pulse center frequencies
        self.efield_wavevectors = []
        self.rhos = dict()

        # Initialize unperturbed wavefunction
        self.set_rho0_auto()

    def set_pulse_delays(self,all_delays):
        """Must be a list of numpy arrays, where each array is a
            list of delay times between pulses
"""
        self.all_pulse_delays = all_delays
        num_delays = len(self.all_pulse_delays)
        num_pulses = len(self.efields)

        if num_delays == num_pulses - 1:
            pass
        elif num_delays == num_pulses - 2 and self.detection_type == 'polarization':
            # If there is a local oscillator, it arrives simultaneously with the last pulse
            self.all_pulse_delays.append(np.array([0]))
        elif num_delays <= num_pulses -2:
            raise Exception('There are not enough delay times')
        elif num_delays >= num_pulses:
            raise Exception('There are too many delay times')

    def calculate_diagrams_all_delays(self,diagrams):
        t0 = time.time()
        num_delays = len(self.all_pulse_delays)
        num_pulses = len(self.efields)

        all_delay_combinations = list(itertools.product(*self.all_pulse_delays))

        signal_shape = [delays.size for delays in self.all_pulse_delays]
        if self.detection_type == 'polarization':
            signal = np.zeros((len(all_delay_combinations),self.w.size),dtype='complex')
            if len(signal_shape) == self.pdc.shape[0]:
                # get rid of the "delay" between the last pulse and the local oscillator
                signal_shape[-1] = self.w.size
            elif len(signal_shape) == self.pdc.shape[0] - 1:
                # append the shape of the polariation-detection axis
                signal_shape.append(self.w.size)
            else:
                raise Exception('Cannot automatically determine final signal shape')
        else:
            signal = np.zeros((len(all_delay_combinations)),dtype='complex')

        counter = 0
        for delays in all_delay_combinations:
            arrival_times = [0]
            for delay in delays:
                arrival_times.append(arrival_times[-1]+delay)

            if self.detection_type == 'polarization':
                signal[counter,:] = self.calculate_diagrams(diagrams,arrival_times)
            else:
                signal[counter] = self.calculate_diagrams(diagrams,arrival_times)
            counter += 1

        self.signal = signal.reshape(signal_shape)
        self.calculation_time = time.time() - t0

        return self.signal

    def save_timing(self):
        save_dict = {'RKE_calculation_time':self.calculation_time}
        np.savez(os.path.join(self.base_path,'RKE_calculation_time.npz'),**save_dict)

    def calculate_signal_all_delays(self):
        t0 = time.time()
        num_delays = len(self.all_pulse_delays)
        num_pulses = len(self.efields)

        all_delay_combinations = list(itertools.product(*self.all_pulse_delays))

        signal_shape = [delays.size for delays in self.all_pulse_delays]
        if self.detection_type == 'polarization':
            signal = np.zeros((len(all_delay_combinations),self.w.size),dtype='complex')
            if len(signal_shape) == self.pdc.shape[0]:
                # get rid of the "delay" between the last pulse and the local oscillator
                signal_shape[-1] = self.w.size
            elif len(signal_shape) == self.pdc.shape[0] - 1:
                # append the shape of the polariation-detection axis
                signal_shape.append(self.w.size)
            else:
                raise Exception('Cannot automatically determine final signal shape')
        else:
            signal = np.zeros((len(all_delay_combinations)),dtype='complex')

        counter = 0
        for delays in all_delay_combinations:
            arrival_times = [0]
            for delay in delays:
                arrival_times.append(arrival_times[-1]+delay)

            if self.detection_type == 'polarization':
                signal[counter,:] = self.calculate_signal(arrival_times)
            else:
                signal[counter] = self.calculate_signal(arrival_times)
            counter += 1

        self.signal = signal.reshape(signal_shape)
        self.calculation_time = time.time() - t0
        return self.signal

    def set_t(self,optical_dephasing_rate,*,dt='auto'):
        """Sets the time grid upon which all frequency-detected signals will
be calculated on
"""
        max_pos_t = int(self.gamma_res/optical_dephasing_rate)
        if dt == 'auto':
            dt = self.dts[-1] # signal detection bandwidth determined by local oscillator
        self.t = np.arange(-max_pos_t,max_pos_t+dt/2,dt)
        # if self.t.size % 2:
        #     self.t = self.t[:-1]
        self.w = fftshift(fftfreq(self.t.size,d=dt)*2*np.pi)

    def execute_diagram(self,instructions):
        r = self.rho0
        name = ''
        for i in range(len(instructions)):
            key, num = instructions[i]
            name += key+str(num)
            # Try to re-use previous calculations, if they exist
            try:
                new_r = self.rhos[name]
            except KeyError:
                new_r = self.KB_dict[key](r,pulse_number=num)
                self.rhos[name] = new_r
            r = new_r
        sig = self.rho_to_signal(r)
        return sig

    def remove_rhos_by_pulse_number(self,pulse_number):
        num = str(pulse_number)
        keys = self.rhos.keys()
        keys_to_remove = []
        for key in keys:
            flag = key.find(num)
            if flag >= 0:
                keys_to_remove.append(key)
        for key in keys_to_remove:
            self.rhos.pop(key)

    def calculate_signal(self,arrival_times):
        t0 = time.time()
        try:
            old_pulse_times = self.pulse_times
            for i in range(len(old_pulse_times)):
                if old_pulse_times[i] != arrival_times[i]:
                    self.remove_rhos_by_pulse_number(i)
        except AttributeError:
            pass

        self.pulse_times = arrival_times
        if self.detection_type == 'polarization':
            times = [self.efield_times[i] + arrival_times[i] for i in range(len(arrival_times)-1)]
        elif self.detection_type == 'integrated_polarization':
            times = [self.efield_times[i] + arrival_times[i] for i in range(len(arrival_times)-1)]
        elif self.detection_type == 'fluorescence':
            times = [self.efield_times[i] + arrival_times[i] for i in range(len(arrival_times))]

        efield_permutations = self.relevant_permutations(times)

        diagram_instructions = []
        for perm in efield_permutations:
            diagram_instructions += self.instructions_from_permutation(perm)
        self.current_instructions = diagram_instructions

        t1 = time.time()
        try:
            instructions = diagram_instructions[0]
            signal = self.execute_diagram(instructions)
            for instructions in diagram_instructions[1:]:
                signal += self.execute_diagram(instructions)
        except IndexError:
            signal = 0

        t2 = time.time()
        self.automation_time += t1-t0
        self.diagram_to_signal_time += t2-t1
        return signal

    def calculate_diagrams(self,diagram_instructions,arrival_times):
        try:
            old_pulse_times = self.pulse_times
            for i in range(len(old_pulse_times)):
                if old_pulse_times[i] != arrival_times[i]:
                    self.remove_rhos_by_pulse_number(i)
        except AttributeError:
            pass

        self.pulse_times = arrival_times

        self.current_instructions = diagram_instructions
        instructions = diagram_instructions[0]
        signal = self.execute_diagram(instructions)
        for instructions in diagram_instructions[1:]:
            signal += self.execute_diagram(instructions)
        return signal

    def polarization_detection_rho_to_signal(self,rho):
        p_of_t = self.dipole_expectation(rho,pulse_number=-1,ket_flag=True)
        return self.polarization_to_signal(p_of_t,local_oscillator_number=-1)

    def integrated_polarization_detection_rho_to_signal(self,rho):
        p = self.integrated_dipole_expectation(rho,ket_flag=True)
        return self.integrated_polarization_to_signal(p,local_oscillator_number=-1)

    # def fluorescence_detection_rho_to_signal(self,rho):
    #     L_size = self.eigenvalues[0].size
    #     H_size = int(np.sqrt(L_size))

    #     # reshape rho into a normal density matrix representation
    #     rho = rho.reshape((H_size,H_size))

    #     fluorescence_yield = np.array([0,1,1,self.f_yield])

    #     signal = np.dot(np.diagonal(rho),fluorescence_yield)

    #     return signal

    def set_efields(self,times_list,efields_list,centers_list,phase_discrimination,*,reset_rhos = True,
                    plot_fields = False):
        self.efield_times = times_list
        self.efields = efields_list
        self.centers = centers_list
        self.set_phase_discrimination(phase_discrimination)
        self.dts = []
        self.efield_frequencies = []
        if reset_rhos:
            self.rhos = dict()
        for t in times_list:
            if t.size == 1:
                dt = 1
                w = np.array([0])
            else:
                dt = t[1] - t[0]
                w = fftshift(fftfreq(t.size,d=dt))*2*np.pi
            self.dts.append(dt)
            self.efield_frequencies.append(w)

        self.dt = self.dts[0]

        if self.detection_type == 'polarization':
            try:
                self.local_oscillator = self.efields[-1].copy()
            except:
                self.local_oscillator = copy.deepcopy(self.efields[-1])

        for field in self.efields:
            if len(field) == 1:
                # M = 1 is the impulsive limit
                pass
            else:
                self.check_efield_resolution(field,plot_fields = plot_fields)

    def check_efield_resolution(self,efield,*,plot_fields = False):
        efield_tail = np.max(np.abs([efield[0],efield[-1]]))

        if efield_tail > np.max(np.abs(efield))/100:
            warnings.warn('Consider using larger num_conv_points, pump does not decay to less than 1% of maximum value in time domain')

        efield_fft = fftshift(fft(ifftshift(efield)))*self.dt
        efield_fft_tail = np.max(np.abs([efield_fft[0],efield_fft[-1]]))

        if efield_fft_tail > np.max(np.abs(efield_fft))/100:
            warnings.warn('''Consider using smaller value of dt, pump does not decay to less than 1% of maximum value in frequency domain''')

        if plot_fields:
            fig, axes = plt.subplots(1,2)
            l1,l2, = axes[0].plot(self.efield_t,np.real(efield),self.efield_t,np.imag(efield))
            plt.legend([l1,l2],['Real','Imag'])
            axes[1].plot(self.efield_w,np.real(efield_fft),self.efield_w,np.imag(efield_fft))

            axes[0].set_ylabel('Electric field Amp')
            axes[0].set_xlabel('Time ($\omega_0^{-1})$')
            axes[1].set_xlabel('Frequency ($\omega_0$)')

            fig.suptitle('Check that efield is well-resolved in time and frequency')
            plt.show()

    def set_local_oscillator_phase(self,phase):
        self.efields[-1] = np.exp(1j*phase) * self.local_oscillator

    def get_closest_index_and_value(self,value,array):
        """Given an array and a desired value, finds the closest actual value
stored in that array, and returns that value, along with its corresponding
array index
"""
        index = np.argmin(np.abs(array - value))
        value = array[index]
        return index, value

    def load_L(self):
        """Load in known eigenvalues. Must be stored as a numpy archive file,
with keys: GSM, SEM, and optionally DEM.  The eigenvalues for each manifold
must be 1d arrays, and are assumed to be ordered by increasing energy. The
energy difference between the lowest energy ground state and the lowest
energy singly-excited state should be set to 0
"""
        L_save_name = os.path.join(self.base_path,'L.npz')
        try:
            with np.load(L_save_name,allow_pickle=True) as L_archive:
                self.L = dict()
                for key in L_archive.keys():
                    L = L_archive[key]
                    if L.dtype == np.dtype('O'):
                        self.L[key] = L[()]
                    else:
                        if self.check_sparsity(L):
                            self.L[key] = csr_matrix(L)
                        else:
                            self.L[key] = L
        except:
            self.L = {'all_manifolds':load_npz(L_save_name)}
        self.manifolds = list(self.L.keys())

    def check_sparsity(self,mat):
        csr_mat = csr_matrix(mat)
        sparsity = csr_mat.nnz / (csr_mat.shape[0]*csr_mat.shape[1])
        if sparsity < self.sparsity_threshold:
            return True
        else:
            return False

    def dL(self,t,rho):
        try:
            L = self.L['all_manifolds']
        except KeyError:
            L = self.L[rho.manifold_key]
        return L.dot(rho)

    def get_dL_manual(self,manifold_key):
        try:
            L = self.L['all_manifolds']
        except KeyError:
            L = self.L[manifold_key]

        def L_fun(t,rho):
            return L.dot(rho)

        return L_fun

    def one_time_step_function(self,rho0,t0,tf,*,manifold_key = None):
        num_steps = 0
        if manifold_key == None:
            rk45 = RK45(self.dL,t0,rho0,tf,atol=self.atol,rtol=self.rtol)
        else:
            dL = self.get_dL_manual(manifold_key)
            rk45 = RK45(dL,t0,rho0,tf,atol=self.atol,rtol=self.rtol)
        while rk45.t < tf:
            rk45.step()
            num_steps += 1
        rho_final = rk45.y
        return rho_final

    def get_bottom_eigenvector(self):
        try:
            L = self.L['all_manifolds']
        except KeyError:
            L = self.L['00']
        if L.shape == (1,1):
            e = L[0,0]
            ev = np.array([[1]])
        else:
            e, ev = eigs(L,k=1,which='SM',maxiter=10000)
        if e.size == 1 and np.allclose(e,0):
            pass
        else:
            raise Exception('Smallest magnitude eigenvalue of L is {}. L must have a single stationary state for this code to work'.format(e))
        v = ev[:,0]
        H_size = int(np.sqrt(v.size))
        rho = v.reshape((H_size,H_size))
        trace = rho.trace()
        v = v/trace # Need to start with a trace 1 object
        return v

    def set_rho0_auto(self):
        try:
            rho0 = np.load(os.path.join(self.base_path,'rho0.npy'))
        except FileNotFoundError:
            rho0 = self.get_bottom_eigenvector()
        t = np.array([-np.inf,0,np.inf])
        rho0 = rho0[:,np.newaxis] * np.ones((rho0.size,t.size))
        pulse_number = None
        manifold_key = '00'
        self.rho0 = RK_rho_container(t,rho0,pulse_number,manifold_key,
                                     interp_kind = 'zero',optical_gap = self.optical_gap)

    def set_rho_shapes(self):
        self.rho_shapes = dict()
        if 'all_manifolds' in self.manifolds:
            L_size = self.L['all_manifolds'].size
            H_size = int(np.sqrt(L_size))
            self.rho_shapes['all_manifolds'] = (H_size,H_size)
        else:
            H_sizes = dict()
            for key in self.manifolds:
                ket_key, bra_key = key
                if ket_key == bra_key:
                    L_size = self.L[key].shape[0]
                    H_size = int(np.sqrt(L_size))
                    H_sizes[ket_key] = H_size
            for key in self.manifolds:
                ket_key, bra_key = key
                ket_size = H_sizes[ket_key]
                bra_size = H_sizes[bra_key]
                self.rho_shapes[key] = (ket_size,bra_size)

    def load_mu(self):
        """Load the precalculated dipole overlaps.  The dipole operator must
            be stored as a .npz file, and must contain at least one array, each with three
            indices: (new manifold index, old manifold eigenfunction,
            cartesian coordinate)."""
        try:
            file_name = os.path.join(self.base_path,'mu_site_basis.npz')
            with np.load(file_name) as mu_archive:
                self.mu = {key:mu_archive[key] for key in mu_archive.keys()}
        except FileNotFoundError:
            try:
                file_name = os.path.join(self.base_path,'mu_original_L_basis.npz')
                with np.load(file_name) as mu_archive:
                    self.mu = {key:mu_archive[key] for key in mu_archive.keys()}
            except FileNotFoundError:
                file_name = os.path.join(self.base_path,'mu.npz')
                with np.load(file_name) as mu_archive:
                    self.mu = {key:mu_archive[key] for key in mu_archive.keys()}
        sparse_flags = []
        for key in self.mu.keys():
            mu_2D = np.sum(np.abs(self.mu[key])**2,axis=-1)
            sparse_flags.append(self.check_sparsity(mu_2D))
        sparse_flags = np.array(sparse_flags)
        if np.allclose(sparse_flags,True):
            self.sparse_mu_flag = True
        else:
            self.sparse_mu_flag = False

        for key in self.mu.keys():
            mu_x = self.mu[key][...,0]
            mu_y = self.mu[key][...,1]
            mu_z = self.mu[key][...,2]

            if self.sparse_mu_flag:
                self.mu[key] = [csr_matrix(mu_x),csr_matrix(mu_y),csr_matrix(mu_z)]
            else:
                self.mu[key] = [mu_x,mu_y,mu_z]

        print('RKE_sparse_mu_flag',self.sparse_mu_flag)

    ### Setting the electric field to be used

    def set_polarization_sequence(self,polarization_list,*,reset_rhos=True):
        """Sets the sequences used for either parallel or crossed pump and probe

        Args:
            polarization_list (list): list of four strings, can be 'x','y' or 'z'
        Returns:
            None: sets the attribute polarization sequence
"""

        x = np.array([1,0,0])
        y = np.array([0,1,0])
        z = np.array([0,0,1])
        pol_options = {'x':x,'y':y,'z':z}

        self.polarization_sequence = [pol_options[pol] for pol in polarization_list]

        if reset_rhos:
            self.rhos = dict()

    ### Tools for recursively calculating perturbed density maatrices using TDPT

    def dipole_matrix(self,pulse_number,key,ket_flag=True,up_flag=True):
        """Calculates the dipole matrix given the electric field polarization vector,
            if ket_flag = False then uses the bra-interaction"""
        t0 = time.time()
        pol = self.polarization_sequence[pulse_number]

        x = np.array([1,0,0])
        y = np.array([0,1,0])
        z = np.array([0,0,1])
        try:
            mu = self.mu[key]
        except KeyError:
            if ket_flag:
                key = 'ket'
            else:
                key = 'bra'
            if up_flag:
                key += '_up'
            else:
                key += '_down'
            mu = self.mu[key]

        if np.all(pol == x):
            overlap_matrix = mu[0]#.copy()
        elif np.all(pol == y):
            overlap_matrix = mu[1]#.copy()
        elif np.all(pol == z):
            overlap_matrix = mu[2]#.copy()
        else:
            overlap_matrix = mu[0]*pol[0] + mu[1]*pol[1] + mu[2]*pol[2]

        # if self.sparse_mu_flag:
        #     to_return = csr_matrix(overlap_matrix)
        # else:
        #     to_return = overlap_matrix

        t1 = time.time()
        self.dipole_time += t1-t0

        return overlap_matrix

    def manifold_key_to_array(self,key):
        """Key must be a string of exactly 2 integers, the first describing
            the ket manifold, the second the bra manifold.  If the density
            matrix is represented in the full space, rather than being divided
            into manifolds, the first integer reperesents the total number of
            excitations to the ket side, and the second integers represents
            the sum of all excitations to the bra side."""
        if len(key) != 2:
            raise Exception('manifold key must be a string of exactly two intgers')
        return np.array([int(char) for char in key],dtype=int)

    def manifold_array_to_key(self,manifold):
        """Inverse of self.manifold_key_to_array"""
        if manifold.size != 2 or manifold.dtype != int:
            raise Exception('manifold array must contain exactly 2 integer')
        return str(manifold[0]) + str(manifold[1])

    def next_order(self,rho_in,*,ket_flag=True,up_flag=True,pulse_number = 0):
        """This function connects psi_p to psi+pj^(*) using the Euler Method.

        Args:
            rho_in (rho_container): input density matrix
            pulse_number (int): index of optical pulse (0,1,2,...)

        Return:
            rho_dict (rho_container): next-order density matrix
"""
        pulse_time = self.pulse_times[pulse_number]
        t = self.efield_times[pulse_number] + pulse_time
        old_manifold_key = rho_in.manifold_key
        if up_flag:
            change = 1
        else:
            change = -1
        if ket_flag:
            manifold_change = np.array([change,0],dtype=int)
        else:
            manifold_change = np.array([0,change],dtype=int)
        old_manifold = self.manifold_key_to_array(old_manifold_key)
        new_manifold = old_manifold + manifold_change
        new_manifold_key = self.manifold_array_to_key(new_manifold)
        mu_key = old_manifold_key + '_to_' + new_manifold_key

        if ket_flag == up_flag:
            # Rotating term excites the ket and de-excites the bra
            conjugate_flag = False
        else:
            # Counter-rotating term
            conjugate_flag = True

        if conjugate_flag:
            center = -self.centers[pulse_number]
        else:
            center = self.centers[pulse_number]

        M = t.size
        old_rho = rho_in(t)

        if self.conserve_memory:
            # move back to the basis the Liouvillian was written in
            if 'all_manifolds' in self.manifolds:
                ket_size,bra_size = self.rho_shapes['all_manifolds']
            else:
                ket_size,bra_size = self.rho_shapes[old_manifold_key]

            old_rho = old_rho.reshape(ket_size,bra_size,M)

            if ket_flag:
                old_ket_key = old_manifold_key[0]
                new_ket_key = new_manifold_key[0]
                if up_flag:
                    H_mu_key = old_ket_key + '_to_' + new_ket_key
                else:
                    H_mu_key = new_ket_key + '_to_' + old_ket_key

                mu_up_flag = up_flag
            else:
                old_bra_key = old_manifold_key[1]
                new_bra_key = new_manifold_key[1]
                if up_flag:
                    H_mu_key = old_bra_key + '_to_' + new_bra_key
                else:
                    H_mu_key = new_bra_key + '_to_' + old_bra_key
                mu_up_flag = not up_flag

            overlap_matrix = self.get_H_mu(pulse_number,H_mu_key,up_flag=mu_up_flag)

            ta = time.time()
            if ket_flag:
                mu_old_rho = np.einsum('ij,jkl',overlap_matrix,old_rho)
            else:
                mu_old_rho = np.einsum('ijl,jk',old_rho,overlap_matrix)
            tb = time.time()

            rho_vec_size = mu_old_rho.shape[0]*mu_old_rho.shape[1]
            mu_old_rho = mu_old_rho.reshape(rho_vec_size,M)

        else:
            overlap_matrix = self.dipole_matrix(pulse_number,mu_key,ket_flag=ket_flag,up_flag=up_flag)

            ta = time.time()
            mu_old_rho = overlap_matrix.dot(old_rho)
            tb = time.time()

        self.dipole_dot_rho_time += tb - ta
        next_rho = np.zeros(mu_old_rho.shape,dtype='complex')

        if M == 1:
            next_rho[:,0] = self.efields[pulse_number] * mu_old_rho
        else:
            if conjugate_flag:
                efield = self.efields[pulse_number]*np.exp(-1j*center*t)
            else:
                efield = np.conjugate(self.efields[pulse_number])*np.exp(-1j*center*t)

            ############
            # This 1j vs -1j needs to be derived!!! #####
            ############
            if ket_flag:
                efield = 1j * efield
            else:
                efield = -1j * efield
            ###########
            ###########

            dt = self.dts[pulse_number]
            next_rho[:,0] = efield[0] * mu_old_rho[:,0] * dt
            for i in range(1,t.size):
                rho0 = next_rho[:,i-1]
                t0 = t[i-1]
                ta = time.time()
                next_rho[:,i] = self.one_time_step_function(rho0,t0,t[i],manifold_key=new_manifold_key)

                tb = time.time()
                self.RK45_step_time += tb - ta

                next_rho[:,i] += efield[i] * mu_old_rho[:,i] * dt

        # # i/hbar Straight from perturbation theory
        # if ket_flag:
        #     rho *= 1j
        # else:
        #     rho *= -1j

        rho_out = RK_rho_container(t,next_rho,pulse_number,new_manifold_key,
                                   optical_gap = self.optical_gap)
        rho_out.one_time_step_function = self.one_time_step_function

        return rho_out

    def ket_up(self,rho_in,*,pulse_number = 0):
        """This method connects psi_p to psi_pj where the next order psi
            is one manifold above the current manifold.

        Args:
            rho_in (rho_container): input density matrix
            pulse_number (int): index of optical pulse (0,1,2,...)

        Returns:
            (rho_container): output from method next_order
"""
        return self.next_order(rho_in,ket_flag=True,up_flag=True,
                               pulse_number = pulse_number)

    def ket_down(self,rho_in,*,pulse_number = 0):
        """This method connects psi_p to psi_pj where the next order psi
            is one manifold above the current manifold.

        Args:
            rho_in (rho_container): input density matrix
            pulse_number (int): index of optical pulse (0,1,2,...)

        Returns:
            (rho_container): output from method next_order
"""
        return self.next_order(rho_in,ket_flag=True,up_flag=False,
                               pulse_number = pulse_number)

    def bra_up(self,rho_in,*,pulse_number = 0):
        """This method connects psi_p to psi_pj where the next order psi
            is one manifold above the current manifold.

        Args:
            rho_in (rho_container): input density matrix
            pulse_number (int): index of optical pulse (0,1,2,...)

        Returns:
            (rho_container): output from method next_order
"""
        return self.next_order(rho_in,ket_flag=False,up_flag=True,
                               pulse_number = pulse_number)

    def bra_down(self,rho_in,*,pulse_number = 0):
        """This method connects psi_p to psi_pj where the next order psi
            is one manifold above the current manifold.

        Args:
            rho_in (rho_container): input density matrix
            pulse_number (int): index of optical pulse (0,1,2,...)

        Returns:
            (rho_container): output from method next_order
"""
        return self.next_order(rho_in,ket_flag=False,up_flag=False,
                               pulse_number = pulse_number)

    ### Tools for taking the expectation value of the dipole operator with perturbed density matrices

    def dipole_down_H_mu(self,rho,manifold_key,*,new_manifold_mask = None,pulse_number = -1,
                    ket_flag=True):
        """This method is similar to the method down, but does not involve
            the electric field shape or convolutions. It is the action of the
            dipole operator on the ket-side without TDPT effects.  It also includes
            the dot product of the final electric field polarization vector."""

        if not ket_flag:
            raise Exception('Not implemented for bra-side')
        old_manifold_key = manifold_key
        old_ket_key = old_manifold_key[0]
        new_ket_key = str(int(old_ket_key)-1)
        mu_key = new_ket_key + '_to_' + old_ket_key

        if ket_flag:
            center = - self.centers[pulse_number]
            conjugate_flag = True
        else:
            center = self.centers[pulse_number]
            conjugate_flag = False

        t_size = rho.shape[-1]

        if 'all_manifolds' in self.L.keys():
            L_size = rho.size
            H_size = int(np.sqrt(L_size))
            rho = rho.reshape(H_size,H_size,t_size)
        else:
            ket_manifold_key = old_manifold_key[0] + old_manifold_key[0]
            ket_L_manifold_size = self.L[ket_manifold_key].shape[0]
            ket_H_size = int(np.sqrt(ket_L_manifold_size))

            bra_manifold_key = old_manifold_key[1] + old_manifold_key[1]
            bra_L_manifold_size = self.L[bra_manifold_key].shape[0]
            bra_H_size = int(np.sqrt(bra_L_manifold_size))

            rho = rho.reshape(ket_H_size,bra_H_size,t_size)

        overlap_matrix = self.get_H_mu(pulse_number,mu_key,ket_flag=ket_flag,up_flag=False)

        t0 = time.time()
        polarization_field = np.einsum('ij,jik',overlap_matrix,rho)

        t1 = time.time()

        return polarization_field

    def dipole_down_L_mu(self,rho,manifold_key,*,new_manifold_mask = None,
                         pulse_number = -1,ket_flag=True):
        """This method is similar to the method down, but does not involve
            the electric field shape or convolutions. It is the action of the
            dipole operator on the ket-side without TDPT effects.  It also includes
            the dot product of the final electric field polarization vector."""
        old_manifold_key = manifold_key
        change = -1
        if ket_flag:
            manifold_change = np.array([change,0])
        else:
            manifold_change = np.array([0,change])

        old_manifold = self.manifold_key_to_array(old_manifold_key)
        new_manifold = old_manifold + manifold_change
        new_manifold_key = self.manifold_array_to_key(new_manifold)
        mu_key = old_manifold_key + '_to_' + new_manifold_key

        if ket_flag:
            center = - self.centers[pulse_number]
            conjugate_flag = True
        else:
            center = self.centers[pulse_number]
            conjugate_flag = False

        rho_in = rho

        overlap_matrix = self.dipole_matrix(pulse_number,mu_key,ket_flag=True,up_flag=False)

        t0 = time.time()
        rho = overlap_matrix.dot(rho_in)

        t1 = time.time()

        L_size = rho.shape[0]
        H_size = int(np.sqrt(L_size))

        # reshape rho into a normal density matrix representation
        rho = rho.reshape((H_size,H_size,rho.shape[-1]))

        polarization_field = np.einsum('iij',rho)

        return polarization_field

    def load_H_mu(self):
        parent_dir = os.path.split(self.base_path)[0]
        file_name = os.path.join(parent_dir,'closed','mu.npz')

        with np.load(file_name) as mu_archive:
            self.H_mu = {key:mu_archive[key] for key in mu_archive.keys()}

    def get_H_mu(self,pulse_number,key,ket_flag=True,up_flag=True):
        """Calculates the dipole matrix given the electric field polarization vector,
            if ket_flag = False then uses the bra-interaction"""
        t0 = time.time()
        pol = self.polarization_sequence[pulse_number]

        x = np.array([1,0,0])
        y = np.array([0,1,0])
        z = np.array([0,0,1])
        try:
            mu = self.H_mu[key]
        except KeyError:
            try:
                key = 'up'
                mu = self.H_mu[key]
            except KeyError:
                key = 'ket_up'
                mu = self.H_mu[key]

        if np.all(pol == x):
            overlap_matrix = mu[:,:,0].copy()
        elif np.all(pol == y):
            overlap_matrix = mu[:,:,1].copy()
        elif np.all(pol == z):
            overlap_matrix = mu[:,:,2].copy()
        else:
            overlap_matrix = np.tensordot(mu,pol,axes=(-1,0))

        if not up_flag:
            overlap_matrix = overlap_matrix.T

        t1 = time.time()
        self.dipole_time += t1-t0

        return overlap_matrix

    def set_undersample_factor(self,frequency_resolution):
        """dt is set by the pulse. However, the system dynamics may not require such a
            small dt.  Therefore, this allows the user to set a requested frequency
            resolution for any spectrally resolved signals."""
        # f = pi/dt
        dt = np.pi/frequency_resolution
        u = int(np.floor(dt/self.dt))
        self.undersample_factor = max(u,1)

    def dipole_expectation(self,rho_in,*,pulse_number = -1,ket_flag=True):
        """Computes the expectation value of the dipole operator"""
        t0 = time.time()

        pulse_number = -1

        pulse_time = self.pulse_times[pulse_number]

        efield_t = self.efield_times[pulse_number] + pulse_time

        if ket_flag:
            center = - self.centers[pulse_number]
        else:
            center = self.centers[pulse_number]

        # The signal is zero before the final pulse arrives, and persists
        # until it decays. Therefore we avoid taking the sum at times
        # where the signal is zero.
        t = self.t + pulse_time

        pulse_start_ind = np.argmin(np.abs(t-efield_t[0]))
        if efield_t[0] < t[pulse_start_ind]:
            pulse_start_ind -= 1
        pulse_end_ind = np.argmin(np.abs(t-efield_t[-1]))
        if efield_t[-1] > t[pulse_end_ind]:
            pulse_end_ind += 1

        t_slice = slice(pulse_start_ind,None,None)
        t1_slice = slice(pulse_start_ind,pulse_end_ind,None)
        u = self.undersample_factor
        t2_slice = slice(pulse_end_ind,None,u)

        t = self.t[t_slice] + pulse_time
        t1 = self.t[t1_slice] + pulse_time
        t2 = self.t[t2_slice] + pulse_time

        rho1 = rho_in(t1)
        rho2 = rho_in(t2)

        rho1 *= np.exp(-1j * center * t1[np.newaxis,:])

        rho2 *= np.exp(-1j * center * t2[np.newaxis,:])

        # _u is an abbreviation for undersampled
        t_u = np.hstack((t1,t2))
        rho_u = np.hstack((rho1,rho2))

        tb = time.time()
        self.slicing_time += tb-t0

        t0 = time.time()
        exp_val_u = self.dipole_down(rho_u,rho_in.manifold_key,pulse_number = pulse_number,
                                     ket_flag = ket_flag)

        tb = time.time()
        self.expectation_time += tb-t0

        t0 = time.time()

        # Interpolate expectation value back onto the full t-grid
        if u != 1:
            # Often must extrapolate the final point
            exp_val_interp = scipy.interpolate.interp1d(t_u,exp_val_u,kind='cubic',fill_value='extrapolate')
            exp_val = exp_val_interp(t)
        else:
            exp_val = exp_val_u
        # print(exp_val.size/exp_val_u.size)
        tb = time.time()
        self.interpolation_time += tb-t0

        # Initialize return array with zeros
        ret_val = np.zeros(self.t.size,dtype='complex')

        # set non-zero values using t_slice
        ret_val[pulse_start_ind:] = exp_val
        return ret_val

    def integrated_dipole_expectation(self,rho_in,*,ket_flag=True):
        """Computes the expectation value of the dipole operator"""

        pulse_number = -1

        pulse_time = self.pulse_times[pulse_number]
        t = pulse_time + self.efield_times[pulse_number]

        if ket_flag:
            center = - self.centers[pulse_number]
        else:
            center = self.centers[pulse_number]

        rho = rho_in(t)

        rho_nonzero = rho_in.bool_mask
        try:
            ev = self.eigenvalues['all_manifolds'][rho_nonzero]
        except KeyError:
            ev = self.eigenvalues[rho_in.manifold_key][rho_nonzero]

        rho = rho * np.exp((ev[:,np.newaxis] - 1j*center)*t)

        rho_dict = {'bool_mask':rho_nonzero,'rho':rho,'manifold_key':rho_in.manifold_key}

        t0 = time.time()
        exp_val = self.dipole_down(rho_dict,pulse_number = pulse_number,
                                     ket_flag = ket_flag)

        tb = time.time()
        self.expectation_time += tb-t0

        return exp_val

    def get_local_oscillator(self):
        local_oscillator_number = -1
        efield_t = self.efield_times[local_oscillator_number]
        efield = self.efields[local_oscillator_number]

        if efield_t.size == 1:
            # Impulsive limit: delta in time is flat in frequency
            efield_ft = np.ones(self.w.size)*efield
            return efield_ft

        e_dt = efield_t[1] - efield_t[0]
        dt = self.t[1] - self.t[0]

        if (np.isclose(e_dt,dt) and efield_t[-1] <= self.t[-1]):
            full_efield = np.zeros(self.t.size,dtype='complex')

            # the local oscillator sets the "zero" on the clock
            pulse_time_ind = np.argmin(np.abs(self.t))

            pulse_start_ind = pulse_time_ind - efield_t.size//2
            pulse_end_ind = pulse_time_ind + efield_t.size//2 + efield_t.size%2

            t_slice = slice(pulse_start_ind, pulse_end_ind,None)

            full_efield[t_slice] = efield
            efield_ft = fftshift(ifft(ifftshift(full_efield)))*full_efield.size * dt
        else:
            efield_ft = fftshift(ifft(ifftshift(efield))) * efield.size * e_dt
            efield_w = fftshift(fftfreq(efield_t.size,d=e_dt)) * 2 * np.pi
            fill_value = (efield_ft[0],efield_ft[-1])
            f = sinterp1d(efield_w,efield_ft,fill_value = fill_value,
                          bounds_error=False,kind='quadratic')
            efield_ft = f(self.w)

        return efield_ft

    def polarization_to_signal(self,P_of_t_in,*,
                                local_oscillator_number = -1,undersample_factor = 1):
        """This function generates a frequency-resolved signal from a polarization field
           local_oscillator_number - usually the local oscillator will be the last pulse
                                     in the list self.efields"""
        undersample_slice = slice(None,None,undersample_factor)
        P_of_t = P_of_t_in[undersample_slice].copy()
        t = self.t[undersample_slice]
        dt = t[1] - t[0]
        pulse_time = self.pulse_times[local_oscillator_number]
        efield_t = self.efield_times[local_oscillator_number]

        pulse_time_ind = np.argmin(np.abs(self.t))

        efield = self.get_local_oscillator()

        halfway = self.w.size//2
        pm = self.w.size//(2*undersample_factor)
        efield_min_ind = halfway - pm
        efield_max_ind = halfway + pm + self.w.size%2
        efield = efield[efield_min_ind:efield_max_ind]

        P_of_w = fftshift(ifft(ifftshift(P_of_t)))*len(P_of_t)*dt#/np.sqrt(2*np.pi)

        signal = P_of_w * np.conjugate(efield)
        if not self.return_complex_signal:
            return np.imag(signal)
        else:
            return 1j*signal

    def integrated_polarization_to_signal(self,P,*,
                                local_oscillator_number = -1):
        """This function generates a frequency-resolved signal from a polarization field
           local_oscillator_number - usually the local oscillator will be the last pulse
                                     in the list self.efields"""
        efield_t = self.efield_times[local_oscillator_number]

        efield = self.efields[local_oscillator_number]

        signal = np.trapz(P * np.conjugate(efield),x=efield_t)
        return np.imag(signal)

    def save(self,file_name,pulse_delay_names,*,use_base_path=True):
        if use_base_path:
            file_name = os.path.join(self.base_path,file_name)
        save_dict = {}
        for name,delays in zip(pulse_delay_names,self.all_pulse_delays):
            save_dict[name] = delays
        if self.detection_type == 'polarization':
            save_dict['wt'] = self.w
        save_dict['signal'] = self.signal
        save_dict['signal_calculation_time'] = self.calculation_time
        np.savez(file_name,**save_dict)


#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Mon Jul 29 10:52:50 2019

@author: daniel
"""
from pylab import *
import os
import numpy as np
import cobra
import networkx as nx
from Env_ball_class import Env_ball

def apply_environment(mdl, env_vec,transporters):
    for i in range(len(transporters)):
        try:
            mdl.reactions.get_by_id(transporters[i]).lower_bound=-env_vec[i]
            mdl.reactions.get_by_id(transporters[i]).upper_bound=1000.
        except KeyError:
            pass
    sol_m=mdl.optimize()
    return sol_m.objective_value

def buid_bipartite_graph(model):
    G=nx.MultiDiGraph()

    for met in model.metabolites:
        G.add_node(met.id, name=met.name, nodeType = 'metabolite')

    for reac in model.reactions:
        if not reac.reversibility:
            G.add_node(reac.id, name=reac.name, nodeType = 'reaction')

            for met in reac.reactants:
                G.add_edge(met.id, reac.id, name = reac.id)
            for met in reac.products:
                G.add_edge(reac.id, met.id, name = reac.id)

        else:
            G.add_node(reac.id + '_r', name=reac.name, nodeType = 'reaction')
            G.add_node(reac.id + '_f', name=reac.name, nodeType = 'reaction')

            for met in reac.reactants:
                G.add_edge(met.id, reac.id + '_f', name = reac.id)
                G.add_edge(reac.id + '_r',met.id, name = reac.id)
            for met in reac.products:
                G.add_edge(reac.id + '_f', met.id, name = reac.id)
                G.add_edge(met.id,reac.id + '_r', name = reac.id)

    return G

class panEFM_run:
    def __init__(self, file_path, file_name):
        print(file_path, file_name)

        self.file_name=file_name
        self.file_path = file_path
        self.file_ = os.path.join(file_path, file_name)

        self.environment=self.__get_environment(file_name)
        self.reactome=None
        self.binary_m=None
        self.freq_t=None

        self.__parse_run_file(self.file_)

    def __get_environment(self, file_name):
        return file_name.split(".")[0].split('_')[-1]

    def __parse_run_file(self, file_):
        with open(file_) as f:

            reactome = np.array(f.readline().strip().split('\t'))

            dt=[]
            sorter = np.argsort(reactome)

            self.reactome = reactome[sorter]

            for line in f:
                a=line.strip().split('\t')
                d = [*map(float, a)]
                d=np.array(d)
                dt.append(d[sorter])
            self.binary_m = np.array(dt).T

            self.freq_t = np.sum(self.binary_m, axis=1)/self.binary_m.shape[1]

class panEFM_family:
    def __init__(self, family_name, file_path, model_path):
        self.family = family_name
        self.file_path = file_path
        self.model_path = model_path

        self.model = cobra.io.read_sbml_model(os.path.join(self.model_path, self.family,self.family+'.ensembl.sbml'))

        self.file_path_fam = os.path.join(self.file_path, self.family)
        self.model_path_fam = os.path.join(self.model_path, self.family, 'gapfilled')

        self.files = os.listdir(self.file_path_fam)
        self.model_files= os.listdir(self.model_path_fam)

        self.environments=None
        self.reactome=None
        self.freq_m=None

        self.model_reac_freq=None
        self.gene_counts = None

        self.panEFM=None

        self.include_reactome = None

        self.model_reactomes = None

        self.__parse_run_files()
        self.__get_model_reac_freq()
        self.__get_include_reactome()

    def __parse_run_files(self):
        panEFM1= panEFM_run(self.file_path_fam, self.files[0])

        #sorter = np.argsort(panEFM1.reactome)
        self.reactome = panEFM1.reactome.copy()

        self.panEFM={}
        self.panEFM[panEFM1.environment]=panEFM1.binary_m.copy()

        self.environments = np.zeros(len(self.files))
        self.freq_m = np.zeros((len(self.files), len(self.reactome)))

        for i in range(len(self.files)):
            print (self.files[i])
            panEFM= panEFM_run(self.file_path_fam, self.files[i])

            self.environments[i] = float(panEFM.environment)
            self.freq_m[i]=panEFM.freq_t

            self.panEFM[panEFM.environment]=panEFM.binary_m.copy()

        env_sorter = np.argsort(self.environments)
        self.environments=self.environments[env_sorter]
        self.freq_m = self.freq_m[env_sorter]

    def __get_model_reac_freq(self):
        '''
        Obtain binary vectors indicating the presence of reactions in specific models.
        Summarize these vectors by their frequency.

        Returns
        -------
        None.

        '''
        models_n = len(self.model_files)
        reactions_n = len(self.reactome)

        v =np.zeros((models_n, reactions_n))
        freq= np.zeros(reactions_n)
        gene_counts = np.zeros(reactions_n)
        for i,name in enumerate(self.model_files):
            mod=cobra.io.read_sbml_model(os.path.join(self.model_path_fam, name))
            for r in range(reactions_n):
                if mod.reactions.has_id(self.reactome[r]):
                    freq[r]+=(1.0/models_n)
                    gene_counts[r]+=len(mod.reactions.get_by_id(self.reactome[r]).genes)
                    v[i][r] = 1.0
        self.model_reac_freq=freq
        self.gene_counts = gene_counts
        self.model_reactomes = v

    def __get_include_reactome(self):
        '''
        reactions that are compared. The rules are:
            1) Are associated to a gene in any of the models;
            2) Are connected to the biomass reaction in a bipartite graph;
            3) Are not lethal to the model, as defined by flux variability analysis.
        '''

        index = np.zeros(len(self.reactome))

        fva = cobra.flux_analysis.flux_variability_analysis(self.model)

        g = buid_bipartite_graph(self.model)

        for i, name in enumerate(self.reactome):
            if self.gene_counts[i]>0:#condition 1)

                #condition 3)
                if fva.loc[name]['minimum']<0:
                    index[i]=1
                elif fva.loc[name]['maximum']>0:
                    index[i]=1

                #condition 2)
                if self.model.reactions.get_by_id(name).reversibility:
                    if nx.has_path(g, source=name+'_f', target='bio1'):
                        index[i]=1
                    elif nx.has_path(g, source=name+'_r', target='bio1'):
                        index[i]=1
                else:
                    if nx.has_path(g, source=name, target='bio1'):
                        index[i]=1

        self.include_reactome = index.astype(np.bool)

ev =Env_ball(1000)

#families = os.listdir(pathToFamilyEFMs)

#fam_mfs={}

#for i in [families[0]]:
#    fam_panEFM[i]=panEFM_family(i, pathToFamilyEFMs,pathToModels)