# -*- coding: utf-8 -*-
"""
Created on Thu Sep 28 09:28:15 2017

@author: tkoller
"""
import warnings
import casadi as cas
import numpy as np
from casadi import MX, mtimes, vertcat, sum2, sqrt
from casadi import reshape as cas_reshape
from .gp_reachability_casadi import lin_ellipsoid_safety_distance
from .gp_reachability_casadi import multi_step_reachability as cas_multistep
from .uncertainty_propagation_casadi import mean_equivalent_multistep, \
    multi_step_taylor_symbolic
from .utils import dlqr, feedback_ctrl, array_of_vec_to_array_of_mat

ATTR_NAMES_PERF = ['type_perf_traj', 'n_perf', 'r', 'perf_has_fb']
DEFAULT_OPT_PERF = {'type_perf_traj': 'mean_equivalent', 'n_perf': 5, 'r': 1,
                    'perf_has_fb': True}

ATTR_NAMES_ENV = ['l_mu', 'l_sigma', 'h_mat_safe', 'h_safe', 'lin_model', 'ctrl_bounds',
                  'safe_policy', 'h_mat_obs', 'h_obs']
DEFAULT_OPT_ENV = {'ctrl_bounds': None, 'safe_policy': None, 'lin_model': None,
                   'h_mat_obs': None, 'h_obs': None}

