def transfer_op_python(As, Bs, direction, x):
    """
    (mixed) transfer operator for a list of mps tensors

    Parameters:
    ----------------------
    As,Bs:        list of tf.Tensor
                  the mps tensors (Bs are on the conjugated side)
    direction:    int or str 
                  can be (1,'l','left') or (-1,'r','right) for left or right 
                  operation
    x:            tf.Tensor 
                  input matrix
    Returns:
    ----------------------
    tf.Tensor:  the evolved matrix

    """

    if direction in ('l', 'left', 1):
        for n in range(len(As)):
            x = ncon([x, As[n], tf.conj(Bs[n])], [(0, 1), (0, 2, -1),
                                                       (1, 2, -2)])
    elif direction in ('r', 'right', -1):
        for n in reversed(range(len(As))):
            x = ncon([x, As[n], tf.conj(Bs[n])], [(0, 1), (-1, 2, 0),
                                                       (-2, 2, 1)])
    else:
        raise ValueError("Invalid direction: {}".format(direction))

    return x

def add_layer_python(B, mps_tensor, mpo, conj_mps_tensor, direction):
    """
    adds an mps-mpo-mps layer to a left or right block "E"; used in dmrg to calculate the left and right
    environments
    Parameters:
    ---------------------------
    B:               Tensor object  
                     a tensor of shape (D1,D1',M1) (for direction>0) or (D2,D2',M2) (for direction>0)
    mps_tensor:      Tensor object of shape =(Dl,Dr,d)
    mpo_tensor:      Tensor object of shape = (Ml,Mr,d,d')
    conj_mps_tensor: Tensor object of shape =(Dl',Dr',d')
                     the mps tensor on the conjugated side
                     this tensor will be complex conjugated inside the routine; usually, the user will like to pass 
                     the unconjugated tensor
    direction:       int or str
                     direction in (1,'l','left'): add a layer to the right of ```B```
                     direction in (-1,'r','right'): add a layer to the left of ```B```
    Return:
    -----------------
    Tensor of shape (Dr,Dr',Mr) for direction in (1,'l','left')
    Tensor of shape (Dl,Dl',Ml) for direction in (-1,'r','right')
    """
    if direction in ('l', 'left', 1):
        return ncon(
            [B, mps_tensor, mpo, tf.conj(conj_mps_tensor)],
            [[1, 4, 3], [1, 2, -1], [3, -3, 5, 2], [4, 5, -2]])
    
    if direction in ('r', 'right', -1):
        return ncon(
            [B, mps_tensor, mpo, tf.conj(conj_mps_tensor)],
            [[1, 4, 3], [-1, 2, 1], [-3, 3, 5, 2], [-2, 5, 4]])

    def __init__(self,
                 mps,
                 mpo,
                 name='InfiniteDMRG',
                 precision=1E-12,
                 precision_canonize=1E-12,
                 nmax=1000,
                 nmax_canonize=1000,
                 ncv=40,
                 numeig=1,
                 pinv=1E-20,
                 power_method=False):

        # if not isinstance(mps, InfiniteMPSCentralGauge):
        #     raise TypeError(
        #         'in InfiniteDMRGEngine.__init__(...): mps of type InfiniteMPSCentralGauge expected, got {0}'
        #         .format(type(mps)))

        mps.restore_form(
            precision=precision_canonize,
            ncv=ncv,
            nmax=nmax_canonize,
            numeig=numeig,
            power_method=power_method,
            pinv=pinv)  #this leaves state in left-orthogonal form

        lb, hl = misc_mps.compute_steady_state_Hamiltonian_GMRES(
            'l',
            mps,
            mpo,
            left_dominant=tf.diag(tf.ones(mps.D[-1], dtype=mps.dtype)),
            right_dominant=ncon.ncon([mps.mat, tf.conj(mps.mat)],
                                     [[-1, 1], [-2, 1]]),
            precision=precision,
            nmax=nmax)

        rmps = mps.get_right_orthogonal_imps(
            precision=precision_canonize,
            ncv=ncv,
            nmax=nmax_canonize,
            numeig=numeig,
            pinv=pinv,
            restore_form=False)

        rb, hr = misc_mps.compute_steady_state_Hamiltonian_GMRES(
            'r',
            rmps,
            mpo,
            right_dominant=tf.diag(tf.ones(mps.D[0], dtype=mps.dtype)),
            left_dominant=ncon.ncon([mps.mat, tf.conj(mps.mat)],
                                    [[1, -1], [1, -2]]),
            precision=precision,
            nmax=nmax)

        left_dominant = ncon.ncon([mps.mat, tf.conj(mps.mat)],
                                  [[1, -1], [1, -2]])
        out = mps.unitcell_transfer_op('l', left_dominant)

        super().__init__(mps=mps, mpo=mpo, lb=lb, rb=rb, name=name)

