# Copyright 2021 D-Wave Systems Inc.
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

r"""Contains implementations for applying MPO's to MPS's with subsequent
compression.
"""

#####################################
## Load libraries/packages/modules ##
#####################################

# For deep copies of objects.
import copy

# For general array handling.
import numpy as np

# For calculating the eigenspectra of Hermitian matrices.
import scipy.linalg

# For creating tensor networks and performing contractions.
import tensornetwork as tn

# Assign an alias to the ``spinbosonchain`` library.
import spinbosonchain as sbc

# For performing SVD truncation sweeps, shifting orthogonal centers,
# and single-node SVD.
import spinbosonchain._svd

# For performing QR sweeps, and single-node QR factorizations.
import spinbosonchain._qr

# For finding the dominant eigenvectors of specially structured tensor networks.
import spinbosonchain._arnoldi

############################
## Authorship information ##
############################

__author__     = "D-Wave Systems Inc."
__copyright__  = "Copyright 2021"
__credits__    = ["Matthew Fitzpatrick"]
__maintainer__ = "D-Wave Systems Inc."
__email__      = "support@dwavesys.com"
__status__     = "Development"

##################################
## Define classes and functions ##
##################################

def apply_mpo_to_mps_and_compress(mpo_nodes,
                                  mps_nodes,
                                  compress_params,
                                  is_infinite):
    kwargs = {"mpo_nodes": mpo_nodes,
              "mps_nodes": mps_nodes,
              "compress_params": compress_params}
    if is_infinite:
        # The Schmidt spectrum for bonds between unit cells is calculated.
        schmidt_spectra = \
            apply_infinite_mpo_to_infinite_mps_and_compress(**kwargs)
    else:
        # The Schmidt spectra is calculated for all bonds if no variational
        # compression is to be performed, otherwise schmidt_spectra is set to
        # `None`.
        schmidt_spectra = \
            apply_finite_mpo_to_finite_mps_and_compress(**kwargs)

    return schmidt_spectra

def apply_infinite_mpo_to_infinite_mps_and_compress(mpo_nodes,
                                                    mps_nodes,
                                                    compress_params):
    apply_directly_mpo_to_mps(mpo_nodes, mps_nodes)

    # See comments in function spinbosonchain._svd.Lambda_Theta_form for a
    # description of the Lambda_Theta object.
    Lambda_Theta = sbc._svd.Lambda_Theta_form(mps_nodes)

    canonicalize_and_compress_infinite_mps(Lambda_Theta, compress_params)

    Lambda = Lambda_Theta[0]
    Theta_nodes = Lambda_Theta[1]

    L = len(mps_nodes)

    nodes_to_contract = [Lambda, Theta_nodes[0]]
    network_struct = [(-1, 1), (1, -2, -3)]
    mps_nodes[0] = tn.ncon(nodes_to_contract, network_struct)

    for r in range(1, L):
        mps_nodes[r] = Theta_nodes[r]

    schmidt_spectra = [Lambda]

    return schmidt_spectra

