from __future__ import division
import numpy as np
import pandas as pd
from matplotlib.colors import to_rgb

from logomaker.src.error_handling import check
from logomaker.src.matrix import ALPHABET_DICT

# Sets default color schemes specified sets of characters
CHARS_TO_COLORS_DICT = {
    tuple('ACGT'): 'classic',
    tuple('ACGU'): 'classic',
    tuple('ACDEFGHIKLMNPQRSTVWY'): 'weblogo_protein',
}

weblogo_blue = [.02, .09, .74]
weblogo_pink = [.83, .11, .75]
weblogo_green = [.13, .83, .15]
weblogo_red = [.83, .04, .08]
weblogo_black = [0, 0, 0]

# COLOR_SCHEME_DICT provides a default set of logo colorschemes
# that can be passed to the 'color_scheme' argument of Logo()
three_ones = np.ones(3)
COLOR_SCHEME_DICT = {
    'classic': {
        'G': [1, .65, 0],
        'TU': [1, 0, 0],
        'C': [0, 0, 1],
        'A': [0, .5, 0]
    },

    'grays': {
        'A': .2 * three_ones,
        'C': .4 * three_ones,
        'G': .6 * three_ones,
        'TU': .8 * three_ones
    },

    'base_pairing': {
        'TAU': [1, .55, 0],
        'GC': [0, 0, 1]
    },

    # Suggested by Ryan Z. Friedman
    # https://twitter.com/rfriedman22/status/1085722502649786368
    'colorblind_safe': {
        'A': '#009980',
        'C': '#59B3E6',
        'G': '#E69B04',
        'TU': '#1A1A1A'
    },

    # Weblogo: http://weblogo.berkeley.edu/examples.html
    # BlockLogo: http://research4.dfci.harvard.edu/cvc/blocklogo/HTML/examples.php
    'weblogo_protein': {
        'RHK' : weblogo_blue, # positive charge
        'DE' : weblogo_red, # negative charge
        'QN' : weblogo_pink, # polar uncharged long
        'GCSTY' : weblogo_green, # polar uncharged short and special cases (???)
        'ILMAFVPW' : weblogo_black # hydrophobic
    },

    # Skylign: http://skylign.org/logo/example
    'skylign_protein': {
        'F': [.16, .99, .18],
        'Y': [.04, .40, .05],
        'L': [.99, .60, .25],
        'V': [1.0, .80, .27],
        'I': [.80, .60, .24],
        'H': [.40, .02, .20],
        'W': [.42, .79, .42],
        'A': [.99, .60, .42],
        'S': [.04, .14, .98],
        'T': [.17, 1.0, 1.0],
        'M': [.80, .60, .80],
        'N': [.21, .40, .40],
        'Q': [.40, .41, .79],
        'R': [.59, .02, .04],
        'K': [.40, .20, .03],
        'E': [.79, .04, .22],
        'G': [.95, .94, .22],
        'D': [.99, .05, .11],
        'P': [.10, .61, .99],
        'C': [.09, .60, .60]
    },

    # dmslogo: https://jbloomlab.github.io/dmslogo/dmslogo.colorschemes.html
    'dmslogo_charge': {
        'A': '#000000',
        'C': '#000000',
        'D': '#0000FF',
        'E': '#0000FF',
        'F': '#000000',
        'G': '#000000',
        'H': '#FF0000',
        'I': '#000000',
        'K': '#FF0000',
        'L': '#000000',
        'M': '#000000',
        'N': '#000000',
        'P': '#000000',
        'Q': '#000000',
        'R': '#FF0000',
        'S': '#000000',
        'T': '#000000',
        'V': '#000000',
        'W': '#000000',
        'Y': '#000000'
    },

    # dmslogo: https://jbloomlab.github.io/dmslogo/dmslogo.colorschemes.html
    'dmslogo_funcgroup': {
        'A': '#f76ab4',
        'C': '#ff7f00',
        'D': '#e41a1c',
        'E': '#e41a1c',
        'F': '#84380b',
        'G': '#f76ab4',
        'H': '#3c58e5',
        'I': '#12ab0d',
        'K': '#3c58e5',
        'L': '#12ab0d',
        'M': '#12ab0d',
        'N': '#972aa8',
        'P': '#12ab0d',
        'Q': '#972aa8',
        'R': '#3c58e5',
        'S': '#ff7f00',
        'T': '#ff7f00',
        'V': '#12ab0d',
        'W': '#84380b',
        'Y': '#84380b'
    },

    'hydrophobicity': {
        'RKDENQ': [0, 0, 1],
        'SGHTAP': [0, .5, 0],
        'YVMCLFIW': [0, 0, 0]
    },

    'chemistry': {
        'GSTYC': [0, .5, 0],
        'QN': [.5, 0, .5],
        'KRH': [0, 0, 1],
        'DE': [1, 0, 0],
        'AVLIPWFM': [0, 0, 0]
    },

    'charge': {
        'KRH': [0, 0, 1],
        'DE': [1, 0, 0],
        'GSTYCQNAVLIPWFM': [.5, .5, .5]
    },

    'NajafabadiEtAl2017': {
        'DEC': [.42, .16, .42],
        'PG': [.47, .47, 0.0],
        'MIWALFV': [.13, .35, .61],
        'NTSQ': [.25, .73, .28],
        'RK': [.74, .18, .12],
        'HY': [.09, .47, .46],
    },
}

def list_color_schemes():
    """
    Provides user with a list of valid color_schemes built into Logomaker.

    returns
    -------
    colors_df: (dataframe)
        A pandas dataframe listing each color_scheme and the corresponding
        set of characters for which colors are specified.
    """

    names = list(COLOR_SCHEME_DICT.keys())
    colors_df = pd.DataFrame()
    for i, name in enumerate(names):
        color_scheme = COLOR_SCHEME_DICT[name]
        characters = list(''.join(list(color_scheme.keys())))
        characters.sort()
        colors_df.loc[i, 'color_scheme'] = name
        colors_df.loc[i, 'characters'] = ''.join(characters)

    return colors_df

def _restrict_dict(in_dict, keys_to_keep):
    """ Restricts a in_dict to keys that fall within keys_to_keep. """
    return dict([(k, v) for k, v in in_dict.items() if k in keys_to_keep])

def _expand_color_dict(color_dict):
    """ Expands the string keys in color_dict, returning new_dict that has
    the same values but whose keys are single characters. These single
    characters are both uppercase and lowercase versions of the characters
    in the color_dict keys. """
    new_dict = {}
    for key in color_dict.keys():
        value = color_dict[key]
        for char in key:
            new_dict[char.upper()] = value
            new_dict[char.lower()] = value
    return new_dict

def get_rgb(color_spec):
    """
    Safely returns an RGB np.ndarray given a valid color specification
    """

    # TODO: the following code should be reviewed for edge-cases:
    # initalizing rgb to handle referenced before assignment type error
    rgb = None

    # If color specification is a string
    if isinstance(color_spec, str):
        try:
            rgb = np.array(to_rgb(color_spec))

        # This will trigger if to_rgb does not recognize color_spec.
        # In this case, raise an error to user.
        except:
            check(False, 'invalid choice: color_spec=%s' % color_spec)

    # Otherwise, if color_specification is array-like, it should
    # specify RGB values; convert to np.ndarray
    elif isinstance(color_spec, (list, tuple, np.ndarray)):

        # color_spec must have length 3 to be RGB
        check(len(color_spec) == 3,
              'color_scheme, if array, must be of length 3.')

        # color_spec must only contain numbers between 0 and 1
        check(all(0 <= x <= 1 for x in color_spec),
              'Values of color_spec must be between 0 and 1 inclusive.')

        # Cast color_spec as RGB
        rgb = np.array(color_spec)

    # Otherwise, throw error
    else:
        check(False, 'type(color_spec) = %s is invalid.' % type(color_spec))

    # Return RGB as an np.ndarray
    return rgb

def get_color_dict(color_scheme, chars):
    """
    Return a color_dict constructed from a user-specified color_scheme and
    a list of characters
    """

    # Check that chars is a list
    check(isinstance(chars, (str, list, tuple, np.ndarray)),
          "chars must be a str or be array-like")

    # Check that chars has length of at least 1
    check(len(chars) >= 1, 'chars must have length >= 1')

    # Sort characters
    chars = list(chars)
    chars.sort()

    # Check that all entries in chars are strings of length 1
    for i, c in enumerate(chars):
        c = str(c) # convert from unicode to string to work with python 2
        check(isinstance(c, str) and len(c)==1,
              'entry number %d in chars is %s; ' % (i, repr(c)) +
              'must instead be a single character')

    # if color_scheme is None, choose default based on chars
    if color_scheme is None:
        key = tuple(chars)
        color_scheme = CHARS_TO_COLORS_DICT.get(key, 'gray')
        color_dict = get_color_dict(color_scheme, chars)

    # otherwise, if color_scheme is a dictionary
    elif isinstance(color_scheme, dict):

        # make sure all the keys are strings
        for key in color_scheme.keys():
            check(isinstance(key, str),
                  'color_scheme dict contains a key (%s) ' % repr(key) +
                  'that is not of type str.')

        # expand the dictionary
        color_dict = _expand_color_dict(color_scheme)

        # set all values to rgb
        for key in color_dict.keys():
            color_dict[key] = to_rgb(color_dict[key])

    # otherwise, if color_scheme is a string, it must either be a valid key in
    # COLOR_SCHEME_DICT or a named matplotlib color
    elif isinstance(color_scheme, str):

        # If a valid key, get the color scheme dict and expand
        if color_scheme in COLOR_SCHEME_DICT.keys():
            tmp_dict = COLOR_SCHEME_DICT[color_scheme]
            color_dict = _expand_color_dict(tmp_dict)

            # Force each color to rgb
            for c in color_dict.keys():
                color = color_dict[c]
                rgb = to_rgb(color)
                color_dict[c] = np.array(rgb)

        # Otherwise, try to convert color_scheme to RGB value, then create
        # color_dict using keys from chars and setting all values to RGB value.
        else:
            try:
                rgb = to_rgb(color_scheme)
                color_dict = dict([(c, rgb) for c in chars])

            # This will trigger if to_rgb does not recognize color_scheme.
            # In this case, raise an error to user.
            except:
                check(False, 'invalid choice: color_scheme=%s' % color_scheme)

    # Otherwise, if color_scheme is array-like, it should be an RGB value
    elif isinstance(color_scheme, (list, tuple, np.ndarray)):

        # color_scheme must have length 3 to be RGB
        check(len(color_scheme) == 3,
              'color_scheme, if array, must be of length 3.')

        # Cast color_scheme as RGB
        rgb = np.ndarray(color_scheme)

        # Construct color_dict with rgb as value for every character in chars
        color_dict = dict([(c, rgb) for c in chars])

    # Otherwise, raise error
    else:
        check(False,
              'Error: color_scheme has invalid type %s'
              % type(color_scheme))

    # If all the characters in chars are not also within the keys of color_dict,
    # add them with color 'black'
    if not set(chars) <= set(color_dict.keys()):
        for c in chars:
            if not c in color_dict:
                print(" Warning: Character '%s' is not in color_dict. " % c +
                      "Using black.")
                color_dict[c] = to_rgb('black')

    return color_dict


"""
django-sage-painless - UWSGI Config Generator

:author: Mehran Rahmanzadeh (mrhnz13@gmail.com)
"""
import os
import time

from django.conf import settings

# Base
from sage_painless.services.abstract import AbstractUWSGIGenerator

# Helpers
from sage_painless.utils.git_service import GitSupport
from sage_painless.utils.jinja_service import JinjaHandler
from sage_painless.utils.json_service import JsonHandler
from sage_painless.utils.timing_service import TimingService

from sage_painless import templates

class UwsgiGenerator(AbstractUWSGIGenerator, JinjaHandler, JsonHandler, TimingService, GitSupport):
    """generate uwsgi config"""
    UWSGI_TEMPLATE = 'uwsgi.jinja'

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def generate(self, diagram_path, git_support=False):
        """generate uwsgi.ini
        template:
            sage_painless/templates/uwsgi.jinja
        """
        start_time = time.time()

        diagram = self.load_json(diagram_path)

        config = self.extract_uwsgi_config(diagram)  # get uwsgi config from diagram

        if git_support:
            self.init_repo(settings.BASE_DIR)

        # generate uwsgi.ini
        self.stream_to_template(
            output_path=f'{settings.BASE_DIR}/uwsgi.ini',
            template_path=os.path.abspath(templates.__file__).replace('__init__.py', self.UWSGI_TEMPLATE),
            data={
                'config': config
            }
        )

        if git_support:
            self.commit_file(
                f'{settings.BASE_DIR}/uwsgi.ini',
                f'deploy (uwsgi): Create uwsgi config file'
            )

        end_time = time.time()
        return True, 'uwsgi config generated ({:.3f} ms)'.format(self.calculate_execute_time(start_time, end_time))