class SimpleSafeMPC:
    """ Simplified implementation of the SafeMPC algorithm in Casadi

    """

    def __init__(self, n_safe, ssm, opt_env, wx_cost, wu_cost, beta_safety=2.5,
                 rhc=True,
                 safe_policy=None, opt_perf_trajectory={}, lin_trafo_gp_input=None, opts_solver=None, verbosity=0):
        """ Initialize the SafeMPC object with dynamic model information

        Parameters
        ----------
        n_safe: int
            Length of the safety trajectory (number of safe controls)
        ssm: StateSpaceModel
            The underlying statistical model
        opt_env: dict
            Dictionary of environment options. List of accepted attributes are given in ATTR_NAMES_ENV.
            All values that are NOT specified in DEFAULT_OPT_ENV are mandatory.
        wx_cost: n_x x n_x np.ndarray[float]
            State cost matrix for the LQR
        wu_cost: n_u x n_u np.ndarray[float]
            Control cost matrix for the LQR
        beta_safety: float, optional
            The safety coefficient for the confidence intervals (Denoted by \beta in paper)
        rhc: boolean, optional
            True, if we want to reinitialize the SafeMPC problem in a receding horizon fashion
        lin_model: Tuple, optional
            The linear prior model. Consists of tuple (a,b) that defines the linear system.
            Default is: x_{t+1} = x_t, i.e. (I,0)
        ctrl_bounds: n_u x 2 np.ndarray[float], optional
            Upper and lower bound per control input
        safe_policy: function, optional
            Custom safe policy u_t^safe = safe_policy(x_t)
            Default: LQR based on prior model and wx_cost,wu_cost
        opt_perf_trajectory: dict, optional
            Dictionary of environment options. List of accepted attributes are given in ATTR_NAMES_PERF.
            All values that are NOT specified in DEFAULT_OPT_PERF are mandatory.
        lin_trafo_gp_input: n_x_gp_in x n_x np.ndarray[float], optional
            Allows for a linear transformation of the gp input (e.g. removing an input)

        """
        self.rhc = rhc
        self.ssm = ssm
        self.ssm_forward = ssm.get_forward_model_casadi(True)
        self.n_safe = n_safe
        self.n_fail = self.n_safe  # initialize s.t. there is no backup strategy
        self.n_s = self.ssm.num_states
        self.n_u = self.ssm.num_actions
        self.has_openloop = False
        self.opts_solver = opts_solver

        self.safe_policy = safe_policy

        self.cost_func = None  # This is updated wheenver the solver is newly initialized (possibly again with None)

        self._set_attributes_from_dict(ATTR_NAMES_ENV, DEFAULT_OPT_ENV, opt_env)

        self.lin_trafo_gp_input = lin_trafo_gp_input
        if self.lin_trafo_gp_input is None:
            self.lin_trafo_gp_input = np.eye(self.n_s)

        if self.h_mat_obs is None:
            m_obs_mat = 0
        else:
            m_obs_mat, n_s_obs = np.shape(self.h_mat_obs)
            assert n_s_obs == self.n_s, " Wrong shape of obstacle matrix"
            assert np.shape(self.h_obs) == (m_obs_mat,
                                            1), " Shapes of obstacle linear inequality matrix/vector must match "
        self.m_obs = m_obs_mat

        m_safe_mat, n_s_safe = np.shape(self.h_mat_safe)
        assert n_s_safe == self.n_s, " Wrong shape of safety matrix"
        assert np.shape(self.h_safe) == (
            m_safe_mat,
            1), " Shapes of safety linear inequality matrix/vector must match "
        self.m_safe = m_safe_mat

        # init safety constraints evaluator
        p_cas = MX.sym('p', (self.n_s, self.n_u))
        q_cas = MX.sym('q', (self.n_s, self.n_s))
        g_val_term_cas = lin_ellipsoid_safety_distance(p_cas, q_cas, self.h_mat_safe,
                                                       self.h_safe)
        self.g_term_cas = cas.Function("g_term", [p_cas, q_cas], [g_val_term_cas])

        if not self.h_mat_obs is None:
            g_val_interm_cas = lin_ellipsoid_safety_distance(p_cas, q_cas,
                                                             self.h_mat_obs, self.h_obs)
            self.g_interm_cas = cas.Function("g_interm", [p_cas, q_cas],
                                             [g_val_term_cas])

        self.has_ctrl_bounds = False

        if not self.ctrl_bounds is None:
            self.has_ctrl_bounds = True
            assert np.shape(self.ctrl_bounds) == (self.n_u, 2), """control bounds need
            to be of shape n_u x 2 with i,0 lower bound and i,1 upper bound per dimension"""

        self.wx_cost = wx_cost
        self.wu_cost = wu_cost
        self.wx_feedback = wx_cost
        self.wu_feedback = 1 * wu_cost

        self.do_shift_solution = True
        self.solver_initialized = False

        self.beta_safety = beta_safety
        self.verbosity = verbosity

        # SET ALL ATTRIBUTES FOR THE ENVIRONMENT

        self.lin_prior = False
        self.a = np.eye(self.n_s)
        self.b = np.zeros((self.n_s, self.n_u))
        if not self.lin_model is None:
            self.a, self.b = self.lin_model
            self.lin_prior = True
            if self.safe_policy is None:
                # no safe policy specified? Use lqr as safe policy
                K = self.get_lqr_feedback()
                self.safe_policy = lambda x: np.dot(K, x)

        # if self.performance_trajectory is None:
        #    self.performance_trajectory = mean_equivalent
        self._set_attributes_from_dict(ATTR_NAMES_PERF, DEFAULT_OPT_PERF,
                                       opt_perf_trajectory)
        self._set_perf_trajectory(self.type_perf_traj)

        self.k_fb_all = None
        if self.safe_policy is None:
            warnings.warn("No SafePolicy!")

        # init safe

    def init_solver(self, cost_func=None, opt_x0=False, init_uncertainty=False):
        """ Initialize a casadi solver object corresponding to the SafeMPC optimization problem

        Parameters:
        -----------
        cost_func: Function
            A function which admits casadi.SX type inputs
            and returns a scalar function
            If performance controls exist function has to be of the form:
                cost_func(p_all,k_ff_all,x_perf,u_perf)
            otherwise:
                cost_func(p_all,k_ff_all)

        """
        self.cost_func = cost_func
        self.opt_x0 = opt_x0
        self.init_uncertainty = init_uncertainty

        u_0 = MX.sym("init_control", (self.n_u, 1))
        k_ff_all = MX.sym("feed-forward control", (self.n_safe - 1, self.n_u))
        g = []
        lbg = []
        ubg = []
        g_name = []

        p_0 = MX.sym("initial state", (self.n_s, 1))
        q_0 = None
        k_fb_0 = None
        if init_uncertainty:
            q_0 = MX.sym("init uncertainty", (self.n_s, self.n_s))
            k_fb_0 = MX.sym("init feddback control matrix", (self.n_u, self.n_s))

        k_fb_safe = MX.sym("feedback matrices",
                        (self.n_safe - 1, self.n_s * self.n_u))

        p_all, q_all, gp_sigma_pred_safe_all = cas_multistep(p_0, u_0, k_fb_safe, k_ff_all,
                                                             self.ssm_forward, self.l_mu,
                                                             self.l_sigma,
                                                             self.beta_safety, self.a,
                                                             self.b,
                                                             self.lin_trafo_gp_input, q_0, k_fb_0)

        # generate open_loop trajectory function [vertcat(x_0,u_0)],[f_x])

        if init_uncertainty:
            self._f_multistep_eval = cas.Function("safe_multistep",
                                                  [p_0, u_0, k_fb_safe, k_ff_all, q_0, k_fb_0],
                                                  [p_all, q_all, gp_sigma_pred_safe_all])
        else:
            self._f_multistep_eval = cas.Function("safe_multistep",
                                                  [p_0, u_0, k_fb_safe, k_ff_all],
                                                  [p_all, q_all, gp_sigma_pred_safe_all])

        g_safe, lbg_safe, ubg_safe, g_names_safe = self.generate_safety_constraints(
            p_all, q_all, u_0, k_fb_safe, k_ff_all, q_0, k_fb_0)
        g = vertcat(g, g_safe)
        lbg += lbg_safe
        ubg += ubg_safe
        g_name += g_names_safe

        # Generate performance trajectory
        if self.n_perf > 1:
            k_ff_perf, k_fb_perf, k_ff_perf_traj, k_fb_perf_traj, mu_perf, sigma_perf, gp_sigma_pred_perf_all, g_perf, lbg_perf, ubg_perf, g_names_perf = self._generate_perf_trajectory_casadi(
                p_0, u_0, k_ff_all, k_fb_safe, self.a, self.b, self.lin_trafo_gp_input)
            g = vertcat(g, g_perf)
            lbg += lbg_perf
            ubg += ubg_perf
            g_name += g_names_perf
        else:
            k_ff_perf = np.array([])
            k_fb_perf = np.array([])
            k_fb_perf_traj = np.array([])
            k_ff_perf_traj = np.array([])
            mu_perf = np.array([])
            sigma_perf = np.array([])
            gp_sigma_pred_perf_all = None

        cost = self.generate_cost_function(p_0, u_0, p_all, q_all, mu_perf, sigma_perf,
                                           k_ff_all, k_fb_safe, gp_sigma_pred_safe_all,
                                           k_fb_perf=k_fb_perf_traj,
                                           k_ff_perf=k_ff_perf_traj,
                                           gp_pred_sigma_perf=gp_sigma_pred_perf_all,
                                           custom_cost_func=cost_func)

        if self.opt_x0:
            opt_vars = vertcat(p_0, u_0, k_ff_perf, k_ff_all.reshape((-1, 1)))
            opt_params = vertcat(k_fb_safe.reshape((-1, 1)), k_fb_perf.reshape((-1, 1)))
        else:
            opt_vars = vertcat(u_0, k_ff_perf, k_ff_all.reshape((-1, 1)))
            opt_params = vertcat(p_0, k_fb_safe.reshape((-1, 1)), k_fb_perf.reshape((-1, 1)))

        if self.init_uncertainty:
            opt_params = vertcat(opt_params, q_0.reshape((-1, 1)), k_fb_0.reshape((-1, 1)))

        prob = {'f': cost, 'x': opt_vars, 'p': opt_params, 'g': g}

        opt = self.opts_solver
        if opt is None:
            opt = {'error_on_fail': False,
                   'ipopt': {'hessian_approximation': 'limited-memory', "max_iter": 100,
                             "expect_infeasible_problem": "no", \
                             'acceptable_tol': 1e-4, "acceptable_constr_viol_tol": 1e-5,
                             "bound_frac": 0.5, "start_with_resto": "no",
                             "required_infeasibility_reduction": 0.85,
                             "acceptable_iter": 8}}  # ipopt

            # opt = {'max_iter':120,'hessian_approximation':'limited-memory'}#,"c1":5e-4} #sqpmethod #,
        # opt = {'max_iter':120,'qpsol':'qpoases'}

        solver = cas.nlpsol('solver', 'ipopt', prob, opt)
        # solver = cas.nlpsol('solver','sqpmethod',prob,opt)
        # solver = cas.nlpsol('solver','blocksqp',prob,opt)

        self.solver = solver
        self.lbg = lbg
        self.ubg = ubg
        self.solver_initialized = True
        self.g = g
        self.g_name = g_name

    def generate_cost_function(self, p_0, u_0, p_all, q_all, mu_perf, sigma_perf,
                               k_ff_safe, k_fb_safe, sigma_safe, k_fb_perf=None,
                               k_ff_perf=None, gp_pred_sigma_perf=None,
                               custom_cost_func=None, eps_noise=0.0):
        # Generate cost function
        if custom_cost_func is None:
            cost = 0
            if self.n_perf > 1:

                n_cost_deviation = np.minimum(self.n_perf, self.n_safe)
                for i in range(1, n_cost_deviation):
                    cost += mtimes(mu_perf[i, :] - p_all[i, :],
                                   mtimes(.1 * self.wx_cost,
                                          (mu_perf[i, :] - p_all[i, :]).T))

                for i in range(self.n_perf):
                    cost -= sqrt(sum2(gp_pred_sigma_perf[i, :] + eps_noise))
            else:
                for i in range(self.n_safe):
                    cost -= sqrt(sum2(sigma_safe[i, :] + eps_noise))
        else:
            if self.n_perf > 1:
                cost = custom_cost_func(p_0, u_0, p_all, q_all, k_ff_safe, k_fb_safe,
                                        sigma_safe, mu_perf, sigma_perf,
                                        gp_pred_sigma_perf, k_fb_perf, k_ff_perf)
            else:
                cost = custom_cost_func(p_0, u_0, p_all, q_all, k_ff_safe, k_fb_safe,
                                        sigma_safe)

        return cost

    def generate_safety_constraints(self, p_all, q_all, u_0, k_fb, k_ff_all, q_0=None, k_fb_0=None):
        """ Generate all safety constraints

        Parameters
        ----------
        p_all: n_safe x n_s casadi.SX
            The centers of the safe trajctory ellipsoids
        q_all: n_safe x n_s x n_s ndarray[float]

        u_0 The initial
            The shape matrices of the safe trajectory ellipsoids
        k_fb: (n_safe-1) x (n_x * n_u) casadi.SX
            Feedback control matrices
        k_ff_all: (n_safe-1) x n_u casadi.SX
            Feed-forward controls

        Returns
        -------
        g: list[casadi.SX]
            The constraint functions
        lbg: list[casadi.SX]
            Lower bounds for the constraints
        ubg: list[casadi.SX]
            Upper bounds for the constraints
        """
        g = []
        lbg = []
        ubg = []
        g_name = []

        H = np.shape(p_all)[0]
        # control constraints
        if self.has_ctrl_bounds:
            g_u_0, lbg_u_0, ubg_u_0 = self._generate_control_constraint(u_0, q_0, k_fb_0)
            g = vertcat(g, g_u_0)
            lbg += lbg_u_0
            ubg += ubg_u_0
            g_name += ["u_0_ctrl_constraint"]

            for i in range(H - 1):
                p_i = p_all[i, :].T
                q_i = q_all[i, :].reshape((self.n_s, self.n_s))
                k_ff_i = k_ff_all[i, :].reshape((self.n_u, 1))
                k_fb_i = k_fb[i, :].reshape((self.n_u, self.n_s))

                g_u_i, lbg_u_i, ubg_u_i = self._generate_control_constraint(k_ff_i, q_i,
                                                                            k_fb_i)
                g = vertcat(g, g_u_i)
                lbg += lbg_u_i
                ubg += ubg_u_i
                g_name += ["ellipsoid_ctrl_constraint_{}".format(i)] * len(lbg_u_i)

        # intermediate state constraints
        if not self.h_mat_obs is None:
            for i in range(H - 1):
                p_i = p_all[i, :].T
                q_i = q_all[i, :].reshape((self.n_s, self.n_s))
                g_state = lin_ellipsoid_safety_distance(p_i, q_i, self.h_mat_obs,
                                                        self.h_obs)
                g = vertcat(g, g_state)
                lbg += [-cas.inf] * self.m_obs
                ubg += [0] * self.m_obs
                g_name += ["obstacle_avoidance_constraint{}".format(i)] * self.m_obs

        # terminal state constraint
        p_T = p_all[-1, :].T
        q_T = q_all[-1, :].reshape((self.n_s, self.n_s))

        g_terminal = lin_ellipsoid_safety_distance(p_T, q_T, self.h_mat_safe,
                                                   self.h_safe)
        g = vertcat(g, g_terminal)
        g_name += ["terminal constraint"] * self.m_safe
        lbg += [-cas.inf] * self.m_safe
        ubg += [0] * self.m_safe

        return g, lbg, ubg, g_name

    def _generate_perf_trajectory_casadi(self, mu_0, u_0, k_ff_ctrl, k_fb_safe, a=None,
                                         b=None, lin_trafo_gp_input=None,
                                         safety_constr=False):
        """ Generate the performance trajectory variables for the casadi solver

        Parameters:
        mu_0: n_x x 1 casadi.SX
            Initial state
        u_0: n_u x 1 casadi.SX
            Initial control
        k_ff_ctrl: (n_safe-1) x n_u casadi.SX
            Safe feed-forward controls
        k_fb_safe: (n_safe-1) x (n_x * n_u) casadi.SX
            Safe feedback control matrices
        a: n_x x n_x np.ndarray[float], optional
            The A-matrix of the prior linear model
        b: n_x x n_u np.ndarray[float], optional
            The B-matrix of the prior linear model
        lin_trafo_gp_input: n_x_gp_in x n_x np.ndarray[float], optional
            Allows for a linear transformation of the gp input (e.g. removing an input)
        safety_constr: boolean, optional
            True, if we want to put a constraint after (n_safe+1)th performance trajectory state to
            be inside the terminal safe set. Can potentially help with (recursive) feasibility
        """
        if self.r > 1:
            warnings.warn(
                "Coupling performance and safety trajectory for more than one step is UNTESTED")

        # we don't have a performance trajectory, so nothing to do here. Might wanna catch this even before
        if self.n_perf <= 1:
            return np.array([]), np.array([]), np.array([]), np.array(
                []), None, np.array([]), [], [], [], [], []
        else:
            k_ff_perf = MX.sym("k_ff_perf", (self.n_perf - self.r, self.n_u))
            k_ff_perf_traj = vertcat(k_ff_ctrl[:self.r - 1, :], k_ff_perf)

            k_fb_perf_traj = np.array([])
            for i in range(self.r - 1):
                k_fb_perf_traj = np.append(k_fb_perf_traj,
                                           [k_fb_safe[i, :].reshape((self.n_u, self.n_s))])
            if self.perf_has_fb and self.n_perf - self.r > 0:
                k_fb_perf = MX.sym("k_fb_perf", (self.n_u, self.n_s))
                for i in range(self.n_perf - self.r):
                    k_fb_perf_traj = np.append(k_fb_perf_traj, [k_fb_perf])

            mu_perf_all, sigma_perf_all, gp_sigma_pred_perf_all = self.perf_trajectory(
                mu_0, self.ssm_forward, vertcat(u_0, k_ff_perf_traj), k_fb_perf_traj, None, a, b,
                lin_trafo_gp_input)

            # evaluation trajectory (mainly for verbosity)
            mu_0_eval = MX.sym("mu_0", (self.n_s, 1))
            u_0_eval = MX.sym("u_0", (self.n_u, 1))
            k_fb_perf_all_eval = MX.sym("k_fb_perf",
                                        (self.n_perf - 1, self.n_u * self.n_s))
            k_ff_perf_all_eval = MX.sym("k_ff_perf", (self.n_perf - 1, self.n_u))

            list_kfb_perf = [cas_reshape(k_fb_perf_all_eval[i, :], (self.n_u, self.n_s))
                             for i in range(self.n_perf - 1)]
            mu_perf_eval_all, sigma_perf_eval_all, gp_sigma_pred_perf_all_eval = self.perf_trajectory(
                mu_0_eval, self.ssm_forward, vertcat(u_0_eval, k_ff_perf_all_eval),
                list_kfb_perf, None, a, b, lin_trafo_gp_input)
            self._f_multistep_perf_eval = cas.Function("f_multistep_perf_eval",
                                                       [mu_0_eval, u_0_eval,
                                                        k_fb_perf_all_eval,
                                                        k_ff_perf_all_eval],
                                                       [mu_perf_eval_all,
                                                        gp_sigma_pred_perf_all_eval])

        # generate (approxiamte) constraints for the performance trajectory
        g_name = []
        g = []
        lbg = []
        ubg = []
        if safety_constr and self.n_perf > self.n_safe:
            g_name += ["Terminal safety performance"]
            g_term = lin_ellipsoid_safety_distance(
                cas_reshape(mu_perf_all[self.n_safe + 1, :], (self.n_s, 1)),
                cas_reshape(sigma_perf_all[self.n_safe + 1, :], (self.n_s, self.n_s)),
                self.h_mat_safe, self.h_safe)
            g = vertcat(g, g_term)
            lbg += [-np.inf] * self.m_safe
            ubg += [0.] * self.m_safe

        if self.has_ctrl_bounds:
            for i in range(self.n_perf - self.r):
                g_u_i, lbu_i, ubu_i = self._generate_control_constraint(
                    k_ff_perf[i, :].T)
                g = vertcat(g, g_u_i)
                lbg += lbu_i
                ubg += ubu_i
                g_name += ["ctrl_constr_performance_{}".format(i)]

        return k_ff_perf, k_fb_perf, k_ff_perf_traj, k_fb_perf_traj, mu_perf_all, sigma_perf_all, gp_sigma_pred_perf_all, g, lbg, ubg, g_name

    def _generate_control_constraint(self, k_ff, q=None, k_fb=None, ctrl_bounds=None):
        """ Build control constraints from state ellipsoids and linear feedback controls

        k_ff: n_u x 1 ndarray[casadi.SX]
            The feed-forward control gain
        q: n_s x n_s ndarray[casadi.SX]
            The shape matrix of the state ellipsoid
        k_fb: n_u x n_s ndarray[casadi.SX]
            The feedback gain
        ctrl_bounds: n_u x 2 ndarray[float], optional

        Returns
        -------
        g: 2*n_u x 1 ndarray[casadi.SX]
            The control constraints (symbollicaly) evaluated at the current
            state/controls
        lbg: 2*n_u x 0 list[float]
            Lower bounds for the control constraints
        ubg: 2*n_u x 0 list[float]
            Upper bounds for the control constraints
        """
        if ctrl_bounds is None:
            if not self.has_ctrl_bounds:
                raise ValueError("""Either ctrl_bounds has to be specified or
                the objects' ctrl_bounds has to be specified """)
            ctrl_bounds = self.ctrl_bounds

        # no feedback term. Reduces to simple feed-forward control bounds

        n_u, _ = np.shape(ctrl_bounds)
        u_min = ctrl_bounds[:, 0]
        u_max = ctrl_bounds[:, 1]

        if k_fb is None:
            return k_ff, u_min.tolist(), u_max.tolist()

        h_vec = np.vstack((u_max[:, None], -u_min[:, None]))
        h_mat = np.vstack((np.eye(n_u), -np.eye(n_u)))

        p_u = k_ff
        q_u = mtimes(k_fb, mtimes(q, k_fb.T))

        g = lin_ellipsoid_safety_distance(p_u, q_u, h_mat, h_vec)

        return g, [-cas.inf] * 2 * n_u, [0] * 2 * n_u

    def _eval_prior_casadi(self, state, action):
        """ symbolically evaluate the prior

        Parameters
        ----------
        state: n x n_s array[casadi.SX]
            Symbolic array of states
        action: n x 1 array[casadi.SX]
            Symbolic array of actions

        Returns
        -------
        x_prior: n x n_s array[casadi.SX]
            The (state,action) pairs evaluated at the prior
        """

        return mtimes(self.a, state.T) + mtimes(self.b, action.T)

    def eval_prior(self, state, action):
        """ Evaluate the prior numerically

        Parameters
        ----------
        state: n x n_s array[float]
            Array of states
        action: n x n_u array[float]

        Returns
        -------
        x_prior: n x n_s array[float]
            The (state,action) pairs evaluated at the prior
        """

        return np.dot(state, self.a.T) + np.dot(action, self.b.T)

    def get_lqr_feedback(self, x_0=None, u_0=None):
        """ Get the initial feedback controller k_fb

        x_0: n_s x 1 ndarray[float], optional
            Current state of the system
        u_0: n_u x 1 ndarray[float], optional
            Initialization of the control input

        """
        q = self.wx_feedback
        r = self.wu_feedback

        if x_0 is None:
            x_0 = np.zeros((self.n_s, 1))
        if u_0 is None:
            u_0 = np.zeros((self.n_u, 1))

        if self.lin_prior:
            a = self.a
            b = self.b

            k_lqr, _, _ = dlqr(a, b, q, r)
            k_fb = -k_lqr
        else:

            raise NotImplementedError(
                "Cannot compute feed-back matrices without prior model")

        return k_fb.reshape((1, self.n_s * self.n_u))

    def get_safety_trajectory_openloop(self, x_0, u_0, k_fb=None, k_ff=None, q_0=None, k_fb_0=None):
        """ Compute a trajectory of ellipsoids based on an initial state and a set of controls

        Parameters
        ----------
        x_0: n_s x 1 2darray[float]
            The initial state
        u_0: n_u x 1 2darray[float]
            The initial action
        k_fb: (n_safe-1) x n_u x n_s  or n_u x n_s ndarray[float], optional
            The feedback controls. Uses the most recent solution to the
            MPC Problem (when calling solve()) if this parameter is not set
        k_ff: (n_safe-1) x n_u, optional
            The feed-forward controls. Uses the most recent solution to the
            MPC Problem (when calling solve()) if this parameter is not set
        get_controls: bool, optional
            Additionally returns the applied controls if this flag is set to TRUE
        Returns
        -------
        p_all: T x n_s ndarray[float]
            The centers of the trajctory ellipsoids
        q_all: T x n_s x n_s ndarray[float]
            The shape matrices of the trajectory ellipsoids
        """
        if not self.has_openloop:
            return None, None

        if k_fb is None:
            k_fb = np.array(self.k_fb_all)

        if k_ff is None:
            k_ff = np.array(self.k_ff_all)

        if self.init_uncertainty:
            p_all, q_all, gp_sigma_pred_safe_all = self._f_multistep_eval(x_0, u_0, k_fb, k_ff, q_0, k_fb_0)
        else:
            p_all, q_all, gp_sigma_pred_safe_all = self._f_multistep_eval(x_0, u_0, k_fb, k_ff)

        return p_all, q_all, gp_sigma_pred_safe_all

    def get_action(self, x0_mu, lqr_only=False, sol_verbose=False):
        """ Wrapper around the solve function

        Parameters
        ----------
        x0_mu: n_s x 0 1darray[float]
            The current state of the system

        Returns
        -------
        u_apply: n_u x 0 1darray[float]
            The action to be applied to the system
        success: bool
            The control was not successful if we are outside the safezone
            AND we have to revert to the safe controller.
        """
        safety_failure = False
        if lqr_only:
            u_apply = self.safe_policy(x0_mu)

            return u_apply, safety_failure

        if sol_verbose:
            _, u_apply, feasible, success, k_fb_apply, k_ff_all, p_all, q_all = self.solve(
                x0_mu[:, None], sol_verbose=True)
            return u_apply.reshape(
                self.n_u, ), feasible, success, k_fb_apply, k_ff_all, p_all, q_all

        else:
            _, u_apply, success = self.solve(x0_mu[:, None])

            return u_apply.reshape(self.n_u, ), success

    def solve(self, p_0, u_0=None, k_ff_all_0=None, k_fb_safe=None, u_perf_0=None,
              k_fb_perf_0=None, sol_verbose=False, q_0=None, k_fb_0=None):
        """ Solve the MPC problem for a given set of input parameters

        Parameters
        ----------
        p_0: n_s x 1 array[float]
            The initial (current) state
        k_ff_all_0: n_safe x n_u  array[float], optional
            The initialization of the feed-forward controls
        k_fb_all_0: n_safe x (n_s * n_u) array[float], optional
            The initialization of the feedback controls

        Returns
        -------
        k_fb_apply: n_u x n_s array[float]
            The feedback control term to be applied to the system
        k_ff_apply: n_u x 1 array[float]
            The feed-forward control term to be applied to the system
        k_fb_all: n_safe x n_u x n_s
            The feedback control terms for all time steps
        k_ff_all: n_safe x n_u x 1
        """
        assert self.solver_initialized, "Need to initialize the solver first!"

        u_0_init, k_ff_all_0_init, k_fb_safe_init, u_perf_0_init, k_fb_perf_0_init = self._get_init_controls()

        if u_0 is None:
            u_0 = u_0_init
        if k_ff_all_0 is None:
            k_ff_all_0 = k_ff_all_0_init
        if k_fb_safe is None:
            k_fb_safe = k_fb_safe_init
        if u_perf_0 is None:
            u_perf_0 = u_perf_0_init
        if k_fb_perf_0 is None:
            k_fb_perf_0 = k_fb_perf_0_init
        if q_0 is not None:
            if k_fb_0 is None:
                k_fb_0 = self.get_lqr_feedback()

        if self.opt_x0:
            params = np.vstack(
                (cas_reshape(k_fb_safe, (-1, 1)), cas_reshape(k_fb_perf_0, (-1, 1))))

            opt_vars_init = vertcat(cas_reshape(p_0, (-1, 1)), cas_reshape(u_0, (-1, 1)), u_perf_0, \
                               cas_reshape(k_ff_all_0, (-1, 1)))
        else:
            params = np.vstack(
                (p_0, cas_reshape(k_fb_safe, (-1, 1)), cas_reshape(k_fb_perf_0, (-1, 1))))

            opt_vars_init = vertcat(cas_reshape(u_0, (-1, 1)), u_perf_0, \
                             cas_reshape(k_ff_all_0, (-1, 1)))

        if self.init_uncertainty:
            params = vertcat(params, cas_reshape(q_0, (-1, 1)), cas_reshape(k_fb_0, (-1, 1)))

        crash = False
        sol = self.solver(x0=opt_vars_init, lbg=self.lbg, ubg=self.ubg, p=params)
        try:
            # pass
            sol = self.solver(x0=opt_vars_init, lbg=self.lbg, ubg=self.ubg, p=params)
        except:
            crash = True
            warnings.warn("NLP solver crashed, solution infeasible")
            sol = None

        return self._get_solution(p_0, sol, k_fb_safe, k_fb_perf_0, sol_verbose, crash, q_0=q_0, k_fb_0=k_fb_0)

    def _get_solution(self, x_0, sol, k_fb, k_fb_perf_0, sol_verbose=False,
                      crashed=False, feas_tol=1e-6, q_0=None, k_fb_0=None):
        """ Process the solution dict of the casadi solver

        Processes the solution dictionary of the casadi solver and
        (depending on the chosen mode) saves the solution for reuse in the next
        time step. Depending on the chosen verbosity level, it also prints
        some statistics.

        Parameters
        ----------
        sol: dict
            The solution dictionary returned by the casadi solver
        sol_verbose: boolean, optional
            Return additional solver results such as the constraint values

        Returns
        -------
        k_fb_apply: n_u x n_s array[float]
            The feedback control term to be applied to the system
        k_ff_apply: n_u x 1 array[float]
            The feed-forward control term to be applied to the system
        k_fb_all: n_safe x n_u x n_s
            The feedback control terms for all time steps
        k_ff_all: n_safe x n_u x 1

        h_values: (m_obs*n_safe + m_safe) x 0 array[float], optional
            The values of the constraint evaluation (distance to obstacle)
        """

        success = True
        feasible = True
        if crashed:
            feasible = False

            if self.verbosity > 1:
                print("Optimization crashed, infeasible soluion!")
        else:
            g_res = np.array(sol["g"]).squeeze()

            # This is not sufficient, since casadi gives out wrong feasibility values
            if np.any(np.array(self.lbg) - feas_tol > g_res) or np.any(
                    np.array(self.ubg) + feas_tol < g_res):
                feasible = False

            x_opt = sol["x"]
            self.has_openloop = True

            if self.opt_x0:
                x_0 = x_opt[:self.n_s]
                x_opt = x_opt[self.n_s:, :]

            # get indices of the respective variables
            n_u_0 = self.n_u
            n_u_perf = 0
            if self.n_perf > 1:
                n_u_perf = (self.n_perf - self.r) * self.n_u
            n_k_ff = (self.n_safe - 1) * self.n_u

            c = 0
            idx_u_0 = np.arange(n_u_0)
            c += n_u_0
            idx_u_perf = np.arange(c, c + n_u_perf)
            c += n_u_perf
            idx_k_ff = np.arange(c, c + n_k_ff)
            c += n_k_ff

            u_apply = np.array(cas_reshape(x_opt[idx_u_0], (1, self.n_u)))
            k_ff_perf = np.array(
                cas_reshape(x_opt[idx_u_perf], (self.n_perf - self.r, self.n_u)))

            k_ff_safe = np.array(
                cas_reshape(x_opt[idx_k_ff], (self.n_safe - 1, self.n_u)))
            k_ff_safe_all = np.vstack((u_apply, k_ff_safe))

            k_fb_safe_output = array_of_vec_to_array_of_mat(np.copy(k_fb), self.n_u,
                                                            self.n_s)

            p_safe, q_safe, gp_sigma_pred_safe_all = self.get_safety_trajectory_openloop(x_0, u_apply,
                                                                 np.copy(k_fb),
                                                                 k_ff_safe, q_0, k_fb_0)

            p_safe = np.array(p_safe)
            q_safe = np.array(q_safe)

            if self.verbosity > 1:
                print("=== Safe Trajectory: ===")
                print("Centers:")
                print(p_safe)
                print("Shape matrices:")
                print(q_safe)
                print("Safety controls:")
                print(u_apply)
                print(k_ff_safe)

            k_fb_perf_traj_eval = np.empty((0, self.n_s * self.n_u))
            k_ff_perf_traj_eval = np.empty((0, self.n_u))
            if self.n_safe > 1:
                k_fb_perf_traj_eval = np.vstack(
                    (k_fb_perf_traj_eval, k_fb[:self.r - 1, :]))
                k_ff_perf_traj_eval = np.vstack(
                    (k_ff_perf_traj_eval, k_ff_safe[:self.r - 1, :]))
            if self.n_perf > self.r:
                k_fb_perf_traj_eval = np.vstack((k_fb_perf_traj_eval,
                                                 np.matlib.repmat(k_fb_perf_0,
                                                                  self.n_perf - self.r,
                                                                  1)))
            k_ff_perf_traj_eval = np.vstack((k_ff_perf_traj_eval, k_ff_perf))

            if self.n_perf > 1:
                mu_perf, sigma_perf = self._f_multistep_perf_eval(x_0.squeeze(),
                                                                  u_apply,
                                                                  k_fb_perf_traj_eval,
                                                                  k_ff_perf_traj_eval)

                if self.verbosity > 1:
                    print("=== Performance Trajectory: ===")
                    print("Mu perf:")
                    print(mu_perf)
                    print("Peformance controls:")
                    print(k_ff_perf_traj_eval)

            feasible, _ = self.eval_safety_constraints(p_safe, q_safe)

            if self.rhc and feasible:
                self.k_ff_safe = k_ff_safe
                self.k_ff_perf = k_ff_perf
                self.p_safe = p_safe
                self.k_fb_safe_all = np.copy(k_fb)
                self.u_apply = u_apply
                self.k_fb_perf_0 = k_fb_perf_0

        if feasible:
            self.n_fail = 0

        if not feasible:
            self.n_fail += 1
            q_all = None
            k_fb_safe_output = None
            k_ff_all = None
            p_safe = None
            q_safe = None
            g_res = None

            if self.n_fail >= self.n_safe:
                # Too many infeasible solutions -> switch to safe controller
                if self.verbosity > 1:
                    print(
                        "Infeasible solution. Too many infeasible solutions, switching to safe controller")
                u_apply = self.safe_policy(x_0)
                k_ff_safe_all = u_apply
            else:
                # can apply previous solution
                if self.verbosity > 1:
                    print((
                        "Infeasible solution. Switching to previous solution, n_fail = {}, n_safe = {}".format(
                            self.n_fail, self.n_safe)))
                if sol_verbose:
                    u_apply, k_fb_safe_output, k_ff_safe_all, p_safe = self.get_old_solution(
                        x_0, get_ctrl_traj=True)
                else:
                    u_apply = self.get_old_solution(x_0)
                    k_ff_safe_all = u_apply

        if sol_verbose:
            return x_0, u_apply, feasible, success, k_fb_safe_output, k_ff_safe_all, p_safe, q_safe, sol, gp_sigma_pred_safe_all

        return x_0, u_apply, success

    def eval_safety_constraints(self, p_all, q_all, ubg_term=0., lbg_term=-np.inf,
                                ubg_interm=0., lbg_interm=-np.inf, terminal_only=False,
                                eps_constraints=1e-5):
        """ Evaluate the safety constraints """
        g_term_val = self.g_term_cas(p_all[-1, :, None],
                                     cas_reshape(q_all[-1, :], (self.n_s, self.n_s)))

        feasible_term = np.all(lbg_term - eps_constraints < g_term_val) and np.all(
            g_term_val < ubg_term + eps_constraints)

        feasible = feasible_term
        if terminal_only or self.h_mat_obs is None:

            if self.verbosity > 1:
                print((
                    "\n===== Evaluated terminal constraint values: FEASIBLE = {} =====".format(
                        feasible)))
                print(g_term_val)
                print("\n===== ===== ===== ===== ===== ===== ===== =====")

            return feasible, g_term_val

        g_interm_val = self.g_interm_cas(p_all[-1, :, None], cas_reshape(q_all[-1, :], (
            self.n_s, self.n_s)))

        feasible_interm = np.all(
            lbg_interm - eps_constraints < g_interm_val) and np.all(
            g_interm_val < ubg_interm + eps_constraints)

        feasible = feasible_term and feasible_interm

        return feasible, np.vstack((g_term_val, g_interm_val))

    def get_old_solution(self, x, k=None, get_ctrl_traj=False):
        """ Shift previously obtained solutions in time and return solution to be applied

        Prameters
        ---------
        k: int, optional
            The number of steps to shift back in time. This is number is
            already tracked by the algorithm, so a custom value should be used with caution
        get_ctrl_traj: bool, optional
            Return the safety trajectory state feedback ctrl laws in terms of k_fb,k_ff,p_ctrl

        Returns
        -------
        u_apply: n_s x 0 1darray[float]
            The controls to be applied at the current time step

        if get_ctrl_traj:
            k_fb_safe_traj:
                The feedback ctrls of the remaining safety trajectory
            k_ff_safe_traj:
                The current ff ctrl and the remaining ff ctrls of the safety trajectory
            p_ctrl_safe_traj:
                The

        """
        if self.n_fail > self.n_safe:
            warnings.warn(
                "There are no previous solution to be applied. Returning None")
            return None
        if k is None:
            k = self.n_fail

        if k < 1:
            warnings.warn("Have to shift at least one timestep back")
            return None

        k_fb_old = self.k_fb_safe_all[k - 1]
        k_ff = self.k_ff_safe[k - 1, :, None]
        p_safe = self.p_safe[k - 1, :, None]

        u_apply = feedback_ctrl(x, k_ff, k_fb_old, p_safe)
        if get_ctrl_traj:
            k_fb_safe_traj = None
            k_ff_safe_traj = u_apply
            p_ctrl_safe_traj = None

            if k < self.n_safe:
                k_fb_safe_traj = self.k_fb_safe_all[k:, :]
                # in accordance to the structure current ctrl u_apply is part of the k_ff ctrl trajectory
                k_ff_safe_traj = np.vstack((u_apply, self.k_ff_safe[k + 1:, :]))
                p_ctrl_safe_traj = self.p_safe[k:, :]

            return u_apply, k_fb_safe_traj, \
                   k_ff_safe_traj, p_ctrl_safe_traj

        return u_apply

    def _set_attributes_from_dict(self, attrib_names, default_attribs={},
                                  custom_attrib={}):
        """ Set class attributes from a list of keys,values """
        for attr_name in attrib_names:
            if attr_name in custom_attrib:
                attr_val = custom_attrib[attr_name]
            elif attr_name in default_attribs:
                attr_val = default_attribs[attr_name]
            else:
                raise ValueError(
                    "Neither a custom nor a default value is given for the requried attribute {}".format(
                        attr_name))

            setattr(self, attr_name, attr_val)

    def _set_perf_trajectory(self, name):
        """ Get the peformance trajectory function from identifier"""
        if name == 'mean_equivalent':
            self.perf_trajectory = mean_equivalent_multistep
        elif name == 'taylor':
            self.perf_trajectory = multi_step_taylor_symbolic
        else:
            raise NotImplementedError("Unknown uncertainty propagation method")

    def _get_init_controls(self):
        """ Initialize the controls for the MPC step

        Returns
        -------
        u_0: n_u x 0 np.array[float]
            Initialization of the first (shared) input
        k_ff_safe_new:  (n_safe-1) x n_u np.ndarray[float]
            Initialization of the safety feed-forward control inputs
        k_fb_new: n_u x n_x np.ndarray[float]
            Initialization of the safety feed-back control inputs
        k_ff_perf_new: (n_perf - r_1) x n_u np.ndarray[float]
            Initialization of the performance feed-forward control inputs.
            Not including the shared controls (if r > 1)
        k_fb_perf_0: n_u x n_x np.ndarray[float]
            Initialization of the safety feed-back control inputs

        """

        u_perf_0 = None
        k_fb_perf_0 = None
        k_fb_lqr = self.get_lqr_feedback()

        if self.do_shift_solution and self.n_fail == 0:
            if self.n_safe > 1:
                k_fb_safe = np.copy(self.k_fb_safe_all)

                # Shift the safe controls
                k_ff_safe = np.copy(self.k_ff_safe)

                u_0 = k_ff_safe[0, :]

                if self.n_safe > self.r and self.n_perf > self.n_safe:  # the first control after the shared controls
                    k_ff_perf = np.copy(self.k_ff_perf)
                    k_ff_r_last = (k_ff_perf[0, :] + k_ff_safe[self.r - 1,
                                                     :]) / 2  # mean of first perf ctrl and safe ctrl after shared
                else:
                    k_ff_r_last = k_ff_safe[-1, :]  # just the last safe control

                k_ff_safe_new = np.vstack((k_ff_safe[1:self.r, :], k_ff_r_last))

                if self.n_safe > self.r + 1:
                    k_ff_safe_new = np.vstack((k_ff_safe_new, k_ff_safe[self.r:, :]))
            else:
                u_0 = self.u_apply
                k_ff_safe_new = np.array([])

            if self.n_perf - self.r > 0:
                k_ff_perf = np.copy(self.k_ff_perf)
                k_ff_perf_new = np.vstack((k_ff_perf[1:, :], k_ff_perf[-1, :]))

                if self.perf_has_fb:
                    k_fb_perf_0 = np.copy(self.k_fb_perf_0)
                else:
                    k_fb_perf_0 = np.array([])
            else:
                k_ff_perf_new = np.array([])
                k_fb_perf_0 = np.array([])
        else:
            k_fb_safe = np.empty((self.n_safe - 1, self.n_s * self.n_u))
            for i in range(self.n_safe - 1):
                k_fb_safe[i] = cas_reshape(k_fb_lqr, (1, -1))

            k_ff_safe_new = np.zeros((self.n_safe - 1, self.n_u))
            u_0 = np.zeros((self.n_u, 1))

            k_ff_perf_new = np.array([])
            if self.n_perf > 1:
                k_ff_perf_new = np.zeros((self.n_perf - self.r, self.n_u))

                if self.perf_has_fb:
                    k_fb_perf_0 = k_fb_lqr
            else:
                k_fb_perf_0 = np.array([])

        if self.n_safe > 1:
            k_fb_safe_new = np.vstack((k_fb_safe[1:, :], k_fb_safe[-1, :]))

        else:
            k_fb_safe_new = np.array([])

        return u_0, k_ff_safe_new, k_fb_safe, k_ff_perf_new, k_fb_perf_0

    def update_model(self, x, y, opt_hyp=False, replace_old=True,
                     reinitialize_solver=True):
        """ Update the model of the dynamics

        Parameters
        ----------
        x: n x (n_s+n_u) array[float]
            The raw training input (state,action) pairs
        y: n x (n_s) array[float]
            The raw training targets (observations of next state)
        opt_hyp: bool
            True, if we want to re-optimize the GP hyperparameters
        replace_old: bool
            True, if we replace the current training set of the GP with x,y
        reinitialize_solver:
            True, if we re-initialize the solver (otherwise the MPC will not be updated with the new GP)

        """
        n_train = np.shape(x)[0]
        x_s = x[:, :self.n_s].reshape((n_train, self.n_s))
        x_u = x[:, self.n_s:].reshape((n_train, self.n_u))
        y_prior = self.eval_prior(x_s, x_u)
        x_trafo = mtimes(x_s, self.lin_trafo_gp_input.T)

        x = np.hstack((x_trafo, x_u))

        self.ssm.update_model(x, y - y_prior, opt_hyp, replace_old)
        self.ssm_forward = self.ssm.get_forward_model_casadi(True)

        if reinitialize_solver:
            self.init_solver(self.cost_func)
        else:
            warnings.warn("""Updating gp without reinitializing the solver! \n
                This is potentially dangerous, since the new GP is not incorporated in the MPC""")


