from pytorch_pretrained_biggan import (BigGAN, truncated_noise_sample)
from imagenet_class_picker import ImageNetClassPicker
from tqdm import tqdm
from scipy.misc import toimage
import librosa
import argparse
import numpy as np
import moviepy.editor as mpy
import random
import torch
torch.cuda.empty_cache()

# get input arguments
parser = argparse.ArgumentParser()
parser.add_argument("--song", required=True)
parser.add_argument("--resolution", default='512')
parser.add_argument("--duration", type=int)
parser.add_argument("--pitch_sensitivity", type=int, default=220)
parser.add_argument("--tempo_sensitivity", type=float, default=0.25)
parser.add_argument("--depth", type=float, default=1)
parser.add_argument("--classes", nargs='+', type=int)
parser.add_argument("--num_classes", type=int, default=12)
parser.add_argument("--class_picker", type=int, default=1)
parser.add_argument("--image_net_classes_file",
                    default="imagenet1000_clsidx_to_labels.txt")
parser.add_argument("--lyrics_file", default="lyrics.txt")
parser.add_argument("--class_list_file", default="class_list_file.txt")
parser.add_argument("--classes_per_lyric_line", type=int, default=1)
parser.add_argument("--sort_classes_by_power", type=int, default=0)
parser.add_argument("--jitter", type=float, default=0.5)
parser.add_argument("--frame_length", type=int, default=512)
parser.add_argument("--truncation", type=float, default=1)
parser.add_argument("--smooth_factor", type=int, default=20)
parser.add_argument("--batch_size", type=int, default=30)
parser.add_argument("--use_previous_classes", type=int, default=0)
parser.add_argument("--use_previous_vectors", type=int, default=0)
parser.add_argument("--output_file", default="output.mp4")
args = parser.parse_args()

# read song
if args.song:
    song = args.song
    print('\nReading audio \n')
    y, sr = librosa.load(song)
else:
    raise ValueError(
        "you must enter an audio file name in the --song argument")

# set model name based on resolution
model_name = 'biggan-deep-' + args.resolution

# set ImageNet classes file
image_net_classes_file = args.image_net_classes_file

# initialize ImageNetClassPicker
image_net_class_picker = ImageNetClassPicker()
image_net_class_picker.load_imagenet_class_labels(image_net_classes_file)

frame_length = args.frame_length

# set pitch sensitivity
pitch_sensitivity = (300-args.pitch_sensitivity) * 512 / frame_length

# set tempo sensitivity
tempo_sensitivity = args.tempo_sensitivity * frame_length / 512

# set depth
depth = args.depth

# set number of classes
num_classes = args.num_classes

# set how classes will be picked
class_picker = args.class_picker

# set the lyrics file name
lyrics_file = args.lyrics_file

# set class list file
class_list_file = args.class_list_file

# set the number of classes to pick per lyrical line
classes_per_lyric_line = args.classes_per_lyric_line

# set sort_classes_by_power
sort_classes_by_power = args.sort_classes_by_power

# set jitter
jitter = args.jitter

# set truncation
truncation = args.truncation

# set batch size
batch_size = args.batch_size

# set use_previous_classes
use_previous_vectors = args.use_previous_vectors

# set use_previous_vectors
use_previous_classes = args.use_previous_classes

# set output name
outname = args.output_file

# set smooth factor
if args.smooth_factor > 1:
    smooth_factor = int(args.smooth_factor * 512 / frame_length)
else:
    smooth_factor = args.smooth_factor

# set duration
if args.duration:
    seconds = args.duration
    frame_lim = int(np.floor(seconds*22050/frame_length/batch_size))
else:
    frame_lim = int(np.floor(len(y)/sr*22050/frame_length/batch_size))

# Load pre-trained model
model = BigGAN.from_pretrained(model_name)

# set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

########################################
########################################
########################################
########################################
########################################

# create spectrogram
spec = librosa.feature.melspectrogram(
    y=y, sr=sr, n_mels=128, fmax=8000, hop_length=frame_length)

# get mean power at each time point
specm = np.mean(spec, axis=0)

# compute power gradient across time points
gradm = np.gradient(specm)

# set max to 1
gradm = gradm/np.max(gradm)

# set negative gradient time points to zero
gradm = gradm.clip(min=0)

# normalize mean power between 0-1
specm = (specm-np.min(specm))/np.ptp(specm)

# create chromagram of pitches X time points
chroma = librosa.feature.chroma_cqt(y=y, sr=sr, hop_length=frame_length)

# sort pitches by overall power
chromasort = np.argsort(np.mean(chroma, axis=1))[::-1]

########################################
########################################
########################################
########################################
########################################

if args.classes:
    classes = args.classes
    if len(classes) not in [12, num_classes]:
        raise ValueError(
            "The number of classes entered in the --class argument must equal 12 or [num_classes] if specified")

elif args.use_previous_classes == 1:
    cvs = np.load('class_vectors.npy')
    classes = list(np.where(cvs[0] > 0)[0])

elif class_picker == 2:  # sentiment analysis
    image_net_class_picker.load_song_lyrics(lyrics_file)
    classes = image_net_class_picker.get_classes_from_sentiment_analysis(
        num_classes)[:num_classes]
    print("List of classes determined based off sentiment analysis:",
          classes)

elif class_picker == 3:  # contextual similarity
    image_net_class_picker.load_song_lyrics(lyrics_file)
    classes = image_net_class_picker.get_classes_from_semantic_similarity(
        classes_per_lyric_line, num_classes)
    image_net_class_picker.save_list("stay_classes.txt", classes)
    print("List of classes determined based off contextual similarity:",
          classes)

elif class_picker == 4:  # sentiment + contextual
    image_net_class_picker.load_song_lyrics(lyrics_file)
    sentiment_classes_list = image_net_class_picker.get_classes_from_sentiment_analysis()
    print("List of classes determined based off sentiment analysis:",
          sentiment_classes_list)
    semantic_classes_list = image_net_class_picker.get_classes_from_semantic_similarity(
        classes_per_lyric_line, 1000)
    print("List of classes determined based off contextual similarity:",
          semantic_classes_list)

    classes = []
    for class_id in semantic_classes_list:
        if class_id in sentiment_classes_list:
            classes.append(class_id)

    print("List of classes determined based off sentiment analysis & contextual similarity:",
          classes)

    classes = classes[:num_classes]

    print("Final list of classes:", classes)

elif class_picker == 5:  # load in classes from pickle file
    classes = list(set([int(x)
                        for x in image_net_class_picker.load_list(class_list_file)]))
    random.shuffle(classes)
    classes = classes[:num_classes]
    print("List of loaded classes:", classes)

else:  # select 12 random classes
    cls1000 = list(range(1000))
    random.shuffle(cls1000)
    classes = cls1000[:12]

if sort_classes_by_power == 1:

    classes = [classes[s] for s in np.argsort(chromasort[:num_classes])]

# initialize first class vector
cv1 = np.zeros(1000)
for pi, p in enumerate(chromasort[:num_classes]):

    if num_classes < 12:
        cv1[classes[pi]] = chroma[p][np.min(
            [np.where(chrow > 0)[0][0] for chrow in chroma])]
    else:
        cv1[classes[p]] = chroma[p][np.min(
            [np.where(chrow > 0)[0][0] for chrow in chroma])]

# initialize first noise vector
nv1 = truncated_noise_sample(truncation=truncation)[0]

# initialize list of class and noise vectors
class_vectors = [cv1]
noise_vectors = [nv1]

# initialize previous vectors (will be used to track the previous frame)
cvlast = cv1
nvlast = nv1

# initialize the direction of noise vector unit updates
update_dir = np.zeros(128)
for ni, n in enumerate(nv1):
    if n < 0:
        update_dir[ni] = 1
    else:
        update_dir[ni] = -1

# initialize noise unit update
update_last = np.zeros(128)

########################################
########################################
########################################
########################################
########################################

# get new jitters
def new_jitters(jitter):
    jitters = np.zeros(128)
    for j in range(128):
        if random.uniform(0, 1) < 0.5:
            jitters[j] = 1
        else:
            jitters[j] = 1-jitter
    return jitters

# get new update directions
def new_update_dir(nv2, update_dir):
    for ni, n in enumerate(nv2):
        if n >= 2*truncation - tempo_sensitivity:
            update_dir[ni] = -1

        elif n < -2*truncation + tempo_sensitivity:
            update_dir[ni] = 1
    return update_dir