"""
This file contains methods to evaluate the trained model
"""
from read_args import read_command

import numpy as np
import tensorflow as tf
from sklearn.metrics import confusion_matrix, recall_score, precision_score, accuracy_score

import model, read_input

tf.app.flags.DEFINE_integer('e', 1, 'eval measure to use')
tf.app.flags.DEFINE_string('c', '', 'directory to store checkpoints')
tf.app.flags.DEFINE_boolean('no-train', True, 'whether to train the inside weights')

def eval_measure_1(sess, logits, labels, total_examples):
	true_count = 0

	top_k_op = tf.nn.in_top_k(logits, labels, 1)

	iter_n = 0
	while True:
		try:
			print(iter_n)
			trues = sess.run([top_k_op])
			true_count += np.sum(trues)
			iter_n += 1
		except tf.errors.OutOfRangeError:
			break

	accuracy = true_count / total_examples
	print(accuracy)

def eval_measure_2(sess, logits, labels):

	iter_n = 0
	predictions_all = np.array([])
	labels_all = np.array([])
	while True:
		try:
			print(iter_n)
			predictions, labels_r = sess.run([tf.argmax(logits, axis=1), labels])

			predictions_all = np.append(predictions_all, predictions)
			labels_all = np.append(labels_all, labels_r)

			iter_n += 1
		except tf.errors.OutOfRangeError:
			break

	print(labels_all.shape, predictions_all.shape)
	conf_matrix = confusion_matrix(labels_all, predictions_all)
	print(conf_matrix)

	precision_macro = precision_score(labels_all, predictions_all, average='macro')
	precision_weighted = precision_score(labels_all, predictions_all, average='weighted')
	print(precision_macro, precision_weighted)

	recall_macro = recall_score(labels_all, predictions_all, average='macro')
	recall_weighted = recall_score(labels_all, predictions_all, average='weighted')
	print(recall_macro, recall_weighted)

	accuracy = accuracy_score(labels_all, predictions_all)
	print(accuracy)

def eval_once(saver, eval_measure, logits, labels, args, checkpoint_dir):

	with tf.Session() as sess:
		ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
		if ckpt and ckpt.model_checkpoint_path:
			saver.restore(sess, ckpt.model_checkpoint_path)
		else:
			print('No checkpoint file found')
			return

		eval_measure(sess, logits, labels, *args)

def eval(args):
	with tf.Graph().as_default():

		images, labels = args['inputs'](train=False)

		# logits = args['inference'](images, args['n_classes'])
		logits = model.inference(images, args['n_classes'], args['relu'],
								 args['lrn'], args['pool_type'], args['trainable'])

		# top_k_op = tf.nn.in_top_k(logits, labels, 1)

		saver = tf.train.Saver()

		eval_once(saver, args['eval_measure'], logits, labels,
				  args['eval_measure_args'], args['checkpoint_dir'])

def mnist_eval(args):

	args['inputs'] = read_input.mnist_input
	args['n_classes'] = 10

	if args['eval_measure_id'] == 1:
		args['eval_measure'] = eval_measure_1
		args['eval_measure_args'] = (10000,) # examples per epoch
	else:
		args['eval_measure'] = eval_measure_2
		args['eval_measure_args'] = tuple()

	eval(args)

def caltech_101_eval(args):
	args['inputs'] = read_input.caltech_101_input
	args['n_classes'] = 102

	if args['eval_measure_id'] == 1:
		args['eval_measure'] = eval_measure_1
		args['eval_measure_args'] = (30*102,) # examples per epoch
	else:
		args['eval_measure'] = eval_measure_2
		args['eval_measure_args'] = tuple()

	eval(args)

if __name__ == '__main__':
	args = read_command()

	if args['dataset'] == 'mnist':
		mnist_eval(args)
	elif args['dataset'] == 'caltech_101':
		caltech_101_eval(args)


import interfacesDist as interfaces
import numpy as np
import time
from sixDir_UI import UI
import pygame as pg

interface = interfaces.DiscreteActionsRobot()
interface.mode = 3
interface.angle = 0
interface.debugLines = 1
interface.open()
robot_open = 1
target_pos = np.array([0.0, 300.0])

logCursor = 0
updateRate = 0.125
pg.init()
pg.display.set_mode((500,500))
pg.display.set_caption("Control Interface")
runUI = UI(logCursor, updateRate)

y = -0.1
target_pos = np.array([[0.25, y, -0.2], [0.05, y, -0.2] , [-0.15,y, -.2], [-0.35,y, -.2]])
interface.updateRefresh(0.125);

i = 0
runUI.mode = interface.mode

while True:

	time.sleep(10.0)

	while i < 4:
		current_target = target_pos[i]
		print(current_target)
		i = i+1

		interface.create_target3D_Dist(current_target,1)

		time.sleep(3.0)
		interface.reset()

	# print(current_target)
# 	interface.reset()

# 	# for i in range(6):
# 	# 	interface.create_target3D(target_pos[i])

# 	pos = np.array([0., 0., 0.])
# 	pos[0] = interface.robotenv.pos[0] - interface.robotenv.center[0]
# 	pos[1] = interface.robotenv.pos[1] - interface.robotenv.center[1]
# 	pos[2] = interface.robotenv.pos[2]
# 	dist = np.linalg.norm(pos - current_target)
# 	print("B", current_target)

# 	if interface.mode == 1:
# 		notdone = 1
# 		while notdone == 1:
# 			time.sleep(0.125)

# 			runUI.update()
# 			key = runUI.state
# 			interface.update_joystick(key)
# 			interface.render()

# 			if interface.target == 1:

# 				if interface.robotenv.fing < .01:
# 					notdone = 0
# 			if interface.target == 3:
# 				if interface.robotenv.fing > 1.34:
# 					notdone = 0
# 			if interface.target == 2:
# 				if interface.robotenv.pos[2] > 0.4:
# 					notdone = 0
# 			if interface.target == 4:
# 				if interface.robotenv.pos[2] < 0.05:
# 					notdone = 0
# 	else:
# 		while dist > 0.04:
# 			pos[0] = interface.robotenv.pos[0] - interface.robotenv.center[0]
# 			pos[1] = interface.robotenv.pos[1] - interface.robotenv.center[1]
# 			pos[2] = interface.robotenv.pos[2]
# 			# print(pos)
# 			print(current_target)
# 			dist = np.linalg.norm(pos - current_target)
# 			# print(dist)
# 			time.sleep(0.125)

# 			runUI.update()
# 			key = runUI.state
# 			# print(key)
# 			interface.update_joystick(key)
# 			interface.render()
# 	time.sleep(1.0)

	# if command == 0:
	# 	if val1 == 0:		# Open Robot
	# 		if robot_open == 0:
	# 			interface.open()
	# 			robot_open = 1
	# 	if val1 == 1:		# Reset
	# 		interface.reset()
	# 	if val1 == 2:		# Change hold time on debug lines
	# 		updateRate = 1.0 /val2;
	# 		interface.updateRefresh(updateRate);
	# if command == 1:	# Set Target
	# 	# interface.reset()
	# 	target_pos[0] = (val1 - 128) / 100
	# 	target_pos[1] = (val2 - 128) / 100
	# 	interface.create_target(target_pos)
	# 	print(val1, val2)
	# 	print(target_pos)
	# if command == 2:	# Set Dirr
	# 	key = val1


import os
import sys
import json
import logging as logger
import boto3
from botocore.config import Config
from botocore.exceptions import ClientError
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes

import time

logger.getLogger().setLevel(logger.INFO)

def get_aws_client(name):
    return boto3.client(
        name,
        config=Config(retries={"max_attempts": 10, "mode": "standard"}),
    )

def create_resources(
    thing_name: str,
    iot_policy: str,
    iot_policy_name: str,
    stack_name: str,
    encryption_algo: str,
):
    """Create AWS IoT thing, certificate and attach certificate with policy and thing.
    Returns the Arns for the values and a Parameter Store Arn for the private key
    """
    c_iot = get_aws_client("iot")
    c_ssm = get_aws_client("ssm")

    result = {}

    # Create thing
    try:
        response = c_iot.create_thing(thingName=thing_name)
        result["ThingArn"] = response["thingArn"]
        result["ThingName"] = response["thingName"]
    except ClientError as e:
        logger.error(f"Error creating thing {thing_name}, {e}")
        sys.exit(1)

    # Create certificate and private key
    if encryption_algo == "ECC":
        key = ec.generate_private_key(curve=ec.SECP256R1(), backend=default_backend())
    elif encryption_algo == "RSA":
        key = rsa.generate_private_key(
            public_exponent=65537, key_size=4096, backend=default_backend()
        )
    else:
        logger.error(
            f"Should not get here. Encryption algorithm of 'ECC' or 'RSA' expected, received {encryption_algo}. Exiting"
        )
        sys.exit(1)

    private_key = key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption(),
    ).decode("utf-8")

    # Generate a CSR and set subject (CN=dispenserId)
    csr = (
        x509.CertificateSigningRequestBuilder()
        .subject_name(
            x509.Name(
                [
                    # Provide various details about who we are.
                    x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
                    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "CO"),
                    x509.NameAttribute(NameOID.LOCALITY_NAME, "Denver"),
                    x509.NameAttribute(
                        NameOID.ORGANIZATION_NAME, "Greengrass Accelerator Testing"
                    ),
                    x509.NameAttribute(NameOID.COMMON_NAME, thing_name),
                ]
            )
        )
        .sign(key, hashes.SHA256(), default_backend())
    )
    try:
        response = c_iot.create_certificate_from_csr(
            certificateSigningRequest=str(
                csr.public_bytes(serialization.Encoding.PEM), "utf-8"
            ),
            setAsActive=True,
        )
        certificate_pem = response["certificatePem"]
        result["CertificateArn"] = response["certificateArn"]
    except ClientError as e:
        logger.error(f"Error creating certificate, {e}")
        sys.exit(1)

    # policy
    try:
        response = c_iot.create_policy(
            policyName=iot_policy_name, policyDocument=iot_policy
        )
        result["IotPolicyArn"] = response["policyArn"]
    except ClientError as e:
        logger.error(f"Error creating policy {iot_policy_name}, {e}")
        sys.exit(1)

    # attach cert-pol
    try:
        c_iot.attach_policy(policyName=iot_policy_name, target=result["CertificateArn"])
    except ClientError as e:
        logger.error(
            f"Error attaching certificate {result['CertificateArn']} to policy {iot_policy_name}, {e}"
        )
        sys.exit(1)

    # attach cert-thing
    try:
        c_iot.attach_thing_principal(
            thingName=thing_name,
            principal=result["CertificateArn"],
        )
    except ClientError as e:
        logger.error(
            f"Error attaching certificate {result['CertificateArn']} to thing {thing_name}, {e}"
        )
        sys.exit(1)

    # store certificate and private key in SSM param store
    try:
        parameter_private_key = f"/{stack_name}/{thing_name}/private_key"
        parameter_certificate_pem = f"/{stack_name}/{thing_name}/certificate_pem"
        # private key
        response = c_ssm.put_parameter(
            Name=parameter_private_key,
            Description=f"Certificate private key for IoT thing {thing_name}",
            Value=private_key,
            Type="SecureString",
            Tier="Advanced",
        )
        result["PrivateKeySecretParameter"] = parameter_private_key
        # certificate pem
        response = c_ssm.put_parameter(
            Name=parameter_certificate_pem,
            Description=f"Certificate PEM for IoT thing {thing_name}",
            Value=certificate_pem,
            Type="String",
            Tier="Advanced",
        )
        result["CertificatePemParameter"] = parameter_certificate_pem
    except ClientError as e:
        logger.error(f"Error creating secure string parameters, {e}")
        sys.exit(1)

    # Additional data - these calls and responses are used in other constructs or external applications

    # Get the IoT-Data endpoint
    try:
        response = c_iot.describe_endpoint(endpointType="iot:Data-ATS")
        result["DataAtsEndpointAddress"] = response["endpointAddress"]
    except ClientError as e:
        logger.error(f"Could not obtain iot:Data-ATS endpoint, {e}")
        result["DataAtsEndpointAddress"] = "stack_error: see log files"

    # Get the Credential Provider endpoint
    try:
        response = c_iot.describe_endpoint(endpointType="iot:CredentialProvider")
        result["CredentialProviderEndpointAddress"] = response["endpointAddress"]
    except ClientError as e:
        logger.error(f"Could not obtain iot:CredentialProvider endpoint, {e}")
        result["CredentialProviderEndpointAddress"] = "stack_error: see log files"

    return result