def steady_state_density_matrices(nsteps, rhoAB, rhoBA, w_isometry, v_isometry, unitary, refsym):
    for n in range(nsteps):
        rhoAB_, rhoBA_ = descending_super_operator(rhoAB, rhoBA, w_isometry, v_isometry, unitary,
                                                   refsym)
        rhoAB = 1/2 * (rhoAB_ + tf.conj(tf.transpose(rhoAB_,(2,3,0,1))))/ncon.ncon([rhoAB_],[[1,2,1,2]])
        rhoBA = 1/2 * (rhoBA_ + tf.conj(tf.transpose(rhoBA_,(2,3,0,1))))/ncon.ncon([rhoBA_],[[1,2,1,2]])
        
    return rhoAB, rhoBA

    def _inference(self, x, dropout):
        with tf.name_scope('conv1'):
            # Transform to Fourier domain
            x_2d = tf.reshape(x, [-1, 28, 28])
            x_2d = tf.complex(x_2d, 0)
            xf_2d = tf.fft2d(x_2d)
            xf = tf.reshape(xf_2d, [-1, NFEATURES])
            xf = tf.expand_dims(xf, 1)  # NSAMPLES x 1 x NFEATURES
            xf = tf.transpose(xf)  # NFEATURES x 1 x NSAMPLES
            # Filter
            Wreal = self._weight_variable([int(NFEATURES/2), self.F, 1])
            Wimg = self._weight_variable([int(NFEATURES/2), self.F, 1])
            W = tf.complex(Wreal, Wimg)
            xf = xf[:int(NFEATURES/2), :, :]
            yf = tf.matmul(W, xf)  # for each feature
            yf = tf.concat([yf, tf.conj(yf)], axis=0)
            yf = tf.transpose(yf)  # NSAMPLES x NFILTERS x NFEATURES
            yf_2d = tf.reshape(yf, [-1, 28, 28])
            # Transform back to spatial domain
            y_2d = tf.ifft2d(yf_2d)
            y_2d = tf.real(y_2d)
            y = tf.reshape(y_2d, [-1, self.F, NFEATURES])
            # Bias and non-linearity
            b = self._bias_variable([1, self.F, 1])
#            b = self._bias_variable([1, self.F, NFEATURES])
            y += b  # NSAMPLES x NFILTERS x NFEATURES
            y = tf.nn.relu(y)
        with tf.name_scope('fc1'):
            W = self._weight_variable([self.F*NFEATURES, NCLASSES])
            b = self._bias_variable([NCLASSES])
            y = tf.reshape(y, [-1, self.F*NFEATURES])
            y = tf.matmul(y, W) + b
        return y

def sparse_dot_product0(emb, tuples, use_matmul=True, output_type='real'):
    """
    Compute the dot product of complex vectors.
    It uses complex vectors but tensorflow does not optimize in the complex space (or there is a bug in the gradient
    propagation with complex numbers...)
    :param emb: embeddings
    :param tuples: indices at which we compute dot products
    :return: scores (dot products)
    """
    n_t = tuples.get_shape()[0].value
    rk = emb.get_shape()[1].value
    emb_sel_a = tf.gather(emb, tuples[:, 0])
    emb_sel_b = tf.gather(emb, tuples[:, 1])
    if use_matmul:
        pred_cplx = tf.squeeze(tf.batch_matmul(
                tf.reshape(emb_sel_a, [n_t, rk, 1]),
                tf.reshape(emb_sel_b, [n_t, rk, 1]), adj_x=True))
    else:
        pred_cplx = tf.reduce_sum(tf.mul(tf.conj(emb_sel_a), emb_sel_b), 1)
    if output_type == 'complex':
        return pred_cplx
    elif output_type == 'real':
        return tf.real(pred_cplx) + tf.imag(pred_cplx)
    elif output_type == 'real':
        return tf.abs(pred_cplx)
    elif output_type == 'angle':
        raise NotImplementedError('No argument or inverse-tanh function for complex number in Tensorflow')
    else:
        raise NotImplementedError()

def tridiag_tensorflow(vecs, alpha, beta):
    Heff=tf.contrib.distributions.tridiag(beta, alpha, tf.conj(beta))
    eta, u = tf.linalg.eigh(Heff)  #could use tridiag
    out=ncon.ncon([vecs,u],[[1,-1,-2,-3],[1,-4]])
    out=out[:,:,:,0]
    out=tf.math.divide(out,tf.linalg.norm(out))
    return eta[0], out