#!/home/jue/.conda/envs/ampere/bin/python

import os
import mock
import numpy as np
import sys
import datetime

from typing import Any, Mapping, Optional, Sequence, Tuple
import collections
from collections import OrderedDict

from timeit import default_timer as timer
import argparse

import io
from Bio import PDB
from Bio.PDB.Polypeptide import PPBuilder
from Bio.PDB import PDBParser
from Bio.PDB.mmcifio import MMCIFIO

sys.path.insert( 0, '/home/nrbennet/software/dl/af2/alphafold' )

import scipy

import jax
import jax.numpy as jnp

from alphafold.common import residue_constants
from alphafold.common import protein
from alphafold.common import confidence
from alphafold.data import pipeline
from alphafold.data import templates
from alphafold.data import mmcif_parsing
from alphafold.model import data
from alphafold.model import config
from alphafold.model import model
from alphafold.data.tools import hhsearch

sys.path.append( '/home/nrbennet/software/silent_tools' )
import silent_tools

# Run with Nate's ampere environment
from pyrosetta import *
from rosetta import *
init( '-in:file:silent_struct_type binary' )

def get_args():
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-silent", dest="silent", required=True,
                        help="silent file to predict")
    parser.add_argument("-batch", dest="batch_size", type=int, default=1,
                        help="Number of structures to batch together")
    parser.add_argument("-recycle", dest="recycle", type=int, default=3,
                        help="Number of recycle iterations to perform")

    args = parser.parse_args()
    return args