def delete_resources(
    thing_name: str, certificate_arn: str, iot_policy_name: str, stack_name: str
):
    """Delete thing, certificate, and policy in reverse order. Check for modifications
    since create (policy versions, etc)"""

    c_iot = get_aws_client("iot")
    c_ssm = get_aws_client("ssm")

    # delete ssm param store
    try:
        parameter_private_key = f"/{stack_name}/{thing_name}/private_key"
        parameter_certificate_pem = f"/{stack_name}/{thing_name}/certificate_pem"
        c_ssm.delete_parameters(
            Names=[parameter_private_key, parameter_certificate_pem]
        )
    except ClientError as e:
        logger.error(f"Unable to delete parameter store values, {e}")

    # delete policy (prune versions, detach from targets)
    # delete all non active policy versions
    try:
        response = c_iot.list_policy_versions(policyName=iot_policy_name)
        for version in response["policyVersions"]:
            if not version["isDefaultVersion"]:
                c_iot.delete_policy_version(
                    policyName=iot_policy_name, policyVersionId=version["versionId"]
                )
    except ClientError as e:
        logger.error(
            f"Unable to delete policy versions for policy {iot_policy_name}, {e}"
        )
    # Detach any principals
    try:
        response = c_iot.list_targets_for_policy(policyName=iot_policy_name)
        for target in response["targets"]:
            c_iot.detach_policy(policyName=iot_policy_name, target=target)
    except ClientError as e:
        logger.error(f"Unable to detach targets from policy {iot_policy_name}, {e}")
    # delete policy
    try:
        c_iot.delete_policy(policyName=iot_policy_name)
    except ClientError as e:
        logger.error(f"Unable to delete policy {iot_policy_name}, {e}")

    # delete cert
    # detach all policies and things from cert
    try:
        response = c_iot.list_principal_things(principal=certificate_arn)
        for thing in response["things"]:
            c_iot.detach_thing_principal(thingName=thing, principal=certificate_arn)
        response = c_iot.list_attached_policies(target=certificate_arn)
        for policy in response["policies"]:
            c_iot.detach_policy(policyName=policy["policyName"], target=certificate_arn)
    except ClientError as e:
        logger.error(
            f"Unable to list or detach things or policies from certificate {certificate_arn}, {e}"
        )
    try:
        c_iot.update_certificate(
            certificateId=certificate_arn.split("/")[-1], newStatus="REVOKED"
        )
        c_iot.delete_certificate(certificateId=certificate_arn.split("/")[-1])
    except ClientError as e:
        logger.error(f"Unable to delete certificate {certificate_arn}, {e}")

    # delete thing
    # Check and detach principals attached to thing
    try:
        response = c_iot.list_thing_principals(thingName=thing_name)
        for principal in response["principals"]:
            c_iot.detach_thing_principal(thingName=thing_name, principal=principal)
    except ClientError as e:
        logger.error(f"Unable to list or detach principals from {thing_name}, {e}")
    try:
        c_iot.delete_thing(thingName=thing_name)
    except ClientError as e:
        logger.error(f"Error calling iot.delete_thing() for thing: {thing_name}, {e}")

def handler(event, context):
    logger.info("Received event: %s", json.dumps(event, indent=2))
    logger.info("Environment: %s", dict(os.environ))
    props = event["ResourceProperties"]
    physical_resource_id = ""
    try:
        # Check if this is a Create and we're failing Creates
        if event["RequestType"] == "Create" and event["ResourceProperties"].get(
            "FailCreate", False
        ):
            raise RuntimeError("Create failure requested, logging")
        elif event["RequestType"] == "Create":
            logger.info("Request CREATE")
            resp = create_resources(
                thing_name=props["ThingName"],
                iot_policy=props["IotPolicy"],
                iot_policy_name=props["IoTPolicyName"],
                encryption_algo=props["EncryptionAlgorithm"],
                stack_name=props["StackName"],
            )

            # set response data (PascalCase key)
            response_data = {
                "ThingArn": resp["ThingArn"],
                "CertificateArn": resp["CertificateArn"],
                "IotPolicyArn": resp["IotPolicyArn"],
                "PrivateKeySecretParameter": resp["PrivateKeySecretParameter"],
                "CertificatePemParameter": resp["CertificatePemParameter"],
                "DataAtsEndpointAddress": resp["DataAtsEndpointAddress"],
                "CredentialProviderEndpointAddress": resp[
                    "CredentialProviderEndpointAddress"
                ],
            }
            # Certificate arn is resource id since it is created with an unknown
            # value, which is needed for the delete operation
            physical_resource_id = response_data["CertificateArn"]
        elif event["RequestType"] == "Update":
            logger.info("Request UPDATE")
            response_data = {}
        elif event["RequestType"] == "Delete":
            logger.info("Request DELETE")
            resp = delete_resources(
                thing_name=props["ThingName"],
                # resource id contains the cert arn
                certificate_arn=event["PhysicalResourceId"],
                iot_policy_name=props["IoTPolicyName"],
                stack_name=props["StackName"],
            )
            response_data = {}
            physical_resource_id = event["PhysicalResourceId"]
        else:
            logger.info("Should not get here in normal cases - could be REPLACE")
        output = {"PhysicalResourceId": physical_resource_id, "Data": response_data}
        logger.info("Output from Lambda: %s", json.dumps(output, indent=2))
        return output
    except Exception as e:
        logger.exception(e)
        sys.exit(1)


# Dependencies
import networkx as nx
import matplotlib.pyplot as plt
from operator import itemgetter

# Number of nodes and m is the connection with the last node (in the center)
n = 5
m = 4

# A graph of n nodes is grown by attaching new nodes each with m
# edges that are preferentially attached to existing nodes with high degree.
# the node with the highest degree is generally the one with the last label.
G3 = nx.generators.barabasi_albert_graph(n, m)

# Finds the node basically with the highest degree associated, and makes it the
# center,typically it is the last node.
node_and_degree = G3.degree()

# Get the node with the most degree (the label), and the degree associated with
# it
(largest_hub, degree) = sorted(node_and_degree.items(), key=itemgetter(1))[-1]

# Make an ego graph with the node with the highest degree right in the center
hub_ego = nx.ego_graph(G3, largest_hub)

# This is basically used to layout the graph in the image, like for example
# at what position to center it, and which nodes to keep fixed etc.
pos = nx.spring_layout(hub_ego)

# Set the labels for each of the nodes which will be show in the graph
labels = {0: 'harsh', 1: 'deep', 2: 'singh', 3: 'chauhan', 4: 'gino'}

# Draw the other nodes in blue color with the size associated and the labels
# and the position associated with them
nx.draw(hub_ego, pos, node_color='b', node_size=50, labels=labels, with_labels=True)

# Draw the network's nodes with the nodelist (node at the center with a much more larger size)
# and with a different color and position (which we get from spring layout)
nx.draw_networkx_nodes(hub_ego, pos, nodelist=[largest_hub], node_size=300, node_color='r')

# Save the image
plt.savefig('ego_graph.png')

# Show the image
plt.show()


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations
import django.utils.timezone
import django.db.models.deletion
from django.conf import settings
import djangobmf.numbering.validators

class Migration(migrations.Migration):

    dependencies = [
        ('contenttypes', '__first__'),
        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
    ]

    operations = [
        migrations.CreateModel(
            name='Activity',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('topic', models.CharField(max_length=100, null=True, verbose_name='Topic', blank=True)),
                ('text', models.TextField(null=True, verbose_name='Text', blank=True)),
                ('action', models.PositiveSmallIntegerField(default=1, verbose_name='Action', null=True, editable=False, choices=[(1, 'Comment'), (2, 'Created'), (3, 'Updated'), (4, 'Workflow'), (5, 'File')])),
                ('template', models.CharField(verbose_name='Template', max_length=100, null=True, editable=False)),
                ('parent_id', models.PositiveIntegerField()),
                ('modified', models.DateTimeField(auto_now=True, verbose_name='Modified')),
                ('parent_ct', models.ForeignKey(to='contenttypes.ContentType', related_name='bmf_history_parent')),
                ('user', models.ForeignKey(blank=True, to=settings.AUTH_USER_MODEL, null=True)),
            ],
            options={
                'ordering': ('-modified',),
                'get_latest_by': 'modified',
                'verbose_name': 'History',
                'verbose_name_plural': 'History',
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Configuration',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('app_label', models.CharField(verbose_name='Application', max_length=100, null=True, editable=False)),
                ('field_name', models.CharField(verbose_name='Fieldname', max_length=100, null=True, editable=False)),
                ('value', models.TextField(null=True, verbose_name='Value')),
            ],
            options={
                'default_permissions': ('change',),
                'verbose_name': 'Configuration',
                'verbose_name_plural': 'Configurations',
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Dashboard',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('name', models.CharField(max_length=100, null=True, verbose_name='Name')),
                ('user', models.ForeignKey(blank=True, to=settings.AUTH_USER_MODEL, null=True)),
            ],
            options={
                'ordering': ('name', 'id'),
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Notification',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('obj_id', models.PositiveIntegerField()),
                ('unread', models.NullBooleanField(default=True)),
                ('created', models.DateTimeField(auto_now_add=True, verbose_name='Created')),
                ('changed', models.DateTimeField(auto_now=True, verbose_name='Changed')),
                ('activity', models.ForeignKey(to='djangobmf.Activity', null=True)),
                ('created_by', models.ForeignKey(on_delete=django.db.models.deletion.SET_NULL, blank=True, editable=False, to=settings.AUTH_USER_MODEL, null=True, related_name='+')),
                ('obj_ct', models.ForeignKey(to='contenttypes.ContentType')),
                ('user', models.ForeignKey(to=settings.AUTH_USER_MODEL, null=True)),
            ],
            options={
                'ordering': ('-created',),
                'default_permissions': (),
                'get_latest_by': 'modified',
                'verbose_name': 'Notification',
                'verbose_name_plural': 'Notifications',
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='NumberCycle',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('name_template', models.CharField(max_length=64, null=True, verbose_name='Template', validators=[djangobmf.numbering.validators.template_name_validator])),
                ('counter_start', models.PositiveIntegerField(default=1, null=True)),
                ('current_period', models.DateField(default=django.utils.timezone.now, null=True)),
                ('ct', models.OneToOneField(null=True, editable=False, to='contenttypes.ContentType')),
            ],
            options={
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Report',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('reporttype', models.CharField(max_length=20, verbose_name='Reporttype')),
                ('mimetype', models.CharField(default='pdf', verbose_name='Mimetype', max_length=20, editable=False)),
                ('options', models.TextField(help_text='Options for the renderer. Empty this field to get all available options with default values', verbose_name='Options', blank=True)),
                ('modified', models.DateTimeField(auto_now=True, verbose_name='Modified')),
                ('contenttype', models.ForeignKey(blank=True, to='contenttypes.ContentType', help_text='Connect a Report to an BMF-Model', null=True)),
            ],
            options={
                'get_latest_by': 'modified',
                'verbose_name': 'Report',
                'verbose_name_plural': 'Reports',
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='View',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('category', models.CharField(max_length=100, null=True, blank=True)),
                ('name', models.CharField(max_length=100, null=True)),
                ('url', models.CharField(max_length=80, null=True)),
                ('kwargs', models.CharField(max_length=255, null=True, blank=True)),
                ('search', models.CharField(max_length=255, null=True, blank=True)),
                ('dashboard', models.ForeignKey(to='djangobmf.Dashboard', null=True)),
            ],
            options={
                'ordering': ('name',),
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Watch',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('watch_id', models.PositiveIntegerField()),
                ('active', models.BooleanField(default=False, verbose_name='Active', db_index=True, editable=False)),
                ('new_entry', models.BooleanField(default=False, db_index=True, verbose_name='New entry')),
                ('comment', models.BooleanField(default=False, db_index=True, verbose_name='Comment written')),
                ('file', models.BooleanField(default=False, db_index=True, verbose_name='File added')),
                ('changed', models.BooleanField(default=False, db_index=True, verbose_name='Object changed')),
                ('workflow', models.BooleanField(default=False, db_index=True, verbose_name='Workflowstate changed')),
                ('modified', models.DateTimeField(auto_now=True, verbose_name='Modified')),
                ('user', models.ForeignKey(blank=True, to=settings.AUTH_USER_MODEL, null=True)),
                ('watch_ct', models.ForeignKey(to='contenttypes.ContentType')),
            ],
            options={
                'get_latest_by': 'modified',
                'ordering': ('-modified',),
                'verbose_name_plural': 'Watched activities',
                'default_permissions': (),
                'verbose_name': 'Watched activity',
            },
            bases=(models.Model,),
        ),
        migrations.AlterUniqueTogether(
            name='watch',
            unique_together=set([('user', 'watch_ct', 'watch_id')]),
        ),
    ]


# pylint: disable=invalid-name

import sys
import os
from functools import partial
import numpy as np
import opt_einsum as oem
import pytest
from pytest_lazyfixture import lazy_fixture as lf

# see https://stackoverflow.com/questions/56307329/how-can-i-parametrize-tests-to-run-with-different-fixtures-in-pytest

thisfile = os.path.abspath(__file__)
modulepath = os.path.dirname(os.path.dirname(thisfile))

sys.path.insert(0, modulepath)
import tensorcircuit as tc