# smooth class vectors
def smooth(class_vectors, smooth_factor):

    if smooth_factor == 1:
        return class_vectors

    class_vectors_terp = []
    for c in range(int(np.floor(len(class_vectors)/smooth_factor)-1)):
        ci = c*smooth_factor
        cva = np.mean(class_vectors[int(ci):int(ci)+smooth_factor], axis=0)
        cvb = np.mean(
            class_vectors[int(ci)+smooth_factor:int(ci)+smooth_factor*2], axis=0)

        for j in range(smooth_factor):
            cvc = cva*(1-j/(smooth_factor-1)) + cvb*(j/(smooth_factor-1))
            class_vectors_terp.append(cvc)

    return np.array(class_vectors_terp)

# normalize class vector between 0-1
def normalize_cv(cv2):
    min_class_val = min(i for i in cv2 if i != 0)
    for ci, c in enumerate(cv2):
        if c == 0:
            cv2[ci] = min_class_val
    cv2 = (cv2-min_class_val)/np.ptp(cv2)

    return cv2

print('\nGenerating input vectors \n')

for i in tqdm(range(len(gradm))):

    # print progress
    pass

    # update jitter vector every 100 frames by setting ~half of noise vector units to lower sensitivity
    if i % 200 == 0:
        jitters = new_jitters(jitter)

    # get last noise vector
    nv1 = nvlast

    # set noise vector update based on direction, sensitivity, jitter, and combination of overall power and gradient of power
    update = np.array([tempo_sensitivity for k in range(128)]) * \
        (gradm[i]+specm[i]) * update_dir * jitters

    # smooth the update with the previous update (to avoid overly sharp frame transitions)
    update = (update+update_last*3)/4

    # set last update
    update_last = update

    # update noise vector
    nv2 = nv1+update

    # append to noise vectors
    noise_vectors.append(nv2)

    # set last noise vector
    nvlast = nv2

    # update the direction of noise units
    update_dir = new_update_dir(nv2, update_dir)

    # get last class vector
    cv1 = cvlast

    # generate new class vector
    cv2 = np.zeros(1000)
    for j in range(num_classes):

        cv2[classes[j]] = (cvlast[classes[j]] + ((chroma[chromasort[j]][i]) /
                           (pitch_sensitivity)))/(1+(1/((pitch_sensitivity))))

    # if more than 6 classes, normalize new class vector between 0 and 1, else simply set max class val to 1
    if num_classes > 6:
        cv2 = normalize_cv(cv2)
    else:
        cv2 = cv2/np.max(cv2)

    # adjust depth
    cv2 = cv2*depth

    # this prevents rare bugs where all classes are the same value
    if np.std(cv2[np.where(cv2 != 0)]) < 0.0000001:
        cv2[classes[0]] = cv2[classes[0]]+0.01

    # append new class vector
    class_vectors.append(cv2)

    # set last class vector
    cvlast = cv2

# interpolate between class vectors of bin size [smooth_factor] to smooth frames
class_vectors = smooth(class_vectors, smooth_factor)

# check whether to use vectors from last run
if use_previous_vectors == 1:
    # load vectors from previous run
    class_vectors = np.load('class_vectors.npy')
    noise_vectors = np.load('noise_vectors.npy')
else:
    # save record of vectors for current video
    np.save('class_vectors.npy', class_vectors)
    np.save('noise_vectors.npy', noise_vectors)

########################################
########################################
########################################
########################################
########################################

# convert to Tensor
noise_vectors = torch.Tensor(np.array(noise_vectors))
class_vectors = torch.Tensor(np.array(class_vectors))

# Generate frames in batches of batch_size

print('\n\nGenerating frames \n')

# send to CUDA if running on GPU
model = model.to(device)
noise_vectors = noise_vectors.to(device)
class_vectors = class_vectors.to(device)

frames = []

for i in tqdm(range(frame_lim)):

    # print progress
    pass

    if (i+1)*batch_size > len(class_vectors):
        torch.cuda.empty_cache()
        break

    # get batch
    noise_vector = noise_vectors[i*batch_size:(i+1)*batch_size]
    class_vector = class_vectors[i*batch_size:(i+1)*batch_size]

    # Generate images
    with torch.no_grad():
        output = model(noise_vector, class_vector, truncation)

    output_cpu = output.cpu().data.numpy()

    # convert to image array and add to frames
    for out in output_cpu:
        im = np.array(toimage(out))
        frames.append(im)

    # empty cuda cache
    torch.cuda.empty_cache()

image_net_class_picker.save_list('frames.pkl', frames)

# Save video
aud = mpy.AudioFileClip(song, fps=44100)

if args.duration:
    aud.duration = args.duration

clip = mpy.ImageSequenceClip(frames, fps=22050/frame_length)
clip = clip.set_audio(aud)
clip.write_videofile(outname, audio_codec='aac')


import os
#import sys
import math
import numpy as np
import time
from config import *

lo0p=1
while lo0p:
    lo0p=0
    start_time = time.time()

    # fill rainbow pattern
    rgb = np.zeros((32, 3))
    rgb[0]=[3, 3, 3]
    rgb[1]=[0, 0, 0]
    rgb[2]=[0, 0, 1]
    rgb[3]=[0, 0, 2]
    rgb[4]=[0, 0, 3]
    rgb[5]=[0, 1, 3]
    rgb[6]=[0, 2, 3]
    rgb[7]=[0, 3, 3]
    rgb[8]=[0, 3, 2]
    rgb[9]=[0, 3, 1]
    rgb[10]=[0, 3, 0]
    rgb[11]=[1, 3, 0]
    rgb[12]=[2, 3, 0]
    rgb[13]=[3, 3, 0]
    rgb[14]=[3, 2, 0]
    rgb[15]=[3, 1, 0]
    rgb[16]=[3, 0, 0]
    rgb[17]=[3, 0, 1]
    rgb[18]=[3, 0, 2]
    rgb[19]=[3, 0, 3]
    rgb[20]=[2, 0, 3]
    rgb[21]=[1, 0, 3]
    rgb[22]=[0, 0, 3]
    rgb[23]=[3, 2, 1]
    rgb[24]=[0, 0, 0]
    rgb[25]=[1, 1, 1]
    rgb[26]=[2, 2, 2]
    rgb[27]=[3, 3, 3]
    rgb[28]=[0, 0, 0]
    rgb[29]=[3, 0, 0]
    rgb[30]=[0, 3, 0]
    rgb[31]=[0, 0, 3]

    # generate image (LED matrix), based on given pattern
    arr1 = np.zeros((64, 32, 8, 3))

    width=64
    height=32
    for y in range(height):
        for x in range (width):
            if (y==0 or y>28): # generate gradient for given color, alter orientation for some lines
                b=int(rgb[y][2]*(63-x))
                g=int(rgb[y][1]*(63-x))
                r=int(rgb[y][0]*(63-x))
            else:
                b=int(rgb[y][2]*x)
                g=int(rgb[y][1]*x)
                r=int(rgb[y][0]*x)

            for bit in range(0,8): # split pixel into 8 bitfields
                arr1[x][y][bit][0]= ((r)&(1<<bit))
                arr1[x][y][bit][1]= ((g)&(1<<bit))
                arr1[x][y][bit][2]= ((b)&(1<<bit))

    xs=int(xf/2)
    # create black line, empty line.
    black_line =bytearray(0)
    for x in range(0,xs):
        black_line.extend((0).to_bytes(4, byteorder='little')) #
        black_line.extend((0).to_bytes(4, byteorder='little')) #

    bit0=bytearray(0)
    bit1=bytearray(0)
    bit2=bytearray(0)
    bit3=bytearray(0)
    bit4=bytearray(0)
    bit5=bytearray(0)
    bit6=bytearray(0)
    bit7=bytearray(0)

    # translate bitplanes into strings of data
    for bit in range(0,8):
        dat=bytearray(0)
        for y in range (16):
            count=0
            for x in range(0,xs):
                count+=1
                d2=0
                if (count <=64):
                    d2=enable
                    if (count>90):
                        d2=0
                    if (arr1[x][y][bit][2]): #b
                        d2+=B1
                    if (arr1[x][y][bit][1]): #g
                        d2+=G1
                    if (arr1[x][y][bit][0]): #r
                        d2+=R1

                    if (arr1[x][y+16][bit][2]): #b
                        d2+=B2
                    if (arr1[x][y+16][bit][1]): #g
                        d2+=G2
                    if (arr1[x][y+16][bit][0]): #r
                        d2+=R2
                    low=(d2).to_bytes(4, byteorder='little') #
                    hig=(d2+clk).to_bytes(4, byteorder='little') #

                    dat.extend(low)
                    dat.extend(hig)
                else:
                    dat.extend((0).to_bytes(4, byteorder='little')) #
                    dat.extend((0).to_bytes(4, byteorder='little')) #
            #f.write(dat) #
        if(bit==0):
            bit0=dat
        if(bit==1):
            bit1=dat
        if(bit==2):
            bit2=dat
        if(bit==3):
            bit3=dat
        if(bit==4):
            bit4=dat
        if(bit==5):
            bit5=dat
        if(bit==6):
            bit6=dat
        if(bit==7):
            bit7=dat

    with open('/dev/fb1','wb') as f:
        for n in range(13):
            f.write(black_line) # padding at the beginning

        # display data in order,
        #f.write(bit7)
        #f.write(bit7)
        #f.write(bit7)
        #f.write(bit7)
        #f.write(bit7)
        #f.write(bit7)
        #f.write(bit7)
        #f.write(bit7)
        #f.write(bit6)
        #f.write(bit6)
        #f.write(bit6)
        #f.write(bit6)
        #f.write(bit5)
        #f.write(bit5)
        #f.write(bit4)

        # shuffle data around to get less flicker
        f.write(bit7)
        f.write(bit6)
        f.write(bit7)
        f.write(bit5)
        f.write(bit7)
        f.write(bit6)
        f.write(bit7)
        f.write(bit4)
        f.write(bit7)
        f.write(bit6)
        f.write(bit7)
        f.write(bit5)
        f.write(bit7)
        f.write(bit6)
        f.write(bit7)
        #f.write(bit3)
        for n in range(30):
            f.write(black_line) # paddung at the end
        f.close()
    print("fps: ", 1.0/(time.time()-start_time))


# Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
import glob
import json
import logging
import math
import os
import pickle
import random
import uuid
from dataclasses import dataclass, field
from functools import lru_cache
from itertools import cycle, product
from pathlib import Path
from typing import Any, Dict, Sequence, Tuple

import numpy as np
from cached_property import cached_property

from smarts.core.coordinates import Heading, BoundingBox
from smarts.core.data_model import SocialAgent
from smarts.core.route import ShortestRoute
from smarts.core.sumo_road_network import SumoRoadNetwork
from smarts.core.traffic_history import TrafficHistory
from smarts.core.utils.file import file_md5_hash, make_dir_in_smarts_log_dir, path2hash
from smarts.core.utils.id import SocialAgentId
from smarts.core.utils.math import radians_to_vec, vec_to_radians
from smarts.sstudio import types as sstudio_types
from smarts.sstudio.types import CutIn, EntryTactic, UTurn
from smarts.sstudio.types import Via as SSVia

from smarts.core import VEHICLE_CONFIGS

# XXX: consider using smarts.core.coordinates.Pose for this
@dataclass(frozen=True)
class Start:
    position: Tuple[int, int]
    heading: Heading

@dataclass(frozen=True)
class Goal:
    def is_endless(self):
        return True

    def is_reached(self, vehicle):
        return False

@dataclass(frozen=True)
class EndlessGoal(Goal):
    pass

@dataclass(frozen=True)
class PositionalGoal(Goal):
    position: Tuple[int, int]
    # target_heading: Heading
    radius: float

    @classmethod
    def fromedge(cls, edge_id, road_network, lane_index=0, lane_offset=None, radius=1):
        edge = road_network.edge_by_id(edge_id)
        lane = edge.getLanes()[lane_index]

        if lane_offset is None:
            # Default to the midpoint safely ensuring we are on the lane and not
            # bordering another
            lane_offset = lane.getLength() * 0.5

        position = road_network.world_coord_from_offset(lane, lane_offset)
        return cls(position=position, radius=radius)

    def is_endless(self):
        return False

    def is_reached(self, vehicle):
        a = vehicle.position
        b = self.position
        dist = math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)
        return dist <= self.radius

@dataclass(frozen=True)
class TraverseGoal(Goal):
    """A TraverseGoal is satisfied whenever an Agent-driven vehicle
    successfully finishes traversing a non-closed (acyclical) map
    It's a way for the vehicle to exit the simulation successfully,
    for example, driving across from one side to the other on a
    straight road and then continuing off the map.  This goal is
    non-specific about *where* the map is exited, save for that
    the vehicle must be going the correct direction in its lane
    just prior to doing so."""

    def __init__(self, road_network: SumoRoadNetwork):
        super().__init__()
        self._road_network = road_network

    def is_endless(self):
        return True

    def is_reached(self, vehicle):
        return self._road_network.drove_off_map(vehicle.position, vehicle.heading)

def default_entry_tactic(default_entry_speed: float = None) -> EntryTactic:
    return sstudio_types.TrapEntryTactic(
        wait_to_hijack_limit_s=0,
        exclusion_prefixes=tuple(),
        zone=None,
        default_entry_speed=default_entry_speed,
    )

@dataclass(frozen=True)
class Via:
    lane_id: str
    edge_id: str
    lane_index: int
    position: Tuple[float, float]
    hit_distance: float
    required_speed: float

@dataclass(frozen=True)
class VehicleSpec:
    veh_id: str
    veh_type: str
    dimensions: BoundingBox

@dataclass(frozen=True)
class Mission:
    start: Start
    goal: Goal
    # An optional list of edge IDs between the start and end goal that we want to
    # ensure the mission includes
    route_vias: Tuple[str] = field(default_factory=tuple)
    start_time: float = 0.1
    entry_tactic: EntryTactic = None
    task: Tuple[CutIn, UTurn] = None
    via: Tuple[Via, ...] = ()
    # if specified, will use vehicle_spec to build the vehicle (for histories)
    vehicle_spec: VehicleSpec = None

    @property
    def has_fixed_route(self):
        return not self.goal.is_endless()

    def is_complete(self, vehicle, distance_travelled):
        return self.goal.is_reached(vehicle)

@dataclass(frozen=True)
class LapMission:
    start: Start
    goal: Goal
    route_length: float
    num_laps: int = None  # None means infinite # of laps
    # An optional list of edge IDs between the start and end goal that we want to
    # ensure the mission includes
    route_vias: Tuple[str] = field(default_factory=tuple)
    start_time: float = 0.1
    entry_tactic: EntryTactic = None
    via_points: Tuple[Via, ...] = ()

    @property
    def has_fixed_route(self):
        return True

    def is_complete(self, vehicle, distance_travelled):
        return (
            self.goal.is_reached(vehicle)
            and distance_travelled > self.route_length * self.num_laps
        )