args = get_args()

model_name = "model_1_ptm"

model_config = config.model_config(model_name)
model_config.data.eval.num_ensemble = 1

model_config.data.common.num_recycle = args.recycle
model_config.model.num_recycle = args.recycle

model_config.model.embeddings_and_evoformer.initial_guess = True

model_config.model.global_config.mixed_precision = False

model_config.data.common.max_extra_msa = 5
model_config.data.eval.max_msa_clusters = 5

model_params = data.get_model_haiku_params(model_name=model_name, data_dir="/projects/ml/alphafold") # CHANGE THIS (directory where "params/ folder is")
model_runner = model.RunModel(model_config, model_params)

def get_seq_from_pdb( pdb_fn ):
  to1letter = {
    "ALA":'A', "ARG":'R', "ASN":'N', "ASP":'D', "CYS":'C',
    "GLN":'Q', "GLU":'E', "GLY":'G', "HIS":'H', "ILE":'I',
    "LEU":'L', "LYS":'K', "MET":'M', "PHE":'F', "PRO":'P',
    "SER":'S', "THR":'T', "TRP":'W', "TYR":'Y', "VAL":'V' }

  seq = []
  seqstr = ''
  with open(pdb_fn) as fp:
    for line in fp:
      if line.startswith("TER"):
          seq.append(seqstr)
          seqstr = ''
      if not line.startswith("ATOM"):
        continue
      if line[12:16].strip() != "CA":
        continue
      resName = line[17:20]
      #
      seqstr += to1letter[resName]
  return seq