def visualize_data_transformations():
    records = glob.glob(os.path.join(utils.working_dir, 'train_fragment_*.tfrecords'))
    dataset = tf.data.TFRecordDataset(records)
    dataset = dataset.map(parse_tfrecord_raw)
    dataset = dataset.repeat()
    dataset = dataset.shuffle(buffer_size=10)
    dataset = dataset.prefetch(2)
    it = dataset.make_one_shot_iterator()

    data_x = tf.placeholder(tf.float32, shape=(utils.sample_rate * utils.audio_clip_len,))
    data_y = tf.placeholder(tf.float32, shape=(utils.timesteps,))
    stfts = tf.contrib.signal.stft(data_x, frame_length=utils.frame_length, frame_step=utils.frame_step,
                                   fft_length=4096)
    power_stfts = tf.real(stfts * tf.conj(stfts))
    magnitude_spectrograms = tf.abs(stfts)
    power_magnitude_spectrograms = tf.abs(power_stfts)

    num_spectrogram_bins = magnitude_spectrograms.shape[-1].value

    # scale frequency to mel scale and put into bins to reduce dimensionality
    lower_edge_hertz, upper_edge_hertz = 30.0, 17000.0
    num_mel_bins = utils.mel_bins_base * 4
    linear_to_mel_weight_matrix = tf.contrib.signal.linear_to_mel_weight_matrix(
        num_mel_bins, num_spectrogram_bins, utils.sample_rate, lower_edge_hertz,
        upper_edge_hertz)
    mel_spectrograms = tf.tensordot(magnitude_spectrograms, linear_to_mel_weight_matrix, 1)
    mel_spectrograms.set_shape(
        magnitude_spectrograms.shape[:-1].concatenate(linear_to_mel_weight_matrix.shape[-1:]))

    # log scale the mel bins to better represent human loudness perception
    log_offset = 1e-6
    log_mel_spectrograms = tf.log(mel_spectrograms + log_offset)

    # compute first order differential and concat. "It indicates a raise or reduction of the energy for each
    # frequency bin at a frame relative to its predecessor"
    first_order_diff = tf.abs(
        tf.subtract(log_mel_spectrograms, tf.manip.roll(log_mel_spectrograms, shift=1, axis=1)))
    mel_fod = tf.concat([log_mel_spectrograms, first_order_diff], 1)

    with tf.Session() as sess:
        while True:
            try:
                raw_x, raw_y = sess.run(it.get_next())
                np_stfts = sess.run(power_stfts, feed_dict={data_x: raw_x})
                np_magnitude_spectrograms = sess.run(power_magnitude_spectrograms, feed_dict={data_x: raw_x})
                np_mel_spectrograms = sess.run(mel_spectrograms, feed_dict={data_x: raw_x})
                np_log_mel_spectrograms = sess.run(log_mel_spectrograms, feed_dict={data_x: raw_x})
                np_mel_fod = sess.run(mel_fod, feed_dict={data_x: raw_x})

                utils.plot_signal_transforms(raw_x,
                                            np_stfts,
                                            np_magnitude_spectrograms,
                                            np_mel_spectrograms,
                                            np_log_mel_spectrograms,
                                            np_mel_fod)
                print('wank')

            except tf.errors.OutOfRangeError:
                break

 def _compareConj(self, cplx, use_gpu):
   np_ans = np.conj(cplx)
   with self.test_session(use_gpu=use_gpu):
     inx = tf.convert_to_tensor(cplx)
     tf_conj = tf.conj(inx)
     tf_ans = tf_conj.eval()
   self.assertAllEqual(np_ans, tf_ans)
   self.assertShapeEqual(np_ans, tf_conj)

 def gram_schmidt_step(j, basis, v):
     """Makes v orthogonal to the j'th vector in basis."""
     #v_shape = v.get_shape()
     basis_vec = basis.read(j)
     v -=  ncon.ncon([tf.reshape(tf.conj(basis_vec), [basis_vec.shape[0] * basis_vec.shape[1] * basis_vec.shape[2]]),
                      tf.reshape(v, [v.shape[0] * v.shape[1] * v.shape[2]])], [[1], [1]])* basis_vec
     #v.set_shape(v_shape)
     return j + 1, basis, v