def apply_finite_mpo_to_finite_mps_and_compress(mpo_nodes,
                                                mps_nodes,
                                                compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.
    # [2]: New J. Phys. 12, 055026 (2010).

    # Put MPO in left-canonical form. MPS is assumed to be already in said form.
    kwargs = {"nodes": mpo_nodes, "normalize": False}
    sbc._qr.left_to_right_sweep(**kwargs)

    # See docs of class spinbosonchain.compress.Params for descriptions of the
    # various compression parameters used below.
    if compress_params.max_num_var_sweeps > 0:
        norm_of_mps_to_compress = \
            norm_of_mpo_mps_network_wo_compression(mpo_nodes, mps_nodes)
        initial_mps_nodes = copy.deepcopy(mps_nodes)
    else:
        initial_mps_nodes = None

    if compress_params.method == "direct":
        # Here we essentially follow the discussion in the paragraph containing
        # Eq. (16) of [2], and the subsequent paragraph, with a few changes:
        # we first perform a sweep from right to left without compression,
        # followed by a return sweep from left to right. Since there is no
        # compression in first sweep, we can use QR instead of SVD.
        apply_directly_mpo_to_mps(mpo_nodes, mps_nodes)
        kwargs["nodes"] = mps_nodes
        sbc._qr.right_to_left_sweep(**kwargs)
    else:
        # For a discussion on the zip-up method, see the paragraph above that
        # containing Eq. (17) of [2], and read through until the end of that
        # section. Note that mps_nodes is updated in place.
        zip_up(mpo_nodes, mps_nodes, compress_params)

    # Perform a return SVD truncation sweep. If the SVD truncation sweep is
    # followed by variational compression, then L-nodes are calculated and
    # cached while performing the SVD truncation sweep for efficiency. See
    # Sec. 4.5.2 of [1] for a discussion on the L-nodes and variational
    # compression. Note that mps_nodes is updated in place.
    kwargs = {"mpo_nodes": mpo_nodes,
              "mps_nodes": mps_nodes,
              "initial_mps_nodes": initial_mps_nodes,
              "compress_params": compress_params}
    L_cache, schmidt_spectra = return_svd_sweep(**kwargs)

    if compress_params.max_num_var_sweeps > 0:
        # If variational compression is performed, then the Schmidt spectra is
        # not calculated.
        schmidt_spectra = None
        kwargs["norm_of_mps_to_compress"] = norm_of_mps_to_compress
        kwargs["L_cache"] = L_cache
        variational_compression(**kwargs)

    return schmidt_spectra

def norm_of_mpo_mps_network_wo_compression(mpo_nodes, mps_nodes):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for some helpful context. In this scenario, the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2 of
    # [1]. The R-nodes are discussed in Sec. 4.5.2 of [1].
    imax = len(mps_nodes) - 1
    R = trivial_R(mpo_nodes, mps_nodes, num_legs=4)  # Right-most R-node.

    for i in range(imax, -1, -1):
        MWWR = contract_MWWR_network(mps_nodes[i], mpo_nodes[i], R)
        conj_mps_node = tn.conj(mps_nodes[i])
        conj_mps_node[1] ^ MWWR[3]
        conj_mps_node[2] ^ MWWR[4]
        output_edge_order = (MWWR[0], MWWR[1], MWWR[2], conj_mps_node[0])
        R = tn.contract_between(node1=conj_mps_node,
                                node2=MWWR,
                                output_edge_order=output_edge_order)

    nodes_to_contract = [R]
    network_struct = [(1, 2, 2, 1)]
    result = tn.ncon(nodes_to_contract, network_struct)
    result = float(np.sqrt(np.real(result.tensor)))

    return result

def trivial_R(mpo_nodes, mps_nodes, num_legs):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for a discussion on the R-nodes. The 'trivial'
    # R-node is essentially a multi-dimensional Kronecker delta.

    w_r = mpo_nodes[-1].shape[-1]  # w_r: right dangling mpo bond dimension.
    chi_r = mps_nodes[-1].shape[-1]  # chi_r: right dangling mps bond dimension.
    if num_legs == 3:
        tensor = np.zeros([chi_r, w_r, chi_r*w_r], dtype=np.complex128)
        for m2 in range(chi_r*w_r):
            m1 = 0 if chi_r == 1 else m2
            w = 0 if w_r == 1 else m2
            tensor[m1, w, m2] = 1
        R = tn.Node(tensor)
    else:
        tensor = np.zeros([chi_r, w_r, w_r, chi_r], dtype=np.complex128)
        for m in range(chi_r):
            for w in range(w_r):
                tensor[m, w, w, m] = 1
        R = tn.Node(tensor)

    return R

def trivial_L(mpo_nodes, mps_nodes):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for a discussion on the L-nodes. The 'trivial'
    # L-node is essentially a multi-dimensional Kronecker delta.

    if mps_nodes is None:
        L = None
        return L

    w_l = mpo_nodes[0].shape[0]  # w_l: left dangling mpo bond dimension.
    chi_l = mps_nodes[0].shape[0]  # chi_l: left dangling mps bond dimension.
    tensor = np.zeros([chi_l, w_l, chi_l*w_l], dtype=np.complex128)
    for m2 in range(chi_l*w_l):
        m1 = 0 if chi_l == 1 else m2
        w = 0 if w_l == 1 else m2
        tensor[m1, w, m2] = 1
    L = tn.Node(tensor)

    return L

def contract_MR_network(M, R, num_R_legs=3):
    # See comments in function contract_MWWR_network for context.

    nodes_to_contract = [M, R]
    if num_R_legs == 3:
        network_struct = [(-1, -2, 1), (1, -3, -4)]
    else:
        network_struct = [(-1, -2, 1), (1, -3, -4, -5)]
    MR = tn.ncon(nodes_to_contract, network_struct)

    return MR

def contract_MWR_network(M, W, R, num_R_legs=3):
    # See comments in function contract_MWWR_network for context.

    MR = contract_MR_network(M, R, num_R_legs)

    W[2] ^ MR[1]
    W[3] ^ MR[2]
    if num_R_legs == 3:
        output_edge_order = (MR[0], W[0], W[1], MR[3])

    else:
        output_edge_order = (MR[0], W[0], W[1], MR[3], MR[4])
    MWR = tn.contract_between(node1=W,
                              node2=MR,
                              output_edge_order=output_edge_order)

    return MWR

def contract_MWWR_network(M, W, R):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # M is the MPS node at the current site of interest. W is the MPO node being
    # applied to M. R is the R-node just to the right of the current site of
    # interest. See Sec. 4.5.2 of [1] for a discussion on the R-nodes.

    MWR = contract_MWR_network(M, W, R, num_R_legs=4)
    W = tn.conj(W)

    W[2] ^ MWR[2]
    W[3] ^ MWR[3]
    output_edge_order = (MWR[0], MWR[1], W[0], W[1], MWR[4])
    MWWR = tn.contract_between(node1=W,
                               node2=MWR,
                               output_edge_order=output_edge_order)

    return MWWR

def contract_LM_network(L, M):
    # See comments in function contract_LMW_network for context.

    nodes_to_contract = [L, M]
    network_struct = [(1, -3, -4), (1, -2, -1)]
    LM = tn.ncon(nodes_to_contract, network_struct)

    return LM

def contract_LMW_network(L, M, W):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # M is the MPS node at the current site of interest to which a MPO is to be
    # applied at some point. L is the L-node just to the left of the current
    # site of interest. See Sec. 4.5.2 of [1] for a discussion on the L-nodes.

    LM = contract_LM_network(L, M)

    LM[1] ^ W[2]
    LM[2] ^ W[0]
    output_edge_order = (LM[0], W[3], W[1], LM[3])
    LMW = tn.contract_between(node1=LM,
                              node2=W,
                              output_edge_order=output_edge_order)

    return LMW

def zip_up(mpo_nodes, mps_nodes, compress_params):
    # [2]: New J. Phys. 12, 055026 (2010).

    # For a discussion on the zip-up method, see the paragraph above that
    # containing Eq. (17) of [2], and read through until the end of that
    # section.

    imax = len(mps_nodes) - 1
    U = None
    S = None

    # See second last paragraph of Sec. 2.1.3 of [2] for the reasoning behind
    # the code block below.
    compress_params = copy.deepcopy(compress_params)
    if compress_params.max_num_singular_values is not None:
        compress_params.max_num_singular_values *= 2
    if compress_params.max_trunc_err is not None:
        compress_params.max_trunc_err /= 10
    if compress_params.svd_rel_tol is not None:
        compress_params.svd_rel_tol /= 10

    for i in range(imax, -1, -1):
        node_idx = i
        U, S = update_mps_node_in_zip_up(node_idx,
                                         mpo_nodes,
                                         mps_nodes,
                                         U,
                                         S,
                                         compress_params)

    return None

def update_mps_node_in_zip_up(node_idx,
                              mpo_nodes,
                              mps_nodes,
                              U,
                              S,
                              compress_params):
    # [2]: New J. Phys. 12, 055026 (2010).

    # For a discussion on the zip-up method, see the paragraph above that
    # containing Eq. (17) of [2], and read through until the end of that
    # section.

    i = node_idx
    imax = len(mps_nodes) - 1
    mpo_node = mpo_nodes[i]
    mps_node = mps_nodes[i]

    if i == imax:
        nodes_to_contract = (mpo_node, mps_node)
        network_struct = [(-1, -3, 1, -4), (-2, 1, -5)]
        temp_node_2 = tn.ncon(nodes_to_contract, network_struct)
        tn.flatten_edges([temp_node_2[3], temp_node_2[4]])
    elif 0 <= i < imax:
        nodes_to_contract = (mps_node, U, S)
        network_struct = [(-1, -2, 2), (-3, 2, 1), (1, -4)]
        temp_node_1 = tn.ncon(nodes_to_contract, network_struct)
        temp_node_1[1] ^ mpo_node[2]
        temp_node_1[2] ^ mpo_node[3]
        output_edge_order = \
            (mpo_node[0], temp_node_1[0], mpo_node[1], temp_node_1[3])
        temp_node_2 = tn.contract_between(node1=temp_node_1,
                                          node2=mpo_node,
                                          output_edge_order=output_edge_order)

    if 0 < i <= imax:
        left_edges = (temp_node_2[0], temp_node_2[1])
        right_edges = (temp_node_2[2], temp_node_2[3])
        U, S, V_dagger = sbc._svd.split_node_full(temp_node_2,
                                                  left_edges,
                                                  right_edges,
                                                  compress_params)
        mps_nodes[i] = V_dagger
    elif i == 0:
        tn.flatten_edges([temp_node_2[0], temp_node_2[1]])
        temp_node_2.reorder_edges([temp_node_2[2],
                                   temp_node_2[0],
                                   temp_node_2[1]])
        mps_nodes[i] = temp_node_2

    return U, S

def return_svd_sweep(mpo_nodes, mps_nodes, initial_mps_nodes, compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # Perform a return SVD truncation sweep. If the SVD truncation sweep is
    # followed by variational compression, then L-nodes are calculated and
    # cached while performing the SVD truncation sweep for efficiency. See
    # Sec. 4.5.2 of [1] for a discussion on the L-nodes and variational
    # compression. Readers should note that the mpo-mps network without
    # compression is denoted by |psi> in Sec. 4.5.2 and should keep this in
    # mind when reading about the L-nodes.

    # Note further that initial_mps_nodes represents the MPS before the MPO
    # application, whereas at the end of this function, mps_nodes represents
    # the MPS obtained after applying the MPO.

    imin = 0
    imax = len(mps_nodes) - 2
    L = trivial_L(mpo_nodes, initial_mps_nodes)
    L_cache = [L]
    schmidt_spectra = []

    # Normalize MPS now if no variational compression is to be performed.
    # Otherwise, normalization is performed after variational compression.
    normalize_schmidt_spectra = (compress_params.max_num_var_sweeps == 0)

    for i in range(imin, imax+1):
        kwargs = {"nodes": mps_nodes,
                  "current_orthogonal_center_idx": i,
                  "compress_params": compress_params,
                  "normalize_schmidt_spectra": normalize_schmidt_spectra}
        U, S, V_dagger = sbc._svd.shift_orthogonal_center_to_the_right(**kwargs)
        schmidt_spectra.insert(0, S)

        if compress_params.max_num_var_sweeps > 0:
            M = initial_mps_nodes[i]
            W = mpo_nodes[i]
            LMW = contract_LMW_network(L, M, W)
            conj_mps_node = tn.conj(mps_nodes[i])
            LMW[2] ^ conj_mps_node[1]
            LMW[3] ^ conj_mps_node[0]
            output_edge_order = (LMW[0], LMW[1], conj_mps_node[2])
            L = tn.contract_between(node1=LMW,
                                    node2=conj_mps_node,
                                    output_edge_order=output_edge_order)
            L_cache.append(L)

    if compress_params.max_num_var_sweeps == 0:
        # Normalize MPS such that its 'norm' equal unity.
        mps_nodes[-1] /= tn.norm(mps_nodes[-1])

    return L_cache, schmidt_spectra

def variational_compression(mpo_nodes,
                            mps_nodes,
                            initial_mps_nodes,
                            norm_of_mps_to_compress,
                            L_cache,
                            compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # This function essentially implements the variational compression procedure
    # described in Sec. 4.5.2 of [1], see therein for discussion on the L- and
    # R-nodes. Readers should note that the mpo-mps network without compression
    # is denoted by |psi> in Sec. 4.5.2 and should keep this in mind when
    # reading about the L- and R-nodes.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    imax = len(mps_nodes) - 1
    R_cache = [trivial_R(mpo_nodes, initial_mps_nodes, num_legs=3)]

    # A full sweep goes right to left, then left to right.
    sweep_count = 0
    while sweep_count < compress_params.max_num_var_sweeps:
        kwargs = {"mpo_nodes": mpo_nodes,
                  "mps_nodes": mps_nodes,
                  "initial_mps_nodes": initial_mps_nodes,
                  "norm_of_mps_to_compress": norm_of_mps_to_compress,
                  "L_cache": L_cache,
                  "R_cache": R_cache,
                  "compress_params": compress_params}

        if variational_compression_has_converged(**kwargs):
            break

        kwargs = {"mpo_nodes": mpo_nodes,
                  "mps_nodes": mps_nodes,
                  "initial_mps_nodes": initial_mps_nodes,
                  "L_cache": L_cache,
                  "R_cache": R_cache}

        for i in range(imax, 0, -1):
            update_node_and_shift_left_in_variational_compression(**kwargs)
        for i in range(imax):
            update_node_and_shift_right_in_variational_compression(**kwargs)

        sweep_count += 1

    # Normalize MPS such that its 'norm' equal unity.
    mps_nodes[-1] /= tn.norm(mps_nodes[-1])

    return None

def variational_compression_has_converged(mpo_nodes,
                                          mps_nodes,
                                          initial_mps_nodes,
                                          norm_of_mps_to_compress,
                                          L_cache,
                                          R_cache,
                                          compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for a discussion on the variational compression
    # procedure implemented in this module. Readers should note that the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2.

    # This function checks whether the variational compression procedure has
    # achieved convergence. This is done by evaluating the left hand side of
    # Eq. (150) of [1].

    # 'compressed' and 'uncompressed' MPS's are refering to |psi'> and |psi>
    # in Sec. 4.5.2 [see above note on |psi>]. Hence, in both cases these are
    # MPS's obtained after applying the MPO.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    overlap = overlap_btwn_compressed_and_uncompressed_mps(L_cache,
                                                           initial_mps_nodes,
                                                           mpo_nodes,
                                                           mps_nodes,
                                                           R_cache)

    nodes_to_contract = (mps_nodes[-1], tn.conj(mps_nodes[-1]))
    network_struct = [(1, 2, 3), (1, 2, 3)]
    norm_sq_of_compressed_mps = tn.ncon(nodes_to_contract, network_struct)
    norm_of_compressed_mps = \
        float(np.sqrt(np.real(norm_sq_of_compressed_mps.tensor)))

    var_rel_err_sq = (np.abs(norm_of_mps_to_compress * norm_of_mps_to_compress
                             - 2 * np.real(overlap)
                             + norm_of_compressed_mps * norm_of_compressed_mps)
                      / norm_of_mps_to_compress)
    var_rel_err = np.sqrt(var_rel_err_sq)

    if var_rel_err < compress_params.var_rel_tol:
        result = True
    else:
        result = False

    return result

def overlap_btwn_compressed_and_uncompressed_mps(L_cache,
                                                 initial_mps_nodes,
                                                 mpo_nodes,
                                                 mps_nodes,
                                                 R_cache):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for a discussion on the variational compression
    # procedure implemented in this module. Readers should note that the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2.

    # 'compressed' and 'uncompressed' MPS's are refering to |psi'> and |psi>
    # in Sec. 4.5.2 [see above note on |psi>]. Hence, in both cases these are
    # MPS's obtained after applying the MPO.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    L = L_cache[-1]
    M = initial_mps_nodes[-1]
    W = mpo_nodes[-1]
    R = R_cache[0]

    MWR = contract_MWR_network(M, W, R)

    L[0] ^ MWR[0]
    L[1] ^ MWR[1]
    output_edge_order = (L[2], MWR[2], MWR[3])
    LMWR = tn.contract_between(node1=L,
                               node2=MWR,
                               output_edge_order=output_edge_order)

    conj_mps_node = tn.conj(mps_nodes[-1])

    for edge_idx in range(3):
        conj_mps_node[edge_idx] ^ LMWR[edge_idx]

    overlap = tn.contract_between(node1=conj_mps_node, node2=LMWR)
    overlap = complex(overlap.tensor)

    return overlap

def update_node_and_shift_left_in_variational_compression(mpo_nodes,
                                                          mps_nodes,
                                                          initial_mps_nodes,
                                                          L_cache,
                                                          R_cache):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # This function essentially implements a single shift to the left in the
    # variational compression procedure described in Sec. 4.5.2 of [1], see
    # therein for discussion on the L- and R-nodes. Readers should note that the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2 and
    # should keep this in mind when reading about the L- and R-nodes.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    orthogonal_center_idx = len(L_cache) - 1
    L = L_cache[-1]
    M = initial_mps_nodes[orthogonal_center_idx]
    W = mpo_nodes[orthogonal_center_idx]
    R = R_cache[0]

    MWR = contract_MWR_network(M, W, R)

    L[0] ^ MWR[0]
    L[1] ^ MWR[1]
    output_edge_order = (L[2], MWR[2], MWR[3])
    LMWR = tn.contract_between(node1=L,
                               node2=MWR,
                               output_edge_order=output_edge_order)

    L_cache.pop()

    QR = tn.conj(LMWR)
    left_node, right_node = sbc._qr.split_node(node=QR,
                                               left_edges=(QR[1], QR[2]),
                                               right_edges=(QR[0],))
    new_edge_order = [left_node[2], left_node[0], left_node[1]]
    B = tn.conj(left_node.reorder_edges(new_edge_order))
    mps_nodes[orthogonal_center_idx] = B
    conj_B = tn.conj(B)

    conj_B[1] ^ MWR[2]
    conj_B[2] ^ MWR[3]
    output_edge_order = (MWR[0], MWR[1], conj_B[0])
    next_R = tn.contract_between(node1=conj_B,
                                 node2=MWR,
                                 output_edge_order=output_edge_order)
    R_cache.insert(0, next_R)

    return None

def update_node_and_shift_right_in_variational_compression(mpo_nodes,
                                                           mps_nodes,
                                                           initial_mps_nodes,
                                                           L_cache,
                                                           R_cache):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # This function essentially implements a single shift to the right in the
    # variational compression procedure described in Sec. 4.5.2 of [1], see
    # therein for discussion on the L- and R-nodes. Readers should note that the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2 and
    # should keep this in mind when reading about the L- and R-nodes.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    orthogonal_center_idx = len(L_cache) - 1
    L = L_cache[-1]
    M = initial_mps_nodes[orthogonal_center_idx]
    W = mpo_nodes[orthogonal_center_idx]
    R = R_cache[0]

    LMW = contract_LMW_network(L, M, W)

    LMW[0] ^ R[0]
    LMW[1] ^ R[1]
    output_edge_order = (LMW[3], LMW[2], R[2])
    LMWR = tn.contract_between(node1=LMW,
                               node2=R,
                               output_edge_order=output_edge_order)
    R_cache.pop(0)

    QR = LMWR
    left_node, right_node = sbc._qr.split_node(node=QR,
                                               left_edges=(QR[0], QR[1]),
                                               right_edges=(QR[2],))
    A = left_node
    mps_nodes[orthogonal_center_idx] = A
    conj_A = tn.conj(A)

    if orthogonal_center_idx == len(mps_nodes) - 2:
        nodes_to_contract = [right_node, mps_nodes[-1]]
        network_struct = [(-1, 1), (1, -2, -3)]
        mps_nodes[-1] = tn.ncon(nodes_to_contract, network_struct)

    LMW[2] ^ conj_A[1]
    LMW[3] ^ conj_A[0]
    output_edge_order = (LMW[0], LMW[1], conj_A[2])
    next_L = tn.contract_between(node1=LMW,
                                 node2=conj_A,
                                 output_edge_order=output_edge_order)
    L_cache.append(next_L)

    return None

def apply_directly_mpo_to_mps(mpo_nodes, mps_nodes):
    for idx, (mpo_node, mps_node) in enumerate(zip(mpo_nodes, mps_nodes)):
        mps_nodes[idx] = apply_directly_mpo_node_to_mps_node(mpo_node, mps_node)

    return None

def apply_directly_mpo_node_to_mps_node(mpo_node, mps_node):
    nodes_to_contract = (mpo_node, mps_node)
    network_struct = [(-1, -3, 1, -4), (-2, 1, -5)]
    new_mps_node = tn.ncon(nodes_to_contract, network_struct)

    # A single call to tn.flatten_edges will take the flatten edge set and move
    # it to the right end of the list of edges, hence the necessary call to
    # reorder_edges.
    tn.flatten_edges([new_mps_node[0], new_mps_node[1]])
    tn.flatten_edges([new_mps_node[1], new_mps_node[2]])
    new_mps_node.reorder_edges([new_mps_node[1],
                                new_mps_node[0],
                                new_mps_node[2]])

    return new_mps_node

def canonicalize_and_compress_infinite_mps(Lambda_Theta, compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.
    # [3]: Phys. Rev. B 78, 155117 (2008)
    # [4]: arXiv.0804.2509 (2008)

    # This function compresses a given infinite MPS. The procedure implemented
    # here draws from [1], [3], and [4]. Notation-wise we follow mostly [1].
    # In [1], the relevant content is found in Sec. 10.5.

    # See comments in function spinbosonchain._svd.Lambda_Theta_form for a
    # description of the Lambda_Theta object.

    # V_L, V_R, X, X_inv, Y, and Y_inv are the same as those defined in
    # Sec. 10.5 of [1] except that we have generalized to a L-site unit cell
    # where L>0.
    V_L = calc_V_L(Lambda_Theta)
    V_R = calc_V_R(Lambda_Theta)
    X, X_inv = calc_X_and_X_inv(V_L)
    Y, Y_inv = calc_Y_and_Y_inv(V_R)

    old_Lambda = Lambda_Theta[0]
    old_Theta_nodes = Lambda_Theta[1]

    # Here we obtain the new singular value matrix for the bond between unit
    # cells. Note that the singular value spectrum is truncated.
    X_old_Lambda_Y = tn.ncon([X, old_Lambda, Y], [(-1, 1), (1, 2), (2, -2)])
    U, S, V_dagger = sbc._svd.split_node_full(node=X_old_Lambda_Y,
                                              left_edges=(X_old_Lambda_Y[0],),
                                              right_edges=(X_old_Lambda_Y[1],),
                                              compress_params=compress_params)

    new_Lambda = S / tn.norm(S)

    L = len(old_Theta_nodes)

    # Here we construct the Z_L and Z_R that are used to transform the Theta
    # node, which is a part of the canonicalization procedure.
    if L == 1:
        # Here we are essentially following Fig. 2.ii of [3]. Note that our
        # V_dagger is their V; our Y_inv is their X_inv; and our X_inv is their
        # Y_T_inv.
        Z_L = tn.ncon([V_dagger, Y_inv], [(-1, 1), (1, -2)])
        Z_R = tn.ncon([X_inv, U], [(-1, 1), (1, -2)])
    else:
        # Here we are essentially following the paragraph containing Eq. (355)
        # of [1] with a few minor changes.
        new_Lambda_inv_tensor = \
            np.diag(1 / np.diag(np.array(new_Lambda.tensor)))
        new_Lambda_inv = tn.Node(new_Lambda_inv_tensor)

        conj_U = tn.conj(U)
        V_T = tn.conj(V_dagger)

        nodes_to_contract = [new_Lambda_inv, conj_U, X, old_Lambda]
        network_struct = [(-1, 3), (1, 3), (1, 2), (2, -2)]
        Z_L = tn.ncon(nodes_to_contract, network_struct)

        nodes_to_contract = [old_Lambda, Y, V_T, new_Lambda_inv]
        network_struct = [(-1, 1), (1, 2), (3, 2), (3, -2)]
        Z_R = tn.ncon(nodes_to_contract, network_struct)

    new_Theta_nodes = old_Theta_nodes

    nodes_to_contract = [Z_L, new_Theta_nodes[0]]
    network_struct = [(-1, 1), (1, -2, -3)]
    new_Theta_nodes[0] = tn.ncon(nodes_to_contract, network_struct)

    nodes_to_contract = [new_Theta_nodes[-1], Z_R]
    network_struct = [(-1, -2, 1), (1, -3)]
    new_Theta_nodes[-1] = tn.ncon(nodes_to_contract, network_struct)

    compress_and_normalize_Theta(new_Lambda, new_Theta_nodes, compress_params)

    Lambda_Theta[0] = new_Lambda
    Lambda_Theta[1] = new_Theta_nodes

    return None

def compress_and_normalize_Theta(Lambda, Theta_nodes, compress_params):
    # [3]: Phys. Rev. B 78, 155117 (2008)

    # See comments in function spinbosonchain._svd.Lambda_Theta_form for a
    # description of the Lambda_Theta object.

    # Essentially following [3], here we compress and normalize the Theta node
    # of the Lambda-Theta pair. We apply a direct SVD compression procedure.

    kwargs = {"nodes": Theta_nodes, "normalize": False}
    sbc._qr.left_to_right_sweep(**kwargs)
    kwargs["compress_params"] = compress_params
    kwargs["normalize"] = True
    sbc._svd.right_to_left_sweep(**kwargs)

    # Normalize the Theta node such that the MPS it represents has a state
    # 'norm' equal to unity.
    M = tn.ncon([Lambda, Theta_nodes[0]], [(-1, 1), (1, -2, -3)])
    Theta_nodes[0] /= tn.norm(M)

    return None

def calc_V_L(Lambda_Theta):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # V_L is the same as that defined in Sec. 10.5 of [1] except that we have
    # generalized to a L-site unit cell where L>0.

    kwargs = {"Lambda_Theta": Lambda_Theta,
              "unit_cell_type": "left",
              "krylov_dim": 10,
              "epsilon_D": 1.0e-14}

    mu_L, V_L = sbc._arnoldi.dominant_eigpair_of_transfer_matrix(**kwargs)

    return V_L

def calc_V_R(Lambda_Theta):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # V_R is the same as that defined in Sec. 10.5 of [1] except that we have
    # generalized to a L-site unit cell where L>0.

    kwargs = {"Lambda_Theta": Lambda_Theta,
              "unit_cell_type": "right",
              "krylov_dim": 10,
              "epsilon_D": 1.0e-14}

    mu_R, V_R = sbc._arnoldi.dominant_eigpair_of_transfer_matrix(**kwargs)

    return V_R

def calc_X_and_X_inv(V_L):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # V_L, X, and X_inv are the same as those defined in Sec. 10.5 of [1] except
    # that we have generalized to a L-site unit cell where L>0.

    if V_L.backend.name != "numpy":
        sbc._backend.tf_to_np(V_L)

    try:
        D, W = scipy.linalg.eigh(V_L.tensor, driver='ev')
    except TypeError as err:  # Occurs with some versions scipy.
        D, W = scipy.linalg.eigh(V_L.tensor)
    D = np.abs(D)
    tol = 1.0e-14
    D[D<tol] = tol

    X_dagger = W @ np.diag(np.sqrt(D))
    X = np.conj(np.transpose(X_dagger))
    X = tn.Node(X)

    # Note that X_inv is really a pseudo-inverse of X.
    X_inv = W @ np.diag(1/np.sqrt(D))
    X_inv = tn.Node(X_inv)

    return X, X_inv

def calc_Y_and_Y_inv(V_R):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # V_R, Y, and Y_inv are the same as those defined in Sec. 10.5 of [1] except
    # that we have generalized to a L-site unit cell where L>0.

    if V_R.backend.name != "numpy":
        sbc._backend.tf_to_np(V_R)

    try:
        D, W = scipy.linalg.eigh(V_R.tensor, driver='ev')
    except TypeError as err:  # Occurs with some versions scipy.
        D, W = scipy.linalg.eigh(V_R.tensor)
    D = np.abs(D)
    tol = 1.0e-14
    D[D<tol] = tol

    Y = W @ np.diag(np.sqrt(D))
    Y = tn.Node(Y)  # Convert from numpy array to tensornetwork node.

    # Note that Y_inv is really a pseudo-inverse of Y.
    W_dagger = np.conj(np.transpose(W))
    Y_inv = np.diag(1/np.sqrt(D)) @ W_dagger
    Y_inv = tn.Node(Y_inv)

    return Y, Y_inv


import gin
import lime
import pandas as pd
import pytest
import numpy as np
import numpy.testing as npt
import tensorflow as tf
import lime.for_biologists.architectures.vcharge

def test_mini_fit():
    ds = gin.i_o.from_sdf.to_ds('data/mols.sdf', True)

    ds = ds.map(lambda atoms, adjacency_map, coordinates, charges:
        (
            tf.py_function(
                lambda atoms, adjacency_map: \
                    tf.reshape(
                        lime.for_biologists.architectures.vcharge\
                        .VChargeTyping(
                        [atoms, adjacency_map]).get_assignment(),
                    [-1]),
                [atoms, adjacency_map],
                tf.int64),
            adjacency_map,
            coordinates,
            charges
        ))

    charge_model = lime.for_biologists.architectures.vcharge.VCharge()
    lime.for_biologists.architectures.vcharge.train(
        ds,
        charge_model)

test_mini_fit()

'''
df = pd.read_csv('data/SAMPL.csv')
df = df[~df['smiles'].str.contains('B')]
# df = df[~df['smiles'].str.contains('P')]
df = df[~df['smiles'].str.contains('\+')]
df = df[~df['smiles'].str.contains('\-')]
smiles_array = df[['smiles']].values.flatten()

@pytest.mark.parametrize('smiles', smiles_array)
def test_VCharge_typing_mutually_exclusive(smiles):
    mol = gin.i_o.from_smiles.smiles_to_mol(smiles)
    mol = gin.deterministic.hydrogen.add_hydrogen(mol)
    typing = lime.for_biologists.architectures.vcharge.VChargeTyping(mol)
    gaff_typing = [
        tf.expand_dims(
            getattr(typing, 'is_' + str(idx)).__call__(),
            0) \
        for idx in range(1, 40)]

    gaff_typing = tf.concat(
        gaff_typing,
        axis=0)

    print(tf.where(
        tf.logical_not(
            tf.equal(
                tf.math.count_nonzero(
                    gaff_typing,
                    axis=0),
                tf.constant(1, dtype=tf.int64)))))

    print(tf.boolean_mask(
        gaff_typing,
        tf.logical_not(
            tf.equal(
                tf.math.count_nonzero(
                    gaff_typing,
                    axis=0),
                tf.constant(1, dtype=tf.int64))),
            axis=1))

    tf.debugging.Assert(
        tf.reduce_all(
            tf.equal(
                tf.math.count_nonzero(
                    gaff_typing,
                    axis=0),
                tf.constant(1, dtype=tf.int64))),
        data=[smiles])
'''


#!/usr/bin/env python

import rospy
import actionlib
import sys
import numpy as np
import math
import time
from geometry_msgs.msg import Pose, PoseWithCovarianceStamped, Point, Quaternion, Twist, PoseStamped
from std_msgs.msg import String, Bool
from kobuki_msgs.msg import Sound, Led
from move_base_msgs.msg import MoveBaseAction, MoveBaseGoal
from std_srvs.srv import Empty
from sensor_msgs.msg import Joy, LaserScan
from smach import State,StateMachine
import smach_ros
import cv2
import costmap_2d
from sound_play.msg import SoundRequest
from sound_play.libsoundplay import SoundClient
from map_msgs.msg import OccupancyGridUpdate
from nav_msgs.msg import OccupancyGrid

def global_localize():
	rospy.wait_for_service('global_localization')
	localize_service = rospy.ServiceProxy('global_localization', Empty)
	try:
		res = localize_service()
	except:
		print("Service did not process request")

def joy_callback(msg):
	global x_pressed
	global y_pressed
	global move_base
	x_pressed = msg.buttons[0]
	y_pressed = msg.buttons[2]

	if y_pressed == 1:
		move_base.cancel_goal()
		rospy.signal_shutdown('Y Pressed')

	if x_pressed == 1:
		print 'X Pressed'
		# play a sound? Why do we have this?

def vel_callback(msg):
	global velocity, vel_pub
	velocity.x = msg.linear.x
	velocity.y = msg.linear.x
	velocity.z = msg.linear.x
	vel_pub.publish(velocity)

def ar_callback(msg):
	global ar_found, ar_location
	# print msg.data
	if len(msg.data) > 0:
		# maybe if we get 10 msgs within a certain time limit set to true?
		ar_found = True
		ar_location = msg.data.split()
		ar_location[0] = float(ar_location[0])
		ar_location[1] = float(ar_location[1])
		# print 'ar location: ', ar_location
	else:
		ar_found = False

def pose_callback(data):
	global current_pose, pose_pub
	current_pose = data.pose
	pose_pub.publish(data)
	# print current_pose

def scan_callback(msg):
	global g_range_ahead
	largeDistance = 100000
	arrayLen = 201
	triPoint = [i for i in range (0,arrayLen)]
	for i in range(0,arrayLen):
		triPoint[i] = msg.ranges[(len(msg.ranges)/2) - ((arrayLen-1)/2) + i]
	g_range_ahead = largeDistance
	for i in range (0, len(triPoint)):
		if (triPoint[i] < g_range_ahead and not math.isnan(triPoint[i])):
			g_range_ahead = triPoint[i]
	if g_range_ahead == largeDistance:
		g_range_ahead = 0

def original_cost_map_callback(data):
	global costmap_og, map_height, map_width
	# print data
	# print('og costmap')
	print(len(data.data))
	# print('----------------------------------------------------------------------------------------------------')
	costmap_og = data

def cost_map_callback(data):
	global costmap, map_width, map_height, map_initialized
# 	header:
#   seq: 5
#   stamp:
#     secs: 1523406448
#     nsecs: 109183886
#   frame_id: odom
# x: 24
# y: 6
# width: 56
# height: 56
# data: [0,
	# print (data)
	print('costmap update: x: ', data.x, ' y:', data.y, ' height: ', data.height, ' width: ', data.width)
	# print(len(data.data))
	# print('___________________________________________________________________________________________________________')
	if not map_initialized:
		map_width = data.width
		map_height = data.height
		map_initialized = True
	costmap = data

def goal_pose(return_pose):

	goal_pose = MoveBaseGoal()
	goal_pose.target_pose.header.frame_id = 'map'
	goal_pose.target_pose.header.stamp = rospy.Time.now()
	goal_pose.target_pose.pose.position.x = return_pose.pose.position.x
	goal_pose.target_pose.pose.position.y = return_pose.pose.position.y
	goal_pose.target_pose.pose.position.z = return_pose.pose.position.z
	goal_pose.target_pose.pose.orientation.w = return_pose.pose.orientation.w
	goal_pose.target_pose.pose.orientation.x = return_pose.pose.orientation.x
	goal_pose.target_pose.pose.orientation.y = return_pose.pose.orientation.y
	goal_pose.target_pose.pose.orientation.z = return_pose.pose.orientation.z

	return goal_pose

class Start(State):
	def __init__(self):
		State.__init__(self, outcomes=['success'])

	def execute(self, userdata):
		global x_pressed
		print 'START'
		while x_pressed == 0:
			time.sleep(0.01)
		global_localize()
		return 'success'

class Localize(State):
	def __init__(self):
		State.__init__(self, outcomes=['success'])

	def execute(self, userdata):
		global g_range_ahead
		print 'LOCALIZE'
		rotateTime = rospy.Time.now() + rospy.Duration(10)
		forwardTime = rospy.Time.now() + rospy.Duration(12)
		count = 0

		return 'success'

		while True:
			if g_range_ahead > 1.4 and (rospy.Time.now() < rotateTime):
				twist = Twist()
				twist.angular.z = 0.4
				twist.linear.x = 0
				cmd_vel_pub.publish(twist)
			elif g_range_ahead < 1.51:
				twist = Twist()
				twist.angular.z = 0.3
				twist.linear.x = 0
				cmd_vel_pub.publish(twist)
			elif g_range_ahead > 1.5 and (rospy.Time.now() < forwardTime):
				twist = Twist()
				twist.angular.z = 0
				twist.linear.x = 0.3
				cmd_vel_pub.publish(twist)
			elif count > 0:
				break
			else:
				rotateTime = rospy.Time.now() + rospy.Duration(10)
				forwardTime = rospy.Time.now() + rospy.Duration(12)
				count += 1

		return 'success'

class Seek(State):
	def __init__(self):
		State.__init__(self, outcomes=['success', 'failure'])
		global move_base, ar_found, ar_location, waypoints, map_height, map_width, costmap

	def execute(self, data):
		print 'SEEK'
		while not ar_found:
			self.find_lowest_cost_area()
			continue
			# move around according to cost map

		return 'success'

	def move(self, goal):
		global move_base
		move_base.send_goal(goal)
		cap_time = move_base.wait_for_result(rospy.Duration(60))
		if not cap_time:
			move_base.cancel_goal()
		else:
			state = move_base.get_state()
			if state == actionlib.simple_action_client.GoalStatus.SUCCEEDED:
				self.atPoint = True
				print 'goal succeeded'
				return

	def find_lowest_cost_area(self):
		# current position x, y -> width*y + x = current index in list

		temp_costmap = costmap

		# top
		top_sum = 0
		if (temp_costmap.y + 15) < temp_costmap.height and (temp_costmap.x + 5) < temp_costmap.width and (temp_costmap.x - 5) > 0:
			for i in range(10, 15):
				for j in range(-5, 5):
					if temp_costmap.data[(temp_costmap.y * temp_costmap.width) + (i * temp_costmap.width) + (temp_costmap.x + j)] >= 0:
						top_sum += temp_costmap.data[(temp_costmap.y * temp_costmap.width) + (i * temp_costmap.width) + (temp_costmap.x + j)]

		print 'top sum: ', top_sum

		# bot
		bot_sum = 0
		if (temp_costmap.y - 15) > 0 and (temp_costmap.x + 5) < temp_costmap.width and (temp_costmap.x - 5) > 0:
			for i in range(-10, -15):
				for j in range(-5, 5):
					if temp_costmap.data[(temp_costmap.y * temp_costmap.width) + (i * temp_costmap.width) + (temp_costmap.x + j)] >= 0:
						bot_sum += temp_costmap.data[(temp_costmap.y * temp_costmap.width) + (i * temp_costmap.width) + (temp_costmap.x + j)]

		print 'bot sum: ', bot_sum

class CatchHider(State):
	def __init__(self):
		State.__init__(self, outcomes=['success', 'failure'])
		global move_base, ar_found, ar_location, waypoints, cmd_vel_pub
		self.twist = Twist()

	def execute(self, userdata):
		print 'CATCH HIDER'
		caught = False

		while ar_found and not caught:
			self.pursue(ar_location)

		if caught:
				return 'success'
		else:
			return 'failure'

	def pursue(self, location):
		# lin_kp = 10.0
		ang_kp = 0.005
		# x = lin_kp * (1/float(lin_diff))

		ideal = 315
		c = location[0]
		ang_diff = ideal - float(c)
		# if ang_diff < -10:
		#     z = -0.1
		# elif ang_diff > 10:
		#     z = 0.1
		# else:
		#     z = 0
		z = ang_diff * ang_kp

		# if x > 0.4:
		#     x = 0.4
		# # print x
		# if x < 0.03:
		#     return True

		# print "z: ", z

		self.twist.angular.z = z
		self.twist.linear.x = 0.0 #should be 0.7 when full speed chasing
		cmd_vel_pub.publish(self.twist)

class Win(State):
	def __init__(self):
		State.__init__(self, outcomes=['success'])

	def execute(self, userdata):
		print 'WIN'
		return 'success'

class Lose(State):
	def __init__(self):
		State.__init__(self, outcomes=['failure'])

	def execute(self, userdata):
		print 'LOSE'
		return 'failure'

if __name__ == '__main__':
	global x_pressed, y_pressed, ar_found, ar_location, current_pose, move_base, cmd_vel_pub
	global g_range_ahead, waypoints, pose_pub, velocity, map_height, map_width, map_initialized
	# global soundhandle
	velocity = Point()
	x_pressed = 0
	y_pressed = 0
	g_range_ahead = 0
	ar_found = False
	ar_location = []
	current_pose = None
	move_base = actionlib.SimpleActionClient('move_base', MoveBaseAction)
	map_initialized = False
	# move_base.wait_for_server(rospy.Duration(60))
	# from comp 4 for testing
	waypoints = [[(4.4884057045, 0.0779395103455, 0.0), (0.0, 0.0, 0.991586551372, 0.129445398293)],
			[(1.91507136822, 0.765387296677, 0.0), (0.0, 0.0, 0.977148376444, -0.212558345904)],
			[(0.121204137802, 0.593794822693, 0.0), (0.0, 0.0, 0.997421420608, 0.0717670517161)],
			[(-2.08100700378, 1.54591023922, 0.0), (0.0, 0.0, 0.997192362885, 0.074882517348)],
			[(-4.89103031158, 1.8487932682, 0.0), (0.0, 0.0, -0.761332749671, 0.648361353166)],
			[(-5.474401474, -1.43959307671, 0.0), (0.0, 0.0, -0.119438998482, 0.992841541054)],
			[(-3.64519500732, -1.72245955467, 0.0), (0.0, 0.0, -0.081600949758, 0.996665081659)],
			[(-1.02599549294, -2.39323973656, 0.0), (0.0, 0.0, -0.0681311425473, 0.997676374089)],
			[(1.60786616802, -2.74715876579, 0.0), (0.0, 0.0, -0.102695379529, 0.994712852548)],
			[(3.8603272438, -2.69449019432, 0.0), (0.0, 0.0, 0.611738147471, 0.791060325721)]]
	# soundhandle = SoundClient()

	# # soundhandle.say('oaoaoaoaoaoaoaoaoaoaoaoa')
	# # soundhandle.say('wawawawawawawawawawawawa')
	# # soundhandle.playWave('/home/hegberg/Music/mlg-airhorn.mp3')
	scan_sub = rospy.Subscriber('/scan', LaserScan, scan_callback)
	pose_sub = rospy.Subscriber('/amcl_pose', PoseWithCovarianceStamped, pose_callback, queue_size = 100)
	ar_sub = rospy.Subscriber('/ar_location', String, ar_callback)
	joy_sub = rospy.Subscriber('/joy', Joy, joy_callback)
	vel_sub = rospy.Subscriber('/cmd_vel', Twist, vel_callback)
	cmd_vel_pub = rospy.Publisher('/cmd_vel_mux/input/teleop', Twist, queue_size=1)
	pose_pub = rospy.Publisher('/seeker_pose', PoseWithCovarianceStamped, queue_size=100)
	vel_pub = rospy.Publisher('/seeker_velocity', Point, queue_size=100)

	# cost_sub = rospy.Subscriber('/move_base/global_costmap/costmap_updates', OccupancyGridUpdate, cost_map_callback, queue_size = 100)
	# og_cost_sub = rospy.Subscriber('/move_base/global_costmap/costmap', OccupancyGrid, original_cost_map_callback, queue_size = 100)
	cost_sub = rospy.Subscriber('/move_base/local_costmap/costmap_updates', OccupancyGridUpdate, cost_map_callback, queue_size = 100)
	og_cost_sub = rospy.Subscriber('/move_base/local_costmap/costmap', OccupancyGrid, original_cost_map_callback, queue_size = 100)

	rospy.init_node('Seeker')

	sm = StateMachine(outcomes=['success', 'failure'])

	sis = smach_ros.IntrospectionServer('server_name', sm, '/SM_SEEK_ROOT')
	sis.start()

	with sm:
		StateMachine.add('Start', Start(), transitions={'success':'Localize'})
		StateMachine.add('Localize', Localize(), transitions={'success':'Seek'})
		StateMachine.add('Seek', Seek(), transitions={'success':'CatchHider', 'failure':'Lose'})
		StateMachine.add('CatchHider', CatchHider(), transitions={'success':'Win', 'failure':'Seek'})
		StateMachine.add('Win', Win(), transitions={'success':'Win'})
		StateMachine.add('Lose', Lose(), transitions={'failure':'Lose'})

outcome = sm.execute()


from re import I
from research.wrappers import preproc_vec_env
import yaml
import sys
from collections import defaultdict
import numpy as np
import matplotlib.pyplot as plt
from torch.optim import Adam
import torch as th
from torch import distributions as thd
from torch import nn
import torch.nn.functional as F
from research.nets.common import TransformerBlock, BinaryHead
from research import utils
from research.nets.autoencoders.bvae import BVAE
from ._base import VideoModel

class FBT(VideoModel):
  def __init__(self, env, G):
    super().__init__(env, G)
    # <LOAD BVAE>
    sd = th.load(G.weightdir / 'BVAE.pt')
    bvaeC = sd.pop('G')
    self.bvae = BVAE(env, bvaeC)
    self.bvae.load(G.weightdir)
    for p in self.bvae.parameters():
      p.requires_grad = False
    #self.bvae.eval()
    print('LOADED BVAE', G.weightdir)
    # </LOAD BVAE>

    self.zW = int(self.bvae.G.wh_ratio*4)
    self.size = self.bvae.G.vqD * 4 * self.zW
    self.z_size = self.bvae.z_size
    self.block_size = self.G.window
    # GPT STUFF
    self.pos_emb = nn.Parameter(th.zeros(1, self.block_size, G.n_embed))
    self.cond_in = nn.Linear(self.act_n, G.n_embed // 2, bias=False)
    # input embedding stem
    self.embed = nn.Linear(self.size, G.n_embed // 2, bias=False)
    # transformer
    self.blocks = nn.Sequential(*[TransformerBlock(self.block_size, G) for _ in range(G.n_layer)])
    # decoder head
    self.ln_f = nn.LayerNorm(G.n_embed)
    self.dist_head = BinaryHead(G.n_embed, self.size, G)
    self._init()

  def train(self, *args):
    super().train(*args)
    self.bvae.eval()

  def forward(self, z, action):
    x = self.embed(z)
    # forward the GPT model
    BS, T, E = x.shape
    # SHIFT RIGHT (add a padding on the left)
    x = th.cat([th.zeros(BS, 1, E).to(self.G.device), x[:, :-1]], dim=1)
    action = th.cat([th.zeros(BS, 1, action.shape[-1]).to(self.G.device), action[:, :-1]], dim=1)
    cin = self.cond_in(action)
    if action.ndim == 2:
      x = th.cat([x, cin[:, None].repeat_interleave(self.block_size, 1)], -1)
    else:
      x = th.cat([x, cin], -1)
    x += self.pos_emb  # each position maps to a (learnable) vector
    x = self.blocks(x)
    logits = self.ln_f(x)
    return logits

  def loss(self, batch):
    z = self.bvae.encode(batch, noise=False).detach()
    logits = self.forward(z, batch['action'])
    dist = self.dist_head(logits)
    loss = -dist.log_prob(z).mean()
    #import ipdb; ipdb.set_trace()
    return loss, {'loss/total': loss}

  def onestep(self, batch, i, temp=1.0):
    z = self.bvae.encode(batch, noise=False)
    logits = self.forward(z, batch['action'])
    dist = self.dist_head(logits[:,i:i+1] / temp)
    z_sample = dist.sample().reshape([-1, self.bvae.G.vqD, 4, self.zW])
    dist = self.bvae.decoder(z_sample)
    batch['lcd'][:, i] = 1.0*(dist['lcd'].probs > 0.5)[:, 0]
    batch['proprio'][:, i] = dist['proprio'].mean
    return batch

  def latent_onestep(self, z, a, i, temp=1.0):
    logits = self.forward(z, a)
    dist = self.dist_head(logits / temp)
    z[:, i] = dist.sample()[:, i]
    return z

  def latent_sample(self, z, a, start, temp=1.0):
    for i in range(start, self.block_size):
      z = self.latent_onestep(z, a, i, temp=temp)
    return z

  def sample(self, n, action=None, prompts=None, prompt_n=10, temp=1.0):
    # TODO: feed act_n
    with th.no_grad():
      # CREATE PROMPT
      if action is None:
        action = (th.rand(n, self.block_size, self.act_n) * 2 - 1).to(self.G.device)
      else:
        n = action.shape[0]
      batch = {}
      batch['lcd'] = th.zeros(n, self.block_size, self.G.lcd_h, self.G.lcd_w).to(self.G.device)
      batch['proprio'] = th.zeros(n, self.block_size, self.proprio_n).to(self.G.device)
      start = 0
      if prompts is not None:
        batch['lcd'][:, :prompt_n] = prompts['lcd'][:, :prompt_n]
        batch['proprio'][:, :prompt_n] = prompts['proprio'][:, :prompt_n]
        start = prompt_n
      z = self.bvae.encode(batch, noise=False).detach()
      z_sample = th.zeros(n, self.block_size, self.bvae.G.vqD * 4 * self.zW).to(self.G.device)
      z_sample[:, :prompt_n] = z[:, :prompt_n]
      # SAMPLE FORWARD IN LATENT SPACE, ACTION CONDITIONED
      z_sample = self.latent_sample(z_sample, action, start, temp)
      z_sample = z_sample.reshape([n * self.block_size, self.bvae.G.vqD, 4, self.zW])

      # DECODE
      dist = self.bvae.decoder(z_sample)
      batch['lcd'] = (1.0 * (dist['lcd'].probs > 0.5)).reshape([n, self.block_size, 1, self.G.lcd_h, self.G.lcd_w])
      batch['proprio'] = (dist['proprio'].mean).reshape([n, self.block_size, -1])
    return batch


import os
import usb.core
import usb.util
import skimage.draw
import skimage.morphology
import numpy as np
import time
import threading
from projector.pycrafter6500 import dmd
from models import devices_db, devices_list, projected_frames

RES_Y = 1080  # projectors' resolution
RES_X = 1920

class ProjectorThread(threading.Thread):
    def __init__(self, name, address):
        # Init thread
        threading.Thread.__init__(self)

        # Initialize projector
        self.name = name
        self.address = address
        self.dmd = dmd(address)

        devices = list(usb.core.find(idVendor=0x0451,           # just for testing
                       idProduct=0xc900, find_all=True))
        print("addresses available")
        for dev in devices:
            print(dev.address)

    def run(self):
        while(True):
            # check if projector is connected
            self.dmd.CheckConnection()

            # update connected status in db
            db_connected_status = devices_db.search(
                devices_list.name == self.name)[0]["connected"]
            if self.dmd.connected is not db_connected_status and db_connected_status is not None:
                devices_db.update({'connected': self.dmd.connected},
                                  devices_list.name == self.name)

            # if (not self.dmd.connected):
            #     print(self.name, "disconnected")

            time.sleep(1)

    def SendPattern(self, pattern):
        # Check if projector is connected
        connected = self.dmd.CheckConnection()
        if (not connected):
            return False

        # Init empty array
        frames_array = []
        for frame in range(pattern['number-of-measurements']):
            frames_array.append(np.zeros((RES_Y, RES_X)).astype(np.uint8))

        if self.name == 'green-projector':
            # Draw each green dot trajectory in the frames array
            for trajectory in pattern['trajectories']:
                # Get every point in the line
                rr, cc = skimage.draw.line(int(trajectory['start'][1] * RES_Y), int(trajectory['start'][0] * RES_X),
                                           int(trajectory['end'][1] * RES_Y), int(trajectory['end'][0] * RES_X))

                # Calculate distance between dots
                dot_step = (len(rr) - 1) / \
                    (pattern['number-of-measurements'] - 1)

                # Go through each frame, draw corresponding points
                for frame_idx, frame in enumerate(frames_array):
                    # Get point in the line for this frame
                    dot_idx = round(frame_idx * dot_step)

                    # Draw point in current frame
                    rr_disk, cc_disk = skimage.draw.disk(
                        (rr[dot_idx], cc[dot_idx]), pattern['green-point-diameter'])
                    frame[rr_disk, cc_disk] = 1

        if self.name == 'laser-projector':
            # Draw the middle point of each trajectory for the laser beam
            for trajectory in pattern['trajectories']:
                # Get the middle point of the trajectory
                middle_point = (int((trajectory['start'][1] + trajectory['end'][1]) * RES_Y / 2),
                                int((trajectory['start'][0] + trajectory['end'][0]) * RES_X / 2))

                # Draw the middle point
                rr_disk, cc_disk = skimage.draw.disk(
                    middle_point, pattern['laser-point-diameter'])

                # Draw point in each frame
                for frame in frames_array:
                    frame[rr_disk, cc_disk] = 1

        self.dmd.stopsequence()
        self.dmd.changemode(3)

        # Set secondary parameters
        exposure = [pattern['measurement-time'] * 1000]*len(frames_array)
        dark_time = [0]*len(frames_array)
        trigger_in = [0]*len(frames_array)
        trigger_out = [0]*len(frames_array)
        repetitions = 0  # infinite loop

        # Start sequence
        self.dmd.defsequence(frames_array, 'green', exposure, trigger_in,
                             dark_time, trigger_out, repetitions)
        self.dmd.startsequence()

        return True

    def SendBox(self, box):
        # Check if projector is connected
        connected = self.dmd.CheckConnection()
        if (not connected):
            return False

        # Init empty array
        frame = np.zeros((RES_Y, RES_X)).astype(np.uint8)

        # Draw 1-pixel-thick box
        start = ((0.5 + box['x-offset']) * RES_X - box['side-length'] * RES_Y / 2,
                 (0.5 + box['y-offset']) * RES_Y - box['side-length'] * RES_Y / 2)
        rr_disk, cc_disk = skimage.draw.rectangle_perimeter(start=(start[1], start[0]),
                                                            extent=(box['side-length'] * RES_Y, box['side-length'] * RES_Y))
        frame[rr_disk, cc_disk] = 1

        # Dilate to box thickness using morphology
        frame = skimage.morphology.dilation(
            frame, skimage.morphology.disk(box['thickness']))

        self.dmd.stopsequence()
        self.dmd.changemode(3)

        # Set secondary parameters
        exposure = [0]
        dark_time = [0]
        trigger_in = [0]
        trigger_out = [0]
        repetitions = 0  # infinite loop

        # Start sequence
        self.dmd.defsequence([frame], 'green', exposure, trigger_in,
                             dark_time, trigger_out, repetitions)
        self.dmd.startsequence()

        return True


import numpy as np
import os, sys
import pickle, h5py

ROOT_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')

GT_TEST_PATH = os.path.join(ROOT_PATH, 'dataset/h36m_test.pkl')
GT_TRAIN_PATH = os.path.join(ROOT_PATH, 'dataset/h36m_train.pkl')

def flip_data(data):
    """
    horizontal flip
        data: [N, 17*k] or [N, 17, k], i.e. [x, y], [x, y, confidence] or [x, y, z]
    Return
        result: [2N, 17*k] or [2N, 17, k]
    """
    left_joints = [4, 5, 6, 11, 12, 13]
    right_joints = [1, 2, 3, 14, 15, 16]

    flipped_data = data.copy().reshape((len(data), 17, -1))
    flipped_data[:, :, 0] *= -1  # flip x of all joints
    flipped_data[:, left_joints+right_joints] = flipped_data[:, right_joints+left_joints]
    flipped_data = flipped_data.reshape(data.shape)

    result = np.concatenate((data, flipped_data), axis=0)

    return result

def unflip_data(data):
    """
    Average original data and flipped data
        data: [2N, 17*3]
    Return
        result: [N, 17*3]
    """
    left_joints = [4, 5, 6, 11, 12, 13]
    right_joints = [1, 2, 3, 14, 15, 16]

    data = data.copy().reshape((2, -1, 17, 3))
    data[1, :, :, 0] *= -1  # flip x of all joints
    data[1, :, left_joints+right_joints] = data[1, :, right_joints+left_joints]
    data = np.mean(data, axis=0)
    data = data.reshape((-1, 17*3))

    return data

class DataReader(object):
    def __init__(self):
        self.gt_trainset = None
        self.gt_testset = None
        self.dt_dataset = None

    def real_read(self, subset):
        file_name = 'h36m_%s.pkl' % subset
        print('loading %s' % file_name)
        file_path = os.path.join(ROOT_PATH, 'dataset', file_name)
        with open(file_path, 'rb') as f:
            gt = pickle.load(f)
        return gt

    def read_2d(self, which='scale', mode='dt_ft', read_confidence=True):
        if self.gt_trainset is None:
            self.gt_trainset = self.real_read('train')
        if self.gt_testset is None:
            self.gt_testset = self.real_read('test')

        if mode == 'gt':
            trainset = np.empty((len(self.gt_trainset), 17, 2))  # [N, 17, 2]
            testset = np.empty((len(self.gt_testset), 17, 2))  # [N, 17, 2]
            for idx, item in enumerate(self.gt_trainset):
                trainset[idx] = item['joint_3d_image'][:, :2]
            for idx, item in enumerate(self.gt_testset):
                testset[idx] = item['joint_3d_image'][:, :2]
            if read_confidence:
                train_confidence = np.ones((len(self.gt_trainset), 17, 1))  # [N, 17, 1]
                test_confidence = np.ones((len(self.gt_testset), 17, 1))  # [N, 17, 1]
        elif mode == 'dt_ft':
            file_name = 'h36m_sh_dt_ft.pkl'
            file_path = os.path.join(ROOT_PATH, 'dataset', file_name)
            print('loading %s' % file_name)
            with open(file_path, 'rb') as f:
                self.dt_dataset = pickle.load(f)

            trainset = self.dt_dataset['train']['joint3d_image'][:, :, :2].copy()  # [N, 17, 2]
            testset = self.dt_dataset['test']['joint3d_image'][:, :, :2].copy()  # [N, 17, 2]
            if read_confidence:
                train_confidence = self.dt_dataset['train']['confidence'].copy()  # [N, 17, 1]
                test_confidence = self.dt_dataset['test']['confidence'].copy()  # [N, 17, 1]
        else:
            assert 0, 'not supported type %s' % mode

        # normalize
        if which == 'scale':
            # map to [-1, 1]
            for idx, item in enumerate(self.gt_trainset):
                camera_name = item['camera_param']['name']
                if camera_name == '54138969' or camera_name == '60457274':
                    res_w, res_h = 1000, 1002
                elif camera_name == '55011271' or camera_name == '58860488':
                    res_w, res_h = 1000, 1000
                else:
                    assert 0, '%d data item has an invalid camera name' % idx
                trainset[idx, :, :] = trainset[idx, :, :] / res_w * 2 - [1, res_h / res_w]
            for idx, item in enumerate(self.gt_testset):
                camera_name = item['camera_param']['name']
                if camera_name == '54138969' or camera_name == '60457274':
                    res_w, res_h = 1000, 1002
                elif camera_name == '55011271' or camera_name == '58860488':
                    res_w, res_h = 1000, 1000
                else:
                    assert 0, '%d data item has an invalid camera name' % idx
                testset[idx, :, :] = testset[idx, :, :] / res_w * 2 - [1, res_h / res_w]
        else:
            assert 0, 'not support normalize type %s' % which

        if read_confidence:
            trainset = np.concatenate((trainset, train_confidence), axis=2)  # [N, 17, 3]
            testset = np.concatenate((testset, test_confidence), axis=2)  # [N, 17, 3]

        # reshape
        trainset, testset = trainset.reshape((len(trainset), -1)), testset.reshape((len(testset), -1))

        return trainset, testset

    def read_3d(self, which='scale', mode='dt_ft'):
        if self.gt_trainset is None:
            self.gt_trainset = self.real_read('train')
        if self.gt_testset is None:
            self.gt_testset = self.real_read('test')

        # normalize
        train_labels = np.empty((len(self.gt_trainset), 17, 3))
        test_labels = np.empty((len(self.gt_testset), 17, 3))
        if which == 'scale':
            # map to [-1, 1]
            for idx, item in enumerate(self.gt_trainset):
                camera_name = item['camera_param']['name']
                if camera_name == '54138969' or camera_name == '60457274':
                    res_w, res_h = 1000, 1002
                elif camera_name == '55011271' or camera_name == '58860488':
                    res_w, res_h = 1000, 1000
                else:
                    assert 0, '%d data item has an invalid camera name' % idx
                train_labels[idx, :, :2] = item['joint_3d_image'][:, :2] / res_w * 2 - [1, res_h / res_w]
                train_labels[idx, :, 2:] = item['joint_3d_image'][:, 2:] / res_w * 2
            for idx, item in enumerate(self.gt_testset):
                camera_name = item['camera_param']['name']
                if camera_name == '54138969' or camera_name == '60457274':
                    res_w, res_h = 1000, 1002
                elif camera_name == '55011271' or camera_name == '58860488':
                    res_w, res_h = 1000, 1000
                else:
                    assert 0, '%d data item has an invalid camera name' % idx
                test_labels[idx, :, :2] = item['joint_3d_image'][:, :2] / res_w * 2 - [1, res_h / res_w]
                test_labels[idx, :, 2:] = item['joint_3d_image'][:, 2:] / res_w * 2
        else:
            assert 0, 'not support normalize type %s' % which

        # reshape
        train_labels, test_labels = train_labels.reshape((-1, 17*3)), test_labels.reshape((-1, 17*3))

        return train_labels, test_labels

    def denormalize(self, data, which='scale'):
        if self.gt_testset is None:
            self.gt_testset = self.real_read('test')

        if which == 'scale':
            data = data.reshape((-1, 17, 3)).copy()
            # denormalize (x,y,z) coordiantes for results
            for idx, item in enumerate(self.gt_testset):
                camera_name = item['camera_param']['name']
                if camera_name == '54138969' or camera_name == '60457274':
                    res_w, res_h = 1000, 1002
                elif camera_name == '55011271' or camera_name == '58860488':
                    res_w, res_h = 1000, 1000
                else:
                    assert 0, '%d data item has an invalid camera name' % idx

                data[idx, :, :2] = (data[idx, :, :2] + [1, res_h / res_w]) * res_w / 2
                data[idx, :, 2:] = data[idx, :, 2:] * res_w / 2
        else:
            assert 0
        return data

if __name__ == '__main__':
    datareader = DataReader()
    train_data, test_data, train_2d_mean, train_2d_std = datareader.read_2d(which='scale', mode='dt_ft', read_confidence=False)
    train_labels, test_labels, train_3d_mean, train_3d_std = datareader.read_3d(which='scale', mode='dt_ft')


import pandas as pd
import numpy as np
import pkg_resources
import torch
from glycowork.glycan_data.loader import lib, unwrap
from glycowork.ml.processing import dataset_to_dataloader

try:
  from torch_geometric.data import Data
  from torch_geometric.loader import DataLoader
except ImportError:
  raise ImportError('<torch_geometric missing; cannot do deep learning>')

io = pkg_resources.resource_stream(__name__,
                                   "glycowork_lectinoracle_background_correction.csv")
df_corr = pd.read_csv(io)

#try:
#  import esm
#except ImportError:
#  print('<esm missing; cannot use get_esm1b_representations>')

def glycans_to_emb(glycans, model, libr = None, batch_size = 32, rep = True,
                   class_list = None):
    """Returns a dataframe of learned representations for a list of glycans\n
    | Arguments:
    | :-
    | glycans (list): list of glycans in IUPAC-condensed as strings
    | model (PyTorch object): trained graph neural network (such as SweetNet) for analyzing glycans
    | libr (list): sorted list of unique glycoletters observed in the glycans of our dataset
    | batch_size (int): change to batch_size used during training; default:32
    | rep (bool): True returns representations, False returns actual predicted labels; default is True
    | class_list (list): list of unique classes to map predictions\n
    | Returns:
    | :-
    | Returns dataframe of learned representations (columns) for each glycan (rows)
    """
    if libr is None:
      libr = lib
    glycan_loader = dataset_to_dataloader(glycans, range(len(glycans)),
                                          libr = libr, batch_size = batch_size,
                                          shuffle = False)
    res = []
    for data in glycan_loader:
        x, y, edge_index, batch = data.x, data.y, data.edge_index, data.batch
        x = x.cuda()
        y = y.cuda()
        edge_index = edge_index.cuda()
        batch = batch.cuda()
        model = model.eval()
        pred, out = model(x, edge_index, batch, inference = True)
        if rep:
            res.append(out)
        else:
            res.append(pred)
    res2 = [res[k].detach().cpu().numpy() for k in range(len(res))]
    res2 = pd.DataFrame(np.concatenate(res2))
    if rep:
      return res2
    else:
      idx = res2.idxmax(axis = "columns").values.tolist()
      preds = [class_list[k] for k in idx]
      return preds

def get_multi_pred(prot, glycans, model, prot_dic,
                   background_correction = False, correction_df = None,
                   batch_size = 128, libr = None):
  """Inner function to actually get predictions for lectin-glycan binding from LectinOracle-type model\n
  | Arguments:
  | :-
  | prot (string): protein amino acid sequence
  | glycans (list): list of glycans in IUPACcondensed
  | model (PyTorch object): trained LectinOracle-type model
  | prot_dic (dictionary): dictionary of type protein sequence:ESM1b representation
  | background_correction (bool): whether to correct predictions for background; default:False
  | correction_df (dataframe): background prediction for (ideally) all provided glycans; default:None
  | batch_size (int): change to batch_size used during training; default:128
  | libr (list): sorted list of unique glycoletters observed in the glycans of our dataset\n
  | Returns:
  | :-
  | Returns dataframe of glycan sequences and predicted binding to prot
  """
  if libr is None:
      libr = lib
  try:
    rep = prot_dic[prot]
  except:
    print('new protein, no stored embedding')
  train_loader = dataset_to_dataloader(glycans, [0.99]*len(glycans),
                                       libr = libr, batch_size = batch_size,
                                       shuffle = False, extra_feature = [rep]*len(glycans))
  model = model.eval()
  res = []
  for k in train_loader:
    x, y, edge_index, prot, batch = k.x, k.y, k.edge_index, k.train_idx, k.batch
    x = x.cuda()
    y = y.cuda()
    prot = prot.view(max(batch)+1, -1).cuda()
    edge_index = edge_index.cuda()
    batch = batch.cuda()
    pred = model(prot, x, edge_index, batch)
    res.append(pred)
  #res = unwrap([res[k].detach().cpu().numpy() for k in range(len(res))][0].tolist())
  res = unwrap([res[k].detach().cpu().numpy() for k in range(len(res))])
  res = [k.tolist()[0] for k in res]
  if background_correction:
    correction_df = pd.Series(correction_df.pred.values,
                              index = correction_df.motif).to_dict()
    bg_res = [correction_df[j] if j in list(correction_df.keys()) else 0 for j in glycans]
    if 0 in bg_res:
      print("Warning: not all glycans are in the correction_df; consider adding their background to correction_df")
    res = [a_i - b_i for a_i, b_i in zip(res, bg_res)]
  return res

def get_lectin_preds(prot, glycans, model, prot_dic, background_correction = False,
                     correction_df = None, batch_size = 128, libr = None, sort = True):
  """Wrapper that uses LectinOracle-type model for predicting binding of protein to glycans\n
  | Arguments:
  | :-
  | prot (string): protein amino acid sequence
  | glycans (list): list of glycans in IUPACcondensed
  | model (PyTorch object): trained LectinOracle-type model
  | prot_dic (dictionary): dictionary of type protein sequence:ESM1b representation
  | background_correction (bool): whether to correct predictions for background; default:False
  | correction_df (dataframe): background prediction for (ideally) all provided glycans; default:V4 correction file
  | batch_size (int): change to batch_size used during training; default:128
  | libr (list): sorted list of unique glycoletters observed in the glycans of our dataset
  | sort (bool): whether to sort prediction results descendingly; default:True\n
  | Returns:
  | :-
  | Returns dataframe of glycan sequences and predicted binding to prot
  """
  if libr is None:
      libr = lib
  if correction_df is None:
    correction_df = df_corr
  preds = get_multi_pred(prot, glycans, model, prot_dic,
                         batch_size = batch_size, libr = libr)
  df_pred = pd.DataFrame(glycans, columns = ['motif'])
  df_pred['pred'] = preds
  if background_correction:
    correction_df = pd.Series(correction_df.pred.values,
                              index = correction_df.motif).to_dict()
    for j in df_pred.motif.values.tolist():
      motif_idx = df_pred.motif.values.tolist().index(j)
      try:
        df_pred.at[motif_idx, 'pred'] = df_pred.iloc[motif_idx, 1] - correction_df[j]
      except:
        pass
  if sort:
    df_pred.sort_values('pred', ascending = True, inplace = True)
  return df_pred

def get_esm1b_representations(prots, model, alphabet):
  """Retrieves ESM1b representations of protein for using them as input for LectinOracle\n
  | Arguments:
  | :-
  | prots (list): list of protein sequences (strings) that should be converted
  | model (ESM1b object): trained ESM1b model; from running esm.pretrained.esm1b_t33_650M_UR50S()
  | alphabet (ESM1b object): used for converting sequences; from running esm.pretrained.esm1b_t33_650M_UR50S()\n
  | Returns:
  | :-
  | Returns dictionary of the form protein sequence:ESM1b representation
  """
  #model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
  batch_converter = alphabet.get_batch_converter()
  prots = list(set(prots))
  data_list = []
  for k in range(0,len(prots)):
    if len(prots[k])<1000:
      data_list.append(('protein'+str(k), prots[k][:np.min([len(prots[k]),
                                                               1000])]))
    else:
      data_list.append(('protein'+str(k), prots[k][:1000]))
  batch_labels, batch_strs, batch_tokens = batch_converter(data_list)
  with torch.no_grad():
      results = model(batch_tokens,
                      repr_layers = [33], return_contacts = False)
  token_representations = results["representations"][33]
  sequence_representations = []
  for i, (_, seq) in enumerate(data_list):
      sequence_representations.append(token_representations[i, 1 : len(seq) + 1].mean(0))
  prot_dic =  {prots[k]:sequence_representations[k].tolist() for k in range(len(sequence_representations))}
  return prot_dic


"""
Functions for explaining classifiers that use tabular data (matrices).
"""
import collections
import json
import copy
import numpy as np
import sklearn
import sklearn.preprocessing
from . import lime_base
from . import explanation

class TableDomainMapper(explanation.DomainMapper):
    """Maps feature ids to names, generates table views, etc"""
    def __init__(self, feature_names, feature_values, scaled_row,
                 categorical_features, discretized_feature_names=None):
        """Init.

        Args:
            feature_names: list of feature names, in order
            feature_values: list of strings with the values of the original row
            scaled_row: scaled row
            categorical_features: list of categorical features ids (ints)
        """
        self.exp_feature_names = feature_names
        self.discretized_feature_names = discretized_feature_names
        self.feature_names = feature_names
        self.feature_values = feature_values
        self.scaled_row = scaled_row
        self.all_categorical = len(categorical_features) == len(scaled_row)
        self.categorical_features = categorical_features

    def map_exp_ids(self, exp):
        """Maps ids to feature names.

        Args:
            exp: list of tuples [(id, weight), (id,weight)]

        Returns:
            list of tuples (feature_name, weight)
        """
        names = self.exp_feature_names
        if self.discretized_feature_names is not None:
            names = self.discretized_feature_names
        return [(names[x[0]], x[1]) for x in exp]

    def visualize_instance_html(self,
                                exp,
                                label,
                                div_name,
                                exp_object_name,
                                show_table=True,
                                show_all=False):
        """Shows the current example in a table format.

        Args:
             exp: list of tuples [(id, weight), (id,weight)]
             label: label id (integer)
             div_name: name of div object to be used for rendering(in js)
             exp_object_name: name of js explanation object
             show_table: if False, don't show table visualization.
             show_all: if True, show zero-weighted features in the table.
        """
        if not show_table:
            return ''
        weights = [0] * len(self.feature_names)
        for x in exp:
            weights[x[0]] = x[1]
        out_list = list(zip(self.exp_feature_names, self.feature_values,
                        weights))
        if not show_all:
            out_list = [out_list[x[0]] for x in exp]
        ret = u'''
            %s.show_raw_tabular(%s, %d, %s);
        ''' % (exp_object_name, json.dumps(out_list), label, div_name)
        return ret

class LimeTabularExplainer(object):
    """Explains predictions on tabular (i.e. matrix) data.
    For numerical features, perturb them by sampling from a Normal(0,1) and
    doing the inverse operation of mean-centering and scaling, according to the
    means and stds in the training data. For categorical features, perturb by
    sampling according to the training distribution, and making a binary
    feature that is 1 when the value is the same as the instance being
    explained."""
    def __init__(self, training_data, feature_names=None,
                 categorical_features=None, categorical_names=None,
                 kernel_width=None, verbose=False, class_names=None,
                 feature_selection='auto', discretize_continuous=True):
        """Init function.

        Args:
            training_data: numpy 2d array
            feature_names: list of names (strings) corresponding to the columns
                in the training data.
            categorical_features: list of indices (ints) corresponding to the
                categorical columns. Everything else will be considered
                continuous. Values in these columns MUST be integers.
            categorical_names: map from int to list of names, where
                categorical_names[x][y] represents the name of the yth value of
                column x.
            kernel_width: kernel width for the exponential kernel.
            If None, defaults to sqrt(number of columns) * 0.75
            verbose: if true, print local prediction values from linear model
            class_names: list of class names, ordered according to whatever the
                classifier is using. If not present, class names will be '0',
                '1', ...
            feature_selection: feature selection method. can be
                'forward_selection', 'lasso_path', 'none' or 'auto'.
                See function 'explain_instance_with_data' in lime_base.py for
                details on what each of the options does.
            discretize_continuous: if True, all non-categorical features will
                be discretized into quartiles.
        """
        self.categorical_names = categorical_names
        self.categorical_features = categorical_features
        if self.categorical_names is None:
            self.categorical_names = {}
        if self.categorical_features is None:
            self.categorical_features = []
        self.discretizer = None
        if discretize_continuous:
            self.discretizer = QuartileDiscretizer(training_data,
                                                   self.categorical_features,
                                                   feature_names)
            self.categorical_features = range(training_data.shape[1])
            discretized_training_data = self.discretizer.discretize(
                training_data)

        if kernel_width is None:
            kernel_width = np.sqrt(training_data.shape[1]) * .75
        kernel_width = float(kernel_width)

        def kernel(d): return np.sqrt(np.exp(-(d**2) / kernel_width ** 2))
        self.feature_selection = feature_selection
        self.base = lime_base.LimeBase(kernel, verbose)
        self.scaler = None
        self.class_names = class_names
        self.feature_names = feature_names
        self.scaler = sklearn.preprocessing.StandardScaler(with_mean=False)
        self.scaler.fit(training_data)
        self.feature_values = {}
        self.feature_frequencies = {}

        for feature in self.categorical_features:
            feature_count = collections.defaultdict(lambda: 0.0)
            column = training_data[:, feature]
            if self.discretizer is not None:
                column = discretized_training_data[:, feature]
                feature_count[0] = 0.
                feature_count[1] = 0.
                feature_count[2] = 0.
                feature_count[3] = 0.
            for value in column:
                feature_count[value] += 1
            values, frequencies = map(list, zip(*(feature_count.items())))
            self.feature_values[feature] = values
            self.feature_frequencies[feature] = (np.array(frequencies) /
                                                 sum(frequencies))
            self.scaler.mean_[feature] = 0
            self.scaler.scale_[feature] = 1

    def explain_instance(self, data_row, classifier_fn, labels=(1,),
                         top_labels=None, num_features=10, num_samples=5000,
                         distance_metric='euclidean', model_regressor=None):
        """Generates explanations for a prediction.

        First, we generate neighborhood data by randomly perturbing features
        from the instance (see __data_inverse). We then learn locally weighted
        linear models on this neighborhood data to explain each of the classes
        in an interpretable way (see lime_base.py).

        Args:
            data_row: 1d numpy array, corresponding to a row
            classifier_fn: classifier prediction probability function, which
                takes a string and outputs prediction probabilities.  For
                ScikitClassifiers , this is classifier.predict_proba.
            labels: iterable with labels to be explained.
            top_labels: if not None, ignore labels and produce explanations for
                the K labels with highest prediction probabilities, where K is
                this parameter.
            num_features: maximum number of features present in explanation
            num_samples: size of the neighborhood to learn the linear model
            distance_metric: the distance metric to use for weights.
            model_regressor: sklearn regressor to use in explanation. Defaults
            to Ridge regression in LimeBase. Must have model_regressor.coef_
            and 'sample_weight' as a parameter to model_regressor.fit()

        Returns:
            An Explanation object (see explanation.py) with the corresponding
            explanations.
        """
        data, inverse = self.__data_inverse(data_row, num_samples)
        scaled_data = (data - self.scaler.mean_) / self.scaler.scale_

        distances = sklearn.metrics.pairwise_distances(
            scaled_data,
            scaled_data[0].reshape(1, -1),
            metric=distance_metric
        ).ravel()

        yss = classifier_fn(inverse)
        if self.class_names is None:
            self.class_names = [str(x) for x in range(yss[0].shape[0])]
        else:
            self.class_names = list(self.class_names)
        feature_names = copy.deepcopy(self.feature_names)
        if feature_names is None:
            feature_names = [str(x) for x in range(data_row.shape[0])]

        def round_stuff(x): return ['%.2f' % a for a in x]
        values = round_stuff(data_row)
        for i in self.categorical_features:
            if self.discretizer is not None and i in self.discretizer.lambdas:
                continue
            name = int(data_row[i])
            if i in self.categorical_names:
                name = self.categorical_names[i][name]
            feature_names[i] = '%s=%s' % (feature_names[i], name)
            values[i] = 'True'
        categorical_features = self.categorical_features
        discretized_feature_names = None
        if self.discretizer is not None:
            categorical_features = range(data.shape[1])
            discretized_instance = self.discretizer.discretize(data_row)
            discretized_feature_names = copy.deepcopy(feature_names)
            for f in self.discretizer.names:
                discretized_feature_names[f] = self.discretizer.names[f][int(
                    discretized_instance[f])]

        domain_mapper = TableDomainMapper(
            feature_names, values, scaled_data[0],
            categorical_features=categorical_features,
            discretized_feature_names=discretized_feature_names)
        ret_exp = explanation.Explanation(domain_mapper=domain_mapper,
                                          class_names=self.class_names)
        ret_exp.predict_proba = yss[0]
        if top_labels:
            labels = np.argsort(yss[0])[-top_labels:]
            ret_exp.top_labels = list(labels)
            ret_exp.top_labels.reverse()
        for label in labels:
            (ret_exp.intercept[label],
             ret_exp.local_exp[label],
             ret_exp.score) = self.base.explain_instance_with_data(
                scaled_data, yss, distances, label, num_features,
                model_regressor=model_regressor,
                feature_selection=self.feature_selection)
        return ret_exp

    def __data_inverse(self,
                       data_row,
                       num_samples):
        """Generates a neighborhood around a prediction.

        For numerical features, perturb them by sampling from a Normal(0,1) and
        doing the inverse operation of mean-centering and scaling, according to
        the means and stds in the training data. For categorical features,
        perturb by sampling according to the training distribution, and making
        a binary feature that is 1 when the value is the same as the instance
        being explained.

        Args:
            data_row: 1d numpy array, corresponding to a row
            num_samples: size of the neighborhood to learn the linear model

        Returns:
            A tuple (data, inverse), where:
                data: dense num_samples * K matrix, where categorical features
                are encoded with either 0 (not equal to the corresponding value
                in data_row) or 1. The first row is the original instance.
                inverse: same as data, except the categorical features are not
                binary, but categorical (as the original data)
        """
        data = np.zeros((num_samples, data_row.shape[0]))
        categorical_features = range(data_row.shape[0])
        if self.discretizer is None:
            data = np.random.normal(
                0, 1, num_samples * data_row.shape[0]).reshape(
                num_samples, data_row.shape[0])
            data = data * self.scaler.scale_ + self.scaler.mean_
            categorical_features = self.categorical_features
            first_row = data_row
        else:
            first_row = self.discretizer.discretize(data_row)
        data[0] = data_row.copy()
        inverse = data.copy()
        for column in categorical_features:
            values = self.feature_values[column]
            freqs = self.feature_frequencies[column]
            inverse_column = np.random.choice(values, size=num_samples,
                                              replace=True, p=freqs)
            binary_column = np.array([1 if x == first_row[column]
                                      else 0 for x in inverse_column])
            binary_column[0] = 1
            inverse_column[0] = data[0, column]
            data[:, column] = binary_column
            inverse[:, column] = inverse_column
        if self.discretizer is not None:
            inverse[1:] = self.discretizer.undiscretize(inverse[1:])
        inverse[0] = data_row
        return data, inverse

class QuartileDiscretizer:
    """Discretizes data into quartiles."""
    def __init__(self, data, categorical_features, feature_names):
        """Initializer

        Args:
            data: numpy 2d array
            categorical_features: list of indices (ints) corresponding to the
                categorical columns. These features will not be discretized.
                Everything else will be considered continuous, and will be
                discretized.
            categorical_names: map from int to list of names, where
                categorical_names[x][y] represents the name of the yth value of
                column x.
            feature_names: list of names (strings) corresponding to the columns
                in the training data.
        """
        to_discretize = ([x for x in range(data.shape[1])
                         if x not in categorical_features])
        self.names = {}
        self.lambdas = {}
        self.ranges = {}
        self.means = {}
        self.stds = {}
        self.mins = {}
        self.maxs = {}
        for feature in to_discretize:
            qts = np.percentile(data[:, feature], [25, 50, 75])
            boundaries = np.min(data[:, feature]), np.max(data[:, feature])
            name = feature_names[feature]
            self.names[feature] = (
                ['%s <= %.2f' % (name, qts[0]),
                 '%.2f < %s <= %.2f' % (qts[0], name, qts[1]),
                 '%.2f < %s <= %.2f' % (qts[1], name, qts[2]),
                 '%s > %.2f' % (name, qts[2])])
            self.lambdas[feature] = lambda x, qts=qts: np.searchsorted(qts, x)
            discretized = self.lambdas[feature](data[:, feature])
            self.means[feature] = []
            self.stds[feature] = []
            for x in range(4):
                selection = data[discretized == x, feature]
                mean = 0 if len(selection) == 0 else np.mean(selection)
                self.means[feature].append(mean)
                std = 0 if len(selection) == 0 else np.std(selection)
                std += 0.00000000001
                self.stds[feature].append(std)
            self.mins[feature] = [boundaries[0], qts[0], qts[1], qts[2]]
            self.maxs[feature] = [qts[0], qts[1], qts[2], boundaries[1]]

    def discretize(self, data):
        """Discretizes the data.

        Args:
            data: numpy 2d or 1d array

        Returns:
            numpy array of same dimension, discretized.
        """
        ret = data.copy()
        for feature in self.lambdas:
            if len(data.shape) == 1:
                ret[feature] = int(self.lambdas[feature](ret[feature]))
            else:
                ret[:, feature] = self.lambdas[feature](
                    ret[:, feature]).astype(int)
        return ret

    def undiscretize(self, data):
        ret = data.copy()
        for feature in self.means:
            mins = self.mins[feature]
            maxs = self.maxs[feature]
            means = self.means[feature]
            stds = self.stds[feature]

            def get_inverse(q): return max(
                mins[q],
                min(np.random.normal(means[q], stds[q]), maxs[q]))
            if len(data.shape) == 1:
                q = int(ret[feature])
                ret[feature] = get_inverse(q)
            else:
                ret[:, feature] = (
                    [get_inverse(int(x)) for x in ret[:, feature]])
        return ret


from __future__ import absolute_import

import numpy as np

from .base import DifferentiableModel

class MXNetModel(DifferentiableModel):
    """Creates a :class:`Model` instance from existing `MXNet` symbols and weights.

    Parameters
    ----------
    data : `mxnet.symbol.Variable`
        The input to the model.
    logits : `mxnet.symbol.Symbol`
        The predictions of the model, before the softmax.
    args : `dictionary mapping str to mxnet.nd.array`
        The parameters of the model.
    ctx : `mxnet.context.Context`
        The device, e.g. mxnet.cpu() or mxnet.gpu().
    num_classes : int
        The number of classes.
    bounds : tuple
        Tuple of lower and upper bound for the pixel values, usually
        (0, 1) or (0, 255).
    channel_axis : int
        The index of the axis that represents color channels.
    aux_states : `dictionary mapping str to mxnet.nd.array`
        The states of auxiliary parameters of the model.
    preprocessing: 2-element tuple with floats or numpy arrays
        Elementwises preprocessing of input; we first subtract the first
        element of preprocessing from the input and then divide the input by
        the second element.

    """

    def __init__(
            self,
            data,
            logits,
            args,
            ctx,
            num_classes,
            bounds,
            channel_axis=1,
            aux_states=None,
            preprocessing=(0, 1)):

        super(MXNetModel, self).__init__(
            bounds=bounds,
            channel_axis=channel_axis,
            preprocessing=preprocessing)

        import mxnet as mx

        self._num_classes = num_classes

        self._device = ctx

        self._data_sym = data
        self._batch_logits_sym = logits

        label = mx.symbol.Variable('label')
        self._label_sym = label

        # workaround for https://github.com/apache/incubator-mxnet/issues/6874
        log_softmax = mx.sym.log_softmax(logits)

        loss = mx.sym.sum(
            mx.sym.one_hot(indices=label, depth=num_classes) * log_softmax)

        # loss = mx.symbol.softmax_cross_entropy(logits, label)
        self._loss_sym = loss

        self._args_map = args.copy()
        self._aux_map = aux_states.copy() if aux_states is not None else None

        # move all parameters to correct device
        for k in self._args_map.keys():
            self._args_map[k] = \
                self._args_map[k].as_in_context(ctx)      # pragma: no cover

        if aux_states is not None:
            for k in self._aux_map.keys():                # pragma: no cover
                self._aux_map[k] = \
                    self._aux_map[k].as_in_context(ctx)   # pragma: no cover

    def num_classes(self):
        return self._num_classes

    def batch_predictions(self, images):
        import mxnet as mx
        images, _ = self._process_input(images)
        data_array = mx.nd.array(images, ctx=self._device)
        self._args_map[self._data_sym.name] = data_array
        model = self._batch_logits_sym.bind(
            ctx=self._device, args=self._args_map, grad_req='null',
            aux_states=self._aux_map)
        model.forward(is_train=False)
        logits_array = model.outputs[0]
        logits = logits_array.asnumpy()
        return logits

    def predictions_and_gradient(self, image, label):
        import mxnet as mx
        label = np.asarray(label)
        image, dpdx = self._process_input(image)
        data_array = mx.nd.array(image[np.newaxis], ctx=self._device)
        label_array = mx.nd.array(label[np.newaxis], ctx=self._device)
        self._args_map[self._data_sym.name] = data_array
        self._args_map[self._label_sym.name] = label_array

        grad_array = mx.nd.zeros(image[np.newaxis].shape, ctx=self._device)
        grad_map = {self._data_sym.name: grad_array}

        logits_loss = mx.sym.Group([self._batch_logits_sym, self._loss_sym])
        model = logits_loss.bind(
            ctx=self._device,
            args=self._args_map,
            args_grad=grad_map,
            grad_req='write',
            aux_states=self._aux_map)
        model.forward(is_train=False)
        logits_array = model.outputs[0]
        model.backward([
            mx.nd.zeros(logits_array.shape),
            mx.nd.array(np.array([1]))
        ])
        logits = logits_array.asnumpy()
        gradient = grad_array.asnumpy()
        gradient = self._process_gradient(dpdx, gradient)
        return np.squeeze(logits, axis=0), np.squeeze(gradient, axis=0)

    def _loss_fn(self, image, label):
        import mxnet as mx
        image, _ = self._process_input(image)
        data_array = mx.nd.array(image[np.newaxis], ctx=self._device)
        label_array = mx.nd.array(np.array([label]), ctx=self._device)
        self._args_map[self._data_sym.name] = data_array
        self._args_map[self._label_sym.name] = label_array
        model = self._loss_sym.bind(
            ctx=self._device, args=self._args_map, grad_req='null',
            aux_states=self._aux_map)
        model.forward(is_train=False)
        loss_array = model.outputs[0]
        loss = loss_array.asnumpy()[0]
        return loss

    def backward(self, gradient, image):
        import mxnet as mx

        assert gradient.ndim == 1

        image, dpdx = self._process_input(image)
        data_array = mx.nd.array(image[np.newaxis], ctx=self._device)
        self._args_map[self._data_sym.name] = data_array

        grad_array = mx.nd.zeros(image[np.newaxis].shape, ctx=self._device)
        grad_map = {self._data_sym.name: grad_array}

        logits = self._batch_logits_sym.bind(
            ctx=self._device,
            args=self._args_map,
            args_grad=grad_map,
            grad_req='write',
            aux_states=self._aux_map)

        logits.forward(is_train=False)

        gradient_pre_array = mx.nd.array(
            gradient[np.newaxis], ctx=self._device)
        logits.backward(gradient_pre_array)

        gradient = grad_array.asnumpy()
        gradient = np.squeeze(gradient, axis=0)
        gradient = self._process_gradient(dpdx, gradient)
        return gradient