def af2_get_atom_positions( pdbfilename ) -> Tuple[np.ndarray, np.ndarray]:
  """Gets atom positions and mask from a list of Biopython Residues."""

  with open(pdbfilename, 'r') as pdb_file:
    lines = pdb_file.readlines()

  # indices of residues observed in the structure
  idx_s = [int(l[22:26]) for l in lines if l[:4]=="ATOM" and l[12:16].strip()=="CA"]
  num_res = len(idx_s)

  all_positions = np.zeros([num_res, residue_constants.atom_type_num, 3])
  all_positions_mask = np.zeros([num_res, residue_constants.atom_type_num],
                                dtype=np.int64)

  residues = collections.defaultdict(list)
  # 4 BB + up to 10 SC atoms
  xyz = np.full((len(idx_s), 14, 3), np.nan, dtype=np.float32)
  for l in lines:
    if l[:4] != "ATOM":
        continue
    resNo, atom, aa = int(l[22:26]), l[12:16], l[17:20]

    residues[ resNo ].append( ( atom.strip(), aa, [float(l[30:38]), float(l[38:46]), float(l[46:54])] ) )

  for resNo in residues:

    pos = np.zeros([residue_constants.atom_type_num, 3], dtype=np.float32)
    mask = np.zeros([residue_constants.atom_type_num], dtype=np.float32)

    for atom in residues[ resNo ]:
      atom_name = atom[0]
      x, y, z = atom[2]
      if atom_name in residue_constants.atom_order.keys():
        pos[residue_constants.atom_order[atom_name]] = [x, y, z]
        mask[residue_constants.atom_order[atom_name]] = 1.0
      elif atom_name.upper() == 'SE' and res.get_resname() == 'MSE':
        # Put the coordinates of the selenium atom in the sulphur column.
        pos[residue_constants.atom_order['SD']] = [x, y, z]
        mask[residue_constants.atom_order['SD']] = 1.0

    idx = idx_s.index(resNo) # This is the order they show up in the pdb
    all_positions[idx] = pos
    all_positions_mask[idx] = mask
  # _check_residue_distances(
  #     all_positions, all_positions_mask, max_ca_ca_distance) # AF2 checks this but if we want to allow massive truncations we don't want to check this

  return all_positions, all_positions_mask

def af2_all_atom_from_struct( pdbfilename, seq_list, just_target=False ):
  template_seq = ''.join( seq_list )

  # Parse a residue mask from the chainbreak sequence
  binder_len = len( seq_list[0] )
  residue_mask = [ int( i ) > binder_len for i in range( 1, len( template_seq ) + 1 ) ]

  all_atom_positions, all_atom_mask = af2_get_atom_positions( pdbfilename )

  all_atom_positions = np.split(all_atom_positions, all_atom_positions.shape[0])

  templates_all_atom_positions = []

  # Initially fill will all zero values
  for _ in template_seq:
    templates_all_atom_positions.append(
        jnp.zeros((residue_constants.atom_type_num, 3)))

  for idx, i in enumerate( template_seq ):
    if just_target and not residue_mask[ idx ]: continue

    templates_all_atom_positions[ idx ] = all_atom_positions[ idx ][0] # assign target indices to template coordinates

  return jnp.array(templates_all_atom_positions)

def template_from_struct( pdbfilename, seq_list ):

  template_seq = ''.join(seq_list)

  # Parse a residue mask from the chainbreak sequence
  binder_len = len( seq_list[0] )
  residue_mask = [ int( i ) > binder_len for i in range( 1, len( template_seq ) + 1 ) ]

  ret_all_atom_positions, ret_all_atom_mask = af2_get_atom_positions( pdbfilename )

  all_atom_positions = np.split(ret_all_atom_positions, ret_all_atom_positions.shape[0])
  all_atom_masks = np.split(ret_all_atom_mask, ret_all_atom_mask.shape[0])

  output_templates_sequence = []
  output_confidence_scores = []
  templates_all_atom_positions = []
  templates_all_atom_masks = []

  # Initially fill will all zero values
  for _ in template_seq:
    templates_all_atom_positions.append(
        np.zeros((residue_constants.atom_type_num, 3)))
    templates_all_atom_masks.append(np.zeros(residue_constants.atom_type_num))
    output_templates_sequence.append('-')
    output_confidence_scores.append(-1)

  confidence_scores = []
  for _ in template_seq: confidence_scores.append( 9 )

  for idx, i in enumerate( template_seq ):

    if not residue_mask[ idx ]: continue

    templates_all_atom_positions[ idx ] = all_atom_positions[ idx ][0] # assign target indices to template coordinates
    templates_all_atom_masks[ idx ] = all_atom_masks[ idx ][0]
    output_templates_sequence[ idx ] = template_seq[ idx ]
    output_confidence_scores[ idx ] = confidence_scores[ idx ] # 0-9 where higher is more confident

  output_templates_sequence = ''.join(output_templates_sequence)

  templates_aatype = residue_constants.sequence_to_onehot(
      output_templates_sequence, residue_constants.HHBLITS_AA_TO_ID)

  template_feat_dict = {'template_all_atom_positions': np.array(templates_all_atom_positions)[None],
       'template_all_atom_masks': np.array(templates_all_atom_masks)[None],
       'template_sequence': [output_templates_sequence.encode()],
       'template_aatype': np.array(templates_aatype)[None],
       'template_confidence_scores': np.array(output_confidence_scores)[None],
       'template_domain_names': ['none'.encode()],
       'template_release_date': ["none".encode()]}

  return template_feat_dict, ret_all_atom_positions, ret_all_atom_mask

def insert_chainbreaks( pose, binderlen ):

    conf = pose.conformation()
    conf.insert_chain_ending( binderlen )
    pose.set_new_conformation( conf )

    splits = pose.split_by_chain()

    newpose = splits[1]
    for i in range( 2, len( splits )+1 ):
        newpose.append_pose_by_jump( splits[i], newpose.size() )

    info = core.pose.PDBInfo( newpose, True )
    newpose.pdb_info( info )

    return newpose

    return final_dict

def get_final_dict(score_dict, string_dict):
    print(score_dict)
    final_dict = OrderedDict()
    keys_score = [] if score_dict is None else list(score_dict)
    keys_string = [] if string_dict is None else list(string_dict)

    all_keys = keys_score + keys_string

    argsort = sorted(range(len(all_keys)), key=lambda x: all_keys[x])

    for idx in argsort:
        key = all_keys[idx]

        if ( idx < len(keys_score) ):
            final_dict[key] = "%8.3f"%(score_dict[key])
        else:
            final_dict[key] = string_dict[key]

    return final_dict

def add2scorefile(tag, scorefilename, write_header=False, score_dict=None):
    with open(scorefilename, "a") as f:
        add_to_score_file_open(tag, f, write_header, score_dict)

def add_to_score_file_open(tag, f, write_header=False, score_dict=None, string_dict=None):
    final_dict = get_final_dict( score_dict, string_dict )
    if ( write_header ):
        f.write("SCORE:     %s description\n"%(" ".join(final_dict.keys())))
    scores_string = " ".join(final_dict.values())
    f.write("SCORE:     %s        %s\n"%(scores_string, tag))

def generate_scoredict( outtag, start_time, binderlen, prediction_result, scorefilename ):

  plddt_array = prediction_result['plddt']
  plddt = np.mean( plddt_array )
  plddt_binder = np.mean( plddt_array[:binderlen] )
  plddt_target = np.mean( plddt_array[binderlen:] )

  pae = prediction_result['predicted_aligned_error']
  pae_interaction1 = np.mean( pae[:binderlen,binderlen:] )
  pae_interaction2 = np.mean( pae[binderlen:,:binderlen] )
  pae_binder = np.mean( pae[:binderlen,:binderlen] )
  pae_target = np.mean( pae[binderlen:,binderlen:] )

  pae_interaction_total = ( pae_interaction1 + pae_interaction2 ) / 2

  # Calculate pae_interface which is pae_interaction of only residues within 10A of the binder
  # This is all derived from some AF2 code to calculate pAE from logits
  # This turns logits to probabilites and then gets the expected value at each position

  edges = prediction_result['distogram']['bin_edges']
  edges = np.squeeze(edges)
  probs = scipy.special.softmax(
          prediction_result['distogram']['logits'],
          axis=-1)
  probs = np.squeeze(probs)

  step = (edges[1] - edges[0])
  bin_centers = edges + step / 2

  bin_centers = np.concatenate([bin_centers, [bin_centers[-1] + step]], axis=-1)

  dgram = np.sum(probs * bin_centers, axis=-1)

  # End AF2-derived

  interface_cutoff = 15 # This can be made into a commandline argument eventually
  distance_mask = np.where(dgram < interface_cutoff, 1, 0) # mask all positions that are not within the interface cutoff
  pae_masked = np.squeeze(pae) * distance_mask
  pae_interface = ( np.mean( pae_masked[:binderlen,binderlen:] ) + np.mean( pae_masked[binderlen:,:binderlen] ) ) / 2

  if np.isnan(pae_interface): pae_interface = 25
  if pae_interface == 0: pae_interface = 25 # If there are no residues within the interface we want to set this as a high number, not 0

  time = timer() - start_time

  score_dict = {
          "plddt_total" : plddt,
          "plddt_binder" : plddt_binder,
          "plddt_target" : plddt_target,
          "pae_interaction1" : pae_interaction1,
          "pae_interaction2" : pae_interaction2,
          "pae_binder" : pae_binder,
          "pae_target" : pae_target,
          "pae_interaction" : pae_interaction_total,
          # "pae_interface" : pae_interface, # Benchmarking with this metric suggests that it's not a very good predictor of binder success
          "time" : time
  }

  write_header=False
  if not os.path.isfile(scorefilename): write_header=True
  add2scorefile(outtag, scorefilename, write_header=write_header, score_dict=score_dict)

  print(score_dict)
  print( f"Tag: {outtag} reported success in {time} seconds" )

  return score_dict

def combine_batches(tags, feature_dict_dict, initial_guess_dict):
    print( tags )
    allkeys = feature_dict_dict[tags[0]].keys()
    processed_feature_dict = {}
    for key in allkeys:
        processed_feature_dict[key] = jnp.stack([feature_dict_dict[tag][key] for tag in tags], axis=0)
    processed_initial_guess_dict = jnp.stack([initial_guess_dict[tag] for tag in tags], axis=0)
    return processed_feature_dict, processed_initial_guess_dict

def unpack_batches( tags, binderlen_dict, start, feature_dict_dict, prediction_result, sfd_out, scorefilename ):

    # First unpack the structures
    all_struct_module = prediction_result['structure_module']
    proteins = {}
    for i in range(all_struct_module['final_atom_positions'].shape[0]):
        key = tags[i]

        proteins[key] = protein.Protein(
           aatype=feature_dict_dict[key]['aatype'][0],
           atom_positions=all_struct_module['final_atom_positions'][i,...],
           atom_mask=all_struct_module['final_atom_mask'][i,...],
           residue_index=feature_dict_dict[key]['residue_index'][0] + 1,
           b_factors=np.zeros_like(all_struct_module['final_atom_mask'][i,...]) )

    # Then unpack and compute confidence metrics
    confidence_metrics = {}
    for i in range(prediction_result['predicted_lddt']['logits'].shape[0]):
        key = tags[i]

        curr_metrics = {}
        curr_metrics['distogram'] = prediction_result['distogram']
        curr_metrics['plddt'] = confidence.compute_plddt(
                prediction_result['predicted_lddt']['logits'][i,...])
        if 'predicted_aligned_error' in prediction_result:
            curr_metrics.update(confidence.compute_predicted_aligned_error(
                prediction_result['predicted_aligned_error']['logits'][i,...],
                prediction_result['predicted_aligned_error']['breaks'][i,...]))

        confidence_metrics[key] = curr_metrics

    for tag in tags:

        unrelaxed_pdb_lines = protein.to_pdb(proteins[tag])

        outtag = f'{tag}_af2pred'
        unrelaxed_pdb_path = f'temp_af2output.pdb'
        with open(unrelaxed_pdb_path, 'w') as f: f.write(unrelaxed_pdb_lines)

        score_dict = generate_scoredict( outtag, start, binderlen_dict[tag], confidence_metrics[tag], scorefilename )

        add2silent( outtag, unrelaxed_pdb_path, score_dict, binderlen_dict[tag], sfd_out )
        os.remove( unrelaxed_pdb_path )

def insert_truncations(residue_index, Ls):
    # Minkyung's chainbreak wizardry
    idx_res = residue_index
    for break_i in Ls:
        idx_res[break_i:] += 200
    residue_index = idx_res

    return residue_index

def add2silent( tag, pdb, score_dict, binderlen, sfd_out ):
    pose = pose_from_file( pdb )

    pose = insert_chainbreaks( pose, binderlen )

    struct = sfd_out.create_SilentStructOP()
    struct.fill_struct( pose, tag )

    for scorename, value in score_dict.items():
        struct.add_energy(scorename, value, 1)

    sfd_out.add_structure( struct )
    sfd_out.write_silent_struct( struct, "out.silent" )

def predict_structure(tags, feature_dict_dict, binderlen_dict, initial_guess_dict, sfd_out, scorefilename, random_seed=0):
  """Predicts structure using AlphaFold for the given sequence."""

  start = timer()
  print(f"running {model_name}")
  model_runner.params = model_params

  processed_feature_dict, processed_initial_guess_dict = combine_batches(tags, feature_dict_dict, initial_guess_dict)

  prediction_result = jax.vmap(model_runner.apply, in_axes=(None,None,0,0))(model_runner.params,
          jax.random.PRNGKey(0), processed_feature_dict, processed_initial_guess_dict)

  unpack_batches( tags, binderlen_dict, start, feature_dict_dict, prediction_result, sfd_out, scorefilename )

  print( f"{args.batch_size} predictions made in {timer() - start} seconds" )