def descend_state_1site_R(state_1site, iso_012):  #χ^4
    """Descends a state from the top to the rightmost index of the isometry `iso`.
    Physically, if `iso` has 012 ordering, this is a descent to the right and
    if `iso` has 021 ordering, this is a descent to the left.
    """
    return tensornetwork.ncon(
        [iso_012, state_1site, tf.conj(iso_012)], [(1, 2, -1), (1, 0),
                                                   (0, 2, -2)])

 def CheckUnitary(self, x):
   # Tests that x[...,:,:]^H * x[...,:,:] is close to the identity.
   xx = tf.matmul(tf.conj(x), x, transpose_a=True)
   identity = tf.matrix_band_part(tf.ones_like(xx), 0, 0)
   if is_single:
     tol = 1e-5
   else:
     tol = 1e-14
   self.assertAllClose(identity.eval(), xx.eval(), atol=tol)

def get_env_disentangler(hamAB,hamBA,rhoBA,w,v,u,refsym):

    indList1 = [[7,8,10,-1],[4,3,9,2],[10,-3,9],[7,5,4],[8,-2,5,6],[1,-4,2],[1,6,3]]
    indList2 = [[7,8,-1,-2],[3,6,2,5],[1,-3,2],[1,9,3],[7,8,9,10],[4,-4,5],[4,10,6]]
    indList3 = [[7,8,-2,10],[3,4,2,9],[1,-3,2],[1,5,3],[-1,7,5,6],[10,-4,9],[8,6,4]]

    uEnv = ncon.ncon([hamAB,rhoBA,w,tf.conj(w),tf.conj(u),v,tf.conj(v)],indList1)
    if refsym:
        uEnv = uEnv + tf.transpose(uEnv,(1,0,3,2))
    else:
        uEnv = uEnv + ncon.ncon([hamAB,rhoBA,w,tf.conj(w),tf.conj(u),v,tf.conj(v)],indList3)
    
    uEnv = uEnv + ncon.ncon([hamBA,rhoBA,w,tf.conj(w),tf.conj(u),v,tf.conj(v)],indList2)

    return uEnv

def _mpo_with_state(iso_012, iso_021, h_mpo_2site, state_1site):
    # contract ascended hamiltonian at level `lup` with nearest 1-site descended state
    h2L, h2R = h_mpo_2site

    envL = [
        tensornetwork.ncon(
            [state_1site, iso_021, h, tf.conj(iso_012)],
            [(0, 2), (0, -1, 1), (3, 1), (2, 3, -2)])  # one transpose required
        for h in h2L
    ]

    envR = [
        tensornetwork.ncon(
            [state_1site, iso_012, h, tf.conj(iso_021)],
            [(0, 2), (0, -1, 1), (3, 1), (2, 3, -2)])  # one transpose required
        for h in h2R
    ]

    return envL, envR

 def testComplexConj(self):
   with self.test_session():
     size = ()
     x = tf.constant(11 - 13j, dtype=tf.complex64)
     y = tf.conj(x)
     analytical, numerical = tf.test.compute_gradient(x, size, y, size)
     correct = np.array([[1, 0], [0, -1]])
     self.assertAllEqual(correct, analytical)
     self.assertAllClose(correct, numerical, rtol=3e-6)
     self.assertLess(tf.test.compute_gradient_error(x, size, y, size), 2e-5)

    def compute_spectrograms(self, waveforms, labels=None):
        
        """Computes spectrograms for a batch of waveforms."""
        
        s = self.settings
        
        # Set final dimension of waveforms, which comes to us as `None`.
        self._set_waveforms_shape(waveforms)

        # Compute STFTs.
        waveforms = tf.cast(waveforms, tf.float32)
        stfts = tf.contrib.signal.stft(
            waveforms, self.window_size, self.hop_size,
            fft_length=self.dft_size, window_fn=self.window_fn)
        
        # Slice STFTs along frequency axis.
        stfts = stfts[..., self.freq_start_index:self.freq_end_index]
        
        # Get STFT magnitudes squared, i.e. squared spectrograms.
        grams = tf.real(stfts * tf.conj(stfts))
        # gram = tf.abs(stft) ** 2
        
        # Take natural log of squared spectrograms. Adding an epsilon
        # avoids log-of-zero errors.
        grams = tf.log(grams + s.spectrogram_log_epsilon)
        
        # Clip spectrograms if indicated.
        if s.spectrogram_clipping_enabled:
            grams = tf.clip_by_value(
                grams, s.spectrogram_clipping_min, s.spectrogram_clipping_max)
            
        # Normalize spectrograms if indicated.
        if s.spectrogram_normalization_enabled:
            grams = \
                s.spectrogram_normalization_scale_factor * grams + \
                s.spectrogram_normalization_offset
        
        # Reshape spectrograms for input into Keras neural network.
        grams = self._reshape_grams(grams)
        
        # Create features dictionary.
        features = {self.output_feature_name: grams}
        
        if labels is None:
            
            return features
        
        else:
            # have labels
        
            # Reshape labels into a single 2D column.
            labels = tf.reshape(labels, (-1, 1))
            
            return features, labels