class Scenario:
    """The purpose of the Scenario is to provide an aggregate of all
    code/configuration/assets that is specialized to a scenario using SUMO.

    Args:
        scenario_root:
            The scenario asset folder ie. './scenarios/trigger'.
        route:
            The social vehicle traffic spec.
        missions:
            agent_id to mission mapping.
    """

    def __init__(
        self,
        scenario_root: str,
        route: str = None,
        missions: Dict[str, Mission] = None,
        social_agents: Dict[str, SocialAgent] = None,
        log_dir: str = None,
        surface_patches: list = None,
        traffic_history: str = None,
    ):

        self._logger = logging.getLogger(self.__class__.__name__)
        self._root = scenario_root
        self._route = route
        self._missions = missions or {}
        self._bubbles = Scenario._discover_bubbles(scenario_root)
        self._social_agents = social_agents or {}
        self._surface_patches = surface_patches
        self._log_dir = self._resolve_log_dir(log_dir)
        self._validate_assets_exist()

        if traffic_history:
            self._traffic_history = TrafficHistory(traffic_history)
            default_lane_width = self.traffic_history.lane_width
        else:
            self._traffic_history = None
            default_lane_width = None

        net_file = os.path.join(self._root, "map.net.xml")
        self._road_network = SumoRoadNetwork.from_file(
            net_file, default_lane_width=default_lane_width, lanepoint_spacing=1.0
        )
        self._net_file_hash = file_md5_hash(self._road_network.net_file)
        self._scenario_hash = path2hash(str(Path(self.root_filepath).resolve()))

    def __repr__(self):
        return f"""Scenario(
  _root={self._root},
  _route={self._route},
  _missions={self._missions},
)"""

    @staticmethod
    def scenario_variations(
        scenarios_or_scenarios_dirs: Sequence[str],
        agents_to_be_briefed: Sequence[str],
        shuffle_scenarios: bool = True,
    ):
        """Generate a cycle of the configurations of scenarios.

        Args:
            scenarios_or_scenarios_dirs:
                A sequence of either the scenario to run (see scenarios/ for some samples you
                can use) OR a directory of scenarios to sample from.
            agents_to_be_briefed:
                Agent IDs that will be assigned a mission ("briefed" on a mission).
        """
        scenario_roots = []
        for root in scenarios_or_scenarios_dirs:
            if Scenario.is_valid_scenario(root):
                # This is the single scenario mode, only training against a single scenario
                scenario_roots.append(root)
            else:
                scenario_roots.extend(Scenario.discover_scenarios(root))

        if shuffle_scenarios:
            np.random.shuffle(scenario_roots)

        return Scenario.variations_for_all_scenario_roots(
            cycle(scenario_roots), agents_to_be_briefed, shuffle_scenarios
        )

    @staticmethod
    def variations_for_all_scenario_roots(
        scenario_roots, agents_to_be_briefed, shuffle_scenarios=True
    ):
        for scenario_root in scenario_roots:
            surface_patches = Scenario.discover_friction_map(scenario_root)

            agent_missions = Scenario.discover_agent_missions(
                scenario_root, agents_to_be_briefed
            )
            agent_missions = [dict(zip(agents_to_be_briefed, agent_missions))]
            social_agent_infos = Scenario._discover_social_agents_info(scenario_root)
            social_agents = [
                {
                    agent_id: (agent.to_agent_spec(), (agent, mission))
                    for agent_id, (
                        agent,
                        mission,
                    ) in per_episode_social_agent_infos.items()
                }
                for per_episode_social_agent_infos in social_agent_infos
            ]

            # `or [None]` so that product(...) will not return an empty result
            # but insted a [(..., `None`), ...].
            routes = Scenario.discover_routes(scenario_root) or [None]
            agent_missions = agent_missions or [None]
            social_agents = social_agents or [None]
            traffic_histories = Scenario.discover_traffic_histories(scenario_root) or [
                None
            ]

            roll_routes = 0
            roll_agent_missions = 0
            roll_social_agents = 0
            roll_traffic_histories = 0

            if shuffle_scenarios:
                roll_routes = random.randint(0, len(routes))
                roll_agent_missions = random.randint(0, len(agent_missions))
                roll_social_agents = random.randint(0, len(social_agents))
                roll_traffic_histories = random.randint(0, len(traffic_histories))

            for (
                concrete_route,
                concrete_agent_missions,
                concrete_social_agents,
                concrete_traffic_history,
            ) in product(
                np.roll(routes, roll_routes, 0),
                np.roll(agent_missions, roll_agent_missions, 0),
                np.roll(social_agents, roll_social_agents, 0),
                np.roll(traffic_histories, roll_traffic_histories, 0),
            ):
                concrete_social_agent_missions = {
                    agent_id: mission
                    for agent_id, (_, (_, mission)) in (
                        concrete_social_agents or {}
                    ).items()
                }

                # Filter out mission
                concrete_social_agents = {
                    agent_id: (_agent_spec, social_agent)
                    for agent_id, (_agent_spec, (social_agent, _)) in (
                        concrete_social_agents or {}
                    ).items()
                }

                yield Scenario(
                    scenario_root,
                    route=concrete_route,
                    missions={
                        **(concrete_agent_missions or {}),
                        **concrete_social_agent_missions,
                    },
                    social_agents=concrete_social_agents,
                    surface_patches=surface_patches,
                    traffic_history=concrete_traffic_history,
                )

    @staticmethod
    def discover_agent_missions_count(scenario_root):
        missions_file = os.path.join(scenario_root, "missions.pkl")
        if os.path.exists(missions_file):
            with open(missions_file, "rb") as f:
                return len(pickle.load(f))

        return 0

    @staticmethod
    def discover_agent_missions(scenario_root, agents_to_be_briefed):
        """Returns a sequence of {agent_id: mission} mappings.

        If no missions are discovered we generate random ones. If there is only one
        agent to be briefed we return a list of `{agent_id: mission}` cycling through
        each mission. If there are multiple agents to be briefed we assume that each
        one is intended to get its own mission and that `len(agents_to_be_briefed) ==
        len(missions)`. In this case a list of one dictionary is returned.
        """

        net_file = os.path.join(scenario_root, "map.net.xml")
        road_network = SumoRoadNetwork.from_file(net_file)

        missions = []
        missions_file = os.path.join(scenario_root, "missions.pkl")
        if os.path.exists(missions_file):
            with open(missions_file, "rb") as f:
                missions = pickle.load(f)

            missions = [
                Scenario._extract_mission(actor_and_mission, road_network, agent_id)
                for actor_and_mission, agent_id in zip(missions, agents_to_be_briefed)
            ]

        if not missions:
            missions = [None for _ in range(len(agents_to_be_briefed))]

        if len(agents_to_be_briefed) == 1:
            # single-agent, so we cycle through all missions individually.
            return missions
        elif len(agents_to_be_briefed) > 1:
            # multi-agent, so we assume missions "drive" the agents (i.e. one
            # mission per agent) and we will not be cycling through missions.
            assert not missions or len(missions) == len(agents_to_be_briefed), (
                "You must either provide an equal number of missions ({}) to "
                "agents ({}) or provide no missions at all so they can be "
                "randomly generated.".format(len(missions), len(agents_to_be_briefed))
            )

        return missions

    @staticmethod
    def discover_friction_map(scenario_root):
        """Returns the list of surface patches parameters defined in
        scenario file. Each element of the list contains the
        parameters of the specifiec surface patch.
        """
        surface_patches = []
        friction_map_file = os.path.join(scenario_root, "friction_map.pkl")
        if os.path.exists(friction_map_file):
            with open(friction_map_file, "rb") as f:
                map_surface_patches = pickle.load(f)
            for surface_patch in map_surface_patches:
                surface_patches.append(
                    {
                        "zone": surface_patch.zone,
                        "begin_time": surface_patch.begin_time,
                        "end_time": surface_patch.end_time,
                        "friction coefficient": surface_patch.friction_coefficient,
                    }
                )
        return surface_patches

    @staticmethod
    @lru_cache(maxsize=16)
    def _discover_social_agents_info(
        scenario,
    ) -> Sequence[Dict[str, SocialAgent]]:
        """Loops through the social agent mission pickles, instantiating corresponding
        implementations for the given types. The output is a list of
        {agent_id: (mission, locator)}, where each dictionary corresponds to the
        social agents to run for a given concrete Scenario (which translates to
        "per episode" when swapping).
        """
        scenario_root = (
            scenario.root_filepath if isinstance(scenario, Scenario) else scenario
        )
        net_file = os.path.join(scenario_root, "map.net.xml")
        road_network = SumoRoadNetwork.from_file(net_file)

        social_agents_path = os.path.join(scenario_root, "social_agents")
        if not os.path.exists(social_agents_path):
            return []

        # [ ( missions_file, agent_actor, Mission ) ]
        agent_bucketer = []

        # like dict.setdefault
        def setdefault(l: Sequence[Any], index: int, default):
            while len(l) < index + 1:
                l.append([])
            return l[index]

        file_match = os.path.join(social_agents_path, "*.pkl")
        for missions_file_path in glob.glob(file_match):
            with open(missions_file_path, "rb") as missions_file:
                count = 0
                missions = pickle.load(missions_file)

            for mission_and_actor in missions:
                # Each pickle file will contain a list of actor/mission pairs. The pairs
                # will most likely be generated in an M:N fashion
                # (i.e. A1: M1, A1: M2, A2: M1, A2: M2). The desired behavior is to have
                # a single pair per concrete Scenario (which would translate to
                # "per episode" when swapping)
                assert isinstance(
                    mission_and_actor.actor, sstudio_types.SocialAgentActor
                )
                namespace = os.path.basename(missions_file_path)
                namespace = os.path.splitext(namespace)[0]

                agent_id = SocialAgentId.new(
                    mission_and_actor.actor.name, group=namespace
                )

                actor = mission_and_actor.actor
                extracted_mission = Scenario._extract_mission(
                    mission_and_actor, road_network, agent_id
                )

                setdefault(agent_bucketer, count, []).append(
                    (
                        SocialAgent(
                            id=agent_id,
                            name=actor.name,
                            is_boid=False,
                            is_boid_keep_alive=False,
                            agent_locator=actor.agent_locator,
                            policy_kwargs=actor.policy_kwargs,
                            initial_speed=actor.initial_speed,
                        ),
                        extracted_mission,
                    )
                )
                count += 1

        social_agents_info = []
        for l in agent_bucketer:
            social_agents_info.append(
                {agent.id: (agent, mission) for agent, mission in l}
            )

        return social_agents_info

    @staticmethod
    def discover_scenarios(scenario_or_scenarios_dir):
        if Scenario.is_valid_scenario(scenario_or_scenarios_dir):
            # This is the single scenario mode, only training against a single scenario
            scenario = scenario_or_scenarios_dir
            discovered_scenarios = [scenario]
        else:
            # Find all valid scenarios in the given scenarios directory
            discovered_scenarios = []
            for scenario_file in os.listdir(scenario_or_scenarios_dir):
                scenario_root = os.path.join(scenario_or_scenarios_dir, scenario_file)
                if Scenario.is_valid_scenario(scenario_root):
                    discovered_scenarios.append(scenario_root)
        assert (
            len(discovered_scenarios) > 0
        ), f"No valid scenarios found in {scenario_or_scenarios_dir}"

        return discovered_scenarios

    @staticmethod
    def discover_routes(scenario_root):
        """Discover the route files in the given scenario.

        >>> Scenario.discover_routes("scenarios/intersections/2lane")
        ['all.rou.xml', 'horizontal.rou.xml', 'turns.rou.xml', 'unprotected_left.rou.xml', 'vertical.rou.xml']
        >>> Scenario.discover_routes("scenarios/loop") # loop does not have any routes
        ['basic.rou.xml']
        """
        return sorted(
            [
                os.path.basename(r)
                for r in glob.glob(os.path.join(scenario_root, "traffic", "*.rou.xml"))
            ]
        )

    @staticmethod
    def _discover_bubbles(scenario_root):
        path = os.path.join(scenario_root, "bubbles.pkl")
        if not os.path.exists(path):
            return []

        with open(path, "rb") as f:
            bubbles = pickle.load(f)
            return bubbles

    def set_ego_missions(self, ego_missions: dict):
        self._missions = ego_missions

    def discover_missions_of_traffic_histories(self) -> Dict[str, Mission]:
        vehicle_missions = {}
        map_offset = self._road_network.net_offset
        for row in self._traffic_history.first_seen_times():
            start_time = float(row[1])
            pphs = self._traffic_history.vehicle_pose_at_time(row[0], start_time)
            assert pphs
            pos_x, pos_y, heading, speed = pphs
            entry_tactic = default_entry_tactic(speed)
            v_id = str(row[0])
            veh_type = self._traffic_history.vehicle_type(v_id)
            veh_length, veh_width, veh_height = self._traffic_history.vehicle_size(v_id)
            # missions start from front bumper, but pos is center of vehicle
            hhx, hhy = radians_to_vec(heading) * (0.5 * veh_length)
            vehicle_missions[v_id] = Mission(
                start=Start(
                    (pos_x + map_offset[0] + hhx, pos_y + map_offset[1] + hhy),
                    Heading(heading),
                ),
                entry_tactic=entry_tactic,
                goal=TraverseGoal(self.road_network),
                start_time=start_time,
                vehicle_spec=VehicleSpec(
                    veh_id=v_id,
                    veh_type=veh_type,
                    dimensions=BoundingBox(veh_length, veh_width, veh_height),
                ),
            )
        return vehicle_missions

    @staticmethod
    def discover_traffic_histories(scenario_root: str):
        return [
            entry
            for entry in os.scandir(scenario_root)
            if entry.is_file() and entry.path.endswith(".shf")
        ]

    @staticmethod
    def _extract_mission(actor_and_mission, road_network, agent_id):
        """Takes a sstudio.types.(Mission, EndlessMission, etc.) and converts it to
        the corresponding SMARTS mission types.
        """

        def resolve_offset(offset, lane_length):
            # epsilon to ensure we are within this edge and not the subsequent one
            epsilon = 1e-6
            lane_length -= epsilon
            if offset == "base":
                return epsilon
            elif offset == "max":
                return lane_length
            elif offset == "random":
                return random.uniform(epsilon, lane_length)
            else:
                return float(offset)

        def to_position_and_heading(edge_id, lane_index, offset, road_network):
            edge = road_network.edge_by_id(edge_id)
            lane = edge.getLanes()[lane_index]
            offset = resolve_offset(offset, lane.getLength())
            position = road_network.world_coord_from_offset(lane, offset)
            lane_vector = road_network.lane_vector_at_offset(lane, offset)
            heading = vec_to_radians(lane_vector)
            return tuple(position), Heading(heading)

        def to_scenario_via(
            vias: Tuple[SSVia, ...], sumo_road_network: SumoRoadNetwork
        ) -> Tuple[Via, ...]:
            s_vias = []
            for via in vias:
                lane = sumo_road_network.lane_by_index_on_edge(
                    via.edge_id, via.lane_index
                )
                hit_distance = (
                    via.hit_distance if via.hit_distance > 0 else lane.getWidth() / 2
                )
                via_position = sumo_road_network.world_coord_from_offset(
                    lane,
                    via.lane_offset,
                )

                s_vias.append(
                    Via(
                        lane_id=lane.getID(),
                        lane_index=via.lane_index,
                        edge_id=via.edge_id,
                        position=tuple(via_position),
                        hit_distance=hit_distance,
                        required_speed=via.required_speed,
                    )
                )

            return tuple(s_vias)

        mission = actor_and_mission.mission
        # For now we discard the route and just take the start and end to form our
        # missions.
        if isinstance(mission, sstudio_types.Mission):
            position, heading = to_position_and_heading(
                *mission.route.begin,
                road_network,
            )
            start = Start(position, heading)

            position, _ = to_position_and_heading(
                *mission.route.end,
                road_network,
            )
            goal = PositionalGoal(position, radius=2)

            vehicle_spec = VehicleSpec(
                veh_id=agent_id,
                veh_type=actor_and_mission.actor.vehicle_type,
                dimensions=VEHICLE_CONFIGS[
                    actor_and_mission.actor.vehicle_type
                ].dimensions,
            )

            return Mission(
                start=start,
                route_vias=mission.route.via,
                goal=goal,
                start_time=mission.start_time,
                entry_tactic=mission.entry_tactic,
                task=mission.task,
                via=to_scenario_via(mission.via, road_network),
                vehicle_spec=vehicle_spec,
            )
        elif isinstance(mission, sstudio_types.EndlessMission):
            position, heading = to_position_and_heading(
                *mission.begin,
                road_network,
            )
            start = Start(position, heading)

            vehicle_spec = VehicleSpec(
                veh_id=agent_id,
                veh_type=actor_and_mission.actor.vehicle_type,
                dimensions=VEHICLE_CONFIGS[
                    actor_and_mission.actor.vehicle_type
                ].dimensions,
            )

            return Mission(
                start=start,
                goal=EndlessGoal(),
                start_time=mission.start_time,
                entry_tactic=mission.entry_tactic,
                via=to_scenario_via(mission.via, road_network),
                vehicle_spec=vehicle_spec,
            )
        elif isinstance(mission, sstudio_types.LapMission):
            start_edge_id, start_lane, start_edge_offset = mission.route.begin
            end_edge_id, end_lane, end_edge_offset = mission.route.end

            travel_edge = road_network.edge_by_id(start_edge_id)
            if start_edge_id == end_edge_id:
                travel_edge = list(travel_edge.getOutgoing())[0]

            end_edge = road_network.edge_by_id(end_edge_id)
            via_edges = [road_network.edge_by_id(e) for e in mission.route.via]

            route_length = ShortestRoute(
                road_network,
                edge_constraints=[travel_edge] + via_edges + [end_edge],
                wraps_around=True,
            ).length

            start_position, start_heading = to_position_and_heading(
                *mission.route.begin,
                road_network,
            )
            end_position, _ = to_position_and_heading(
                *mission.route.end,
                road_network,
            )

            return LapMission(
                start=Start(start_position, start_heading),
                goal=PositionalGoal(end_position, radius=2),
                route_vias=mission.route.via,
                num_laps=mission.num_laps,
                route_length=route_length,
                start_time=mission.start_time,
                entry_tactic=mission.entry_tactic,
                via_points=to_scenario_via(mission.via, road_network),
            )

        raise RuntimeError(
            f"sstudio mission={mission} is an invalid type={type(mission)}"
        )

    @staticmethod
    def is_valid_scenario(scenario_root):
        """Checks if the scenario_root directory matches our expected scenario structure

        >>> Scenario.is_valid_scenario("scenarios/loop")
        True
        >>> Scenario.is_valid_scenario("scenarios/non_existant")
        False
        """
        paths = [
            os.path.join(scenario_root, "map.net.xml"),
        ]

        for f in paths:
            if not os.path.exists(f):
                return False

        # make sure we can load the sumo network
        net_file = os.path.join(scenario_root, "map.net.xml")
        net = SumoRoadNetwork.from_file(net_file)
        if net is None:
            return False

        return True

    @staticmethod
    def next(scenario_iterator, log_id=""):
        """Utility to override specific attributes from a scenario iterator"""

        scenario = next(scenario_iterator)
        scenario._log_id = log_id
        return scenario

    @property
    def name(self):
        return os.path.basename(os.path.normpath(self._root))

    @property
    def root_filepath(self):
        return self._root

    @property
    def surface_patches(self):
        return self._surface_patches

    @property
    def net_file_hash(self):
        return self._net_file_hash

    @property
    def plane_filepath(self):
        return os.path.join(self._root, "plane.urdf")

    @property
    def vehicle_filepath(self):
        for fname in os.listdir(self._root):
            if fname.endswith(".urdf") and fname != "plane.urdf":
                return os.path.join(self._root, fname)
        return None

    @property
    def tire_parameters_filepath(self):
        return os.path.join(self._root, "tire_parameters.yaml")

    @property
    def controller_parameters_filepath(self):
        return os.path.join(self._root, "controller_parameters.yaml")

    @property
    def route(self):
        return self._route

    @property
    def route_files_enabled(self):
        return bool(self._route)

    @property
    def route_filepath(self):
        return os.path.join(self._root, "traffic", self._route)

    @property
    def map_glb_filepath(self):
        return os.path.join(self._root, "map.glb")

    def unique_sumo_log_file(self):
        return os.path.join(self._log_dir, f"sumo-{str(uuid.uuid4())[:8]}")

    @property
    def road_network(self):
        return self._road_network

    @property
    def missions(self):
        return self._missions

    @property
    def social_agents(self):
        return self._social_agents

    @property
    def bubbles(self):
        return self._bubbles

    def mission(self, agent_id):
        return self._missions.get(agent_id, None)

    def _resolve_log_dir(self, log_dir):
        if log_dir is None:
            log_dir = make_dir_in_smarts_log_dir("_sumo_run_logs")

        return os.path.abspath(log_dir)

    def _validate_assets_exist(self):
        assert Scenario.is_valid_scenario(self._root)
        os.makedirs(self._log_dir, exist_ok=True)

    @property
    def traffic_history(self):
        return self._traffic_history

    @property
    def scenario_hash(self):
        return self._scenario_hash

    @property
    def map_bounding_box(self):
        # This function returns the following tuple:
        # (bbox length, bbox width, bbox center)
        net_file = os.path.join(self._root, "map.net.xml")
        road_network = SumoRoadNetwork.from_file(net_file)
        # 2D bbox in format (xmin, ymin, xmax, ymax)
        bounding_box = road_network.graph.getBoundary()
        bounding_box_length = bounding_box[2] - bounding_box[0]
        bounding_box_width = bounding_box[3] - bounding_box[1]
        bounding_box_center = [
            (bounding_box[0] + bounding_box[2]) / 2,
            (bounding_box[1] + bounding_box[3]) / 2,
            0,
        ]
        return (bounding_box_length, bounding_box_width, bounding_box_center)