# Mostly taken from af2 source
# Going to use to detect truncation points
def check_residue_distances(all_positions,
                             all_positions_mask,
                             max_amide_distance):
  """Checks if the distance between unmasked neighbor residues is ok."""
  breaks = []

  c_position = residue_constants.atom_order['C']
  n_position = residue_constants.atom_order['N']
  prev_is_unmasked = False
  this_c = None
  for i, (coords, mask) in enumerate(zip(all_positions, all_positions_mask)):
    this_is_unmasked = bool(mask[c_position]) and bool(mask[n_position])
    if this_is_unmasked:
      this_n = coords[n_position]
      if prev_is_unmasked:
        distance = np.linalg.norm(this_n - prev_c)
        if distance > max_amide_distance:
          breaks.append(i)
          print( f'The distance between residues {i} and {i+1} is {distance:.2f} A' +
                     f' > limit {max_amide_distance} A.' )
          print( f"I'm going to insert a chainbreak after residue {i}" )
      prev_c = coords[c_position]
    prev_is_unmasked = this_is_unmasked

  return breaks

def generate_feature_dict( pdbfile ):
  seq_list = get_seq_from_pdb(pdbfile)
  query_sequence = ''.join(seq_list)

  initial_guess = af2_all_atom_from_struct(pdbfile, seq_list, just_target=False)

  template_dict, all_atom_positions, all_atom_masks = template_from_struct(pdbfile, seq_list)

  # Gather features
  feature_dict = {
      **pipeline.make_sequence_features(sequence=query_sequence,
                                        description="none",
                                        num_res=len(query_sequence)),
      **pipeline.make_msa_features(msas=[[query_sequence]],
                                   deletion_matrices=[[[0]*len(query_sequence)]]),
      **template_dict
  }

  max_amide_distance = 3
  breaks = check_residue_distances(all_atom_positions, all_atom_masks, max_amide_distance)

  feature_dict['residue_index'] = insert_truncations(feature_dict['residue_index'], breaks)

  return feature_dict, initial_guess, len(seq_list[0])

def input_check( pdbfile, tag ):
    with open(pdbfile,'r') as f: lines = f.readlines()

    seen_indices = set()
    chain1 = True

    for line in lines:
        line = line.strip()

        if len(line) == 0: continue

        splits = line.split()

        if splits[0] == "TER":
            chain1 = False
            continue

        if not splits[0] == "ATOM": continue

        if splits[2] == 'CA':
            # Only checking residue index at CA atom
            residx = splits[5]
            if residx in seen_indices:
                sys.exit( f"\nNon-unique residue indices detected for tag: {tag}. " +
                "This will cause AF2 to yield garbage outputs. Exiting." )

            seen_indices.add(residx)

        if ( not splits[4] == "A" ) and chain1:
            sys.exit( f"\nThe first chain in the pose must be the binder and it must be chain A. " +
                    f"Tag: {tag} does not satisfy this requirement. Exiting." )

def tag_buffer2features(tag_buffer, sfd_in):

  feature_dict_dict = {}
  initial_guess_dict = {}
  binderlen_dict = {}

  for tag in tag_buffer:
    # dump pdb
    pose = Pose()
    sfd_in.get_structure(tag).fill_pose(pose)
    pose.dump_pdb(tmppdb)

    # Input Checking
    # Must ensure that:
    # - All residue indices are unique
    # - The first chain is "A"

    input_check(tmppdb, tag)

    feature_dict, initial_guess, binderlen = generate_feature_dict(tmppdb)
    feature_dict_dict[tag] = model_runner.process_features(feature_dict, random_seed=0)
    initial_guess_dict[tag] = initial_guess
    binderlen_dict[tag] = binderlen

    os.remove( tmppdb )

  return feature_dict_dict, initial_guess_dict, binderlen_dict

# Checkpointing Functions

def record_checkpoint( tag_buffer, checkpoint_filename ):
    with open( checkpoint_filename, 'a' ) as f:
        for tag in tag_buffer:
            f.write( tag )
            f.write( '\n' )

def determine_finished_structs( checkpoint_filename ):
    done_set = set()
    if not os.path.isfile( checkpoint_filename ): return done_set

    with open( checkpoint_filename, 'r' ) as f:
        for line in f:
            done_set.add( line.strip() )

    return done_set

# End Checkpointing Functions

################## Begin Main Function ##################

sfd_out = core.io.silent.SilentFileData( "out.silent", False, False, "binary", core.io.silent.SilentFileOptions())

sfd_in = rosetta.core.io.silent.SilentFileData(rosetta.core.io.silent.SilentFileOptions())
sfd_in.read_file(args.silent)

alltags = silent_tools.get_silent_index(args.silent)["tags"]
tmppdb = 'tmp.pdb'

checkpoint_filename = "check.point"
scorefilename = "out.sc"
tag_buffer = []

finished_structs = determine_finished_structs( checkpoint_filename )

for idx,tag in enumerate(alltags):

  if tag in finished_structs:
    print( f"SKIPPING {tag}, since it was already run" )
    continue

  tag_buffer.append(tag)
  if (len(tag_buffer) < args.batch_size) and not idx>=(len(alltags)-1) : continue

  feature_dict_dict, initial_guess_dict, binderlen_dict = tag_buffer2features(tag_buffer, sfd_in)
  predict_structure(tag_buffer, feature_dict_dict, binderlen_dict, initial_guess_dict, sfd_out, scorefilename)

  record_checkpoint( tag_buffer, checkpoint_filename )
  tag_buffer = []

print('done predicting')


# Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved.
#
# 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.

"""
此模块包含构造 MBQC 模型的常用类和配套的运算模拟工具。
"""

from numpy import random, pi
from networkx import Graph, spring_layout, draw_networkx
import matplotlib.pyplot as plt
from paddle import t, to_tensor, matmul, conj, real, reshape, multiply
from paddle_quantum.mbqc.utils import plus_state, cz_gate, pauli_gate
from paddle_quantum.mbqc.utils import basis, kron, div_str_to_float
from paddle_quantum.mbqc.utils import permute_to_front, permute_systems, print_progress, plot_results
from paddle_quantum.mbqc.qobject import State, Pattern
from paddle_quantum.mbqc.transpiler import transpile

__all__ = [
    "MBQC",
    "simulate_by_mbqc"
]