@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")])
def test_wavefunction(backend):
    qc = tc.Circuit(2)
    qc.unitary(
        0,
        1,
        unitary=tc.gates.Gate(np.arange(16).reshape(2, 2, 2, 2).astype(np.complex64)),
    )
    assert np.real(qc.wavefunction()[2]) == 8
    qc = tc.Circuit(2)
    qc.unitary(
        1,
        0,
        unitary=tc.gates.Gate(np.arange(16).reshape(2, 2, 2, 2).astype(np.complex64)),
    )
    qc.wavefunction()
    assert np.real(qc.wavefunction()[2]) == 4
    qc = tc.Circuit(2)
    qc.unitary(
        0, unitary=tc.gates.Gate(np.arange(4).reshape(2, 2).astype(np.complex64))
    )
    qc.wavefunction()
    assert np.real(qc.wavefunction()[2]) == 2

@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")])
def test_basics(backend):
    c = tc.Circuit(2)
    c.x(0)
    np.testing.assert_allclose(tc.backend.numpy(c.amplitude("10")), np.array(1.0))
    c.CNOT(0, 1)
    np.testing.assert_allclose(tc.backend.numpy(c.amplitude("11")), np.array(1.0))

@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")])
def test_measure(backend):
    c = tc.Circuit(3)
    c.H(0)
    c.h(1)
    c.toffoli(0, 1, 2)
    assert c.measure(2)[0] in [0, 1]

def test_gates_in_circuit():
    c = tc.Circuit(2, inputs=np.eye(2**2))
    c.iswap(0, 1)
    ans = tc.gates.iswap_gate().tensor.reshape([4, 4])
    np.testing.assert_allclose(c.state().reshape([4, 4]), ans, atol=1e-5)

def test_control_vgate():
    c = tc.Circuit(2)
    c.x(1)
    c.crx(1, 0, theta=0.3)
    np.testing.assert_allclose(
        c.expectation([tc.gates._z_matrix, 0]), 0.95533645, atol=1e-5
    )

def test_adjoint_gate_circuit():
    c = tc.Circuit(1)
    c.X(0)
    c.SD(0)
    np.testing.assert_allclose(c.state(), np.array([0.0, -1.0j]))

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_jittable_measure(backend):
    @partial(tc.backend.jit, static_argnums=(2, 3))
    def f(param, key, n=6, nlayers=3):
        if key is not None:
            tc.backend.set_random_state(key)
        c = tc.Circuit(n)
        for i in range(n):
            c.H(i)
        for j in range(nlayers):
            for i in range(n - 1):
                c.exp1(i, i + 1, theta=param[2 * j, i], hamiltonian=tc.gates._zz_matrix)
            for i in range(n):
                c.rx(i, theta=param[2 * j + 1, i])
        return c.measure_jit(0, 1, 2, with_prob=True)

    if tc.backend.name == "tensorflow":
        import tensorflow as tf

        print(f(tc.backend.ones([6, 6]), None))
        print(f(tc.backend.ones([6, 6]), None))
        print(f(tc.backend.ones([6, 6]), tf.random.Generator.from_seed(23)))
        print(f(tc.backend.ones([6, 6]), tf.random.Generator.from_seed(24)))
    elif tc.backend.name == "jax":
        import jax

        print(f(tc.backend.ones([6, 6]), jax.random.PRNGKey(23)))
        print(f(tc.backend.ones([6, 6]), jax.random.PRNGKey(24)))

    # As seen here, though I have tried the best, the random API is still not that consistent under jit

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_jittable_depolarizing(backend):
    @tc.backend.jit
    def f1(key):
        n = 5
        if key is not None:
            tc.backend.set_random_state(key)
        c = tc.Circuit(n)
        for i in range(n):
            c.H(i)
        for i in range(n):
            c.cnot(i, (i + 1) % n)
        for i in range(n):
            c.unitary_kraus(
                [
                    tc.gates._x_matrix,
                    tc.gates._y_matrix,
                    tc.gates._z_matrix,
                    tc.gates._i_matrix,
                ],
                i,
                prob=[0.2, 0.2, 0.2, 0.4],
            )
        for i in range(n):
            c.cz(i, (i + 1) % n)
        return c.wavefunction()

    @tc.backend.jit
    def f2(key):
        n = 5
        if key is not None:
            tc.backend.set_random_state(key)
        c = tc.Circuit(n)
        for i in range(n):
            c.H(i)
        for i in range(n):
            c.cnot(i, (i + 1) % n)
        for i in range(n):
            c.unitary_kraus(
                tc.channels.depolarizingchannel(0.2, 0.2, 0.2),
                i,
            )
        for i in range(n):
            c.X(i)
        return c.wavefunction()

    @tc.backend.jit
    def f3(key):
        n = 5
        if key is not None:
            tc.backend.set_random_state(key)
        c = tc.Circuit(n)
        for i in range(n):
            c.H(i)
        for i in range(n):
            c.cnot(i, (i + 1) % n)
        for i in range(n):
            c.depolarizing(i, px=0.2, py=0.2, pz=0.2)
        for i in range(n):
            c.X(i)
        return c.wavefunction()

    @tc.backend.jit
    def f4(key):
        n = 5
        if key is not None:
            tc.backend.set_random_state(key)
        c = tc.Circuit(n)
        for i in range(n):
            c.H(i)
        for i in range(n):
            c.cnot(i, (i + 1) % n)
        for i in range(n):
            c.depolarizing2(i, px=0.2, py=0.2, pz=0.2)
        for i in range(n):
            c.X(i)
        return c.wavefunction()

    @tc.backend.jit
    def f5(key):
        n = 5
        if key is not None:
            tc.backend.set_random_state(key)
        c = tc.Circuit(n)
        for i in range(n):
            c.H(i)
        for i in range(n):
            c.cnot(i, (i + 1) % n)
        for i in range(n):
            c.unitary_kraus2(
                tc.channels.depolarizingchannel(0.2, 0.2, 0.2),
                i,
            )
        for i in range(n):
            c.X(i)
        return c.wavefunction()

    for f in [f1, f2, f3, f4, f5]:
        if tc.backend.name == "tensorflow":
            import tensorflow as tf

            np.testing.assert_allclose(tc.backend.norm(f(None)), 1.0, atol=1e-4)
            np.testing.assert_allclose(
                tc.backend.norm(f(tf.random.Generator.from_seed(23))), 1.0, atol=1e-4
            )
            np.testing.assert_allclose(
                tc.backend.norm(f(tf.random.Generator.from_seed(24))), 1.0, atol=1e-4
            )

        elif tc.backend.name == "jax":
            import jax

            np.testing.assert_allclose(
                tc.backend.norm(f(jax.random.PRNGKey(23))), 1.0, atol=1e-4
            )
            np.testing.assert_allclose(
                tc.backend.norm(f(jax.random.PRNGKey(24))), 1.0, atol=1e-4
            )

@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")])
def test_expectation(backend):
    c = tc.Circuit(2)
    c.H(0)
    np.testing.assert_allclose(
        tc.backend.numpy(c.expectation((tc.gates.z(), [0]))), 0, atol=1e-7
    )

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("cpb")])
def test_exp1(backend):
    @partial(tc.backend.jit, jit_compile=True)
    def sf():
        c = tc.Circuit(2)
        xx = np.array(
            [[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]], dtype=np.complex64
        )
        c.exp1(0, 1, unitary=xx, theta=tc.num_to_tensor(0.2))
        s = c.state()
        return s

    @tc.backend.jit
    def s1f():
        c = tc.Circuit(2)
        xx = np.array(
            [[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]], dtype=np.complex64
        )
        c.exp(0, 1, unitary=xx, theta=tc.num_to_tensor(0.2))
        s1 = c.state()
        return s1

    s = tc.backend.numpy(sf())
    s1 = tc.backend.numpy(s1f())
    np.testing.assert_allclose(s, s1, atol=1e-4)

def test_complex128(highp, tfb):
    c = tc.Circuit(2)
    c.H(1)
    c.rx(0, theta=tc.gates.num_to_tensor(1j))
    c.wavefunction()
    np.testing.assert_allclose(c.expectation((tc.gates.z(), [1])), 0)

# def test_qcode():
#     qcode = """
# 4
# x 0
# cnot 0 1
# r 2 theta 1.0 alpha 1.57
# """
#     c = tc.Circuit.from_qcode(qcode)
#     assert c.measure(1)[0] == "1"
#     assert c.to_qcode() == qcode[1:]

def universal_ad():
    @tc.backend.jit
    def forward(theta):
        c = tc.Circuit(2)
        c.R(0, theta=theta, alpha=0.5, phi=0.8)
        return tc.backend.real(c.expectation((tc.gates.z(), [0])))

    gg = tc.backend.grad(forward)
    vg = tc.backend.value_and_grad(forward)
    gg = tc.backend.jit(gg)
    vg = tc.backend.jit(vg)
    theta = tc.gates.num_to_tensor(1.0)
    grad1 = gg(theta)
    v2, grad2 = vg(theta)
    assert grad1 == grad2
    return v2, grad2

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_ad(backend):
    # this amazingly shows how to code once and run in very different AD-ML engines
    print(universal_ad())

def test_single_qubit():
    c = tc.Circuit(1)
    c.H(0)
    w = c.state()[0]
    np.testing.assert_allclose(w, np.array([1, 1]) / np.sqrt(2), atol=1e-4)

@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")])
def test_expectation_between_two_states(backend):
    zp = np.array([1.0, 0.0])
    zd = np.array([0.0, 1.0])
    assert tc.expectation((tc.gates.y(), [0]), ket=zp, bra=zd) == 1j

    c = tc.Circuit(3)
    c.H(0)
    c.ry(1, theta=tc.num_to_tensor(0.8))
    c.cnot(1, 2)

    state = c.wavefunction()
    x1z2 = [(tc.gates.x(), [0]), (tc.gates.z(), [1])]
    e1 = c.expectation(*x1z2)
    e2 = tc.expectation(*x1z2, ket=state, bra=state, normalization=True)
    np.testing.assert_allclose(tc.backend.numpy(e2), tc.backend.numpy(e1))

    c = tc.Circuit(3)
    c.H(0)
    c.ry(1, theta=tc.num_to_tensor(0.8 + 0.7j))
    c.cnot(1, 2)

    state = c.wavefunction()
    x1z2 = [(tc.gates.x(), [0]), (tc.gates.z(), [1])]
    e1 = c.expectation(*x1z2) / tc.backend.norm(state) ** 2
    e2 = tc.expectation(*x1z2, ket=state, normalization=True)
    np.testing.assert_allclose(tc.backend.numpy(e2), tc.backend.numpy(e1))

    c = tc.Circuit(2)
    c.X(1)
    s1 = c.state()
    c2 = tc.Circuit(2)
    c2.X(0)
    s2 = c2.state()
    c3 = tc.Circuit(2)
    c3.H(1)
    s3 = c3.state()
    x1x2 = [(tc.gates.x(), [0]), (tc.gates.x(), [1])]
    e = tc.expectation(*x1x2, ket=s1, bra=s2)
    np.testing.assert_allclose(tc.backend.numpy(e), 1.0)
    e2 = tc.expectation(*x1x2, ket=s3, bra=s2)
    np.testing.assert_allclose(tc.backend.numpy(e2), 1.0 / np.sqrt(2))

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("cpb")])
def test_any_inputs_state(backend):
    c = tc.Circuit(2, inputs=tc.array_to_tensor(np.array([0.0, 0.0, 0.0, 1.0])))
    c.X(0)
    z0 = c.expectation((tc.gates.z(), [0]))
    assert z0 == 1.0
    c = tc.Circuit(2, inputs=tc.array_to_tensor(np.array([0.0, 0.0, 1.0, 0.0])))
    c.X(0)
    z0 = c.expectation((tc.gates.z(), [0]))
    assert z0 == 1.0
    c = tc.Circuit(2, inputs=tc.array_to_tensor(np.array([1.0, 0.0, 0.0, 0.0])))
    c.X(0)
    z0 = c.expectation((tc.gates.z(), [0]))
    assert z0 == -1.0
    c = tc.Circuit(
        2,
        inputs=tc.array_to_tensor(np.array([1 / np.sqrt(2), 0.0, 1 / np.sqrt(2), 0.0])),
    )
    c.X(0)
    z0 = c.expectation((tc.gates.z(), [0]))
    np.testing.assert_allclose(tc.backend.numpy(z0), 0.0, rtol=1e-4, atol=1e-4)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("cpb")])