# Copyright 2021 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
r"""
This module contains the :class:`~.LightningQubit` class, a PennyLane simulator device that
interfaces with C++ for fast linear algebra calculations.
"""
from typing import List
from warnings import warn
from os import getenv
from itertools import islice

import numpy as np
from pennylane import (
    math,
    gradients,
    BasisState,
    QubitStateVector,
    QubitUnitary,
    Projector,
    Hermitian,
    Rot,
    QuantumFunctionError,
    DeviceError,
)
from pennylane.devices import DefaultQubit
from pennylane.operation import Tensor
from pennylane.measurements import Expectation
from pennylane.wires import Wires

# Remove after the next release of PL
# Add from pennylane import matrix
import pennylane as qml

from ._version import __version__

try:
    from .lightning_qubit_ops import (
        MeasuresC64,
        StateVectorC64,
        AdjointJacobianC64,
        VectorJacobianProductC64,
        MeasuresC128,
        StateVectorC128,
        AdjointJacobianC128,
        VectorJacobianProductC128,
        DEFAULT_KERNEL_FOR_OPS,
        EXPORTED_KERNEL_OPS,
    )

    from ._serialize import _serialize_obs, _serialize_ops, _is_lightning_gate

    CPP_BINARY_AVAILABLE = True
except ModuleNotFoundError:
    CPP_BINARY_AVAILABLE = False