class MBQC:
    r"""定义基于测量的量子计算模型 ``MBQC`` 类。

    用户可以通过实例化该类来定义自己的 MBQC 模型。
    """

    def __init__(self):
        r"""MBQC 类的构造函数，用于实例化一个 ``MBQC`` 对象。
        """
        self.__graph = None  # Graph in a MBQC model
        self.__pattern = None  # Measurement pattern in a MBQC model

        self.__bg_state = State()  # Background state of computation
        self.__history = [self.__bg_state]  # History of background states
        self.__status = self.__history[-1] if self.__history != [] else None  # latest history item

        self.vertex = None  # Vertex class to maintain all the vertices
        self.__outcome = {}  # Dictionary to store all measurement outcomes
        self.max_active = 0  # Maximum number of active vertices so far

        self.__draw = False  # Switch to draw the dynamical running process
        self.__track = False  # Switch to track the running progress
        self.__pause_time = None  # Pause time for drawing
        self.__pos = None  # Position for drawing

    class Vertex:
        r"""定义维护点列表，用于实例化一个 ``Vertex`` 对象。

        将 MBQC 算法中图的节点分为三类，并进行动态维护。

        Note:
            这是内部类，用户不需要直接调用到该类。

        Attributes:
            total (list): MBQC 算法中图上的全部节点，不随运算而改变
            pending (list): 待激活的节点，随着运算的执行而逐渐减少
            active (list): 激活的节点，与当前测量步骤直接相关的节点
            measured (list): 已被测量过的节点，随着运算的执行而逐渐增加
        """

        def __init__(self, total=None, pending=None, active=None, measured=None):
            r"""``Vertex`` 类的构造函数，用于实例化一个 ``Vertex`` 对象。

            Args:
                total (list): MBQC 算法中图上的全部节点，不随运算而改变
                pending (list): 待激活的节点，随着运算的执行而逐渐减少
                active (list): 激活的节点，与当前测量步骤直接相关的节点
                measured (list): 已被测量过的节点，随着运算的执行而逐渐增加
            """
            self.total = [] if total is None else total
            self.pending = [] if pending is None else pending
            self.active = [] if active is None else active
            self.measured = [] if measured is None else measured

    def set_graph(self, graph):
        r"""设置 MBQC 模型中的图。

        该函数用于将用户自己构造的图传递给 ``MBQC`` 实例。

        Args:
            graph (list): MBQC 模型中的图，由列表 ``[V, E]`` 给出， 其中 ``V`` 为节点列表，``E`` 为边列表
        """
        vertices, edges = graph

        vertices_of_edges = set([vertex for edge in edges for vertex in list(edge)])
        assert vertices_of_edges.issubset(vertices), "edge must be between the graph vertices."

        self.__graph = Graph()
        self.__graph.add_nodes_from(vertices)
        self.__graph.add_edges_from(edges)

        self.vertex = self.Vertex(total=vertices, pending=vertices, active=[], measured=[])

    def get_graph(self):
        r"""获取图的信息。

        Returns:
            nx.Graph: 图
        """
        return self.__graph

    def set_pattern(self, pattern):
        r"""设置 MBQC 模型的测量模式。

        该函数用于将用户由电路图翻译得到或自己构造的测量模式传递给 ``MBQC`` 实例。

        Warning:
            输入的 pattern 参数是 ``Pattern`` 类型，其中命令列表为标准 ``EMC`` 命令。

        Args:
            pattern (Pattern): MBQC 算法对应的测量模式
        """
        assert isinstance(pattern, Pattern), "please input a pattern of type 'Pattern'."

        self.__pattern = pattern
        cmds = self.__pattern.commands[:]

        # Check if the pattern is a standard EMC form
        cmd_map = {"E": 1, "M": 2, "X": 3, "Z": 4, "S": 5}
        cmd_num_wild = [cmd_map[cmd.name] for cmd in cmds]
        cmd_num_standard = cmd_num_wild[:]
        cmd_num_standard.sort(reverse=False)
        assert cmd_num_wild == cmd_num_standard, "input pattern is not a standard EMC form."

        # Set graph by entanglement commands
        edges = [tuple(cmd.which_qubits) for cmd in cmds if cmd.name == "E"]
        vertices = list(set([vertex for edge in edges for vertex in list(edge)]))
        graph = [vertices, edges]
        self.set_graph(graph)

    def get_pattern(self):
        r"""获取测量模式的信息。

        Returns:
            Pattern: 测量模式
        """
        return self.__pattern

    def set_input_state(self, state=None):
        r"""设置需要替换的输入量子态。

        Warning:
            与电路模型不同，MBQC 模型通常默认初始态为加态。如果用户不调用此方法设置初始量子态，则默认为加态。
            如果用户以测量模式运行 MBQC，则此处输入量子态的系统标签会被限制为从零开始的自然数，类型为整型。

        Args:
            state (State): 需要替换的量子态，默认为加态
        """
        assert self.__graph is not None, "please set 'graph' or 'pattern' before calling 'set_input_state'."
        assert isinstance(state, State) or state is None, "please input a state of type 'State'."
        vertices = list(self.__graph.nodes)

        if state is None:
            vector = plus_state()
            system = [vertices[0]]  # Activate the first vertex, system should be a list
        else:
            vector = state.vector
            # If a pattern is set, map the input state system to the pattern's input
            if self.__pattern is not None:
                assert all(isinstance(label, int) for label in state.system), "please input system labels of type 'int'"
                assert all(label >= 0 for label in state.system), "please input system labels with non-negative values"

                system = [label for label in self.__pattern.input_ if int(div_str_to_float(label[0])) in state.system]
            else:
                system = state.system
        assert set(system).issubset(vertices), "input system labels must be a subset of graph vertices."

        self.__bg_state = State(vector, system)
        self.__history = [self.__bg_state]
        self.__status = self.__history[-1]
        self.vertex = self.Vertex(total=vertices,
                                  pending=list(set(vertices).difference(system)),
                                  active=system,
                                  measured=[])
        self.max_active = len(self.vertex.active)

    def __set_position(self, pos):
        r"""设置动态过程图绘制时节点的位置坐标。

        Note:
            这是内部方法，用户并不需要直接调用到该方法。

        Args:
            pos (dict or bool, optional): 节点坐标的字典数据或者内置的坐标选择，
                                          内置的坐标选择有：``True`` 为测量模式自带的坐标，``False`` 为 ``spring_layout`` 坐标
        """
        assert isinstance(pos, bool) or isinstance(pos, dict), "'pos' should be either bool or dict."
        if isinstance(pos, dict):
            self.__pos = pos
        elif pos:
            assert self.__pattern is not None, "'pos=True' must be chosen after a pattern is set."
            self.__pos = {v: [div_str_to_float(v[1]), - div_str_to_float(v[0])] for v in list(self.__graph.nodes)}
        else:
            self.__pos = spring_layout(self.__graph)  # Use 'spring_layout' otherwise

    def __draw_process(self, which_process, which_qubit):
        r"""根据当前节点状态绘图，用以实时展示 MBQC 模型的模拟计算过程。

        Note:
            这是内部方法，用户并不需要直接调用到该方法。

        Args:
            which_process (str): MBQC 执行的阶段，"measuring", "active" 或者 "measured"
            which_qubit (any): 当前关注的节点，可以是 ``str``, ``tuple`` 等任意数据类型，但需要和图的标签类型匹配
        """
        if self.__draw:
            assert which_process in ["measuring", "active", "measured"]
            assert which_qubit in self.vertex.total, "'which_qubit' must be in the graph."

            vertex_sets = []
            # Find where the 'which_qubit' is
            if which_qubit in self.vertex.pending:
                pending = self.vertex.pending[:]
                pending.remove(which_qubit)
                vertex_sets = [pending, self.vertex.active, [which_qubit], self.vertex.measured]
            elif which_qubit in self.vertex.active:
                active = self.vertex.active[:]
                active.remove(which_qubit)
                vertex_sets = [self.vertex.pending, active, [which_qubit], self.vertex.measured]
            elif which_qubit in self.vertex.measured:
                vertex_sets = [self.vertex.pending, self.vertex.active, [], self.vertex.measured]

            # Indentify ancilla vertices
            ancilla_qubits = []
            if self.__pattern is not None:
                for vertex in list(self.__graph.nodes):
                    row_coordinate = div_str_to_float(vertex[0])
                    col_coordinate = div_str_to_float(vertex[1])
                    # Ancilla vertices do not have integer coordinates
                    if abs(col_coordinate - int(col_coordinate)) >= 1e-15 \
                            or abs(row_coordinate - int(row_coordinate)) >= 1e-15:
                        ancilla_qubits.append(vertex)

            plt.cla()
            plt.title("MBQC Running Process", fontsize=15)
            plt.xlabel("Measuring (RED)  Active (GREEN)  Pending (BLUE)  Measured (GRAY)", fontsize=12)
            plt.grid()
            mngr = plt.get_current_fig_manager()
            mngr.window.setGeometry(500, 100, 800, 600)
            colors = ['tab:blue', 'tab:green', 'tab:red', 'tab:gray']
            for j in range(4):
                for vertex in vertex_sets[j]:
                    options = {
                        "nodelist": [vertex],
                        "node_color": colors[j],
                        "node_shape": '8' if vertex in ancilla_qubits else 'o',
                        "with_labels": False,
                        "width": 3,
                    }
                    draw_networkx(self.__graph, self.__pos, **options)
                    ax = plt.gca()
                    ax.margins(0.20)
                    plt.axis("on")
                    ax.set_axisbelow(True)
            plt.pause(self.__pause_time)

    def draw_process(self, draw=True, pos=False, pause_time=0.5):
        r"""动态过程图绘制，用以实时展示 MBQC 模型的模拟计算过程。

        Args:
            draw (bool, optional): 是否绘制动态过程图的布尔开关
            pos (bool or dict, optional): 节点坐标的字典数据或者内置的坐标选择，内置的坐标选择有：
                                            ``True`` 为测量模式自带的坐标，``False`` 为 `spring_layout` 坐标
            pause_time (float, optional): 绘制动态过程图时每次更新的停顿时间
        """
        assert self.__graph is not None, "please set 'graph' or 'pattern' before calling 'draw_process'."
        assert isinstance(draw, bool), "'draw' must be bool."
        assert isinstance(pos, bool) or isinstance(pos, dict), "'pos' should be either bool or dict."
        assert pause_time > 0, "'pause_time' must be strictly larger than 0."

        self.__draw = draw
        self.__pause_time = pause_time

        if self.__draw:
            plt.figure()
            plt.ion()
            self.__set_position(pos)

    def track_progress(self, track=True):
        r""" 显示 MBQC 模型运行进度的开关。

        Args:
            track (bool, optional): ``True`` 打开进度条显示功能， ``False`` 关闭进度条显示功能
        """
        assert isinstance(track, bool), "the parameter 'track' must be bool."
        self.__track = track

    def __apply_cz(self, which_qubits_list):
        r"""对给定的两个比特作用控制 Z 门。

        Note:
            这是内部方法，用户并不需要直接调用到该方法。

        Warning:
            作用控制 Z 门的两个比特一定是被激活的。

        Args:
            which_qubits_list (list): 作用控制 Z 门的比特对标签列表，例如 ``[(1, 2), (3, 4),...]``
        """
        for which_qubits in which_qubits_list:
            assert set(which_qubits).issubset(self.vertex.active), \
                "vertices in 'which_qubits_list' must be activated first."
            assert which_qubits[0] != which_qubits[1], \
                'the control and target qubits must not be the same.'

            # Find the control and target qubits and permute them to the front
            self.__bg_state = permute_to_front(self.__bg_state, which_qubits[0])
            self.__bg_state = permute_to_front(self.__bg_state, which_qubits[1])

            new_state = self.__bg_state
            new_state_len = new_state.length
            qua_length = int(new_state_len / 4)
            cz = cz_gate()
            # Reshape the state, apply CZ and reshape it back
            new_state.vector = reshape(matmul(cz, reshape(new_state.vector, [4, qua_length])), [new_state_len, 1])

            # Update the order of active vertices and the background state
            self.vertex.active = new_state.system
            self.__bg_state = State(new_state.vector, new_state.system)

    def __apply_pauli_gate(self, gate, which_qubit):
        r"""对给定的单比特作用 Pauli 门。

        Note:
            这是内部方法，用户并不需要直接调用到该方法。

        Args:
            gate (str): Pauli 门的索引字符，"I", "X", "Y", "Z" 分别表示对应的门，在副产品处理时用 "X" 和 "Z" 门
            which_qubit (any): 作用 Pauli 门的系统标签，
                               可以是 ``str``, ``tuple`` 等任意数据类型，但需要和 MBQC 模型中节点的标签类型匹配
        """
        new_state = permute_to_front(self.__bg_state, which_qubit)
        new_state_len = new_state.length
        half_length = int(new_state_len / 2)
        gate_mat = pauli_gate(gate)
        # Reshape the state, apply X and reshape it back
        new_state.vector = reshape(matmul(gate_mat, reshape(new_state.vector, [2, half_length])), [new_state_len, 1])
        # Update the order of active vertices and the background state
        self.vertex.active = new_state.system
        self.__bg_state = State(new_state.vector, new_state.system)

    def __create_graph_state(self, which_qubit):
        r"""以待测量的比特为输入参数，生成测量当前节点所需要的最小的量子图态。

        Note:
            这是内部方法，用户并不需要直接调用到该方法。

        Args:
            which_qubit (any): 待测量比特的系统标签。
                                可以是 ``str``, ``tuple`` 等任意数据类型，但需要和 MBQC 模型中节点的标签类型匹配
        """
        # Find the neighbors of 'which_qubit'
        which_qubit_neighbors = set(self.__graph.neighbors(which_qubit))
        # Exclude the qubits already measured
        neighbors_not_measured = which_qubit_neighbors.difference(set(self.vertex.measured))
        # Create a list of system labels that will be applied to cz gates
        cz_list = [(which_qubit, qubit) for qubit in neighbors_not_measured]
        # Get the qubits to be activated
        append_qubits = {which_qubit}.union(neighbors_not_measured).difference(set(self.vertex.active))
        # Update active and pending lists
        self.vertex.active += list(append_qubits)
        self.vertex.pending = list(set(self.vertex.pending).difference(self.vertex.active))

        # Compute the new background state vector
        new_bg_state_vector = kron([self.__bg_state.vector] + [plus_state() for _ in append_qubits])

        # Update the background state and apply cz
        self.__bg_state = State(new_bg_state_vector, self.vertex.active)
        self.__apply_cz(cz_list)
        self.__draw_process("active", which_qubit)

    def __update(self):
        r"""更新历史列表和量子态信息。
        """
        self.__history.append(self.__bg_state)
        self.__status = self.__history[-1]

    def measure(self, which_qubit, basis_list):
        r"""以待测量的比特和测量基为输入参数，对该比特进行测量。

        Note:
            这是用户在实例化 MBQC 类之后最常调用的方法之一，此处我们对单比特测量模拟进行了最大程度的优化，
            随着用户对该函数的调用，MBQC 类将自动完成激活相关节点、生成所需的图态以及对特定比特进行测量的全过程，
            并记录测量结果和对应测量后的量子态。用户每调用一次该函数，就完成一次对单比特的测量操作。

        Warning:
            当且仅当用户调用 ``measure`` 类方法时，MBQC 模型才真正进行运算。

        Args:
            which_qubit (any): 待测量量子比特的系统标签，
                                可以是 ``str``, ``tuple`` 等任意数据类型，但需要和 MBQC 模型的图上标签匹配
            basis_list (list): 测量基向量构成的列表，列表元素为 ``Tensor`` 类型的列向量

        代码示例：

        .. code-block:: python

            from paddle_quantum.mbqc.simulator import MBQC
            from paddle_quantum.mbqc.qobject import State
            from paddle_quantum.mbqc.utils import zero_state, basis

            G = [['1', '2', '3'], [('1', '2'), ('2', '3')]]
            mbqc = MBQC()
            mbqc.set_graph(G)
            state = State(zero_state(), ['1'])
            mbqc.set_input_state(state)
            mbqc.measure('1', basis('X'))
            mbqc.measure('2', basis('X'))
            print("Measurement outcomes: ", mbqc.get_classical_output())

        ::

            Measurement outcomes:  {'1': 0, '2': 1}
        """
        self.__draw_process("measuring", which_qubit)
        self.__create_graph_state(which_qubit)
        assert which_qubit in self.vertex.active, 'the qubit to be measured must be activated first.'

        new_bg_state = permute_to_front(self.__bg_state, which_qubit)
        self.vertex.active = new_bg_state.system
        half_length = int(new_bg_state.length / 2)

        eps = 10 ** (-10)
        prob = [0, 0]
        state_unnorm = [0, 0]

        # Calculate the probability and post-measurement states
        for result in [0, 1]:
            basis_dagger = t(conj(basis_list[result]))
            # Reshape the state, multiply the basis and reshape it back
            state_unnorm[result] = reshape(matmul(basis_dagger,
                                                  reshape(new_bg_state.vector, [2, half_length])), [half_length, 1])
            probability = matmul(t(conj(state_unnorm[result])), state_unnorm[result])
            is_complex128 = probability.dtype == to_tensor([], dtype='complex128').dtype
            prob[result] = real(probability) if is_complex128 else probability

        # Randomly choose a result and its corresponding post-measurement state
        if prob[0].numpy().item() < eps:
            result = 1
            post_state_vector = state_unnorm[1]
        elif prob[1].numpy().item() < eps:
            result = 0
            post_state_vector = state_unnorm[0]
        else:  # Take a random choice of outcome
            result = random.choice(2, 1, p=[prob[0].numpy().item(), prob[1].numpy().item()]).item()
            # Normalize the post-measurement state
            post_state_vector = state_unnorm[result] / prob[result].sqrt()

        # Write the measurement result into the dict
        self.__outcome.update({which_qubit: int(result)})
        # Update measured, active lists
        self.vertex.measured.append(which_qubit)
        self.max_active = max(len(self.vertex.active), self.max_active)
        self.vertex.active.remove(which_qubit)

        # Update the background state and history list
        self.__bg_state = State(post_state_vector, self.vertex.active)
        self.__update()

        self.__draw_process("measured", which_qubit)

    def sum_outcomes(self, which_qubits, start=0):
        r"""根据输入的量子系统标签，在存储测量结果的字典中找到对应的测量结果，并进行求和。

        Note:
            在进行副产品纠正操作和定义适应性测量角度时，用户可以调用该方法对特定比特的测量结果求和。

        Args:
            which_qubits (list): 需要查找测量结果并求和的比特的系统标签列表
            start (int): 对结果进行求和后需要额外相加的整数

        Returns:
            int: 指定比特的测量结果的和

        代码示例：

        .. code-block:: python

            from paddle_quantum.mbqc.simulator import MBQC
            from paddle_quantum.mbqc.qobject import State
            from paddle_quantum.mbqc.utils import zero_state, basis

            G = [['1', '2', '3'], [('1', '2'), ('2', '3')]]
            mbqc = MBQC()
            mbqc.set_graph(G)
            input_state = State(zero_state(), ['1'])
            mbqc.set_input_state(input_state)
            mbqc.measure('1', basis('X'))
            mbqc.measure('2', basis('X'))
            mbqc.measure('3', basis('X'))
            print("All measurement outcomes: ", mbqc.get_classical_output())
            print("Sum of outcomes of qubits '1' and '2': ", mbqc.sum_outcomes(['1', '2']))
            print("Sum of outcomes of qubits '1', '2' and '3' with an extra 1: ", mbqc.sum_outcomes(['1', '2', '3'], 1))

        ::

            All measurement outcomes:  {'1': 0, '2': 0, '3': 1}
            Sum of outcomes of qubits '1' and '2':  0
            Sum of outcomes of qubits '1', '2' and '3' with an extra 1:  2
        """
        assert isinstance(start, int), "'start' must be of type int."

        return sum([self.__outcome[label] for label in which_qubits], start)

    def correct_byproduct(self, gate, which_qubit, power):
        r"""对测量后的量子态进行副产品纠正。

        Note:
            这是用户在实例化 MBQC 类并完成测量后，经常需要调用的一个方法。

        Args:
            gate (str): ``'X'`` 或者 ``'Z'``，分别表示 Pauli X 或 Z 门修正
            which_qubit (any): 待操作的量子比特的系统标签，可以是 ``str``, ``tuple`` 等任意数据类型，但需要和 MBQC 中图的标签类型匹配
            power (int): 副产品纠正算符的指数

        代码示例：

            此处展示的是 MBQC 模型下实现隐形传态的一个例子。

        .. code-block:: python

            from paddle_quantum.mbqc.simulator import MBQC
            from paddle_quantum.mbqc.qobject import State
            from paddle_quantum.mbqc.utils import random_state_vector, basis, compare_by_vector

            G = [['1', '2', '3'], [('1', '2'), ('2', '3')]]
            state = State(random_state_vector(1), ['1'])
            mbqc = MBQC()
            mbqc.set_graph(G)
            mbqc.set_input_state(state)
            mbqc.measure('1', basis('X'))
            mbqc.measure('2', basis('X'))
            outcome = mbqc.get_classical_output()
            mbqc.correct_byproduct('Z', '3', outcome['1'])
            mbqc.correct_byproduct('X', '3', outcome['2'])
            state_out = mbqc.get_quantum_output()
            state_std = State(state.vector, ['3'])
            compare_by_vector(state_out, state_std)

        ::

            Norm difference of the given states is:
             0.0
            They are exactly the same states.
        """
        assert gate in ['X', 'Z'], "'gate' must be 'X' or 'Z'."
        assert isinstance(power, int), "'power' must be of type 'int'."

        if power % 2 == 1:
            self.__apply_pauli_gate(gate, which_qubit)
        self.__update()

    def __run_cmd(self, cmd):
        r"""执行测量或副产品处理命令。

        Args:
            cmd (Pattern.CommandM / Pattern.CommandX / Pattern.CommandZ): 测量或副产品处理命令
        """
        assert cmd.name in ["M", "X", "Z"], "the input 'cmd' must be CommandM, CommandX or CommandZ."
        if cmd.name == "M":  # Execute measurement commands
            signal_s = self.sum_outcomes(cmd.domain_s)
            signal_t = self.sum_outcomes(cmd.domain_t)
            # The adaptive angle is (-1)^{signal_s} * angle + {signal_t} * pi
            adaptive_angle = multiply(to_tensor([(-1) ** signal_s], dtype="float64"), cmd.angle) \
                             + to_tensor([signal_t * pi], dtype="float64")
            self.measure(cmd.which_qubit, basis(cmd.plane, adaptive_angle))
        else:  # Execute byproduct correction commands
            power = self.sum_outcomes(cmd.domain)
            self.correct_byproduct(cmd.name, cmd.which_qubit, power)

    def __run_cmd_lst(self, cmd_lst, bar_start, bar_end):
        r"""对列表执行测量或副产品处理命令。

        Args:
            cmd_lst (list): 命令列表，包含测量或副产品处理命令
            bar_start (int): 进度条的开始点
            bar_end (int): 进度条的结束点
        """
        for i in range(len(cmd_lst)):
            cmd = cmd_lst[i]
            self.__run_cmd(cmd)
            print_progress((bar_start + i + 1) / bar_end, "Pattern Running Progress", self.__track)

    def __kron_unmeasured_qubits(self):
        r"""该方法将没有被作用 CZ 纠缠的节点初始化为 |+> 态，并与当前的量子态做张量积。

        Warning:
            该方法仅在用户输入测量模式时调用，当用户输入图时，如果节点没有被激活，我们默认用户没有对该节点进行任何操作。
        """
        # Turn off the plot switch
        self.__draw = False
        # As the create_graph_state function would change the measured qubits list, we need to record it
        measured_qubits = self.vertex.measured[:]

        for qubit in list(self.__graph.nodes):
            if qubit not in self.vertex.measured:
                self.__create_graph_state(qubit)
                # Update vertices and backgrounds
                self.vertex.measured.append(qubit)
                self.max_active = max(len(self.vertex.active), self.max_active)
                self.__bg_state = State(self.__bg_state.vector, self.vertex.active)

        # Restore the measured qubits
        self.vertex.measured = measured_qubits

    def run_pattern(self):
        r"""按照设置的测量模式对 MBQC 模型进行模拟。

        Warning:
            该方法必须在 ``set_pattern`` 调用后调用。
        """
        assert self.__pattern is not None, "please use this method after calling 'set_pattern'!"

        # Execute measurement commands and correction commands
        cmd_m_lst = [cmd for cmd in self.__pattern.commands if cmd.name == "M"]
        cmd_c_lst = [cmd for cmd in self.__pattern.commands if cmd.name in ["X", "Z"]]
        bar_end = len(cmd_m_lst + cmd_c_lst)

        self.__run_cmd_lst(cmd_m_lst, 0, bar_end)
        # Activate unmeasured qubits before byproduct corrections
        self.__kron_unmeasured_qubits()
        self.__run_cmd_lst(cmd_c_lst, len(cmd_m_lst), bar_end)

        # The output state's label is messy (e.g. [(2, 0), (0, 1), (1, 3)...]),
        # so we permute the systems in order
        q_output = self.__pattern.output_[1]
        self.__bg_state = permute_systems(self.__status, q_output)
        self.__update()

    @staticmethod
    def __map_qubit_to_row(out_lst):
        r"""将输出比特的标签与行数对应起来，便于查找其对应关系。

        Returns:
            dict: 返回字典，代表行数与标签的对应关系
        """
        return {int(div_str_to_float(qubit[0])): qubit for qubit in out_lst}

    def get_classical_output(self):
        r"""获取 MBQC 模型运行后的经典输出结果。

        Returns:
            str or dict: 如果用户输入是测量模式，则返回测量输出节点得到的比特串，与原电路的测量结果相一致，没有被测量的比特位填充 "？"，如果用户输入是图，则返回所有节点的测量结果
        """
        # If the input is pattern, return the equivalent result as the circuit model
        if self.__pattern is not None:
            width = len(self.__pattern.input_)
            c_output = self.__pattern.output_[0]
            q_output = self.__pattern.output_[1]
            # Acquire the relationship between row number and corresponding output qubit label
            output_lst = c_output + q_output
            row_and_qubit = self.__map_qubit_to_row(output_lst)

            # Obtain the string, with classical outputs denoted as their measurement outcomes
            # and quantum outputs denoted as "?"
            bit_str = [str(self.__outcome[row_and_qubit[i]])
                       if row_and_qubit[i] in c_output else '?'
                       for i in range(width)]
            string = "".join(bit_str)
            return string

        # If the input is graph, return the outcome dictionary
        else:
            return self.__outcome

    def get_history(self):
        r"""获取 MBQC 计算模拟时的中间步骤信息。

        Returns:
            list: 生成图态、进行测量、纠正副产品后运算结果构成的列表
        """
        return self.__history

    def get_quantum_output(self):
        r"""获取 MBQC 模型运行后的量子态输出结果。

        Returns:
            State: MBQC 模型运行后的量子态
        """
        return self.__status

