text: # -*- coding: utf-8 -*-
  """Routines for constructing estimates of dynamical quantities on trajectory
  data using Galerkin expansion.
  
  @author: Erik
  
  """
  from __future__ import absolute_import, division, print_function, unicode_literals
  
  import numpy as np
  import scipy.linalg as spl
  
  from .data_manipulation import get_initial_final_split, tlist_to_flat, flat_to_tlist
  
  
  def compute_mfpt(basis, stateA, lag=1, dt=1.):
      """Calculates the mean first passage time into state A as a function of
      each point.
  
      Parameters
      ----------
      basis : list of trajectories
          Basis for the Galerkin expansion. Must be zero in state A.
      state_A : list of trajectories
          List of trajectories where each element is 1 or 0, corresponding to whether or not the datapoint is in state A.
      lag : int, optional
          Number of timepoints in the future to use for the finite difference in the discrete-time generator.  If not provided, uses value in the generator.
      timestep : scalar, optional
          Time between timepoints in the trajectory data.
  
      Returns
      -------
      mfpt : list of trajectories
          List of trajectories containing the values of the mean first passage time at each timepoint.
  
      """
      complement = [(A_i - 1.) for A_i in stateA]
      soln = compute_FK(basis, complement, lag=lag, dt=dt)
  
      post_processed_soln = []
      for ht_i in soln:
          ht_i[ht_i < 0.] = 0.
          post_processed_soln.append(ht_i)
      return post_processed_soln
  
  
  def compute_committor(basis, guess_committor, lag=1):
      """Calculates the forward committor into state A as a function of
      each point.
  
      Parameters
      ----------
      basis : list of trajectories
          Basis for the Galerkin expansion. Must be zero in state A and B
      guess_committor : list of trajectories, optional
          The value of the guess function obeying the inhomogenous boundary conditions.
      lag : int, optional
          Number of timepoints in the future to use for the finite difference in the discrete-time generator.
  
      Returns
      -------
      committor: dynamical basis object
          List of trajectories containing the values of the forward committor at each point.
  
      """
      h = [np.zeros(gi.shape) for gi in guess_committor]
      soln = compute_FK(basis, h, r=guess_committor, lag=lag, dt=1.)
  
      post_processed_soln = []
      for qi in soln:
          qi[qi > 1.] = 1.
          qi[qi < 0.] = 0.
          post_processed_soln.append(qi)
      return post_processed_soln
  
  
  def compute_bwd_committor(basis, guess_committor, stationary_com, lag=1):
      """Calculates the backward into state A as a function of
      each point.
  
      Parameters
      ----------
      basis : list of trajectories
          Basis for the Galerkin expansion. Must be zero in state A and B
      guess_committor : list of trajectories, optional
          The value of the guess function obeying the inhomogenous boundary conditions.
      stationary_com : list of trajectories
          Values of the change of measure to the stationary distribution.
      lag : int, optional
          Number of timepoints in the future to use for the finite difference in the discrete-time generator.
  
      Returns
      -------
      bwd_committor: dynamical basis object
          List of trajectories containing the values of the backward_committor at each point.
      """
      h = [np.zeros(gi.shape) for gi in guess_committor]
      soln = compute_adj_FK(basis, h, com=stationary_com, r=guess_committor, lag=lag, dt=1.)
  
      post_processed_soln = []
      for qi in soln:
          qi[qi > 1.] = 1.
          qi[qi < 0.] = 0.
          post_processed_soln.append(qi)
      return post_processed_soln
  
  
  def compute_change_of_measure(basis, lag=1):
      """Calculates the value of the change of measure to the stationary distribution for each datapoint.
  
      Parameters
      ----------
      basis : list of trajectories
          Basis for the Galerkin expansion. Must be zero in state A and B
      lag : int, optional
          Number of timepoints in the future to use for the finite difference in the discrete-time generator.
  
      Returns
      -------
      change_of_measure : dynamical basis object
          List of trajectories containing the values of the change of measure to the stationary distribution at each point.
  
      """
      evals, evecs = compute_esystem(basis, lag, left=True, right=False)
      com = [np.real(ev_i[:, 0]) for ev_i in evecs]
  
      # Ensure positivity
      com_sign = np.sign(np.sum(com))
      com = [ci * com_sign for ci in com]
      post_processed_soln = []
      for ci in com:
          ci[ci < 0.] = 0.
          post_processed_soln.append(ci)
      return post_processed_soln
  
  
  def compute_esystem(basis, lag=1, dt=1., left=False, right=True):
      """Calculates the eigenvectors and eigenvalues of the generator through
      Galerkin expansion.
  
      Parameters
      ----------
      basis : list of trajectories
          List of trajectories containing the basis for the Galerkin expansion.
          This method works much better if the basis set is zero on states A and B, however this is not a necessity.
      lag : int, optional
          Number of timepoints in the future to use for the finite difference in the discrete-time generator.
      left : bool, optional
          Whether or not to calculate the left eigenvectors  of the system.
      right : bool, optional
          Whether or not to calculate the right eigenvectors  of the system.
  
      Returns
      -------
      eigenvalues : numpy array
          Numpy array containing the eigenvalues of the generator.
      left_eigenvectors : list of trajectories, optional
          If left was set to true, the left eigenvectors are returned as a list of trajectories.
      right_eigenvectors : list of trajectories, optional
          If right was set to true, the right eigenvectors are returned as a list of trajectories.
  
      """
      if lag is None:
          lag = basis.lag
      basis_list = basis
      basis_flat_traj, traj_edges = tlist_to_flat(basis_list)
      K = compute_correlation_mat(basis_list, lag=lag)
      S = compute_stiffness_mat(basis_list, lag=lag)
      L = (K - S) / (lag * dt)
  
      # Calculate, sort eigensystem
      if (left and right):
          evals, evecs_l, evecs_r = spl.eig(L, b=S, left=True, right=True)
          evals, [evecs_l, evecs_r] = _sort_esystem(evals, [evecs_l, evecs_r])
          evecs_l = flat_to_tlist(np.dot(basis_flat_traj, evecs_l), traj_edges)
          evecs_r = flat_to_tlist(np.dot(basis_flat_traj, evecs_r), traj_edges)
          return evals, evecs_l, evecs_r
      elif (left or right):
          evals, evecs = spl.eig(L, b=S, left=left, right=right)
          evals, [evecs] = _sort_esystem(evals, [evecs])
          evecs = np.dot(basis_flat_traj, evecs)
          evecs = flat_to_tlist(evecs, traj_edges)
          return evals, evecs
      else:
          evals = spl.eig(L, b=S, left=False, right=False)
          return np.sort(evals)[::-1]
  
  
  def _sort_esystem(evals, evec_collection):
      """Utility function that sorts a collection of eigenvetors in desceding
      order, and then sorts the eigenvectors accordingly.
  
      Parameters
      ----------
      evals : numpy array
          The unsorted eigenvalues.
      evec_collection : list of arrays
          List where each element in the list is a collection of eigenvectors.
          Each collection is sorted according to the eigenvalues.
  
      Returns
      -------
      sorted_evals : numpy array
          The sorted eigenvalues
      sorted_evecs : of arrays
          list where each element is a collection of sorted eigenvectors.
  
      """
      idx = evals.argsort()[::-1]
      sorted_evals = evals[idx]
      sorted_evecs = [evecs[:, idx] for evecs in evec_collection]
      return sorted_evals, sorted_evecs
  
  
  def compute_FK(basis, h, r=None, lag=1, dt=1., return_coeffs=False):
      """
      Solves the forward Feynman-Kac problem Lg=h on a domain D, with boundary
      conditions g=b on the complement of D. To account for the boundary
      conditions, we solve the homogeneous problem Lg = h - Lr, where r is the
      provided guess.
  
      Parameters
      ----------
      traj_data : list of arrays OR single numpy array
          Value of the basis functions at every time point.  Should only be nonzero
          for points on the domain.
      h : list of 1d arrays or single 1d array
          Value of the RHS of the FK formula.  This should only be nonzero at
          points on the domain, Domain.
      r : list of 1d arrays or single 1d array, optional
          Value of the guess function.  Should be equal to b every point off of
          the domain.  IF not provided, the boundary conditions are assumed to be
          homogeneous.
      lag : int
          Number of timepoints in the future to use for the finite difference in
          the discrete-time generator.  If not provided, defaults to 1.
      timestep : scalar, optional
          Time between timepoints in the trajectory data.  Defaults to 1.
  
      Returns
      -------
      g : list of arrays
          Estimated solution to the Feynman-Kac problem.
      coeffs : ndarray
          Coefficients for the solution, only returned if return_coeffs is True.
      """
      h = [h_i.reshape(-1, 1).astype('float') for h_i in h]
      L_basis = compute_generator(basis, lag=lag, dt=dt)
      h_i = compute_stiffness_mat(basis, h, lag=lag)
      if r is not None:
          r = [r_i.reshape(-1, 1) for r_i in r]
          L_guess = compute_generator(basis, r, lag=lag, dt=dt)
          h_i -= L_guess
      coeffs = spl.solve(L_basis, h_i.ravel())
  
      # Construct solution vector
      N_traj = len(basis)
      soln = [np.dot(basis[i], coeffs) for i in range(N_traj)]
      if r is not None:
          soln = [soln[i] + r[i].ravel() for i in range(N_traj)]
  
      # Return calculated values
      if return_coeffs:
          return soln, coeffs
      else:
          return soln
  
  
  def compute_adj_FK(basis, h, com=None, r=None, lag=1, dt=1., return_coeffs=False):
      """
      Solves the Feynman-Kac problem L^t dagger g=h on a domain D, with
      boundary conditions g=b on the complement of D. Here L^t is the adjoint of
      the generator with respect to the provided change of measure.  To account
      for the boundary conditions, we solve the homogeneous problem
      L^t g = h - L^t r, where r is the provided guess.
  
      Parameters
      ----------
      traj_data : list of arrays OR single numpy array
          Value of the basis functions at every time point.  Should only be nonzero
          for points on the domain.
      h : list of 1d arrays or single 1d array
          Value of the RHS of the FK formula.  This should only be nonzero at
          points on the domain, Domain.
      com : list of 1d arrays or single 1d array, optional
          Values of the change of measure against which to take the desired
          adjoint.  If not provided, takes the adjoint against the sampled meaure
      r : list of 1d arrays or single 1d array, optional
          Value of the guess function.  Should be equal to b every point off of
          the domain.  If not provided, the boundary conditions are assumed to be
          homogeneous.
      lag : int
          Number of timepoints in the future to use for the finite difference in
          the discrete-time generator.  If not provided, defaults to 1.
      timestep : scalar, optional
          Time between timepoints in the trajectory data.  Defaults to 1.
  
      Returns
      -------
      g : list of arrays
          Estimated solution to the Feynman-Kac problem.
      coeffs : ndarray
          Coefficients for the solution, only returned if return_coeffs is True.
      """
      h = [h_i.reshape(-1, 1).astype('float') for h_i in h]
  
      L_basis = compute_generator(basis, lag=lag, dt=dt, com=com)
      h_i = compute_stiffness_mat(basis, h, lag=lag, com=com)
      if r is not None:
          r = [r_i.reshape(-1, 1) for r_i in r]
          L_guess = compute_generator(r, basis, lag=lag, dt=dt, com=com).T
          h_i -= L_guess
      coeffs = spl.solve(L_basis.T, h_i.ravel())
  
      # Construct solution vector
      N_traj = len(basis)
      soln = [np.dot(basis[i], coeffs) for i in range(N_traj)]
      if r is not None:
          soln = [soln[i] + r[i].ravel() for i in range(N_traj)]
  
      # Return calculated values
      if return_coeffs:
          return soln, coeffs
      else:
          return soln
  
  
  def compute_correlation_mat(Xs, Ys=None, lag=1, com=None):
      """Computes the time-lagged correlation matrix between two sets of observables.
  
      Parameters
      ----------
      Xs : list of trajectories
          List of trajectories for the first set of observables.
      Ys : list of trajectories, optional
          List of trajectories for the second set of observables.  If None, set to be X.
      lag : int, optional
          Lag to use in the correlation matrix.  Default is one step.
      com : list of trajectories
          Values of the change of measure against which to compute the average
  
      Returns
      -------
      K : numpy array
          The time-lagged correlation matrix between X and Y.
      """
      if Ys is None:
          Ys = Xs
      # Move to flat convention for easier processing.
      X_flat, traj_edges = tlist_to_flat(Xs)
      Y_flat = tlist_to_flat(Ys)[0]
      if com is not None:
          com_flat = tlist_to_flat(com)[0]
          X_flat *= com_flat
  
      # Split into initial, final points
      initial_indices, final_indices = get_initial_final_split(traj_edges, lag=lag)
      Y_t_lag = Y_flat[final_indices]
      X_t_0 = X_flat[initial_indices]
  
      # Calculate correlation, stiffness matrices.
      N = X_t_0.shape[0]
      K = X_t_0.T.dot(Y_t_lag) / float(N)
      return K
  
  
  def compute_stiffness_mat(Xs, Ys=None, lag=1, com=None):
      """Computes the stiffness matrix between two sets of observables.
  
      Parameters
      ----------
      Xs : list of trajectories
          List of trajectories for the first set of observables.
      Ys : list of trajectories, optional
          List of trajectories for the second set of observables.  If None, set to be X.
      lag : int, optional
          Lag to use in the correlation matrix.  Default is one step.
          This is required as the stiffness is only evaluated over the initial points.
      com : list of trajectories
          Values of the change of measure against which to compute the average
  
      Returns
      -------
      S : numpy array
          The time-lagged stiffness matrix between X and Y.
      """
      if Ys is None:
          Ys = Xs
      # Move to flat convention for easier processing.
      X_flat, traj_edges = tlist_to_flat(Xs)
      Y_flat = tlist_to_flat(Ys)[0]
      if com is not None:
          com_flat = tlist_to_flat(com)[0]
          X_flat *= com_flat
  
      # Split into initial, final points
      initial_indices, final_indices = get_initial_final_split(traj_edges, lag=lag)
      Y_t_0 = Y_flat[initial_indices]
      X_t_0 = X_flat[initial_indices]
      N = X_t_0.shape[0]
      S = X_t_0.T.dot(Y_t_0) / float(N)
      return S
  
  
  def compute_generator(Xs, Ys=None, lag=1, dt=1., com=None):
      """
      Computes the matrix of inner product elements against the generator.
  
      Parameters
      ----------
      Xs : list of trajectories
          List of trajectories for the first set of observables.
      Ys : list of trajectories, optional
          List of trajectories for the second set of observables.  If None, set to be X.
      lag : int, optional
          Lag to use in the correlation matrix.  Default is one step.
      dt: float, optional
          time per step of dynamics. Default is one time unit.
      com : list of trajectories
          Values of the change of measure against which to compute the average.
  
      Returns
      -------
      L : numpy array
          The approximation to the inner product <X, L Y>.
      """
      if Ys is None:
          Ys = Xs
      K = compute_correlation_mat(Xs, Ys, lag=lag, com=com)
      S = compute_stiffness_mat(Xs, Ys, lag=lag, com=com)
      return (K - S) / (lag * dt)
  