def get_env_v_isometry(hamAB, hamBA, rhoBA, rhoAB, w_isometry, v_isometry, unitary):

    indList1 = [[6,4,1,3],[9,11,8,-3],[1,2,8],[6,7,9],[3,5,2,-2],[4,5,7,10],[-1,10,11]]
    indList2 = [[3,4,1,2],[8,10,9,-3],[5,6,9],[5,7,8],[1,2,6,-2],[3,4,7,11],[-1,11,10]]
    indList3 = [[9,10,11,-1],[3,4,2,-3],[1,8,2],[1,5,3],[7,11,8,-2],[7,9,5,6],[10,6,4]]
    indList4 = [[7,5,-1,4],[6,3,-3,2],[7,-2,6],[4,1,2],[5,1,3]]

    vEnv = ncon.ncon([hamAB,rhoBA,w_isometry,tf.conj(w_isometry),unitary,tf.conj(unitary),tf.conj(v_isometry)],indList1)
    vEnv = vEnv + ncon.ncon([hamBA,rhoBA,w_isometry,tf.conj(w_isometry),unitary,tf.conj(unitary),tf.conj(v_isometry)],indList2)
    vEnv = vEnv + ncon.ncon([hamAB,rhoBA,w_isometry,tf.conj(w_isometry),unitary,tf.conj(unitary),tf.conj(v_isometry)],indList3)
    vEnv = vEnv + ncon.ncon([hamBA,rhoAB,tf.conj(v_isometry),w_isometry,tf.conj(w_isometry)],indList4)

    return vEnv

def prepare_tensor_QR_python(tensor, direction):
    """
    prepares an mps tensor using svd decomposition 
    Parameters:
    ---------------------
    tensor: tf.Tensors of shape(D1,D2,d)
            an mps tensor
    direction: int
               if >0 returns left orthogonal decomposition, if <0 returns right orthogonal decomposition

    Returns:
    ----------------------------
    direction>0:     out,s,v
                     out: a left isometric tf.Tensor of dimension (D1,D,d)
                     s  : the singular values of length D
                     v  : a right isometric tf.Tensor of dimension (D,D2)
    direction<0:     u,s,out
                     u  : a left isometric tf.Tensor of dimension (D1,D)
                     s  : the singular values of length D
                     out: a right isometric tf.Tensor of dimension (D,D2,d)

    """
    l1, d, l2 = tf.unstack(tf.shape(tensor))
    if direction in ('l', 'left', 1):
        temp = tf.reshape(tensor, [d * l1, l2])
        q, r = tf.linalg.qr(temp)
        Z = tf.linalg.norm(r)
        r/=Z        
        size1, size2 = tf.unstack(tf.shape(q))
        out = tf.reshape(q, [l1, d, size2])
        return out, r, Z

    if direction in ('r', 'right', -1):
        temp = tf.reshape(tensor, [l1, d * l2])
        q, r = tf.linalg.qr(tf.transpose(tf.conj(temp)))
        Z = tf.linalg.norm(r)
        r/=Z
        size1, size2 = tf.unstack(tf.shape(q))
        out = tf.reshape(tf.transpose(tf.conj(q)), [size2, d, l2])
        
        return tf.transpose(tf.conj(r)), out, Z

def ascend_op_2site_to_2site(mpo_2site, iso_012, iso_021):

    def _ascend(op, iso, iso_conj):
        return tensornetwork.ncon([iso_conj, op, iso], [(-1, 2, 0), (0, 1), (-2, 2, 1)])

    op2L, op2R = mpo_2site
    dtype = iso_012.dtype

    M = len(op2L)

    iso_021_conj = tf.conj(iso_021)
    op_asc_R = []
    for m in range(M):
        op_asc_R.append(_ascend(op2R[m], iso_021, iso_021_conj))

    iso_012_conj = tf.conj(iso_012)
    op_asc_L = []
    for m in range(M):
        op_asc_L.append(_ascend(op2L[m], iso_012, iso_012_conj))

    return op_asc_L, op_asc_R

 def _inverse(self, y):
   x = y
   if self.shift is not None:
     x -= self.shift
   if self._is_only_identity_multiplier:
     s = (tf.conj(self._scale)
          if self.adjoint and self._scale.dtype.is_complex
          else self._scale)
     return x / s
   # Solve fails if the op is singular so we may safely skip this assertion.
   x = self.scale.solvevec(x, adjoint=self.adjoint)
   return x