def test_postselection(backend):
    c = tc.Circuit(3)
    c.H(1)
    c.H(2)
    c.mid_measurement(1, 1)
    c.mid_measurement(2, 1)
    s = c.wavefunction()
    np.testing.assert_allclose(tc.backend.numpy(s[3]).real, 0.5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")])
def test_unitary(backend):
    c = tc.Circuit(2, inputs=np.eye(4))
    c.X(0)
    c.Y(1)
    answer = tc.backend.numpy(np.kron(tc.gates.x().tensor, tc.gates.y().tensor))
    np.testing.assert_allclose(
        tc.backend.numpy(c.wavefunction().reshape([4, 4])), answer, atol=1e-4
    )

@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")])
def test_expectation_ps(backend):
    c = tc.Circuit(2)
    c.X(0)
    r = c.expectation_ps(z=[0, 1])
    np.testing.assert_allclose(tc.backend.numpy(r), -1, atol=1e-5)

    c = tc.Circuit(2)
    c.H(0)
    r = c.expectation_ps(z=[1], x=[0])
    np.testing.assert_allclose(tc.backend.numpy(r), 1, atol=1e-5)
    r1 = c.expectation_ps(ps=[1, 3])
    np.testing.assert_allclose(tc.backend.numpy(r1), 1, atol=1e-5)
    r1 = c.expectation_ps(z=[1, 2], ps=[1, 3])
    np.testing.assert_allclose(tc.backend.numpy(r1), 1, atol=1e-5)

def test_probability():
    for c_cls in [tc.Circuit, tc.DMCircuit]:
        c = c_cls(2)
        c.h(0)
        c.h(1)
        np.testing.assert_allclose(
            c.probability(), np.array([1, 1, 1, 1]) / 4, atol=1e-5
        )

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_dqas_type_circuit(backend):
    eye = tc.gates.i().tensor
    x = tc.gates.x().tensor
    y = tc.gates.y().tensor
    z = tc.gates.z().tensor

    def f(params, structures):
        paramsc = tc.backend.cast(params, dtype="complex64")
        structuresc = tc.backend.softmax(structures, axis=-1)
        structuresc = tc.backend.cast(structuresc, dtype="complex64")
        c = tc.Circuit(5)
        for i in range(5):
            c.H(i)
        for j in range(2):
            for i in range(4):
                c.cz(i, i + 1)
            for i in range(5):
                c.any(
                    i,
                    unitary=structuresc[i, j, 0]
                    * (
                        tc.backend.cos(paramsc[i, j, 0]) * eye
                        + tc.backend.sin(paramsc[i, j, 0]) * x
                    )
                    + structuresc[i, j, 1]
                    * (
                        tc.backend.cos(paramsc[i, j, 1]) * eye
                        + tc.backend.sin(paramsc[i, j, 1]) * y
                    )
                    + structuresc[i, j, 2]
                    * (
                        tc.backend.cos(paramsc[i, j, 2]) * eye
                        + tc.backend.sin(paramsc[i, j, 2]) * z
                    ),
                )
        return tc.backend.real(c.expectation([tc.gates.z(), (2,)]))

    structures = tc.array_to_tensor(
        np.random.normal(size=[16, 5, 2, 3]), dtype="float32"
    )
    params = tc.array_to_tensor(np.random.normal(size=[5, 2, 3]), dtype="float32")

    vf = tc.backend.vmap(f, vectorized_argnums=(1,))

    np.testing.assert_allclose(vf(params, structures).shape, [16])

    vvag = tc.backend.vvag(f, argnums=0, vectorized_argnums=1)

    vvag = tc.backend.jit(vvag)

    value, grad = vvag(params, structures)

    np.testing.assert_allclose(value.shape, [16])
    np.testing.assert_allclose(grad.shape, [5, 2, 3])

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_mixed_measurement_circuit(backend):
    n = 4

    def f(params, structures):
        structuresc = tc.backend.cast(structures, dtype="complex64")
        c = tc.Circuit(n)
        for i in range(n):
            c.H(i)
        for j in range(2):
            for i in range(n):
                c.cnot(i, (i + 1) % n)
            for i in range(n):
                c.rz(i, theta=params[j, i])
        obs = []
        for i in range(n):
            obs.append(
                [
                    tc.gates.Gate(
                        sum(
                            [
                                structuresc[i, k] * g.tensor
                                for k, g in enumerate(tc.gates.pauli_gates)
                            ]
                        )
                    ),
                    (i,),
                ]
            )
        loss = c.expectation(*obs, reuse=False)
        return tc.backend.real(loss)

    # measure X0 to X3

    structures = tc.backend.cast(tc.backend.eye(n), "int32")
    structures = tc.backend.onehot(structures, num=4)

    f_vvag = tc.backend.jit(tc.backend.vvag(f, vectorized_argnums=1, argnums=0))
    v, g = f_vvag(tc.backend.ones([2, n], dtype="float32"), structures)
    np.testing.assert_allclose(
        v,
        np.array(
            [
                0.157729,
                0.157729,
                0.157728,
                0.085221,
            ]
        ),
        atol=1e-5,
    )
    np.testing.assert_allclose(
        g[0],
        np.array([-0.378372, -0.624019, -0.491295, -0.378372]),
        atol=1e-5,
    )

def test_circuit_add_demo():
    # to be refactored for better API
    c = tc.Circuit(2)
    c.x(0)
    c2 = tc.Circuit(2, mps_inputs=c.quvector())
    c2.X(0)
    answer = np.array([1.0, 0, 0, 0])
    np.testing.assert_allclose(c2.wavefunction(), answer, atol=1e-4)
    c3 = tc.Circuit(2)
    c3.X(0)
    c3.replace_mps_inputs(c.quvector())
    np.testing.assert_allclose(c3.wavefunction(), answer, atol=1e-4)

def test_circuit_replace_inputs():
    n = 3
    c = tc.Circuit(n, inputs=np.zeros([2**n]))
    for i in range(n):
        c.H(i)
    evenstate = np.ones([2**n])
    evenstate /= np.linalg.norm(evenstate)
    c.replace_inputs(evenstate)
    for i in range(n):
        np.testing.assert_allclose(c.expectation_ps(z=[i]), 1.0, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_circuit_matrix(backend):
    c = tc.Circuit(2)
    c.x(1)
    c.cnot(0, 1)
    np.testing.assert_allclose(c.matrix()[3], np.array([0.0, 0.0, 0.0, 1.0]), atol=1e-5)
    np.testing.assert_allclose(c.state(), np.array([0, 1.0, 0, 0]), atol=1e-5)

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_circuit_split(backend):
    n = 4

    def f(param, max_singular_values=None, max_truncation_err=None, fixed_choice=None):
        if (max_singular_values is None) and (max_truncation_err is None):
            split = None
        else:
            split = {
                "max_singular_values": max_singular_values,
                "max_truncation_err": max_truncation_err,
                "fixed_choice": fixed_choice,
            }
        c = tc.Circuit(
            n,
            split=split,
        )
        for i in range(n):
            c.H(i)
        for j in range(2):
            for i in range(n - 1):
                c.exp1(i, i + 1, theta=param[2 * j, i], hermitian=tc.gates._zz_matrix)
            for i in range(n):
                c.rx(i, theta=param[2 * j + 1, i])
        loss = c.expectation(
            (
                tc.gates.z(),
                [1],
            ),
            (
                tc.gates.z(),
                [2],
            ),
        )
        return tc.backend.real(loss)

    s1 = f(tc.backend.ones([4, n]))
    s2 = f(tc.backend.ones([4, n]), max_truncation_err=1e-5)
    s3 = f(tc.backend.ones([4, n]), max_singular_values=2, fixed_choice=1)

    np.testing.assert_allclose(s1, s2, atol=1e-5)
    np.testing.assert_allclose(s1, s3, atol=1e-5)

    f_jit = tc.backend.jit(f, static_argnums=(1, 2, 3))

    s1 = f_jit(tc.backend.ones([4, n]))
    # s2 = f_jit(tc.backend.ones([4, n]), max_truncation_err=1e-5) # doesn't work now
    # this cannot be done anyway, since variable size tensor network will fail opt einsum
    s3 = f_jit(tc.backend.ones([4, n]), max_singular_values=2, fixed_choice=1)

    # np.testing.assert_allclose(s1, s2, atol=1e-5)
    np.testing.assert_allclose(s1, s3, atol=1e-5)

    f_vg = tc.backend.jit(
        tc.backend.value_and_grad(f, argnums=0), static_argnums=(1, 2, 3)
    )

    s1, g1 = f_vg(tc.backend.ones([4, n]))
    s3, g3 = f_vg(tc.backend.ones([4, n]), max_singular_values=2, fixed_choice=1)

    np.testing.assert_allclose(s1, s3, atol=1e-5)
    # DONE(@refraction-ray): nan on jax backend?
    # i see, complex value SVD is not supported on jax for now :)
    # I shall further customize complex SVD, finally it has applications

    # tf 2.6.2 also doesn't support complex valued SVD AD, weird...
    # if tc.backend.name == "tensorflow":
    np.testing.assert_allclose(g1, g3, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_gate_split(backend):
    n = 4

    def f(param, max_singular_values=None, max_truncation_err=None, fixed_choice=None):
        if (max_singular_values is None) and (max_truncation_err is None):
            split = None
        else:
            split = {
                "max_singular_values": max_singular_values,
                "max_truncation_err": max_truncation_err,
                "fixed_choice": fixed_choice,
            }
        c = tc.Circuit(
            n,
        )
        for i in range(n):
            c.H(i)
        for j in range(2):
            for i in range(n - 1):
                c.exp1(
                    i,
                    i + 1,
                    theta=param[2 * j, i],
                    unitary=tc.gates._zz_matrix,
                    split=split,
                )
            for i in range(n):
                c.rx(i, theta=param[2 * j + 1, i])
        loss = c.expectation(
            (
                tc.gates.x(),
                [1],
            ),
        )
        return tc.backend.real(loss)

    s1 = f(tc.backend.ones([4, n]))
    s2 = f(tc.backend.ones([4, n]), max_truncation_err=1e-5)
    s3 = f(tc.backend.ones([4, n]), max_singular_values=2, fixed_choice=1)

    np.testing.assert_allclose(s1, s2, atol=1e-5)
    np.testing.assert_allclose(s1, s3, atol=1e-5)

def test_toqir():
    split = {
        "max_singular_values": 2,
        "fixed_choice": 1,
    }
    c = tc.Circuit(3)
    c.H(0)
    c.rx(1, theta=tc.array_to_tensor(0.7))
    c.exp1(
        0, 1, unitary=tc.gates._zz_matrix, theta=tc.array_to_tensor(-0.2), split=split
    )
    z1 = c.expectation((tc.gates.z(), [1]))
    qirs = c.to_qir()
    c = tc.Circuit.from_qir(qirs, circuit_params={"nqubits": 3})
    assert len(c._nodes) == 7
    z2 = c.expectation((tc.gates.z(), [1]))
    np.testing.assert_allclose(z1, z2, atol=1e-5)
    c.append_from_qir(qirs)
    z3 = c.expectation((tc.gates.z(), [1]))
    assert len(c._nodes) == 11
    np.testing.assert_allclose(z3, 0.202728, atol=1e-5)

def test_vis_tex():
    c = tc.Circuit(3)
    for i in range(3):
        c.H(i)
    for i in range(3):
        c.any(i, (i + 1) % 3, unitary=tc.backend.ones([4, 4]), name="hihi")
    c.any(2, unitary=tc.backend.ones([2, 2]), name="invisible")
    c.cz(1, 2)
    c.any(1, 0, 2, unitary=tc.backend.ones([8, 8]), name="ccha")
    c.z(2)
    c.cnot(0, 1)
    c.cz(2, 1)

    print(c.vis_tex(init=["0", "1", ""], measure=["x", "y", "z"]))

def test_debug_contract():
    n = 10
    d = 4
    try:
        import cotengra  # pylint: disable=unused-import

    except ImportError:
        pytest.skip("cotengra is not installed")

    @tc.set_function_contractor(
        "custom_stateful",
        optimizer=oem.RandomGreedy,
        max_time=10,
        max_repeats=64,
        minimize="size",
        debug_level=2,
        contraction_info=True,
    )
    def small_tn():
        param = tc.backend.ones([2 * d, n])
        c = tc.Circuit(n)
        c = tc.templates.blocks.example_block(c, param, nlayers=d)
        return c.state()

    np.testing.assert_allclose(small_tn(), np.zeros([2**n]), atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_teleportation(backend):
    key = tc.backend.get_random_state(42)

    @tc.backend.jit
    def f(key):
        tc.backend.set_random_state(key)
        c = tc.Circuit(2)
        c.H(0)
        r = c.cond_measurement(0)
        c.conditional_gate(r, [tc.gates.i(), tc.gates.x()], 1)
        return r, c.expectation([tc.gates.z(), [1]])

    keys = []
    for _ in range(6):
        key, subkey = tc.backend.random_split(key)
        keys.append(subkey)
    rs = [f(k) for k in keys]
    for r, e in rs:
        if tc.backend.numpy(r) > 0.5:
            np.testing.assert_allclose(e, -1, atol=1e-5)
        else:
            np.testing.assert_allclose(e, 1, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_append_circuit(backend):
    c = tc.Circuit(2)
    c.cnot(0, 1)
    c1 = tc.Circuit(2)
    c1.x(0)
    c.append(c1)
    np.testing.assert_allclose(c.expectation_ps(z=[1]), 1.0)

    c = tc.Circuit(2)
    c.cnot(0, 1)
    c1 = tc.Circuit(2)
    c1.x(0)
    c.prepend(c1)
    np.testing.assert_allclose(c.expectation_ps(z=[1]), -1.0)

    c = tc.Circuit(2)
    c1 = tc.Circuit(1)
    c1.x(0)
    c.append(c1, [1])
    np.testing.assert_allclose(c.state(), [0, 1, 0, 0])

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_apply_mpo_gate(backend):
    gate = tc.gates.multicontrol_gate(tc.gates._x_matrix, ctrl=[1, 0])
    ans = np.array(
        [
            [1.0, 0, 0, 0, 0, 0, 0, 0],
            [0, 1.0, 0, 0, 0, 0, 0, 0],
            [0, 0, 1.0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1.0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1.0, 0, 0],
            [0, 0, 0, 0, 1.0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1.0, 0],
            [0, 0, 0, 0, 0, 0, 0, 1.0],
        ]
    )
    c = tc.Circuit(3)
    c.X(0)
    c.mpo(0, 1, 2, mpo=gate.copy())
    np.testing.assert_allclose(c.expectation([tc.gates.z(), [2]]), -1, atol=1e-5)
    c = tc.Circuit(3)
    c.X(1)
    c.mpo(0, 1, 2, mpo=gate.copy())
    np.testing.assert_allclose(c.expectation([tc.gates.z(), [2]]), 1, atol=1e-5)
    np.testing.assert_allclose(gate.eval_matrix(), ans, atol=1e-5)

def test_apply_multicontrol_gate():
    c = tc.Circuit(3)
    c.X(2)
    c.multicontrol(0, 2, 1, ctrl=[0, 1], unitary=tc.gates._x_matrix)
    np.testing.assert_allclose(c.expectation([tc.gates.z(), [1]]), -1, atol=1e-5)
    c = tc.Circuit(3)
    c.X(0)
    c.multicontrol(0, 2, 1, ctrl=[0, 1], unitary=tc.gates._x_matrix)
    np.testing.assert_allclose(c.expectation([tc.gates.z(), [1]]), 1, atol=1e-5)
    c = tc.Circuit(4)
    c.X(0)
    c.X(2)
    c.multicontrol(0, 1, 2, 3, ctrl=[1, 0], unitary=tc.gates.swap())
    np.testing.assert_allclose(c.expectation([tc.gates.z(), [3]]), -1, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_circuit_quoperator(backend):
    c = tc.Circuit(3)
    c.x(0)
    c.cnot(0, 1)
    c.cz(1, 2)
    c.y(2)
    c.exp1(0, 2, theta=1.0, unitary=tc.gates._xx_matrix)
    c.H(1)
    c.multicontrol(0, 2, 1, ctrl=[1, 0], unitary=tc.gates.z())
    qo = c.quoperator()
    np.testing.assert_allclose(qo.eval_matrix(), c.matrix(), atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_qir2cirq(backend):
    try:
        import cirq
    except ImportError:
        pytest.skip("cirq is not installed")
    n = 6
    c = tc.Circuit(n)
    for i in range(n):
        c.H(i)
    zz = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
    for i in range(n):
        c.exp(
            i,
            (i + 1) % n,
            theta=tc.array_to_tensor(np.random.uniform()),
            unitary=tc.array_to_tensor(zz),
            name="zz",
        )
    c.exp1(
        1, 3, theta=tc.array_to_tensor(0.0j), unitary=tc.array_to_tensor(zz), name="zz"
    )
    c.fredkin(0, 1, 2)
    c.cswap(1, 2, 3)
    c.ccnot(1, 2, 3)
    c.cx(2, 3)
    c.swap(0, 1)
    c.iswap(0, 1)
    c.iswap(1, 3, theta=-1.9)
    c.toffoli(0, 1, 2)
    c.s(1)
    c.t(1)
    c.sd(1)
    c.td(1)
    c.x(2)
    c.y(2)
    c.z(2)
    c.wroot(3)
    c.cnot(0, 1)
    c.cy(0, 1)
    c.cz(0, 1)
    c.oy(4, 3)
    c.oz(4, 3)
    c.ox(4, 3)
    c.oy(4, 3)
    c.oz(4, 3)
    c.ox(3, 4)
    c.phase(2, theta=0.3)
    c.cphase(1, 0, theta=-1.2)
    c.rxx(0, 2, theta=0.9)
    c.ryy(1, 4, theta=-2.0)
    c.rzz(1, 3, theta=0.5)
    c.u(2, theta=0, lbd=4.6, phi=-0.3)
    c.cu(4, 1, theta=1.2)
    c.rx(1, theta=tc.array_to_tensor(np.random.uniform()))
    c.r(5, theta=tc.array_to_tensor(np.random.uniform()))
    c.cr(
        1,
        2,
        theta=tc.array_to_tensor(np.random.uniform()),
        alpha=tc.array_to_tensor(np.random.uniform()),
        phi=tc.array_to_tensor(np.random.uniform()),
    )
    c.ry(1, theta=tc.array_to_tensor(np.random.uniform()))
    c.rz(1, theta=tc.array_to_tensor(np.random.uniform()))
    c.crz(2, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.crx(5, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.cry(1, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.orx(5, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.ory(5, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.orz(5, 3, theta=tc.array_to_tensor(np.random.uniform()))

    c.any(1, 3, unitary=tc.array_to_tensor(np.reshape(zz, [2, 2, 2, 2])))

    gate = tc.gates.multicontrol_gate(
        tc.array_to_tensor(tc.gates._x_matrix), ctrl=[1, 0]
    )
    c.mpo(0, 1, 2, mpo=gate.copy())
    c.multicontrol(
        0,
        2,
        4,
        1,
        5,
        ctrl=[0, 1, 0],
        unitary=tc.array_to_tensor(tc.gates._zz_matrix),
        name="zz",
    )
    tc_unitary = c.matrix()
    tc_unitary = np.reshape(tc_unitary, [2**n, 2**n])

    cirq = c.to_cirq()
    cirq_unitary = cirq.unitary()
    cirq_unitary = np.reshape(cirq_unitary, [2**n, 2**n])

    np.testing.assert_allclose(tc_unitary, cirq_unitary, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_qir2qiskit(backend):
    try:
        import qiskit.quantum_info as qi
        from tensorcircuit.translation import perm_matrix
    except ImportError:
        pytest.skip("qiskit is not installed")

    n = 6
    c = tc.Circuit(n, inputs=tc.array_to_tensor(np.eye(2**n)))

    for i in range(n):
        c.H(i)
    zz = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
    for i in range(n):
        c.exp(
            i,
            (i + 1) % n,
            theta=tc.array_to_tensor(np.random.uniform()),
            unitary=tc.array_to_tensor(zz),
            name="zz",
        )
    c.exp1(
        1, 3, theta=tc.array_to_tensor(0.0j), unitary=tc.array_to_tensor(zz), name="zz"
    )
    c.fredkin(1, 2, 3)
    c.cswap(1, 2, 3)
    c.ccnot(1, 2, 3)
    c.cx(2, 3)
    c.swap(0, 1)
    c.iswap(0, 1)
    c.iswap(1, 3, theta=-1.9)
    c.toffoli(0, 1, 2)
    c.s(1)
    c.t(1)
    c.sd(1)
    c.td(1)
    c.x(2)
    c.y(2)
    c.z(2)
    c.wroot(3)
    c.cnot(0, 1)
    c.cy(0, 1)
    c.cz(0, 1)
    c.oy(4, 3)
    c.oz(4, 3)
    c.ox(4, 3)
    c.oy(4, 3)
    c.oz(4, 3)
    c.ox(3, 4)
    c.phase(2, theta=0.3)
    c.cphase(1, 0, theta=-1.2)
    c.rxx(0, 2, theta=0.9)
    c.ryy(1, 4, theta=-2.0)
    c.rzz(1, 3, theta=0.5)
    c.u(2, theta=0, lbd=4.6, phi=-0.3)
    c.cu(4, 1, theta=1.2)
    c.rx(1, theta=tc.array_to_tensor(np.random.uniform()))
    c.r(5, theta=tc.array_to_tensor(np.random.uniform()))
    c.cr(
        1,
        2,
        theta=tc.array_to_tensor(np.random.uniform()),
        alpha=tc.array_to_tensor(np.random.uniform()),
        phi=tc.array_to_tensor(np.random.uniform()),
    )
    c.ry(1, theta=tc.array_to_tensor(np.random.uniform()))
    c.rz(1, theta=tc.array_to_tensor(np.random.uniform()))
    c.crz(2, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.crx(5, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.cry(1, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.orx(5, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.ory(5, 3, theta=tc.array_to_tensor(np.random.uniform()))
    c.orz(5, 3, theta=tc.array_to_tensor(np.random.uniform()))

    c.any(1, 3, unitary=tc.array_to_tensor(np.reshape(zz, [2, 2, 2, 2])))
    gate = tc.gates.multicontrol_gate(
        tc.array_to_tensor(tc.gates._x_matrix), ctrl=[1, 0]
    )
    c.mpo(0, 1, 2, mpo=gate.copy())
    c.multicontrol(
        0,
        2,
        4,
        1,
        5,
        ctrl=[0, 1, 0],
        unitary=tc.array_to_tensor(tc.gates._zz_matrix),
        name="zz",
    )
    tc_unitary = c.wavefunction()
    tc_unitary = np.reshape(tc_unitary, [2**n, 2**n])

    qisc = c.to_qiskit()
    qis_unitary = qi.Operator(qisc)
    qis_unitary = np.reshape(qis_unitary, [2**n, 2**n])

    p_mat = perm_matrix(n)
    np.testing.assert_allclose(p_mat @ tc_unitary @ p_mat, qis_unitary, atol=1e-5)

def test_qiskit2tc():
    try:
        from qiskit import QuantumCircuit
        import qiskit.quantum_info as qi
        from qiskit.circuit.library.standard_gates import MCXGate, SwapGate
        from tensorcircuit.translation import perm_matrix
    except ImportError:
        pytest.skip("qiskit is not installed")
    n = 6
    qisc = QuantumCircuit(n)
    for i in range(n):
        qisc.h(i)
    zz = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
    exp_op = qi.Operator(zz)
    for i in range(n):
        qisc.hamiltonian(exp_op, time=np.random.uniform(), qubits=[i, (i + 1) % n])
    qisc.fredkin(1, 2, 3)
    qisc.cswap(1, 2, 3)
    qisc.swap(0, 1)
    qisc.iswap(0, 1)
    qisc.toffoli(0, 1, 2)
    # test Instructions
    qisc2 = QuantumCircuit(1)
    qisc2.h(0)
    qisc.compose(qisc2, qubits=[1], inplace=True, wrap=True)
    qisc.barrier(0, 1, 2)
    qisc.s(1)
    qisc.t(1)
    qisc.sdg(2)
    qisc.tdg(2)
    qisc.x(3)
    qisc.y(3)
    qisc.z(3)
    qisc.cnot(0, 1)
    qisc.cy(0, 1)
    qisc.cz(0, 1, ctrl_state=0)
    qisc.cy(0, 1, ctrl_state=0)
    qisc.cx(0, 1, ctrl_state=0)
    qisc.rxx(0.3, 1, 2)
    qisc.rzz(-0.8, 2, 0)
    qisc.u(0.3, 0.9, -1.2, 2)
    qisc.rx(np.random.uniform(), 1)
    qisc.ry(np.random.uniform(), 2)
    qisc.rz(np.random.uniform(), 3)
    qisc.crz(np.random.uniform(), 2, 3)
    qisc.crz(np.random.uniform(), 2, 3)
    qisc.crz(np.random.uniform(), 2, 3)
    qisc.crz(np.random.uniform(), 2, 3, ctrl_state=0)
    qisc.crz(np.random.uniform(), 2, 3, ctrl_state=0)
    qisc.crz(np.random.uniform(), 2, 3, ctrl_state=0)
    qisc.r(np.random.uniform(), np.random.uniform(), 1)
    qisc.unitary(exp_op, [1, 3])
    mcx_g = MCXGate(3, ctrl_state="010")
    qisc.append(mcx_g, [0, 1, 2, 3])
    qisc.ccx(0, 1, 2, ctrl_state="01")
    CCCRX = SwapGate().control(2, ctrl_state="01")
    qisc.append(CCCRX, [0, 1, 2, 3])

    c = tc.Circuit.from_qiskit(qisc, n, np.eye(2**n))
    tc_unitary = c.wavefunction()
    tc_unitary = np.reshape(tc_unitary, [2**n, 2**n])
    qis_unitary = qi.Operator(qisc)
    qis_unitary = np.reshape(qis_unitary, [2**n, 2**n])
    p_mat = perm_matrix(n)
    np.testing.assert_allclose(p_mat @ tc_unitary @ p_mat, qis_unitary, atol=1e-5)
    qisc_from_tc = c.to_qiskit(enable_instruction=True)
    qis_unitary2 = qi.Operator(qisc_from_tc)
    qis_unitary2 = np.reshape(qis_unitary2, [2**n, 2**n])
    np.testing.assert_allclose(qis_unitary2, qis_unitary, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_qiskit2tc_parameterized(backend):
    try:
        from qiskit.circuit import QuantumCircuit, Parameter
        from qiskit.quantum_info import Operator
        from qiskit.circuit.library import TwoLocal
        from qiskit_nature.second_q.circuit.library import UCCSD
        from qiskit_nature.second_q.mappers import ParityMapper, QubitConverter
    except ImportError:
        pytest.skip("qiskit or qiskit-nature is not installed")
    from tensorcircuit.translation import perm_matrix

    mapper = ParityMapper()
    converter = QubitConverter(mapper=mapper, two_qubit_reduction=True)
    ansatz1 = UCCSD(2, [1, 1], converter)
    ansatz2 = TwoLocal(2, rotation_blocks="ry", entanglement_blocks="cz")
    ansatz3 = QuantumCircuit(1)
    ansatz3_param = Parameter("θ")
    ansatz3.rx(ansatz3_param, 0)
    ansatz_list = [ansatz1, ansatz2, ansatz3]
    for ansatz in ansatz_list:
        n = ansatz.num_qubits
        if ansatz in [ansatz1, ansatz2]:
            params = np.random.rand(ansatz.num_parameters)
        else:
            params = {ansatz3_param: 0.618}
        qisc = ansatz.assign_parameters(params)
        qiskit_unitary = Operator(qisc)
        qiskit_unitary = np.reshape(qiskit_unitary, [2**n, 2**n])

        # test jit
        @tc.backend.jit
        def get_unitary(_params):
            return tc.Circuit.from_qiskit(
                ansatz, inputs=np.eye(2**n), binding_params=_params
            ).state()

        tc_unitary = get_unitary(params)
        tc_unitary = np.reshape(tc_unitary, [2**n, 2**n])
        p_mat = perm_matrix(n)
        np.testing.assert_allclose(
            p_mat @ tc_unitary @ p_mat, qiskit_unitary, atol=1e-5
        )

        # test grad
        def cost_fn(_params):
            return tc.backend.real(tc.backend.sum(get_unitary(_params)))

        if ansatz in [ansatz1, ansatz2]:
            grad = tc.backend.grad(cost_fn)(tc.backend.convert_to_tensor(params))
            assert np.sum(np.isnan(grad)) == 0
        else:
            # tf only supports tf tensor as input
            grad = tc.backend.grad(cost_fn)(
                {ansatz3_param: tc.backend.convert_to_tensor(0.618)}
            )
            assert not np.isnan(grad[ansatz3_param])

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_qiskit_vs_tc_intialization(backend):
    try:
        from qiskit import QuantumCircuit
        import qiskit.quantum_info as qi
    except ImportError:
        pytest.skip("qiskit is not installed")

    n = 3

    qis_c = QuantumCircuit(n)
    qis_c.h(0)
    qis_c.cnot(0, 1)
    qis_c.y(2)
    state = qi.Statevector(qis_c)
    qis_c = QuantumCircuit(n)
    qis_c.initialize(state)
    qis_c.cnot(1, 2)
    c = tc.Circuit.from_qiskit(qis_c)
    c2 = tc.Circuit(n)
    c2.h(0)
    c2.cnot(0, 1)
    c2.y(2)
    c2.cnot(1, 2)
    np.testing.assert_allclose(c.state(), c2.state(), atol=1e-8)
    np.testing.assert_allclose(
        qi.Statevector(c.to_qiskit(enable_inputs=True)),
        qi.Statevector(qis_c),
        atol=1e-8,
    )

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_batch_sample(backend):
    c = tc.Circuit(3)
    c.H(0)
    c.cnot(0, 1)
    print(c.sample())
    print(c.sample(batch=8))
    print(c.sample(random_generator=tc.backend.get_random_state(42)))
    print(c.sample(allow_state=True))
    print(c.sample(batch=8, allow_state=True))
    print(
        c.sample(
            batch=8, allow_state=True, random_generator=tc.backend.get_random_state(42)
        )
    )
    print(
        c.sample(
            batch=8,
            allow_state=True,
            status=np.random.uniform(size=[8]),
            format="sample_bin",
        )
    )

def test_expectation_y_bug():
    c = tc.Circuit(1, inputs=1 / np.sqrt(2) * np.array([-1, 1.0j]))
    m = c.expectation_ps(y=[0])
    np.testing.assert_allclose(m, -1, atol=1e-5)

def test_lightcone_expectation():
    def construct_c(pbc=True):
        n = 4
        ns = n
        if pbc is False:
            ns -= 1
        c = tc.Circuit(n)
        for j in range(2):
            for i in range(n):
                c.rx(i, theta=0.2, name="rx" + str(j) + "-" + str(i))
            for i in range(ns):
                c.cnot(i, (i + 1) % n, name="cnot" + str(j) + "-" + str(i))
        return c

    for b in [True, False]:
        c = construct_c(b)
        m1 = c.expectation_ps(z=[0], enable_lightcone=True)
        m2 = c.expectation_ps(z=[0])
        np.testing.assert_allclose(m1, m2, atol=1e-5)
        nodes = c.expectation_before([tc.gates.z(), 0], reuse=False)
        l1 = len(nodes)
        nodes = tc.simplify._full_light_cone_cancel(nodes)
        l2 = len(nodes)
        if b is False:
            assert l1 == 37 and l2 == 25
        else:
            assert l1 == 41 and l2 == 41

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_circuit_inverse(backend):
    inputs = np.random.uniform(size=[8])
    inputs /= np.linalg.norm(inputs)
    c = tc.Circuit(3, inputs=inputs)
    c.H(1)
    c.rx(0, theta=0.5)
    c.cnot(1, 2)
    c.rzz(0, 2, theta=-0.8)
    c1 = c.inverse()
    c.append(c1)
    np.testing.assert_allclose(c.state(), inputs, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_circuit_inverse_2(backend):
    inputs = np.random.uniform(size=[8])
    inputs /= np.linalg.norm(inputs)
    c = tc.Circuit(3, inputs=inputs)
    c.iswap(0, 1)
    c.iswap(1, 0, theta=0.6)
    c.rxx(1, 2, theta=-0.2)
    c.cu(0, 1, lbd=2.0, theta=-0.7)
    c.r(2, alpha=0.3)
    c.sd(2)
    c.cx(1, 2)
    c.unitary(0, unitary=tc.gates._x_matrix)
    c1 = c.inverse()
    c.append(c1)
    print(c.draw())
    np.testing.assert_allclose(c.state(), inputs, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_jittable_amplitude(backend):
    # @tc.backend.jit
    def amp(s):
        c = tc.Circuit(3)
        c.H(0)
        c.cnot(0, 1)
        c.swap(1, 2)
        return c.amplitude(s)

    np.testing.assert_allclose(
        amp(tc.array_to_tensor([0, 1, 1], dtype="float32")), 0, atol=1e-5
    )
    np.testing.assert_allclose(
        amp(tc.array_to_tensor([0, 0, 0], dtype="float32")), 1 / np.sqrt(2), atol=1e-5
    )

def test_draw_cond_measure():
    c = tc.Circuit(2)
    c.H(0)
    c.cond_measure(0)
    c.cnot(0, 1)
    print("")
    print(c.draw())

def test_minus_index():
    c = tc.Circuit(3)
    c.H(-2)
    c.H(0)
    np.testing.assert_allclose(tc.backend.real(c.expectation_ps(x=[0])), 1, atol=1e-5)
    np.testing.assert_allclose(tc.backend.real(c.expectation_ps(x=[1])), 1, atol=1e-5)
    np.testing.assert_allclose(tc.backend.real(c.expectation_ps(x=[-1])), 0, atol=1e-5)
    np.testing.assert_allclose(tc.backend.real(c.expectation_ps(z=[-2])), 0, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_sexpps(backend):
    c = tc.Circuit(1, inputs=1 / np.sqrt(2) * np.array([1.0, 1.0j]))
    y = c.sample_expectation_ps(y=[0])
    ye = c.expectation_ps(y=[0])
    np.testing.assert_allclose(y, 1.0, atol=1e-5)
    np.testing.assert_allclose(ye, 1.0, atol=1e-5)

    c = tc.Circuit(4)
    c.H(0)
    c.cnot(0, 1)
    c.rx(1, theta=0.3)
    c.rz(2, theta=-1.2)
    c.ccnot(2, 3, 1)
    c.rzz(0, 3, theta=0.5)
    c.ry(3, theta=2.2)
    c.s(1)
    c.td(2)
    y = c.sample_expectation_ps(x=[1], y=[0], z=[2, 3])
    ye = c.expectation_ps(x=[1], y=[0], z=[2, 3])
    np.testing.assert_allclose(ye, y, atol=1e-5)
    y2 = c.sample_expectation_ps(x=[1], y=[0], z=[2, 3], shots=81920)
    assert np.abs(y2 - y) < 0.01

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_sample_format(backend):
    c = tc.Circuit(2)
    c.H(0)
    c.cnot(0, 1)
    key = tc.backend.get_random_state(42)
    for allow_state in [False, True]:
        print("allow_state: ", allow_state)
        for batch in [None, 1, 3]:
            print("  batch: ", batch)
            for format_ in [
                None,
                "sample_int",
                "sample_bin",
                "count_vector",
                "count_tuple",
                "count_dict_bin",
                "count_dict_int",
            ]:
                print("    format: ", format_)
                print(
                    "      ",
                    c.sample(
                        batch=batch,
                        allow_state=allow_state,
                        format_=format_,
                        random_generator=key,
                    ),
                )

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_channel_auto_register(backend, highp):
    c = tc.Circuit(2)
    c.H(0)
    c.reset(0, status=0.8)
    s = c.state()
    np.testing.assert_allclose(s[0], 1.0, atol=1e-9)

@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_circuit_to_json(backend):
    c = tc.Circuit(3)
    c.h(0)
    c.CNOT(1, 2)
    c.rxx(0, 2, theta=0.3)
    c.crx(0, 1, theta=-0.8)
    c.r(1, theta=tc.backend.ones([]), alpha=0.2)
    c.toffoli(0, 2, 1)
    c.ccnot(0, 1, 2)
    c.multicontrol(1, 2, 0, ctrl=[0, 1], unitary=tc.gates._x_matrix)
    s = c.to_json()
    c2 = tc.Circuit.from_json(s)
    print(c2.draw())
    np.testing.assert_allclose(c.state(), c2.state(), atol=1e-5)

def test_gate_count():
    c = tc.Circuit(3)
    c.x(0)
    c.h(0)
    c.rx(1, theta=-0.2)
    c.h(2)
    c.multicontrol(0, 1, 2, ctrl=[0, 1], unitary=tc.gates._x_matrix)
    c.toffoli(0, 2, 1)
    c.ccnot(1, 2, 0)
    c.ccx(1, 2, 0)
    assert c.gate_count() == 8
    assert c.gate_count("h") == 2
    assert c.gate_count(["ccnot"]) == 3
    assert c.gate_count(["rx", "multicontrol"]) == 2
    assert c.gate_count_by_condition(lambda qir: qir["index"] == (0,)) == 2
    assert c.gate_count_by_condition(lambda qir: qir["mpo"]) == 1
    print(c.gate_summary())
    # {'x': 1, 'h': 2, 'rx': 1, 'multicontrol': 1, 'toffoli': 3}

def test_to_openqasm():
    c = tc.Circuit(3)
    c.H(0)
    c.rz(2, theta=0.2)
    c.cnot(2, 1)
    c.rzz(0, 1, theta=-1.0)
    c.ccx(1, 2, 0)
    c.u(2, theta=0.5, lbd=1.3)
    c.measure_instruction(1)
    print(c.to_openqasm(formatted=True))
    s = c.to_openqasm()
    c1 = tc.Circuit.from_openqasm(s)
    print(c1.draw())
    np.testing.assert_allclose(c.state(), c1.state())
    c.to_openqasm(filename="test.qasm")
    c2 = tc.Circuit.from_openqasm_file("test.qasm")
    np.testing.assert_allclose(c.state(), c2.state())

def test_from_qasm_keep_measure_order():
    qasm_str = """OPENQASM 2.0;
include "qelib1.inc";
qreg q[2];
creg c[2];
h q[0];
measure q[1] -> c[1];
measure q[0] -> c[0];"""
    c = tc.Circuit.from_openqasm(qasm_str)
    c.to_openqasm().split("\n")[-2][-3] == "1"
    c = tc.Circuit.from_openqasm(qasm_str, keep_measure_order=True)
    c.to_openqasm().split("\n")[-2][-3] == "0"

def test_initial_mapping():
    c = tc.Circuit(3)
    c.cnot(0, 1)
    c.h(1)
    c.rx(1, theta=0.5)
    c.cz(2, 1)
    c.measure_instruction(2)

    c1 = c.initial_mapping({0: 1, 1: 2, 2: 0})
    print(c1.draw())

    np.testing.assert_allclose(
        c.expectation_ps(z=[1]), c1.expectation_ps(z=[2]), atol=1e-5
    )
    assert c1._extra_qir[0]["index"][0] == 0

    c2 = c1.initial_mapping({1: 0, 2: 1, 0: 2})
    np.testing.assert_allclose(
        c.expectation_ps(z=[1]), c2.expectation_ps(z=[1]), atol=1e-5
    )

    c3 = c.initial_mapping({0: 2, 1: 7, 2: 0}, n=9)
    np.testing.assert_allclose(
        c.expectation_ps(z=[1]), c3.expectation_ps(z=[7]), atol=1e-5
    )
    print(c3.draw())

def test_get_positional_logical_mapping():
    c = tc.Circuit(3)
    c.cx(0, 1)
    c.cz(1, 2)
    c.h(1)
    assert c.get_positional_logical_mapping() == {0: 0, 1: 1, 2: 2}
    c = tc.Circuit(3)
    c.cx(0, 1)
    c.h(1)
    c.measure_instruction(2)
    c.measure_instruction(0)
    assert c.get_positional_logical_mapping() == {0: 2, 1: 0}

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_inverse_jit(backend):
    K = tc.backend

    def simple_ansatz(param):
        c = tc.Circuit(3)
        for i in range(3):
            c.cx(i, (i + 1) % 3)
            c.rzz(i, (i + 1) % 3, theta=param[i])
        c1 = c.inverse()
        c2 = tc.Circuit(3)
        c2.x(1)
        c1.append(c2)
        return tc.backend.real(c1.expectation_ps(z=[1]))

    v_ansatz = K.jit(K.vvag(simple_ansatz))
    vs, gs = v_ansatz(K.ones([2, 3], dtype="float32"))
    assert K.shape_tuple(gs) == (2, 3)
    np.testing.assert_allclose(K.numpy(vs), -1.0 * K.ones([2]), atol=1e-5)

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("npb")])
def test_fancy_circuit_indexing(backend):
    c = tc.Circuit(4)
    c.cx([0, 1], [-1, -2])
    c.h(range(c._nqubits))
    c.rz([0], theta=0.2)
    c.rx([1, 2], theta=[0.3, 0.5])
    c.rzz([2, 3], [0, 1], theta=tc.backend.ones([2]))
    c.rxx([2, 0, 1], [0, 1, 2], theta=tc.backend.ones([1]))
    assert c.gate_count("h") == 4
    assert c.gate_count("rzz") == 2
    assert c.gate_count("rxx") == 3

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("npb")])
def test_general_kraus(backend):
    c = tc.Circuit(2)
    c.h([0, 1])
    p = 0.5
    status = [0.3, 0.8]
    rs = []
    for i in range(2):
        rs.append(
            c.general_kraus(
                [
                    np.sqrt(p) * np.array([[1.0, 0], [0, 0]]),
                    np.sqrt(p) * np.array([[0, 0], [0, 1.0]]),
                    np.sqrt(1 - p) * np.eye(2),
                ],
                i,
                status=status[i],
            )
        )
    np.testing.assert_allclose(rs[0], 1)
    np.testing.assert_allclose(rs[1], 2)
    np.testing.assert_allclose(c.expectation_ps(z=[0]), -1, atol=1e-5)
    np.testing.assert_allclose(c.expectation_ps(z=[1]), 0, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("npb")])
def test_general_kraus_with_prob(backend):
    c = tc.Circuit(2)
    c.h([0, 1])
    p = 0.5
    status = [0.3, 0.8]
    rs = []
    for i in range(2):
        rs.append(
            c.general_kraus(
                [
                    np.sqrt(p) * np.array([[1.0, 0], [0, 0]]),
                    np.sqrt(p) * np.array([[0, 0], [0, 1.0]]),
                    np.sqrt(1 - p) * np.eye(2),
                ],
                i,
                status=status[i],
                with_prob=True,
            )
        )
    np.testing.assert_allclose(rs[0][0], 1)
    np.testing.assert_allclose(rs[1][0], 2)
    np.testing.assert_allclose(c.expectation_ps(z=[0]), -1, atol=1e-5)
    np.testing.assert_allclose(c.expectation_ps(z=[1]), 0, atol=1e-5)
    np.testing.assert_allclose(rs[0][1], [0.25, 0.25, 0.5], atol=1e-5)
    np.testing.assert_allclose(rs[1][1], [0.25, 0.25, 0.5], atol=1e-5)
    np.testing.assert_allclose(tc.backend.norm(c.state()), 1, atol=1e-5)

@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("npb")])
def test_circuit_copy(backend):
    c = tc.Circuit(2)
    c.h(0)
    c1 = c.copy()
    c.rz(0, theta=0.1)
    assert c1.gate_count() == 1


import numpy as np
from numpy import (max,min,pi,exp,log)
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm
gna, gk, gl = 120, 36, 0.3
vna, vk, vl = 50, -77, -54.5

def get_am(v):
    return 0.1*(v+40)/(1-exp(-0.1*(v+40)))

def get_ah(v):
    return 0.07*exp(-0.05*(v+65))

def get_an(v):
    return 0.01*(v+55)/(1-exp(-0.1*(v+55)))

def get_bm(v):
    return 4*exp(-(v+65)/18)

def get_bh(v):
    return 1/(1+exp(-0.1*(v+35)))

def get_bn(v):
    return 0.125*exp(-(v+65)/80)

def get_m(v):
    return get_am(v) / (get_am(v)+get_bm(v))

def get_h(v):
    return get_ah(v) / (get_ah(v) + get_bh(v))

def get_n(v):
    return get_an(v) / (get_an(v) + get_bn(v))

def get_i(v):
    return gna*get_m(v)**3*get_h(v)*(v-vna)+gk*get_n(v)**4*(v-vk)+gl*(v-vl)

def two_divide(left=-100,right=-20,i=0.0):
    nl = get_i(left)
    nr = get_i(right)
    nm = get_i((left+right)/2)
    while abs(nm-i)>1e-6:
        if nm-i > 0:
            right = (left+right)/2
        else:
            left = (left+right)/2
        nm = get_i((left+right)/2)
        # print(nm)
        # print(left,right)
    return (left+right)/2

def get_jac(v):
    m=get_m(v)
    h=get_h(v)
    n=get_n(v)

    df1_dv = -gna*m**3*h-gk*n**4-gl
    df1_dm = -3*m**2*gna*h*(v-vna)
    df1_dh = -gna*m**3*(v-vna)
    df1_dn = -4*gk*n**3*(v-vk)

    dam_dv = (0.1-(0.5+0.01*v)*exp(-0.1*(v+40)))/(1-exp(-0.1*(v+40)))**2
    dbm_dv = -4/18*exp(-(v+65)/18)
    dah_dv = -0.05*0.07*exp(-0.05*(v+65))
    dbh_dv = -0.1*exp(0.1*(v+35))/(1+exp(0.1*(v+35)))**2
    dan_dv = (0.1-(0.65+0.01*v)*exp(-0.1*(v+55)))/(1-exp(-0.1*(v+55)))**2
    dbn_dv = -1/640*exp(-(65+v)/80)

    df2_dv = (1-m)*dam_dv-m*dbm_dv
    df2_dm = -(get_am(v)+get_bm(v))

    df3_dv = (1-h)*dah_dv - h*dbh_dv
    df3_dh = -(get_ah(v)+get_bh(v))

    df4_dv = (1-n)*dan_dv - n*dbn_dv
    df4_dn = -(get_an(v)+get_bn(v))

    jac = np.zeros((4, 4))
    jac[0,0] = df1_dv
    jac[0,1] = df1_dm
    jac[0,2] = df1_dh
    jac[0,3] = df1_dn

    jac[1,0] = df2_dv
    jac[1,1] = df2_dm

    jac[2,0] = df3_dv
    jac[2,2] = df3_dh

    jac[3,0] = df4_dv
    jac[3,3] = df4_dn
    return jac
# df = []
# for k in np.arange(6.24,6.25,0.0001):
#     v0 = two_divide(i=k)
#     if v0 in [-40,-65,-55]:
#         v0 += 0.02
#     jac = get_jac(v0)
#     eig = np.linalg.eigvals(jac)
#     df.append(eig)
# df1 = pd.DataFrame(df)
def simu(i):
    v0 = two_divide(i=0)
    am = get_am(v0)
    bm = get_bm(v0)
    ah = get_ah(v0)
    bh = get_bh(v0)
    an = get_an(v0)
    bn = get_bn(v0)
    m = get_m(v0)
    h = get_h(v0)
    n = get_n(v0)
    delta = 0.01
    v0_l = [v0]
    for j in range(19999):
        # print(v0)
        v0 += (i-(gna*m**3*h*(v0-vna)+gk*n**4*(v0-vk)+gl*(v0-vl)))*delta
        v0_l.append(v0)
        am = get_am(v0)
        bm = get_bm(v0)
        ah = get_ah(v0)
        bh = get_bh(v0)
        an = get_an(v0)
        bn = get_bn(v0)
        m += delta*(-(am+bm)*m+am)
        h += delta*(-(ah+bh)*h+ah)
        n += delta*(-(an+bn)*n+an)

    return v0_l
i = 6.245
l = simu(i)
plt.subplot(2,2,1)
plt.plot(0.01*np.arange(20000),l)
plt.xlabel(f'I={i}mA/cm2')
plt.ylabel('V(mV)')

plt.subplot(2,2,2)
i = 6.255
l = simu(i)
plt.plot(0.01*np.arange(20000),l)
plt.xlabel(f'I={i}mA/cm2')
plt.ylabel('V(mV)')

plt.subplot(2,2,3)
i = 100
l = simu(i)
plt.plot(0.01*np.arange(20000),l)
plt.xlabel(f'I={i}mA/cm2')
plt.ylabel('V(mV)')

plt.subplot(2,2,4)
i = 158
l = simu(i)
plt.plot(0.01*np.arange(20000),l)
plt.xlabel(f'I={i}mA/cm2')
plt.ylabel('V(mV)')
def get_iv():
    v_l=[]
    for k in tqdm(np.arange(0,180,0.5)):
        l = simu(k)
        v_l.append(max(l[-3000:])-min(l[-3000:]))
    return v_l

# v_l = get_iv()


from gym.envs.registration import register
import sys

print("path", sys.argv[0].split('/')[-1], "!!!")
if sys.argv[0].split('/')[-1] in ["train_ddpg.py", "visitation_plot.py", "vis_fetch.py"]:
    from train_ddpg import args
elif sys.argv[0].split('/')[-1] == "train_hier_ddpg.py":
    from train_hier_ddpg import args
elif sys.argv[0].split('/')[-1] == "train_hier_sac.py":
    from train_hier_sac import args
elif sys.argv[0].split('/')[-1] == "train_hier_ppo.py":
    from train_hier_ppo import args
elif sys.argv[0].split('/')[-1] == "train_covering.py":
    from train_covering import args
else:
    raise Exception("Unknown main file !!!")

robots = ['Point', 'Ant', 'Swimmer']
task_types = ['Maze', 'Maze1', 'Push', 'Fall', 'Block', 'BlockMaze']
all_name = [x + y for x in robots for y in task_types]
random_start = False

if args.image:
    top_down = True
else:
    top_down = False

for name_t in all_name:
    # episode length
    if name_t == "AntMaze":
        max_timestep = 1000
    else:
        max_timestep = 500
    for Test in ['', 'Test', 'Test1', 'Test2']:

        if Test in ['Test', 'Test1', 'Test2']:
            fix_goal = True
        else:
            if name_t == "AntBlock":
                fix_goal = True
            else:
                fix_goal = False
        goal_args = [[-5, -5], [5, 5]]

        register(
            id=name_t + Test + '-v0',
            entry_point='goal_env.mujoco.create_maze_env:create_maze_env',
            kwargs={'env_name': name_t, 'goal_args': goal_args, 'maze_size_scaling': 8, 'random_start': random_start},
            max_episode_steps=max_timestep,
        )

        # v1 is the one we use in the main paper
        register(
            id=name_t + Test + '-v1',
            entry_point='goal_env.mujoco.create_maze_env:create_maze_env',
            kwargs={'env_name': name_t, 'goal_args': goal_args, 'maze_size_scaling': 4, 'random_start': random_start,
                    "fix_goal": fix_goal, "top_down_view": top_down, 'test':Test},
            max_episode_steps=max_timestep,
        )

        register(
            id=name_t + Test + '-v2',
            entry_point='goal_env.mujoco.create_maze_env:create_maze_env',
            kwargs={'env_name': name_t, 'goal_args': goal_args, 'maze_size_scaling': 2, 'random_start': random_start},
            max_episode_steps=max_timestep,
        )