def simulate_by_mbqc(circuit, input_state=None):
    r"""使用等价的 MBQC 模型模拟量子电路。

    该函数通过将量子电路转化为等价的 MBQC 模型并运行，从而获得等价于原始量子电路的输出结果。

    Warning:
        与 ``UAnsatz`` 不同，此处输入的 ``circuit`` 参数包含了测量操作。
        另，MBQC 模型默认初始态为加态，因此，如果用户不输入参数 ``input_state`` 设置初始量子态，则默认为加态。

    Args:
        circuit (Circuit): 量子电路图
        input_state (State, optional): 量子电路的初始量子态，默认为 :math:`|+\rangle` 态

    Returns:
        tuple: 包含如下两个元素:

            - str: 经典输出
            - State: 量子输出
    """
    if input_state is not None:
        assert isinstance(input_state, State), "the 'input_state' must be of type 'State'."

    pattern = transpile(circuit)
    mbqc = MBQC()
    mbqc.set_pattern(pattern)
    mbqc.set_input_state(input_state)
    mbqc.run_pattern()
    c_output = mbqc.get_classical_output()
    q_output = mbqc.get_quantum_output()

    # Return the classical and quantum outputs
    return c_output, q_output

def __get_sample_dict(bit_num, mea_bits, samples):
    r"""根据比特数和测量比特索引的列表，统计采样结果。

    Args:
        bit_num (int): 比特数
        mea_bits (list): 测量的比特列表
        samples (list): 采样结果

    Returns:
        dict: 统计得到的采样结果
    """
    sample_dict = {}
    for i in range(2 ** len(mea_bits)):
        str_of_order = bin(i)[2:].zfill(len(mea_bits))
        bit_str = []
        idx = 0
        for j in range(bit_num):
            if j in mea_bits:
                bit_str.append(str_of_order[idx])
                idx += 1
            else:
                bit_str.append('?')
        string = "".join(bit_str)
        sample_dict[string] = 0

    # Count sampling results
    for string in list(set(samples)):
        sample_dict[string] += samples.count(string)
    return sample_dict

def sample_by_mbqc(circuit, input_state=None, plot=False, shots=1024, print_or_not=True):
    r"""将 MBQC 模型重复运行多次，获得经典结果的统计分布。

    Warning:
        与 ``UAnsatz`` 不同，此处输入的 circuit 参数包含了测量操作。
        另，MBQC 模型默认初始态为加态，因此，如果用户不输入参数 `input_state` 设置初始量子态，则默认为加态。

    Args:
        circuit (Circuit): 量子电路图
        input_state (State, optional): 量子电路的初始量子态，默认为加态
        plot (bool, optional): 绘制经典采样结果的柱状图开关，默认为关闭状态
        shots (int, optional): 采样次数，默认为 1024 次
        print_or_not (bool, optional): 是否打印采样结果和绘制采样进度，默认为开启状态

    Returns:
        dict: 经典结果构成的频率字典
        list: 经典测量结果和所有采样结果（包括经典输出和量子输出）的列表
    """
    # Initialize
    if shots == 1:
        print_or_not = False
    if print_or_not:
        print("Sampling " + str(shots) + " times." + "\nWill return the sampling results.\r\n")
    width = circuit.get_width()
    mea_bits = circuit.get_measured_qubits()

    # Sampling for "shots" times
    samples = []
    all_outputs = []
    for shot in range(shots):
        if print_or_not:
            print_progress((shot + 1) / shots, "Current Sampling Progress")
        c_output, q_output = simulate_by_mbqc(circuit, input_state)
        samples.append(c_output)
        all_outputs.append([c_output, q_output])

    sample_dict = __get_sample_dict(width, mea_bits, samples)
    if print_or_not:
        print("Sample count " + "(" + str(shots) + " shots)" + " : " + str(sample_dict))
    if plot:
        dict_lst = [sample_dict]
        bar_labels = ["MBQC sample outcomes"]
        title = 'Sampling results (MBQC)'
        xlabel = "Measurement outcomes"
        ylabel = "Distribution"
        plot_results(dict_lst, bar_labels, title, xlabel, ylabel)

    return sample_dict, all_outputs