UNSUPPORTED_PARAM_GATES_ADJOINT = (
    "SingleExcitation",
    "SingleExcitationPlus",
    "SingleExcitationMinus",
    "DoubleExcitation",
    "DoubleExcitationPlus",
    "DoubleExcitationMinus",
)

def _chunk_iterable(it, num_chunks):
    "Lazy-evaluated chunking of given iterable from https://stackoverflow.com/a/22045226"
    it = iter(it)
    return iter(lambda: tuple(islice(it, num_chunks)), ())

def _remove_snapshot_from_operations(operations):
    operations = operations.copy()
    operations.discard("Snapshot")
    return operations

class LightningQubit(DefaultQubit):
    """PennyLane Lightning device.

    An extension of PennyLane's built-in ``default.qubit`` device that interfaces with C++ to
    perform fast linear algebra calculations.

    Use of this device requires pre-built binaries or compilation from source. Check out the
    :doc:`/installation` guide for more details.

    Args:
        wires (int): the number of wires to initialize the device with
        kernel_for_ops (dict): Optional argument which kernel to run for a gate operation.
            For example, if {'PauliX': 'LM', 'RX': 'PI'} is passed, the less memory (LM) kernel
            is used for PauliX whereas precomputed indices (PI) kernel is used for RX.
        shots (int): How many times the circuit should be evaluated (or sampled) to estimate
            the expectation values. Defaults to ``None`` if not specified. Setting
            to ``None`` results in computing statistics like expectation values and
            variances analytically.
    """

    name = "Lightning Qubit PennyLane plugin"
    short_name = "lightning.qubit"
    pennylane_requires = ">=0.15"
    version = __version__
    author = "Xanadu Inc."
    _CPP_BINARY_AVAILABLE = True
    operations = _remove_snapshot_from_operations(DefaultQubit.operations)

    def __init__(self, wires, *, kernel_for_ops=None, shots=None, batch_obs=False):
        self._kernel_for_ops = DEFAULT_KERNEL_FOR_OPS
        if kernel_for_ops is not None:
            if not isinstance(kernel_for_ops, dict):
                raise ValueError("Argument kernel_for_ops must be a dictionary.")

            for gate_op, kernel in kernel_for_ops.items():
                if (kernel, gate_op) not in EXPORTED_KERNEL_OPS:
                    raise ValueError(
                        f"The given kernel {kernel} does not implement {gate_op} gate."
                    )
                self._kernel_for_ops[gate_op] = kernel

        super().__init__(wires, shots=shots)
        self._batch_obs = batch_obs

    @classmethod
    def capabilities(cls):
        capabilities = super().capabilities().copy()
        capabilities.update(
            model="qubit",
            supports_reversible_diff=False,
            supports_inverse_operations=True,
            supports_analytic_computation=True,
            returns_state=True,
        )
        capabilities.pop("passthru_devices", None)
        return capabilities

    def apply(self, operations, rotations=None, **kwargs):
        # State preparation is currently done in Python
        if operations:  # make sure operations[0] exists
            if isinstance(operations[0], QubitStateVector):
                self._apply_state_vector(operations[0].parameters[0].copy(), operations[0].wires)
                del operations[0]
            elif isinstance(operations[0], BasisState):
                self._apply_basis_state(operations[0].parameters[0], operations[0].wires)
                del operations[0]

        for operation in operations:
            if isinstance(operation, (QubitStateVector, BasisState)):
                raise DeviceError(
                    "Operation {} cannot be used after other Operations have already been "
                    "applied on a {} device.".format(operation.name, self.short_name)
                )

        # Get the Type of self._state
        # as the reference type
        dtype = self._state.dtype

        if operations:
            self._pre_rotated_state = self.apply_lightning(self._state, operations, dtype=dtype)
        else:
            self._pre_rotated_state = self._state

        if rotations:
            if any(isinstance(r, QubitUnitary) for r in rotations):
                super().apply(operations=[], rotations=rotations)
            else:
                self._state = self.apply_lightning(
                    np.copy(self._pre_rotated_state), rotations, dtype=dtype
                )
        else:
            self._state = self._pre_rotated_state

    def apply_lightning(self, state, operations, dtype=np.complex128):
        """Apply a list of operations to the state tensor.

        Args:
            state (array[complex]): the input state tensor
            operations (list[~pennylane.operation.Operation]): operations to apply
            dtype (type): Type of numpy ``complex`` to be used. Can be important
            to specify for large systems for memory allocation purposes.

        Returns:
            array[complex]: the output state tensor
        """
        state_vector = np.ravel(state)

        if dtype == np.complex64:
            # use_csingle
            sim = StateVectorC64(state_vector)
        elif dtype == np.complex128:
            # self.C_DTYPE is np.complex128 by default
            sim = StateVectorC128(state_vector)
        else:
            raise TypeError(f"Unsupported complex Type: {dtype}")

        # Skip over identity operations instead of performing
        # matrix multiplication with the identity.
        skipped_ops = ["Identity"]

        for o in operations:
            if o.base_name in skipped_ops:
                continue
            name = o.name.split(".")[0]  # The split is because inverse gates have .inv appended
            if _is_lightning_gate(name):
                kernel = self._kernel_for_ops[name]
                method = getattr(sim, f"{name}_{kernel}".format(), None)
            else:
                method = None

            wires = self.wires.indices(o.wires)

            if method is None:
                # Inverse can be set to False since qml.matrix(o) is already in inverted form
                method = getattr(sim, "applyMatrix_{}".format(self._kernel_for_ops["Matrix"]))
                try:
                    method(qml.matrix(o), wires, False)
                except AttributeError:  # pragma: no cover
                    # To support older versions of PL
                    method(o.matrix, wires, False)
            else:
                inv = o.inverse
                param = o.parameters
                method(wires, inv, param)

        return np.reshape(state_vector, state.shape)

    def adjoint_diff_support_check(self, tape):
        """Check Lightning adjoint differentiation method support for a tape.

        Raise ``QuantumFunctionError`` if ``tape`` contains not supported measurements,
        observables, or operations by the Lightning adjoint differentiation method.

        Args:
            tape (.QuantumTape): quantum tape to differentiate
        """
        for m in tape.measurements:
            if m.return_type is not Expectation:
                raise QuantumFunctionError(
                    "Adjoint differentiation method does not support"
                    f" measurement {m.return_type.value}"
                )
            if not isinstance(m.obs, Tensor):
                if isinstance(m.obs, Projector):
                    raise QuantumFunctionError(
                        "Adjoint differentiation method does not support the Projector observable"
                    )
                if isinstance(m.obs, Hermitian):
                    raise QuantumFunctionError(
                        "Lightning adjoint differentiation method does not currently support the Hermitian observable"
                    )
            else:
                if any([isinstance(o, Projector) for o in m.obs.non_identity_obs]):
                    raise QuantumFunctionError(
                        "Adjoint differentiation method does not support the Projector observable"
                    )
                if any([isinstance(o, Hermitian) for o in m.obs.non_identity_obs]):
                    raise QuantumFunctionError(
                        "Lightning adjoint differentiation method does not currently support the Hermitian observable"
                    )

        for op in tape.operations:
            if (
                op.num_params > 1 and not isinstance(op, Rot)
            ) or op.name in UNSUPPORTED_PARAM_GATES_ADJOINT:
                raise QuantumFunctionError(
                    f"The {op.name} operation is not supported using "
                    'the "adjoint" differentiation method'
                )

    def adjoint_jacobian(self, tape, starting_state=None, use_device_state=False):
        if self.shots is not None:
            warn(
                "Requested adjoint differentiation to be computed with finite shots."
                " The derivative is always exact when using the adjoint differentiation method.",
                UserWarning,
            )

        # To support np.complex64 based on the type of self._state
        dtype = self._state.dtype
        if dtype == np.complex64:
            use_csingle = True
        elif dtype == np.complex128:
            use_csingle = False
        else:
            raise TypeError(f"Unsupported complex Type: {dtype}")

        if len(tape.trainable_params) == 0:
            return np.array(0)

        # Check adjoint diff support
        self.adjoint_diff_support_check(tape)

        # Initialization of state
        if starting_state is not None:
            ket = np.ravel(starting_state)
        else:
            if not use_device_state:
                self.reset()
                self.execute(tape)
            ket = np.ravel(self._pre_rotated_state)

        if use_csingle:
            adj = AdjointJacobianC64()
            ket = ket.astype(np.complex64)
        else:
            adj = AdjointJacobianC128()

        obs_serialized = _serialize_obs(tape, self.wire_map, use_csingle=use_csingle)
        ops_serialized, use_sp = _serialize_ops(tape, self.wire_map, use_csingle=use_csingle)

        ops_serialized = adj.create_ops_list(*ops_serialized)

        trainable_params = sorted(tape.trainable_params)
        first_elem = 1 if trainable_params[0] == 0 else 0

        tp_shift = (
            trainable_params if not use_sp else [i - 1 for i in trainable_params[first_elem:]]
        )  # exclude first index if explicitly setting sv

        state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)

        # If requested batching over observables, chunk into OMP_NUM_THREADS sized chunks.
        # This will allow use of Lightning with adjoint for large-qubit numbers AND large
        # numbers of observables, enabling choice between compute time and memory use.
        requested_threads = int(getenv("OMP_NUM_THREADS", "1"))

        if self._batch_obs and requested_threads > 1:
            obs_partitions = _chunk_iterable(obs_serialized, requested_threads)
            jac = []
            for obs_chunk in obs_partitions:
                jac_local = adj.adjoint_jacobian(
                    state_vector,
                    obs_chunk,
                    ops_serialized,
                    tp_shift,
                    tape.num_params,
                )
                jac.extend(jac_local)
            jac = np.array(jac)
        else:
            jac = adj.adjoint_jacobian(
                state_vector,
                obs_serialized,
                ops_serialized,
                tp_shift,
                tape.num_params,
            )
        return jac.reshape(-1, tape.num_params)

    def compute_vjp(self, dy, jac, num=None):
        """Convenience function to compute the vector-Jacobian product for a given
        vector of gradient outputs and a Jacobian.
        Args:
            dy (tensor_like): vector of gradient outputs
            jac (tensor_like): Jacobian matrix. For an n-dimensional ``dy``
                vector, the first n-dimensions of ``jac`` should match
                the shape of ``dy``.
        Keyword Args:
        num (int): The length of the flattened ``dy`` argument. This is an
            optional argument, but can be useful to provide if ``dy`` potentially
            has no shape (for example, due to tracing or just-in-time compilation).
        Returns:
            tensor_like: the vector-Jacobian product
        """
        if jac is None:
            return None

        if not isinstance(dy, np.ndarray) or not isinstance(jac, np.ndarray):
            return gradients.compute_vjp(dy, jac)

        dy_row = math.reshape(dy, [-1])

        if num is None:
            num = math.shape(dy_row)[0]

        jac = math.reshape(jac, [num, -1])
        num_params = jac.shape[1]

        if math.allclose(dy, 0):
            return math.convert_like(np.zeros([num_params]), dy)

        # To support np.complex64 based on the type of self._state
        dtype = self._state.dtype
        if dtype == np.complex64:
            VJP = VectorJacobianProductC64()
        elif dtype == np.complex128:
            VJP = VectorJacobianProductC128()
        else:
            raise TypeError(f"Unsupported complex Type: {dtype}")

        vjp_tensor = VJP.compute_vjp_from_jac(
            math.reshape(jac, [-1]),
            dy_row,
            num,
            num_params,
        )
        return vjp_tensor

    def vjp(self, tape, dy, starting_state=None, use_device_state=False):
        """Generate the processing function required to compute the vector-Jacobian products of a tape.
        Args:
            tape (.QuantumTape): quantum tape to differentiate
            dy (tensor_like): Gradient-output vector. Must have shape
                matching the output shape of the corresponding tape.
        Keyword Args:
            starting_state (tensor_like): post-forward pass state to start execution with. It should be
                complex-valued. Takes precedence over ``use_device_state``.
            use_device_state (bool): use current device state to initialize. A forward pass of the same
                circuit should be the last thing the device has executed. If a ``starting_state`` is
                provided, that takes precedence.
        Returns:
            The processing function required to compute the vector-Jacobian
            products of a tape.
        """
        if self.shots is not None:
            warn(
                "Requested adjoint differentiation to be computed with finite shots."
                " The derivative is always exact when using the adjoint differentiation method.",
                UserWarning,
            )

        num_params = len(tape.trainable_params)

        if num_params == 0:
            return lambda _: None

        if math.allclose(dy, 0):
            return lambda _: math.convert_like(np.zeros([num_params]), dy)

        # To support np.complex64 based on the type of self._state
        dtype = self._state.dtype
        if dtype == np.complex64:
            use_csingle = True
        elif dtype == np.complex128:
            use_csingle = False
        else:
            raise TypeError(f"Unsupported complex Type: {dtype}")

        V = VectorJacobianProductC64() if use_csingle else VectorJacobianProductC128()

        fn = V.vjp_fn(math.reshape(dy, [-1]), tape.num_params)

        def processing_fn(tape):
            # Check adjoint diff support
            self.adjoint_diff_support_check(tape)

            # Initialization of state
            if starting_state is not None:
                ket = np.ravel(starting_state)
            else:
                if not use_device_state:
                    self.reset()
                    self.execute(tape)
                ket = np.ravel(self._pre_rotated_state)

            if use_csingle:
                ket = ket.astype(np.complex64)

            obs_serialized = _serialize_obs(tape, self.wire_map, use_csingle=use_csingle)
            ops_serialized, use_sp = _serialize_ops(tape, self.wire_map, use_csingle=use_csingle)

            ops_serialized = V.create_ops_list(*ops_serialized)

            trainable_params = sorted(tape.trainable_params)
            first_elem = 1 if trainable_params[0] == 0 else 0

            tp_shift = (
                trainable_params if not use_sp else [i - 1 for i in trainable_params[first_elem:]]
            )  # exclude first index if explicitly setting sv

            state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)

            return fn(state_vector, obs_serialized, ops_serialized, tp_shift)

        return processing_fn

    def batch_vjp(
        self, tapes, dys, reduction="append", starting_state=None, use_device_state=False
    ):
        """Generate the processing function required to compute the vector-Jacobian products
        of a batch of tapes.
        Args:
            tapes (Sequence[.QuantumTape]): sequence of quantum tapes to differentiate
            dys (Sequence[tensor_like]): Sequence of gradient-output vectors ``dy``. Must be the
                same length as ``tapes``. Each ``dy`` tensor should have shape
                matching the output shape of the corresponding tape.
        Keyword Args:
            reduction (str): Determines how the vector-Jacobian products are returned.
                If ``append``, then the output of the function will be of the form
                ``List[tensor_like]``, with each element corresponding to the VJP of each
                input tape. If ``extend``, then the output VJPs will be concatenated.
            starting_state (tensor_like): post-forward pass state to start execution with. It should be
                complex-valued. Takes precedence over ``use_device_state``.
            use_device_state (bool): use current device state to initialize. A forward pass of the same
                circuit should be the last thing the device has executed. If a ``starting_state`` is
                provided, that takes precedence.
        Returns:
            The processing function required to compute the vector-Jacobian products of a batch of tapes.
        """
        fns = []

        # Loop through the tapes and dys vector
        for tape, dy in zip(tapes, dys):
            fn = self.vjp(
                tape, dy, starting_state=starting_state, use_device_state=use_device_state
            )
            fns.append(fn)

        def processing_fns(tapes):
            vjps = []
            for t, f in zip(tapes, fns):
                vjp = f(t)

                if vjp is None:
                    if reduction == "append":
                        vjps.append(None)
                    continue

                if isinstance(reduction, str):
                    getattr(vjps, reduction)(vjp)
                elif callable(reduction):
                    reduction(vjps, vjp)

            return vjps

        return processing_fns

    def probability(self, wires=None, shot_range=None, bin_size=None):
        """Return the probability of each computational basis state.

        Devices that require a finite number of shots always return the
        estimated probability.

        Args:
            wires (Iterable[Number, str], Number, str, Wires): wires to return
                marginal probabilities for. Wires not provided are traced out of the system.
            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
                to use. If not specified, all samples are used.
            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
                returns the measurement statistic separately over each bin. If not
                provided, the entire shot range is treated as a single bin.

        Returns:
            array[float]: list of the probabilities
        """
        if self.shots is not None:
            return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)

        wires = wires or self.wires
        wires = Wires(wires)

        # translate to wire labels used by device
        device_wires = self.map_wires(wires)

        # To support np.complex64 based on the type of self._state
        dtype = self._state.dtype
        if dtype == np.complex64:
            use_csingle = True
        elif dtype == np.complex128:
            use_csingle = False
        else:
            raise TypeError(f"Unsupported complex Type: {dtype}")

        # Initialization of state
        ket = np.ravel(self._state)

        if use_csingle:
            ket = ket.astype(np.complex64)

        state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)
        M = MeasuresC64(state_vector) if use_csingle else MeasuresC128(state_vector)

        return M.probs(device_wires)

    def generate_samples(self):
        """Generate samples

        Returns:
            array[int]: array of samples in binary representation with shape ``(dev.shots, dev.num_wires)``
        """

        # To support np.complex64 based on the type of self._state
        dtype = self._state.dtype
        if dtype == np.complex64:
            use_csingle = True
        elif dtype == np.complex128:
            use_csingle = False
        else:
            raise TypeError(f"Unsupported complex Type: {dtype}")

        # Initialization of state
        ket = np.ravel(self._state)

        if use_csingle:
            ket = ket.astype(np.complex64)

        state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)
        M = MeasuresC64(state_vector) if use_csingle else MeasuresC128(state_vector)

        return M.generate_samples(len(self.wires), self.shots).astype(int)

    def expval(self, observable, shot_range=None, bin_size=None):
        """Expectation value of the supplied observable.

        Args:
            observable: A PennyLane observable.
            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
                to use. If not specified, all samples are used.
            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
                returns the measurement statistic separately over each bin. If not
                provided, the entire shot range is treated as a single bin.

        Returns:
            Expectation value of the observable
        """
        if isinstance(observable.name, List) or observable.name in [
            "Identity",
            "Projector",
            "Hermitian",
            "Hamiltonian",
            "SparseHamiltonian",
        ]:
            return super().expval(observable, shot_range=shot_range, bin_size=bin_size)

        if self.shots is not None:
            # estimate the expectation value
            # LightningQubit doesn't support sampling yet
            samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
            return np.squeeze(np.mean(samples, axis=0))

        # To support np.complex64 based on the type of self._state
        dtype = self._state.dtype
        if dtype == np.complex64:
            use_csingle = True
        elif dtype == np.complex128:
            use_csingle = False
        else:
            raise TypeError(f"Unsupported complex Type: {dtype}")

        # Initialization of state
        ket = np.ravel(self._pre_rotated_state)

        if use_csingle:
            ket = ket.astype(np.complex64)

        state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)
        M = MeasuresC64(state_vector) if use_csingle else MeasuresC128(state_vector)

        # translate to wire labels used by device
        observable_wires = self.map_wires(observable.wires)

        return M.expval(observable.name, observable_wires)

    def var(self, observable, shot_range=None, bin_size=None):
        """Variance of the supplied observable.

        Args:
            observable: A PennyLane observable.
            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
                to use. If not specified, all samples are used.
            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
                returns the measurement statistic separately over each bin. If not
                provided, the entire shot range is treated as a single bin.

        Returns:
            Variance of the observable
        """
        if isinstance(observable.name, List) or observable.name in [
            "Identity",
            "Projector",
            "Hermitian",
        ]:
            return super().var(observable, shot_range=shot_range, bin_size=bin_size)

        if self.shots is not None:
            # estimate the var
            # LightningQubit doesn't support sampling yet
            samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
            return np.squeeze(np.var(samples, axis=0))

        # To support np.complex64 based on the type of self._state
        dtype = self._state.dtype
        if dtype == np.complex64:
            use_csingle = True
        elif dtype == np.complex128:
            use_csingle = False
        else:
            raise TypeError(f"Unsupported complex Type: {dtype}")

        # Initialization of state
        ket = np.ravel(self._pre_rotated_state)

        if use_csingle:
            ket = ket.astype(np.complex64)

        state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)
        M = MeasuresC64(state_vector) if use_csingle else MeasuresC128(state_vector)

        # translate to wire labels used by device
        observable_wires = self.map_wires(observable.wires)

        return M.var(observable.name, observable_wires)

if not CPP_BINARY_AVAILABLE:

    class LightningQubit(DefaultQubit):  # pragma: no cover
        name = "Lightning Qubit PennyLane plugin"
        short_name = "lightning.qubit"
        pennylane_requires = ">=0.15"
        version = __version__
        author = "Xanadu Inc."
        _CPP_BINARY_AVAILABLE = False
        operations = _remove_snapshot_from_operations(DefaultQubit.operations)

        def __init__(self, *args, **kwargs):
            warn(
                "Pre-compiled binaries for lightning.qubit are not available. Falling back to "
                "using the Python-based default.qubit implementation. To manually compile from "
                "source, follow the instructions at "
                "https://pennylane-lightning.readthedocs.io/en/latest/installation.html.",
                UserWarning,
            )
            super().__init__(*args, **kwargs)