animation.py

from matplotlib.animation import FFMpegWriter, FuncAnimation
from typing import Sequence, Mapping, Any
import matplotlib.pyplot as plt
import json
import matplotlib
import networkx as nx
import argparse
from abc import ABC, abstractmethod


# physics constants
R_B = 6  # rydberg range
AOD_SEP = 2  # min AOD separation
RYD_SEP = 15  # sufficient distance to avoid Rydberg
SITE_SLMS = 2  # number of SLMs in a site
SLM_SEP = AOD_SEP  # separation of SLMs inside a site
SITE_WIDTH = 4  # total width of SLMs in a site
X_SITE_SEP = RYD_SEP + SITE_WIDTH  # separation of sites in X direction
Y_SITE_SEP = RYD_SEP  # separation of sites in Y direction

# padding of the figure
X_LOW_PAD = 2 * AOD_SEP
Y_LOW_PAD = 4 * AOD_SEP
X_HIGH_PAD = 2 * AOD_SEP
Y_HIGH_PAD = 4 * AOD_SEP

# constants for animation
FPS = 24  # frames per second
INIT_FRM = 24  # initial empty frames
PT_MICRON = 8  # scaling factor: points per micron
MUS_PER_FRM = 8  # microseconds per frame
T_RYDBERG = 0.15  # microseconds for Rydberg
T_ACTIVATE = 50  # microseconds for (de)activating AOD


# class for physical entities: qubits and AOD rows/cols

class Qubit():
    def __init__(self, id: int):
        self.array = 'SLM'
        self.id = id
        self.c = -1  # AOD coloumn index
        self.r = -1  # AOD row index
        self.x = -X_LOW_PAD-1  # real X coordinates in um
        self.y = -Y_LOW_PAD-1  # real Y coordinates in um


class Row():
    def __init__(self, id: int):
        self.id = id
        self.active = False
        self.y = -Y_LOW_PAD-1  # real Y coordinates in um


class Col():
    def __init__(self, id: int):
        self.id = id
        self.active = False
        self.x = -X_LOW_PAD-1  # real X coordinates in um


class Inst(ABC):
    """abstract class of DPQA instructions.

    In general, the __init__ of specific instruction classes looks like
        def __init__(self, *):
            super().__init__(*)
            self.verify(*)
            self.operate(*)
            super().write_code(*)
    """

    def __init__(
            self,
            type: str,
            prefix = None,
            stage: int = -1,
            reduced_keys: Sequence[str] = [],
    ):
        """init method for instructions.

        Args:
            type (str): 
            prefix (str | None, optional): provide the big operation.
                this Inst belongs to. Defaults to None.
            stage (int, optional): stage the Inst belongs to. Defaults to -1.
            reduced_keys (Sequence[str], optional): data to keep in emit()
                from emit_full(). Defaults to [].
        """
        self.type = type
        self.name = prefix + ':' + type if prefix else type
        self.stage = stage
        self.reduced_keys = reduced_keys + ['type', ]
        self.duration = -1
        self.code = {'type': self.type, 'name': self.name, }

    def write_code(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            data: Mapping[str, Any],
    ):
        """write self.code with provided info.

        Args:
            col_objs (Sequence[Col]): Col objects used.
            row_objs (Sequence[Row]): Row objects used.
            qubit_objs (Sequence[Qubit]): Qubit objects used.
            data (Mapping[str, Any]): other info for the Inst.
        """
        for k, v in data.items():
            self.code[k] = v

        # get the current state of DPQA
        curr = {}
        curr['qubits'] = [
            {
                'id': q.id,
                'x': q.x,
                'y': q.y,
                'array': q.array,
                'c': q.c,
                'r': q.r
            } for q in qubit_objs
        ]
        curr['cols'] = [
            {
                'id': c.id,
                'active': c.active,
                'x': c.x
            } for c in col_objs
        ]
        curr['rows'] = [
            {
                'id': r.id,
                'active': r.active,
                'y': r.y
            } for r in row_objs
        ]
        self.code['state'] = curr

    @abstractmethod
    def verify(self):
        """verification of instructions. This is abstract because we require
        each child class to provide its own verification method.
        """
        pass

    def operate(self):
        """perform operation of instructions on Col, Row, and Qubit objects."""
        pass

    def emit(self) -> Sequence[Mapping[str, Any]]:
        """emit the minimum code for executing instructions.

        Returns:
            Sequence[Mapping[str, Any]]: code in a dict.
        """
        return ({k: self.code[k] for k in self.reduced_keys}, )

    def emit_full(self) -> Sequence[Mapping[str, Any]]:
        """emit the code with full info for any purpose.

        Returns:
            Sequence[Mapping[str, Any]]: code with full info in a dict.
        """
        return (self.code, )

    def is_trivial(self) -> bool:
        return True if self.duration == 0 else False

    def remove_trivial_insts(self):
        # this is used in ComboInst. Added here for convience.
        pass


# classes for basic instructions: Init, Move, Activate, Deactivate, Rydberg
# todo: add class Raman (single-qubit gates)


class Init(Inst):
    def __init__(self,
                 col_objs: Sequence[Col],
                 row_objs: Sequence[Row],
                 qubit_objs: Sequence[Qubit],
                 slm_qubit_idx: Sequence[int] = [],
                 slm_qubit_xys: Sequence[Sequence[int]] = [],
                 aod_qubit_idx: Sequence[int] = [],
                 aod_qubit_crs: Sequence[Sequence[int]] = [],
                 aod_col_act_idx: Sequence[int] = [],
                 aod_col_xs: Sequence[int] = [],
                 aod_row_act_idx: Sequence[int] = [],
                 aod_row_ys: Sequence[int] = [],
                 data: Mapping[str, Any] = {},):
        super().__init__(
            'Init',
            reduced_keys=[
                'slm_qubit_idx', 'slm_qubit_xys', 'aod_qubit_idx',
                'aod_qubit_crs', 'aod_col_act_idx', 'aod_col_xs',
                'aod_row_act_idx', 'aod_row_ys', 'n_q',
                'x_high', 'y_high', 'c_high', 'r_high'
            ]
        )
        for k, v in data.items():
            self.code[k] = v
        self.all_slms = []
        self.verify(slm_qubit_idx,
                    slm_qubit_xys,
                    aod_qubit_idx,
                    aod_qubit_crs,
                    aod_col_act_idx,
                    aod_col_xs,
                    aod_row_act_idx,
                    aod_row_ys,)
        self.operate(col_objs,
                     row_objs,
                     qubit_objs,
                     slm_qubit_idx,
                     slm_qubit_xys,
                     aod_qubit_idx,
                     aod_qubit_crs,
                     aod_col_act_idx,
                     aod_col_xs,
                     aod_row_act_idx,
                     aod_row_ys,)
        super().write_code(
            col_objs,
            row_objs,
            qubit_objs,
            {
                'duration': INIT_FRM,
                'slm_qubit_idx': slm_qubit_idx,
                'slm_qubit_xys': slm_qubit_xys,
                'aod_qubit_idx': aod_qubit_idx,
                'aod_qubit_crs': aod_qubit_crs,
                'aod_col_act_idx': aod_col_act_idx,
                'aod_col_xs': aod_col_xs,
                'aod_row_act_idx': aod_row_act_idx,
                'aod_row_ys': aod_row_ys,
            })

    def add_slms(self, slms):
        for slm in slms:
            if slm not in self.all_slms:
                self.all_slms.append(slm)

    def verify(
            self,
            slm_qubit_idx: Sequence[int],
            slm_qubit_xys: Sequence[Sequence[int]],
            aod_qubit_idx: Sequence[int],
            aod_qubit_crs: Sequence[Sequence[int]],
            aod_col_act_idx: Sequence[int],
            aod_col_xs: Sequence[int],
            aod_row_act_idx: Sequence[int],
            aod_row_ys: Sequence[int],
    ):
        a = len(slm_qubit_idx)
        b = len(slm_qubit_xys)
        if a != b:
            raise ValueError(
                f'{self.name}: SLM qubit arguments invalid {a} idx, {b} xys.')
        for i in slm_qubit_xys:
            if len(i) != 2:
                raise ValueError(f'{self.name}: SLM qubit xys {i} invalid.')
        for i in range(len(slm_qubit_xys)):
            for j in range(i+1, len(slm_qubit_xys)):
                if slm_qubit_xys[i] == slm_qubit_xys[j]:
                    raise ValueError(
                        f'{self.name}: SLM qubits {slm_qubit_idx[i]} '
                        f'and {slm_qubit_idx[j]} xys are the same.'
                    )
        # todo: the case when not all atoms are in SLM

    def operate(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            slm_qubit_idx: Sequence[int],
            slm_qubit_xys: Sequence[Sequence[int]],
            aod_qubit_idx: Sequence[int],
            aod_qubit_crs: Sequence[Sequence[int]],
            aod_col_act_idx: Sequence[int],
            aod_col_xs: Sequence[int],
            aod_row_act_idx: Sequence[int],
            aod_row_ys: Sequence[int],
    ):
        for i, q_id in enumerate(slm_qubit_idx):
            qubit_objs[q_id].array = 'SLM'
            (qubit_objs[q_id].x, qubit_objs[q_id].y) = slm_qubit_xys[i]
            self.all_slms.append(slm_qubit_xys[i])

        for i, col_id in enumerate(aod_col_act_idx):
            col_objs[col_id].activate = True
            col_objs[col_id].x = aod_col_xs[i]

        for i, row_id in enumerate(aod_col_act_idx):
            row_objs[row_id].activate = True
            row_objs[row_id].y = aod_row_ys[i]

        for i, q_id in enumerate(aod_qubit_idx):
            qubit_objs[q_id].array = 'AOD'
            (qubit_objs[q_id].c, qubit_objs[q_id].r) = aod_qubit_crs[i]
            qubit_objs[q_id].x = col_objs[qubit_objs[q_id].c].x
            qubit_objs[q_id].y = row_objs[qubit_objs[q_id].r].y

    def emit_full(self):
        # all the used SLMs are counted during the whole codegen process,
        # so the emit_full of Init needs to add this info
        self.code['all_slms'] = self.all_slms
        return super().emit_full()


class Move(Inst):
    def __init__(self,
                 s: int,
                 col_objs: Sequence[Col],
                 row_objs: Sequence[Row],
                 qubit_objs: Sequence[Qubit],
                 col_idx: Sequence[int] = [],
                 col_begin: Sequence[int] = [],
                 col_end: Sequence[int] = [],
                 row_idx: Sequence[int] = [],
                 row_begin: Sequence[int] = [],
                 row_end: Sequence[int] = [],
                 prefix: str = ''):
        super().__init__('Move', prefix=prefix, stage=s)
        self.verify(
            col_objs,
            row_objs,
            col_idx,
            col_begin,
            col_end,
            row_idx,
            row_begin,
            row_end,
        )
        data = self.operate(
            col_objs,
            row_objs,
            qubit_objs,
            col_idx,
            col_begin,
            col_end,
            row_idx,
            row_begin,
            row_end,)
        super().write_code(col_objs, row_objs, qubit_objs, data)

    def operate(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            col_idx: Sequence[int],
            col_begin: Sequence[int],
            col_end: Sequence[int],
            row_idx: Sequence[int],
            row_begin: Sequence[int],
            row_end: Sequence[int],
    ) -> Mapping[str, Any]:
        data = {}
        # calculate the max  move distance of columns
        data['cols'] = []
        max_distance = 0
        for i in range(len(col_idx)):
            distance = abs(col_end[i]-col_begin[i])
            if distance > 0:
                data['cols'].append(
                    {'id': col_idx[i],
                     'shift': col_end[i]-col_begin[i],
                     'begin': col_begin[i],
                     'end': col_end[i]})
                col_objs[col_idx[i]].x = col_end[i]
                max_distance = max(max_distance, distance)

        # calculate the max  move distance of rows
        data['rows'] = []
        for i in range(len(row_idx)):
            distance = abs(row_end[i]-row_begin[i])
            if distance > 0:
                data['rows'].append(
                    {'id': row_idx[i],
                     'shift': row_end[i]-row_begin[i],
                     'begin': row_begin[i],
                     'end': row_end[i]})
                row_objs[row_idx[i]].y = row_end[i]
                max_distance = max(max_distance, distance)

        # movement time per Bluvstein et al. units are us and um.
        self.duration = 200*((max_distance/110)**(1/2))
        data['duration'] = self.duration

        for qubit_obj in qubit_objs:
            if qubit_obj.array == 'AOD':
                qubit_obj.x = col_objs[qubit_obj.c].x
                qubit_obj.y = row_objs[qubit_obj.r].y

        return data

    def verify(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            col_idx: Sequence[int],
            col_begin: Sequence[int],
            col_end: Sequence[int],
            row_idx: Sequence[int],
            row_begin: Sequence[int],
            row_end: Sequence[int],
    ):
        a = len(col_idx)
        b = len(col_begin)
        c = len(col_end)
        if not (a == b and a == c):
            raise ValueError(
                f'{self.name}: col arguments invalid'
                f' {a} idx, {b} begin, {c} end.'
            )
        a = len(row_idx)
        b = len(row_begin)
        c = len(row_end)
        if not (a == b and a == c):
            raise ValueError(
                f'{self.name}: row arguments invalid'
                f' {a} idx, {b} begin, {c} end.'
            )

        activated_col_idx = []
        activated_col_xs = []
        for col_obj in col_objs:
            if col_obj.active:
                if (activated_col_idx
                        and col_obj.x < activated_col_xs[-1] + AOD_SEP):
                    raise ValueError(
                        f'{self.name}: col beginning position invalid'
                        f' col {col_obj.id} at x={col_obj.x} while '
                        f'col {activated_col_idx[-1]} at'
                        f' x={activated_col_xs[-1]}.'
                    )
                activated_col_idx.append(col_obj.id)
                activated_col_xs.append(col_obj.x)
        for i, moving_col_id in enumerate(col_idx):
            if moving_col_id not in activated_col_idx:
                raise ValueError(
                    f'{self.name}: col {moving_col_id} to move'
                    f' is not activated.'
                )
            j = activated_col_idx.index(moving_col_id)
            if col_begin[i] != activated_col_xs[j]:
                raise ValueError(
                    f'{self.name}: col {moving_col_id} beginning x not agree.')
            activated_col_xs[j] = col_end[i]
        for i in range(1, len(activated_col_xs)):
            if activated_col_xs[i - 1] + AOD_SEP > activated_col_xs[i]:
                raise ValueError(
                    f'{self.name}: col ending position invalid'
                    f' col {activated_col_idx[i-1]} at '
                    f'x={activated_col_xs[i-1]} while '
                    f'col {activated_col_idx[i]} at x={activated_col_xs[i]}.')

        activated_row_idx = []
        activated_row_ys = []
        for row_obj in row_objs:
            if row_obj.active:
                if (activated_row_idx
                        and row_obj.y < activated_row_ys[-1] + AOD_SEP):
                    raise ValueError(
                        f'{self.name}: row beginning position invalid '
                        f'row {row_obj.id} at y={row_obj.y} while '
                        f'row {activated_row_idx[-1]} at '
                        f'y={activated_row_ys[-1]}.'
                    )
                activated_row_idx.append(row_obj.id)
                activated_row_ys.append(row_obj.y)
        for i, moving_row_id in enumerate(row_idx):
            if moving_row_id not in activated_row_idx:
                raise ValueError(
                    f'{self.name}: row {moving_row_id} to move '
                    f'is not activated.'
                )
            j = activated_row_idx.index(moving_row_id)
            if row_begin[i] != activated_row_ys[j]:
                raise ValueError(
                    f'{self.name}: row {moving_row_id} beginning y not agree.')
            activated_row_ys[j] = row_end[i]
        for i in range(1, len(activated_row_ys)):
            if activated_row_ys[i - 1] + AOD_SEP > activated_row_ys[i]:
                raise ValueError(
                    f'{self.name}: row ending position invalid '
                    f'row {activated_row_idx[i-1]} at '
                    f'y={activated_row_ys[i-1]} while '
                    f'row {activated_row_idx[i]} at y={activated_row_ys[i]}.')

    def emit(self):
        code = {'type': self.type}
        code['cols'] = [{k: v for k, v in tmp.items() if k in [
            'id', 'shift']} for tmp in self.code['cols']]
        code['rows'] = [{k: v for k, v in tmp.items() if k in [
            'id', 'shift']} for tmp in self.code['rows']]
        return (code,)


class Activate(Inst):
    def __init__(self,
                 s: int,
                 col_objs: Sequence[Col],
                 row_objs: Sequence[Row],
                 qubit_objs: Sequence[Qubit],
                 col_idx: Sequence[int] = [],
                 col_xs: Sequence[int] = [],
                 row_idx: Sequence[int] = [],
                 row_ys: Sequence[int] = [],
                 pickup_qs: Sequence[int] = [],
                 prefix: str = '',):
        super().__init__(
            'Activate',
            prefix=prefix,
            stage=s,
            reduced_keys=['col_idx', 'col_xs', 'row_idx', 'row_ys']
        )
        self.verify(
            col_objs,
            row_objs,
            qubit_objs,
            col_idx,
            col_xs,
            row_idx,
            row_ys,
            pickup_qs,
        )
        self.operate(
            col_objs,
            row_objs,
            qubit_objs,
            col_idx,
            col_xs,
            row_idx,
            row_ys,
            pickup_qs,
        )
        super().write_code(
            col_objs,
            row_objs,
            qubit_objs,
            {
                'col_idx': col_idx,
                'col_xs': col_xs,
                'row_idx': row_idx,
                'row_ys': row_ys,
                'pickup_qs': pickup_qs,
                'duration': T_ACTIVATE
            }
        )

    def operate(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            col_idx: Sequence[int],
            col_xs: Sequence[int],
            row_idx: Sequence[int],
            row_ys: Sequence[int],
            pickup_qs: Sequence[int],
    ):
        for i in range(len(col_idx)):
            col_objs[col_idx[i]].active = True
            col_objs[col_idx[i]].x = col_xs[i]
        for i in range(len(row_idx)):
            row_objs[row_idx[i]].active = True
            row_objs[row_idx[i]].y = row_ys[i]
        for q_id in pickup_qs:
            qubit_objs[q_id].array = 'AOD'

    def verify(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            col_idx: Sequence[int],
            col_xs: Sequence[int],
            row_idx: Sequence[int],
            row_ys: Sequence[int],
            pickup_qs: Sequence[int],
    ):
        a = len(col_idx)
        b = len(col_xs)
        if a != b:
            raise ValueError(
                f'{self.name}: col arguments invalid {a} idx, {b} xs')
        a = len(row_idx)
        b = len(row_ys)
        if a != b:
            raise ValueError(
                f'f{self.name}: row arguments invalid {a} idx, {b} ys.')

        for i in range(len(col_idx)):
            if col_objs[col_idx[i]].active:
                raise ValueError(
                    f'{self.name}: col {col_idx[i]} already activated.')
            for j in range(col_idx[i]):
                if (col_objs[j].active
                        and col_objs[j].x > col_xs[i] - AOD_SEP):
                    raise ValueError(
                        f'{self.name}: col {j} at x={col_objs[j].x} is '
                        f'too left for col {col_idx[i]}'
                        f' to activate at x={col_xs[i]}.'
                    )
            for j in range(col_idx[i] + 1, len(col_objs)):
                if (col_objs[j].active
                        and col_objs[j].x - AOD_SEP < col_xs[i]):
                    raise ValueError(
                        f'{self.name}: col {j} at x={col_objs[j].x} is '
                        f'too right for col {col_idx[i]} '
                        f'to activate at x={col_xs[i]}.'
                    )
        for i in range(len(row_idx)):
            if row_objs[row_idx[i]].active:
                raise ValueError(
                    f'{self.name}: row {row_idx[i]} already activated.')
            for j in range(row_idx[i]):
                if (row_objs[j].active
                        and row_objs[j].y > row_ys[i] - AOD_SEP):
                    raise ValueError(
                        f'{self.name}: row {j} at y={row_objs[j].y} is '
                        f'too high for row {row_idx[i]} '
                        f'to activate at y={row_ys[i]}.'
                    )
            for j in range(row_idx[i] + 1, len(row_objs)):
                if (row_objs[j].active
                        and row_objs[j].y-AOD_SEP < row_ys[i]):
                    raise ValueError(
                        f'{self.name}: row {j} at y={col_objs[j].y} is '
                        f'too low for row {row_idx[i]} '
                        f'to activate at y={row_ys[i]}.'
                    )

        active_xys = []  # the traps that are newly activted by this Activate
        active_xs = [col.x for col in col_objs if col.active]
        active_ys = [row.y for row in row_objs if row.active]
        for x in active_xs:
            for y in row_ys:
                active_xys.append((x, y))
        for y in active_ys:
            for x in col_xs:
                active_xys.append((x, y))
        for x in col_xs:
            for y in row_ys:
                active_xys.append((x, y))

        for q_id in range(len(qubit_objs)):
            if q_id in pickup_qs:
                if (qubit_objs[q_id].x, qubit_objs[q_id].y) not in active_xys:
                    raise ValueError(
                        f'{self.name}: q {q_id} not picked up '
                        f'by col {qubit_objs[q_id].c} '
                        f'row {qubit_objs[q_id].r} at '
                        f'x={qubit_objs[q_id].x} y={qubit_objs[q_id].y}.'
                    )
            else:
                if (qubit_objs[q_id].x, qubit_objs[q_id].y) in active_xys:
                    raise ValueError(
                        f'{self.name}: q {q_id} wrongfully picked up by '
                        f'col {qubit_objs[q_id].c} row {qubit_objs[q_id].r}'
                        f' at x={qubit_objs[q_id].x} y={qubit_objs[q_id].y}.')


class Deactivate(Inst):
    def __init__(self,
                 s: int,
                 col_objs: Sequence[Col],
                 row_objs: Sequence[Row],
                 qubit_objs: Sequence[Qubit],
                 col_idx: Sequence[int] = [],
                 col_xs: Sequence[int] = [],
                 row_idx: Sequence[int] = [],
                 row_ys: Sequence[int] = [],
                 dropoff_qs: Sequence[int] = [],
                 prefix: str = '',):
        super().__init__(
            'Deactivate',
            prefix=prefix,
            stage=s,
            reduced_keys=['col_idx', 'row_idx']
        )
        self.verify(
            col_objs,
            row_objs,
            qubit_objs,
            col_idx,
            col_xs,
            row_idx,
            row_ys,
            dropoff_qs
        )
        self.operate(
            col_objs,
            row_objs,
            qubit_objs,
            col_idx,
            row_idx,
            dropoff_qs
        )
        super().write_code(
            col_objs,
            row_objs,
            qubit_objs,
            {
                'col_idx': col_idx,
                'col_xs': col_xs,
                'row_idx': row_idx,
                'row_ys': row_ys,
                'dropoff_qs': dropoff_qs,
                'duration': T_ACTIVATE
            }
        )

    def operate(
        self,
        col_objs: Sequence[Col],
        row_objs: Sequence[Row],
        qubit_objs: Sequence[Qubit],
        col_idx: Sequence[int],
        row_idx: Sequence[int],
        dropoff_qs: Sequence[int],
    ):
        for i in range(len(col_idx)):
            col_objs[col_idx[i]].active = False
        for i in range(len(row_idx)):
            row_objs[row_idx[i]].active = False
        for q_id in dropoff_qs:
            qubit_objs[q_id].array = 'SLM'

    def verify(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            col_idx: Sequence[int],
            col_xs: Sequence[int],
            row_idx: Sequence[int],
            row_ys: Sequence[int],
            dropoff_qs: Sequence[int],
    ):
        a = len(col_idx)
        b = len(col_xs)
        if a != b:
            raise ValueError(
                f'{self.name}: col arguments invalid {a} idx, {b} xs')
        a = len(row_idx)
        b = len(row_ys)
        if a != b:
            raise ValueError(
                f'{self.name}: row arguments invalid {a} idx, {b} ys.')

        for i in range(len(col_idx)):
            if not col_objs[col_idx[i]].active:
                raise ValueError(
                    f'{self.name}: col {col_idx[i]} already dectivated.')
            for j in range(col_idx[i]):
                if (col_objs[j].active
                        and col_objs[j].x > col_xs[i] - AOD_SEP):
                    raise ValueError(
                        f'{self.name}: col {j} at x={col_objs[j].x} is '
                        f'too left for col {col_idx[i]} '
                        f'to deactivate at x={col_xs[i]}.')
            for j in range(col_idx[i]+1, len(col_objs)):
                if (col_objs[j].active
                        and col_objs[j].x - AOD_SEP < col_xs[i]):
                    raise ValueError(
                        f'{self.name}: col {j} at x={col_objs[j].x} is '
                        f'too right for col {col_idx[i]} '
                        f'to deactivate at x={col_xs[i]}.')
        for i in range(len(row_idx)):
            if not row_objs[row_idx[i]].active:
                raise ValueError(
                    f'{self.name}: row {row_idx[i]} already deactivated.')
            for j in range(row_idx[i]):
                if (row_objs[j].active
                        and row_objs[j].y > row_ys[i] - AOD_SEP):
                    raise ValueError(
                        f'{self.name}: row {j} at y={row_objs[j].y} is '
                        f'too high for row {row_idx[i]} '
                        f'to deactivate at y={row_ys[i]}.'
                    )
            for j in range(row_idx[i]+1, len(row_objs)):
                if (row_objs[j].active
                        and row_objs[j].y-AOD_SEP < row_ys[i]):
                    raise ValueError(
                        f'{self.name}: row {j} at y={col_objs[j].y} is '
                        f'too low for row {row_idx[i]} '
                        f'to deactivate at y={row_ys[i]}.'
                    )

        deactive_xys = []
        active_xs = [col.x for col in col_objs if col.active]
        for x in active_xs:
            for y in row_ys:
                deactive_xys.append((x, y))

        for q_id in range(len(qubit_objs)):
            if q_id in dropoff_qs:
                if (qubit_objs[q_id].x, qubit_objs[q_id].y) not in deactive_xys:
                    raise ValueError(
                        f'{self.name}: q {q_id} not dropped off from '
                        f'col {qubit_objs[q_id].c} row {qubit_objs[q_id].r} '
                        f'at x={qubit_objs[q_id].x} y={qubit_objs[q_id].y}.'
                    )
            elif qubit_objs[q_id].array == 'AOD':
                if (qubit_objs[q_id].x, qubit_objs[q_id].y) in deactive_xys:
                    raise ValueError(
                        f'{self.name}: q {q_id} wrongfully dropped off from '
                        f'col {qubit_objs[q_id].c} row {qubit_objs[q_id].r} '
                        f'at x={qubit_objs[q_id].x} y={qubit_objs[q_id].y}.'
                    )


class Rydberg(Inst):
    def __init__(self,
                 s: int,
                 col_objs: Sequence[Col],
                 row_objs: Sequence[Row],
                 qubit_objs: Sequence[Qubit],
                 gates: Sequence[Mapping[str, int]]):
        super().__init__('Rydberg', prefix=f'Rydberg_{s}', stage=s)
        self.verify(gates, qubit_objs)
        super().write_code(
            col_objs,
            row_objs,
            qubit_objs,
            {'gates': gates, 'duration': T_RYDBERG, }
        )

    def verify(
            self,
            gates: Sequence[Mapping[str, int]],
            qubit_objs: Sequence[Qubit]):
        # for g in gates:
        #     q0 = {'id': qubit_objs[g['q0']].id,
        #           'x': qubit_objs[g['q0']].x,
        #           'y': qubit_objs[g['q0']].y}
        #     q1 = {'id': qubit_objs[g['q1']].id,
        #           'x': qubit_objs[g['q1']].x,
        #           'y': qubit_objs[g['q1']].y}
        #     if (q0['x']-q1['x'])**2 + (q0['y']-q1['y'])**2 > R_B**2:
        #         raise ValueError(
        #             f"{self.name}: q{q0['id']} at x={q0['x']} y={q0['y']} "
        #             f"and q{q1['id']} at x={q1['x']} y={q1['y']} "
        #             f"are farther away than Rydberg range."
        #         )
        return


# class for big ops: ReloadRow, Reload, OffloadRow, Offload, SwapPair, Swap
# internally, these are lists of basic operations.

# todo: is there verification needed on the ComboInst lebvel?
class ComboInst:
    pass


class ComboInst():
    """class for combined instructions which is a sequence of combined
    instructions or DPQA instructions.
    """

    def __init__(
            self,
            type: str,
            prefix = None,
            suffix = None,
            stage: int = -1,
    ):
        """init method for combined instructions.

        Args:
            type (str): 
            prefix (str | None, optional): Defaults to None.
            suffix (str | None, optional): Defaults to None.
            stage (int, optional): Defaults to -1.
        """
        self.type = type
        self.name = (f'{prefix}:' if prefix else '') + \
            type + (f'_{suffix}' if suffix else '')
        self.stage = stage
        self.duration = -1
        self.insts = []

    def emit(self) -> Sequence[Mapping[str, Any]]:
        """combine the code of each Inst inside this ComboInst and return."""
        code = []
        for inst in self.insts:
            code += inst.emit()
        return code

    def emit_full(self) -> Sequence[Mapping[str, Any]]:
        code = []
        for inst in self.insts:
            code += inst.emit_full()
        return code

    def append_inst(self, inst):
        self.insts.append(inst)

    def prepend_inst(self, inst):
        self.insts.insert(0, inst)

    def is_trivial(self) -> bool:
        for inst in self.insts:
            if not inst.is_trivial():
                return False
        return True

    def remove_trivial_insts(self):
        nontrivial_insts = []
        for inst in self.insts:
            inst.remove_trivial_insts()
            if not inst.is_trivial():
                nontrivial_insts.append(inst)
        self.insts = nontrivial_insts


class ReloadRow(ComboInst):
    def __init__(self, s: int, r: int, prefix: str = ''):
        super().__init__('ReloadRow', prefix=prefix, suffix=str(r), stage=s)
        self.row_id = r
        self.moving_cols_id = []
        self.moving_cols_begin = []
        self.moving_cols_end = []

    def add_col_shift(self, id: int, begin: int, end: int):
        self.moving_cols_id.append(id)
        self.moving_cols_begin.append(begin)
        self.moving_cols_end.append(end)

    def generate_col_shift(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
    ):
        self.insts.append(
            Move(s=self.stage,
                 col_objs=col_objs,
                 row_objs=row_objs,
                 qubit_objs=qubit_objs,
                 col_idx=self.moving_cols_id,
                 col_begin=self.moving_cols_begin,
                 col_end=self.moving_cols_end,
                 row_idx=[],
                 row_begin=[],
                 row_end=[],
                 prefix=self.name + ':ColShift')
        )

    def generate_row_activate(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            cols: Sequence[int],
            xs: Sequence[int],
            y: int,
            pickup_qs: Sequence[int],
    ):
        self.insts.append(
            Activate(s=self.stage,
                     col_objs=col_objs,
                     row_objs=row_objs,
                     qubit_objs=qubit_objs,
                     col_idx=cols,
                     col_xs=xs,
                     row_idx=[self.row_id, ],
                     row_ys=[y, ],
                     pickup_qs=pickup_qs,
                     prefix=self.name)
        )

    def generate_parking(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            shift_down: int,
            col_idx: Sequence[int] = [],
            col_begin: Sequence[int] = [],
            col_end: Sequence[int] = []
    ):
        self.insts.append(
            Move(s=self.stage,
                 col_objs=col_objs,
                 row_objs=row_objs,
                 qubit_objs=qubit_objs,
                 col_idx=col_idx,
                 col_begin=col_begin,
                 col_end=col_end,
                 row_idx=[self.row_id, ],
                 row_begin=[row_objs[self.row_id].y],
                 row_end=[row_objs[self.row_id].y - shift_down],
                 prefix=self.name + ':Parking')
        )


class Reload(ComboInst):
    def __init__(self, s: int):
        super().__init__('Reload', suffix=str(s), stage=s)

    def add_row_reload(self, r: int):
        self.insts.append(ReloadRow(self.stage, r, prefix=self.name))
        return self.insts[-1]


class OffloadRow(ComboInst):
    def __init__(self, s: int, r: int, prefix: str = ''):
        super().__init__('OffloadRow', prefix=prefix, suffix=str(r), stage=s)
        self.row_id = r
        self.moving_cols_id = []
        self.moving_cols_begin = []
        self.moving_cols_end = []

    def add_col_shift(self, id: int, begin: int, end: int):
        self.moving_cols_id.append(id)
        self.moving_cols_begin.append(begin)
        self.moving_cols_end.append(end)

    def generate_col_shift(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
    ):
        self.insts.append(
            Move(s=self.stage,
                 col_objs=col_objs,
                 row_objs=row_objs,
                 qubit_objs=qubit_objs,
                 col_idx=self.moving_cols_id,
                 col_begin=self.moving_cols_begin,
                 col_end=self.moving_cols_end,
                 row_idx=[],
                 row_begin=[],
                 row_end=[],
                 prefix=self.name + ':ColShift')
        )

    def generate_row_shift(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            site_y: int,
    ):
        self.insts.append(
            Move(s=self.stage,
                 col_objs=col_objs,
                 row_objs=row_objs,
                 qubit_objs=qubit_objs,
                 col_idx=[],
                 col_begin=[],
                 col_end=[],
                 row_idx=[self.row_id, ],
                 row_begin=[row_objs[self.row_id].y],
                 row_end=[site_y*Y_SITE_SEP, ],
                 prefix=self.name + ':RowDownShift')
        )

    def generate_row_deactivate(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            dropoff_qs: Sequence[int],
    ):
        self.insts.append(
            Deactivate(s=self.stage,
                       col_objs=col_objs,
                       row_objs=row_objs,
                       qubit_objs=qubit_objs,
                       col_idx=[],
                       col_xs=[],
                       row_idx=[self.row_id, ],
                       row_ys=[row_objs[self.row_id].y, ],
                       dropoff_qs=dropoff_qs,
                       prefix=self.name)
        )


class Offload(ComboInst):
    def __init__(self, s: int):
        super().__init__('Offload', suffix=str(s), stage=s)

    def add_row_offload(self, r: int):
        self.insts.append(OffloadRow(self.stage, r, prefix=self.name))
        return self.insts[-1]

    def all_cols_deactivate(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
    ):
        self.insts.append(
            Deactivate(
                s=self.stage,
                col_objs=col_objs,
                row_objs=row_objs,
                qubit_objs=qubit_objs,
                col_idx=[c.id for c in col_objs if c.active],
                col_xs=[c.x for c in col_objs if c.active],
                row_idx=[],
                row_ys=[],
                dropoff_qs=[],
                prefix=self.name
            )
        )


class SwapPair(ComboInst):
    """swap a pair of atoms A and B at the same interaction site
       O stands for an empty trap

            A  B

        step1: activate row 0 and col 0 on A to pick it up
            |
         ---A--B---
            |

        step2: move A to below B
               |
            O  B
               |
           ----A-----
               |

        step3: activate row 1 to pick up B
               |
         ---O--B---
               |
         ------A---
               |

        step4: move A and B to the place A used to be
            |
        ----B--O--
            |
        ----A-----
            |

        step5: deactivate row 1 to drop off B
            |
            B  O
            |
        ----A-----
            |

        step6: move A to the place B used to be
               |
         ---B--A---
               |

        step7: deactivate row 0 and col 0 to drop off A

            B  A
    """

    def __init__(self,
                 s: int,
                 col_objs: Sequence[Col],
                 row_objs: Sequence[Row],
                 qubit_objs: Sequence[Qubit],
                 left_q_id: int,
                 right_q_id: int,
                 prefix: str = ''
                 ):
        super().__init__(
            'SwapPair',
            prefix=prefix,
            suffix=f'({left_q_id},{right_q_id})',
            stage=s
        )

        # precondition
        left_x = qubit_objs[left_q_id].x
        right_x = qubit_objs[right_q_id].x
        y = qubit_objs[left_q_id].y
        qubit_objs[left_q_id].c = 0
        qubit_objs[left_q_id].r = 0
        qubit_objs[right_q_id].c = 0
        qubit_objs[right_q_id].r = 1

        self.insts = [
            Activate(s,
                     col_objs=col_objs,
                     row_objs=row_objs,
                     qubit_objs=qubit_objs,
                     col_idx=[0, ],
                     col_xs=[left_x, ],
                     row_idx=[0, ],
                     row_ys=[y, ],
                     pickup_qs=[left_q_id, ],
                     prefix=self.name + f':PickUp_q{left_q_id}'),
            Move(s,
                 col_objs=col_objs,
                 row_objs=row_objs,
                 qubit_objs=qubit_objs,
                 col_idx=[0, ],
                 col_begin=[left_x, ],
                 col_end=[right_x, ],
                 row_idx=[0, ],
                 row_begin=[y, ],
                 row_end=[y-AOD_SEP, ],
                 prefix=self.name + f':tmp<-q@{left_x}'),
            Activate(s,
                     col_objs=col_objs,
                     row_objs=row_objs,
                     qubit_objs=qubit_objs,
                     col_idx=[],
                     col_xs=[],
                     row_idx=[1, ],
                     row_ys=[y, ],
                     pickup_qs=[right_q_id, ],
                     prefix=self.name + f':PickUp_{right_q_id}'),
            Move(s,
                 col_objs=col_objs,
                 row_objs=row_objs,
                 qubit_objs=qubit_objs,
                 col_idx=[0, ],
                 col_begin=[right_x, ],
                 col_end=[left_x, ],
                 row_idx=[],
                 row_begin=[],
                 row_end=[],
                 prefix=self.name + f':q@{left_x}<-q@{right_x}'),
            Deactivate(s,
                       col_objs=col_objs,
                       row_objs=row_objs,
                       qubit_objs=qubit_objs,
                       col_idx=[],
                       col_xs=[],
                       row_idx=[1, ],
                       row_ys=[y, ],
                       dropoff_qs=[right_q_id, ],
                       prefix=self.name + f':DropOff_q{right_q_id}'),
            Move(s,
                 col_objs=col_objs,
                 row_objs=row_objs,
                 qubit_objs=qubit_objs,
                 col_idx=[0, ],
                 col_begin=[left_x, ],
                 col_end=[right_x, ],
                 row_idx=[0, ],
                 row_begin=[y-AOD_SEP, ],
                 row_end=[y, ],
                 prefix=self.name + f':q@{right_x}<-tmp'),
            Deactivate(s,
                       col_objs=col_objs,
                       row_objs=row_objs,
                       qubit_objs=qubit_objs,
                       col_idx=[0, ],
                       col_xs=[right_x, ],
                       row_idx=[0, ],
                       row_ys=[y, ],
                       dropoff_qs=[left_q_id, ],
                       prefix=self.name + f':DropOff_q{left_q_id}'),
        ]


class Swap(ComboInst):
    def __init__(self, s: int):
        super().__init__('Swap', suffix=str(s), stage=s)

    def add_swap_pair(
            self,
            col_objs: Sequence[Col],
            row_objs: Sequence[Row],
            qubit_objs: Sequence[Qubit],
            q_id: int,
            qq_id: int
    ):
        left_q_id = q_id
        right_q_id = qq_id
        if qubit_objs[q_id].x > qubit_objs[qq_id].x:
            left_q_id = qq_id
            right_q_id = q_id
        self.insts.append(
            SwapPair(self.stage, col_objs, row_objs,
                     qubit_objs, left_q_id, right_q_id, self.name)
        )


class CodeGen():
    """Generate code files: json containing a list of dict, each one 
    corresponding to a DPQA instruction defined above. 
    """

    def __init__(
            self,
            file_name: str,
            no_transfer: bool = False,
            dir: str = None
    ):
        self.read_compiled(file_name)
        program = self.builder(no_transfer)

        if not dir:
            dir = './results/code/'
        self.code_full_file = dir + \
            (file_name.split('/')[-1]).replace('.json', '_code_full.json')
        with open(self.code_full_file, 'w') as f:
            json.dump(program.emit_full(), f, indent = 2)
        with open(self.code_full_file.replace('_code_full', '_code'), 'w') as f:
            json.dump(program.emit(), f, indent = 2)

    def read_compiled(self, compiled_file: str):
        with open(compiled_file, 'r') as f:
            data = json.load(f)
        self.n_q = data['n_q']
        self.x_high = int(data['n_x'])
        self.y_high = int(data['n_y'])
        self.c_high = int(data['n_c'])
        self.r_high = int(data['n_r'])
        self.layers = data['layers']
        self.n_t = len(self.layers)

        """change of convention. In solve() and the SMT model, a/c/r_s govern
        the movement from stage s to stage s+1, i.e.,
            ----------- x/y_0 
            | a/c/r_0 |
            ----------- x/y_1
            | a/c/r_1 |
            ----------- x/y_2
            | a/c/r_2 |

        However, the initial stage is very special in codegen and animation,
        so we generate code in this way:
           Rydberg_0  (special)  <------- x/y_0
           
           Swap_1 (optional)  <---. 
           Reload_1  <----------\  \ 
           BigMove_1  <--------- a/c/r_1
           Offload_1  <---------/ 
           Rydberg_1  <----------- x/y_1

               ...

        Thus, the movement between stage s and s+1 should be govened by a/c/r
        with subscript s+1, e.g., BigMove_1 uses a/c/r_1 above.
        So we need to shift the a/c/r variable values here.
        """
        for i in range(self.n_t-1, 0, -1):
            for q in range(self.n_q):
                self.layers[i]['qubits'][q]['a'] =\
                    self.layers[i - 1]['qubits'][q]['a']
                self.layers[i]['qubits'][q]['c'] =\
                    self.layers[i - 1]['qubits'][q]['c']
                self.layers[i]['qubits'][q]['r'] =\
                    self.layers[i - 1]['qubits'][q]['r']

        # infer some info of the AODs
        self.aod_from_compiled()

    def aod_from_compiled(self):
        for s, layer in enumerate(self.layers):
            if s == 0:
                continue
            layer['row'] = [{'id': i, 'qs': []}
                            for i in range(self.r_high-0)]
            layer['col'] = [{'id': i, 'qs': []}
                            for i in range(self.c_high-0)]
            prev_layer = self.layers[s-1]

            # figure out in the movement from stage s-1 to s:
            # - before the movement, where is each row 'y_begin'
            # - after the movement, where is each row 'y_end'
            # - what qubits are in this row 'qs'
            # similar for each AOD column
            for i, q in enumerate(layer['qubits']):
                if layer['qubits'][i]['a']:
                    layer['row'][q['r']]['y_begin'] =\
                        prev_layer['qubits'][i]['y']
                    layer['row'][q['r']]['y_end'] = q['y']
                    layer['row'][q['r']]['qs'].append(q['id'])
                    layer['col'][q['c']]['x_begin'] =\
                        prev_layer['qubits'][i]['x']
                    layer['col'][q['c']]['x_end'] = q['x']
                    layer['col'][q['c']]['qs'].append(q['id'])

            # figure out in the movement from stage s-1 to s:
            # - before the movement, which columns have site coord = X
            # - for all these cols, what is the relevant order from left 'offset'
            # - after the movement, which columns have site coord = X
            # - for all these cols, what is the relevant order from left 'offset'
            # similar for the AOD rows
            for case in ['_begin', '_end']:
                x_cols = []
                for x in range(self.x_high):
                    cols_at_x = []
                    for c in range(self.c_high):
                        if (layer['col'][c]['qs']
                                and layer['col'][c]['x' + case] == x):
                            cols_at_x.append(c)
                    for i, c in enumerate(cols_at_x):
                        layer['col'][c]['offset' + case] = i
                    x_cols.append(cols_at_x)
                layer['x_cols' + case] = x_cols
                y_rows = []
                for y in range(self.y_high):
                    rows_at_y = []
                    for r in range(self.r_high):
                        if (layer['row'][r]['qs']
                                and layer['row'][r]['y' + case] == y):
                            rows_at_y.append(r)
                    for i, r in enumerate(rows_at_y):
                        layer['row'][r]['offset' + case] = i
                    y_rows.append(rows_at_y)
                layer['y_rows' + case] = y_rows

    def builder(self, no_transfer: bool):
        qubits = [Qubit(i) for i in range(self.n_q)]
        rows = [Row(i) for i in range(self.r_high)]
        cols = [Col(i) for i in range(self.c_high)]
        program = ComboInst('Program')

        # read to comment in read_compiled() for structure of this method.
        init = self.builder_init(cols, rows, qubits, program)  # has Rydberg_0

        for s in range(1, len(self.layers)):
            self.builder_swap(s, cols, rows, qubits, program)

            if (not no_transfer) or s == 1:
                # if we know there is not atom transfer, we can simply skip the
                # reload and offload procedures. However, we keep the first one
                # just for convenience of generating animations.
                self.builder_reload(s, cols, rows, qubits, program)

            self.builder_move(s, cols, rows, qubits, program)

            if not no_transfer:
                self.builder_offload(s, cols, rows, qubits, program)
            
            self.builder_rydberg(s, cols, rows, qubits, program, init)

        program.remove_trivial_insts()
        program.prepend_inst(init)
        return program

    def builder_init(
            self,
            cols: Sequence[Col],
            rows: Sequence[Row],
            qubits: Sequence[Qubit],
            program: ComboInst,
    ) -> Inst:
        slm_qubit_idx = list(range(self.n_q))  # put all qubits in SLM
        slm_qubit_xys = [(
            X_SITE_SEP * self.layers[0]['qubits'][i]['x'],
            Y_SITE_SEP * self.layers[0]['qubits'][i]['y']
        ) for i in range(self.n_q)]  # put all qubits in the left trap

        # when there are more than one qubit in a site at the beginning,
        # need to put one of them in the right trap.
        for g in self.layers[0]['gates']:
            a0 = self.layers[1]['qubits'][g['q0']]['a']
            a1 = self.layers[1]['qubits'][g['q1']]['a']
            x_left = X_SITE_SEP * self.layers[0]['qubits'][g['q0']]['x']
            x_right = x_left + SITE_WIDTH
            y = Y_SITE_SEP * self.layers[0]['qubits'][g['q0']]['y']

            # if both atoms are in AOD, use their column indices to decide
            # which one to put in the left trap and which one to the right
            # if they have the same col index, the order does not matter
            # in Reload, we will pick them up in different rows.
            if (a0 == 1
                and a1 == 1
                and self.layers[1]['qubits'][g['q0']]['c'] >
                    self.layers[1]['qubits'][g['q1']]['c']):
                slm_qubit_xys[g['q0']] = (x_right, y)
            else:
                slm_qubit_xys[g['q1']] = (x_right, y)

        init = Init(cols, rows, qubits,
                    slm_qubit_idx=slm_qubit_idx,
                    slm_qubit_xys=slm_qubit_xys,
                    data={
                        'n_q': self.n_q,
                        'x_high': self.x_high,
                        'y_high': self.y_high,
                        'c_high': self.c_high,
                        'r_high': self.r_high,
                    })

        self.builder_rydberg(0, cols, rows, qubits, program, init)

        return init

    def builder_rydberg(
            self,
            s: int,
            cols: Sequence[Col],
            rows: Sequence[Row],
            qubits: Sequence[Qubit],
            program: ComboInst,
            init: Inst
    ):
        program.append_inst(
            Rydberg(s, cols, rows, qubits, self.layers[s]['gates'])
        )
        init.add_slms([(q.x, q.y) for q in qubits])

    def builder_swap(
            self,
            s: int,
            cols: Sequence[Col],
            rows: Sequence[Row],
            qubits: Sequence[Qubit],
            program: ComboInst,
    ):
        swap_obj = Swap(s)
        prev_layer = self.layers[s-1]
        this_layer = self.layers[s]
        for q0_id in range(self.n_q):
            for q1_id in range(q0_id+1, self.n_q):
                q0_a = this_layer['qubits'][q0_id]['a']
                q1_a = this_layer['qubits'][q1_id]['a']
                if q0_a == 1 and q1_a == 1:
                    q0_x = prev_layer['qubits'][q0_id]['x']
                    q1_x = prev_layer['qubits'][q1_id]['x']
                    q0_y = prev_layer['qubits'][q0_id]['y']
                    q1_y = prev_layer['qubits'][q1_id]['y']
                    q0_c = this_layer['qubits'][q0_id]['c']
                    q1_c = this_layer['qubits'][q1_id]['c']
                    q0_r = this_layer['qubits'][q0_id]['r']
                    q1_r = this_layer['qubits'][q1_id]['r']
                    # if two qubits are at the same site and
                    # both being picked up in the same row
                    if (q0_x, q0_y, q0_r) == (q1_x, q1_y, q1_r):
                        # if their position and col indeces are in reverse order
                        if (
                            (q0_c > q1_c
                             and qubits[q0_id].x < qubits[q1_id].x
                             )
                            or (q0_c < q1_c
                                and qubits[q0_id].x > qubits[q1_id].x
                                )
                        ):
                            swap_obj.add_swap_pair(cols, rows, qubits,
                                                   q0_id, q1_id)
        program.append_inst(swap_obj)

    def builder_move(
            self,
            s: int,
            cols: Sequence[Col],
            rows: Sequence[Row],
            qubits: Sequence[Qubit],
            program: ComboInst,
    ):
        """to avoid collision, the big moves will end in a slightly adjusted
        location: 1 um to +X direction and AOD_SEP(2) um to +Y direction
        Suppose the O's are the two traps in an interaction site
                O--------O
        The atoms will finish the big move in
                O--------O
             -----A-----
                  -----B-----
        (+X is ->, +Y is down, 1um is --, and AOD_SEP is one line height)
        There can be other cases, e.g., A and B are in the same column
                O--------O
             -----A-----
             -----B-----
        Or, A and B are in the same row
                O--------O
             -----A---B------
        """

        col_idx = []
        col_begin = []
        col_end = []
        for col_id in range(self.c_high):
            if cols[col_id].active:
                col_idx.append(col_id)
                col_begin.append(cols[col_id].x)
                site_x = self.layers[s]['col'][col_id]['x_end']
                offset = self.layers[s]['col'][col_id]['offset_end']
                col_end.append(1 + site_x * X_SITE_SEP + AOD_SEP * offset)

        row_idx = []
        row_begin = []
        row_end = []
        for row_id in range(self.r_high):
            if rows[row_id].active:
                row_idx.append(row_id)
                row_begin.append(rows[row_id].y)
                site_y = self.layers[s]['row'][row_id]['y_end']
                offset = self.layers[s]['row'][row_id]['offset_end']
                row_end.append(site_y * Y_SITE_SEP + AOD_SEP * (1 + offset))

        program.append_inst(
            Move(s=s,
                 col_objs=cols,
                 row_objs=rows,
                 qubit_objs=qubits,
                 col_idx=col_idx,
                 col_begin=col_begin,
                 col_end=col_end,
                 row_idx=row_idx,
                 row_begin=row_begin,
                 row_end=row_end,
                 prefix=f'BigMove_{s}')
        )

    def builder_reload(
            self,
            s: int,
            cols: Sequence[Col],
            rows: Sequence[Row],
            qubits: Sequence[Qubit],
            program: ComboInst,
    ):
        layer = self.layers[s]
        prev_layer = self.layers[s-1]
        reload_obj = Reload(s)
        # reload row by row
        for row_id in range(self.r_high):
            if layer['row'][row_id]['qs']:  # there are qubits to load
                reloadRow_obj = reload_obj.add_row_reload(row_id)
                pickup_qs = []
                cols_to_active = []
                x_to_activate = []
                # consider the movements in a row of sites
                for site_x in range(self.x_high):
                    site_qs = []
                    for q_id in range(len(qubits)):
                        # find out the qubits with site_x and in row_id
                        if (layer['qubits'][q_id]['a'] == 1
                            and prev_layer['qubits'][q_id]['x'] == site_x
                                and layer['qubits'][q_id]['r'] == row_id):
                            site_qs.append(q_id)
                            qubits[q_id].r = row_id
                            qubits[q_id].c = layer['qubits'][q_id]['c']

                    # shift the 1 or 2 cols that are picking up qubits
                    if len(site_qs) == 2:
                        # which qubit is on the left and which on the right
                        [q_id_left, q_id_right] = site_qs
                        if layer['qubits'][q_id_left]['c'] >\
                                layer['qubits'][q_id_right]['c']:
                            tmp = q_id_left
                            q_id_left = q_id_right
                            q_id_right = tmp

                        # current location of Cols
                        col_id_left = layer['qubits'][q_id_left]['c']
                        col_id_right = layer['qubits'][q_id_right]['c']
                        lower_offset =\
                            layer['col'][col_id_left]['offset_begin']
                        upper_offset =\
                            layer['col'][col_id_right]['offset_begin']

                        # target locations of Cols
                        lower_x = qubits[q_id_left].x
                        upper_x = qubits[q_id_right].x

                        # process the Col on the left
                        if not cols[col_id_left].active:
                            cols_to_active.append(col_id_left)
                            x_to_activate.append(lower_x)
                        else:
                            reloadRow_obj.add_col_shift(
                                id=col_id_left,
                                begin=cols[col_id_left].x,
                                end=lower_x)

                        # process the Col on the right
                        if not cols[col_id_right].active:
                            cols_to_active.append(col_id_right)
                            x_to_activate.append(upper_x)
                        else:
                            reloadRow_obj.add_col_shift(
                                id=col_id_right,
                                begin=cols[col_id_right].x,
                                end=upper_x)

                    elif len(site_qs) == 1:
                        # for convience later on, we still keep the *_left and
                        # *_right vars even if there is one qubit to pick up
                        q_id = site_qs[0]
                        col_id_left = layer['qubits'][q_id]['c']
                        col_id_right = col_id_left
                        lower_offset =\
                            layer['col'][col_id_left]['offset_begin']
                        upper_offset = lower_offset
                        lower_x = qubits[q_id].x
                        upper_x = lower_x
                        if not cols[col_id_left].active:
                            cols_to_active.append(col_id_left)
                            x_to_activate.append(lower_x)
                        else:  # col already active, shift it to align with q_id
                            reloadRow_obj.add_col_shift(
                                id=col_id_left,
                                begin=cols[col_id_left].x,
                                end=lower_x)

                    elif len(site_qs) > 2:
                        raise ValueError(
                            f"builder reload {s} row {row_id} site {site_x}:"
                            f" more than 2 qubits"
                        )
                    else:
                        continue

                    # shift other Cols that are already activated. Those Cols
                    # may not be picking up any qubit in this row, but due to
                    # AOD order constraints, we need to shift them to have the
                    # correct order when we shift the Cols that indeed are
                    # picking up some qubit in this row.
                    for col_id in layer['x_cols_begin'][site_x]:
                        # the Cols at site_x in the beginning between stage s-1
                        # and s (since we are processing Reload_s now)

                        if (cols[col_id].active
                            and col_id != col_id_left
                                and col_id != col_id_right):
                            # if this Col is not the one involved in loading

                            # if there is a col on the right of the cols for loading
                            if layer['col'][col_id]['offset_begin'] >\
                                    upper_offset:
                                reloadRow_obj.add_col_shift(
                                    id=col_id,
                                    begin=cols[col_id].x,
                                    end=upper_x + AOD_SEP *
                                    (layer['col'][col_id]['offset_begin'] -
                                     upper_offset) + 1)

                            # if there is a col on the left of the cols for loading
                            elif layer['col'][col_id]['offset_begin'] <\
                                    lower_offset:
                                reloadRow_obj.add_col_shift(
                                    id=col_id,
                                    begin=cols[col_id].x,
                                    end=lower_x + AOD_SEP *
                                    (layer['col'][col_id]['offset_begin'] -
                                     lower_offset) - 1)
                            # if there is a col in the middle of the cols for loading
                            else:
                                reloadRow_obj.add_col_shift(
                                    id=col_id,
                                    begin=cols[col_id].x,
                                    end=lower_x + AOD_SEP *
                                    (layer['col'][col_id]['offset_begin'] -
                                     lower_offset))

                    pickup_qs += site_qs

                # apply all the col shifts added previously
                reloadRow_obj.generate_col_shift(cols, rows, qubits)
                reloadRow_obj.generate_row_activate(
                    cols,
                    rows,
                    qubits,
                    cols_to_active,
                    x_to_activate,
                    layer['row'][row_id]['y_begin']*Y_SITE_SEP,
                    pickup_qs
                )

                # shift down the finished row because later on, some other row
                # may need to adjust the Cols again and if we keep this row
                # as is, some qubits may collid into each other since at each
                # site, the y of two SLM traps are the same.
                num_rows = len(layer['y_rows_begin']
                               [layer['row'][row_id]['y_begin']])
                shift_down = (
                    num_rows - layer['row'][row_id]['offset_begin'])*AOD_SEP
                # Also shift the Cols by 1 to avoid collisions.
                col_idx = [col_id for col_id,
                           col in enumerate(cols) if col.active]
                col_begin = [cols[col_id].x for col_id in col_idx]
                col_end = [1 + layer['col'][col_id]['x_begin'] * X_SITE_SEP +
                           AOD_SEP * layer['col'][col_id]['offset_begin']
                           for col_id in col_idx]
                reloadRow_obj.generate_parking(
                    cols,
                    rows,
                    qubits,
                    shift_down,
                    col_idx,
                    col_begin,
                    col_end
                )

        program.append_inst(reload_obj)

    def builder_offload(
            self,
            s: int,
            cols: Sequence[Col],
            rows: Sequence[Row],
            qubits: Sequence[Qubit],
            program: ComboInst,
    ):
        offload_obj = Offload(s)
        layer = self.layers[s]
        # the row-by-row processing is quite similar to Reload
        for row_id in range(self.r_high):
            if rows[row_id].active:
                dropoff_qs = []
                offloadRow_obj = offload_obj.add_row_offload(row_id)
                for site_x in range(self.x_high):
                    site_q_slm = []
                    site_q_aod = []
                    for q_id, q in enumerate(layer['qubits']):
                        if (
                            q['x'], q['y']
                        ) == (
                            site_x, layer['row'][row_id]['y_end']
                        ):
                            if (qubits[q_id].array == 'AOD'
                                    and q['r'] == row_id):
                                dropoff_qs.append(q_id)
                                site_q_aod.append(q_id)
                            if qubits[q_id].array == 'SLM':
                                site_q_slm.append(q_id)
                    if len(site_q_aod) == 2:
                        [q_id_left, q_id_right] = site_q_aod
                        if layer['qubits'][q_id_left]['c'] >\
                                layer['qubits'][q_id_right]['c']:
                            tmp = q_id_left
                            q_id_left = q_id_right
                            q_id_right = tmp
                        col_id_left = layer['qubits'][q_id_left]['c']
                        col_id_right = layer['qubits'][q_id_right]['c']
                        lower_offset = layer['col'][col_id_left]['offset_end']
                        upper_offset = layer['col'][col_id_right]['offset_end']
                        lower_x = X_SITE_SEP*site_x
                        upper_x = X_SITE_SEP*site_x + SITE_WIDTH

                        offloadRow_obj.add_col_shift(
                            id=col_id_left,
                            begin=qubits[q_id_left].x,
                            end=lower_x
                        )
                        offloadRow_obj.add_col_shift(
                            id=col_id_right,
                            begin=qubits[q_id_right].x,
                            end=upper_x
                        )
                    elif len(site_q_aod) == 1:
                        q_id = site_q_aod[0]
                        col_id_left = layer['qubits'][q_id]['c']
                        col_id_right = col_id_left
                        lower_offset = layer['col'][col_id_left]['offset_end']
                        upper_offset = lower_offset
                        lower_x = X_SITE_SEP*site_x
                        if site_q_slm:
                            lower_x = 2*X_SITE_SEP*site_x + \
                                SITE_WIDTH - qubits[site_q_slm[0]].x
                        upper_x = lower_x
                        offloadRow_obj.add_col_shift(
                            id=col_id_left,
                            begin=qubits[q_id].x,
                            end=lower_x
                        )
                    elif len(site_q_aod) > 2:
                        raise ValueError(
                            f"builder offload {s} row {row_id} site {site_x}:"
                            f" more than 2 qubits"
                        )
                    else:
                        continue

                    for col_id in layer['x_cols_end'][site_x]:
                        if (cols[col_id].active
                            and col_id != col_id_left
                                and col_id != col_id_right):
                            if layer['col'][col_id]['offset_end'] >\
                                    upper_offset:
                                offloadRow_obj.add_col_shift(
                                    id=col_id,
                                    begin=cols[col_id].x,
                                    end=upper_x + AOD_SEP *
                                    (layer['col'][col_id]['offset_end'] -
                                     upper_offset) + 1
                                )
                            elif layer['col'][col_id]['offset_end'] <\
                                    lower_offset:
                                offloadRow_obj.add_col_shift(
                                    id=col_id,
                                    begin=cols[col_id].x,
                                    end=lower_x + AOD_SEP *
                                    (layer['col'][col_id]['offset_end'] -
                                     lower_offset) - 1
                                )
                            else:
                                offloadRow_obj.add_col_shift(
                                    id=col_id,
                                    begin=cols[col_id].x,
                                    end=lower_x + AOD_SEP *
                                    (layer['col'][col_id]['offset_end'] -
                                     lower_offset)
                                )

                # align the Cols to the correct locatios
                offloadRow_obj.generate_col_shift(cols, rows, qubits)
                # the rows are at the parked location finishing a BigMove,
                # so we need to shift them to align with the SLM traps.
                offloadRow_obj.generate_row_shift(
                    cols,
                    rows,
                    qubits,
                    layer['row'][row_id]['y_end']
                )
                offloadRow_obj.generate_row_deactivate(
                    cols, rows, qubits, dropoff_qs)
        offload_obj.all_cols_deactivate(cols, rows, qubits)
        program.append_inst(offload_obj)


class Animator():
    """generate animation movie from code_full files."""

    def __init__(self,
                 code_file_name: str,
                 scaling_factor: int = PT_MICRON,
                 font: int = 10,
                 ffmpeg: str = 'ffmpeg',
                 real_speed: bool = False,
                 show_graph: bool = False,
                 edges = None,
                 dir = None
                 ):
        """
        Args:
            code_file_name (str): file name of code_full generated by CodeGen.
            scaling_factor (int, optional): the unit scaling factor between the
             animation and um. Defaults to PT_MICRON.
            font (int, optional): font size in the animation. Defaults to 10.
            ffmpeg (str, optional): path of ffmpeg. Defaults to 'ffmpeg'.
            real_speed (bool, optional): whether to use real speed in the
                movements in Relaod and Offload. Defaults to False.
            show_graph (bool, optional): whether show the graph on the right
                side of the DPQA in the animation. Defaults to False.
            edges (Sequence[Sequence[int]] | None, optional): the edges of the
                graph. Defaults to None.
            dir (str | None, optional): dir to save output. Defaults to None.
        """
        matplotlib.use('Agg')
        matplotlib.rcParams.update({'font.size': font})
        plt.rcParams['animation.ffmpeg_path'] = ffmpeg
        self.show_graph = show_graph
        self.graph_edges = edges
        self.read_files(code_file_name)

        self.keyframes = []
        self.setup_canvas(scaling_factor)
        self.create_master_schedule(real_speed)

        anim = FuncAnimation(
            self.fig,
            self.update,
            init_func=self.update_init,
            frames=self.keyframes[-1],
            interval=1000/FPS
        )

        animation_file = dir + (code_file_name.replace(
            '_code_full.json', '.mp4')).split('/')[-1]
        anim.save(animation_file, writer=FFMpegWriter(FPS))

    def read_files(self, code_file: str):
        with open(code_file, 'r') as f:
            self.code = json.load(f)

        self.n_q = self.code[0]['n_q']
        self.x_high = self.code[0]['x_high']
        self.y_high = self.code[0]['y_high']
        self.c_high = self.code[0]['c_high']
        self.r_high = self.code[0]['r_high']

    def create_master_schedule(self, real_speed: bool):
        """create a list of keyframes. We need this to calculate how many
        frames are there in total when we call FuncAnimation."""

        frame = 0  # the beginning frame of the Inst currently considered
        for inst in self.code:
            # if not using real speed, make all the movements in Reload
            # and Offload 1 frame
            if not real_speed:
                if inst['type'] == 'Move' and 'BigMove' not in inst['name']:
                    inst['duration'] = 1

            # Rydberg interaction is much shorter compared to the movements,
            # so using real speed, we will never see Rydberg.
            if inst['type'] == 'Rydberg':
                inst['duration'] = MUS_PER_FRM * 8  # i.e., 8 frames

            # Activate and Deactivate is on par with some movements in terms of
            # duration, but we do not have ramping-up or -down animations yet,
            # so we opt for them taking 1 frame.
            if inst['type'] == 'Activate' or inst['type'] == 'Deactivate':
                inst['duration'] = MUS_PER_FRM

            inst['f_begin'] = frame
            # new_frame = frame + ceil(duration / MUS_PER_FRM)
            new_frame = frame +\
                int((inst['duration'] + MUS_PER_FRM - 1)/MUS_PER_FRM)
            # new_frame is the begining of the next Inst, so end of this one is
            inst['f_end'] = new_frame-1
            self.keyframes.append(new_frame)
            frame = new_frame

    """layout of the DPQA plot:
          ____________________________________________________________________
          |                                  ^                               |
          |                        Y_HIGH_PAD|                    X_HIGH_PAD |
        y |                                  V                  |<---------->|
         -|-         o   o              o   o--             o   o            |
        t |                                                 |<->|            |
        i |                   ...                    .   X_SITE_WIDTH        |
        c |                                         .                        |
        k |                                        .                         |
         -|-         o   o              o   o             --o   o            |
          |          |<---X_SITE_SEP--->|                  ^                 |
          |                                        ...     |Y_SITE_SEP       |
          |                                                V                 |
         -|-         o   o              o   o--           --o   o            |
          |          |                       ^                               |
          | X_LOW_PAD|              Y_LOW_PAD|                               |
          |<-------->|                       |                               |
          |____________|__________________|__V________________|______________|
                       |                  |                   |
                          x tick
        """

    def setup_canvas(self, scaling_factor: int):
        """set up various objects before actually drawing."""

        # unit conversion factor from um to pt
        px = 1/plt.rcParams['figure.dpi'] * scaling_factor

        self.X_LOW = -X_LOW_PAD
        self.X_HIGH = SITE_WIDTH + \
            (self.x_high-1)*X_SITE_SEP + X_HIGH_PAD
        self.Y_LOW = -Y_LOW_PAD
        self.Y_HIGH = (self.y_high-1)*Y_SITE_SEP + Y_HIGH_PAD

        if self.show_graph:
            # 1 row 2 col, DPQA on the left, the graph on the right
            self.fig, (self.ax, self.network_ax) = plt.subplots(
                1,
                2,
                gridspec_kw={'width_ratios': [3, 1]},
                figsize=((self.X_HIGH-self.X_LOW)*px*4/3,
                         (self.Y_HIGH-self.Y_LOW)*px)
            )
            self.fig.tight_layout()
            self.network = nx.Graph()
            self.network_ax.set_title('The 3-regular graph')
        else:
            self.fig, self.ax, = plt.subplots(
                figsize=((self.X_HIGH-self.X_LOW)*px,
                         (self.Y_HIGH-self.Y_LOW)*px))

        self.title = self.ax.set_title('')
        self.ax.set_xlim([self.X_LOW, self.X_HIGH])
        self.ax.set_xticks(
            [SITE_WIDTH/2+X_SITE_SEP*i for i in range(self.x_high)])
        self.ax.set_xticklabels([i for i in range(self.x_high)])
        self.ax.set_ylim(
            [self.Y_LOW, self.Y_HIGH])
        self.ax.set_yticks(
            [Y_SITE_SEP*i for i in range(self.y_high)])
        self.ax.set_yticklabels([i for i in range(self.y_high)])

        # draw all the SLMs used throught out the computation
        # slm_xs = [slm[0] for slm in self.code[0]['all_slms']]
        # slm_ys = [slm[1] for slm in self.code[0]['all_slms']]
        # self.ax.scatter(slm_xs, slm_ys, marker='o', s=50,
        #                 facecolor='none', edgecolor=(0, 0, 1, 0.5))

    def update_init(self):
        # init the Qubits, Cols, and Rows. Later just update their attributes
        self.qubit_scat = self.ax.scatter([0 for _ in range(self.n_q)],
                                          [0 for _ in range(self.n_q)])
        self.qubit_scat.set_color('b')
        self.col_plots = [self.ax.axvline(0, self.Y_LOW, self.Y_HIGH, c=(
            1, 0, 0, 0), ls='--') for _ in range(self.c_high)]
        self.row_plots = [self.ax.axhline(0, self.X_LOW, self.X_HIGH, c=(
            1, 0, 0, 0), ls='--') for _ in range(self.r_high)]

        # draw the graph with black thin edges
        if self.show_graph:
            self.network.add_edges_from(self.graph_edges)
            self.pos = nx.spring_layout(self.network)
            nx.draw(
                self.network,
                pos=self.pos,
                ax=self.network_ax,
                with_labels=True,
                edgecolors='black',
                node_color='lightgray'
            )
        return

    def update(self, f: int):  # f is the frame
        if f < self.keyframes[0]:
            return
        for i, inst in enumerate(self.code):
            if f >= self.keyframes[i-1] and f < self.keyframes[i]:
                if inst['type'] == 'Rydberg':
                    return self.update_rydberg(f, inst)
                elif inst['type'] == 'Move':
                    return self.update_move(f, inst, self.code[i-1]['state'])
                elif inst['type'] == 'Activate':
                    return self.update_activate(f, inst)
                elif inst['type'] == 'Deactivate':
                    return self.update_deactivate(f, inst)
                elif inst['type'] == 'Init':
                    return
                else:
                    raise ValueError(f"unknown inst type {inst['type']}")

    def update_rydberg(self, f: int, inst: dict):
        edges = [(g['q0'], g['q1']) for g in inst['gates']]
        if f == inst['f_begin']:
            self.title.set_text(inst['name'])

            # find the qubits involved in 2Q gates and annotate their ids
            active_qubits = []
            for g in inst['gates']:
                active_qubits += [g['q0'], g['q1']]
            self.texts = [
                self.ax.text(inst['state']['qubits'][q_id]['x'] + 1,
                             inst['state']['qubits'][q_id]['y'] + 1,
                             q_id) for q_id in active_qubits
            ]

            # whole plane lights up (alpha=.2 blue)
            self.ax.set_facecolor((0, 0, 1, 0.2))

            # draw SLM qubits in blue and at the correct locations
            self.qubit_scat.set_offsets(
                [(q['x'], q['y']) for q in inst['state']['qubits']])

        if self.show_graph and f == int((inst['f_begin'] + inst['f_end'])/2):
            # draw some edges in the graph (each corresponding to a 2Q gate
            # executed at this Rydberg stage) in thicker blue
            nx.draw_networkx_edges(
                self.network,
                pos=self.pos,
                edgelist=edges,
                edge_color='blue',
                ax=self.network_ax,
                width=4
            )

        if f == inst['f_end']:
            # clean up the annotations and blue background at finishing frame
            self.ax.set_facecolor('w')
            for text in self.texts:
                text.remove()

    def interpolate(self, progress: int, duration: int, begin: int, end: int):
        """implement cubic interpolation per Bluvstein et al.
        Suppose we want to move from x=begin at t=0 to x=end at t=duration
        if we are at t=progress*duration, return our x. (0<= progress <=1)
        """

        D = end - begin
        if D == 0:
            return begin
        r = (1+progress)/duration
        return begin + 3*D*(r**2) - 2*D*(r**3)

    def update_move(self, f: int, inst: dict, prev_state: dict):
        if f == inst['f_begin']:
            self.title.set_text(inst['name'])

        # needs to change location every frame, as below
        progress = f - inst['f_begin']
        duration = inst['f_end'] - inst['f_begin'] + 1

        col_xs = [col['x'] for col in prev_state['cols']]
        for col in inst['cols']:
            curr_x = self.interpolate(
                progress, duration, col['begin'], col['end'])
            col_xs[col['id']] = curr_x
            self.col_plots[col['id']].set_xdata((col_xs[col['id']], ))

        row_ys = [row['y'] for row in prev_state['rows']]
        for row in inst['rows']:
            curr_y = self.interpolate(
                progress, duration, row['begin'], row['end'])
            row_ys[row['id']] = curr_y
            self.row_plots[row['id']].set_ydata((row_ys[row['id']], ))

        q_xs = [q['x'] for q in prev_state['qubits']]
        q_ys = [q['y'] for q in prev_state['qubits']]
        for q_id in range(self.n_q):
            if prev_state['qubits'][q_id]['array'] == 'AOD':
                q_xs[q_id] = col_xs[prev_state['qubits'][q_id]['c']]
                q_ys[q_id] = row_ys[prev_state['qubits'][q_id]['r']]
        self.qubit_scat.set_offsets([(q_xs[i], q_ys[i])
                                    for i in range(self.n_q)])
        return

    def update_activate(self, f: int, inst: dict):
        if f == inst['f_begin']:
            self.title.set_text(inst['name'])

            # set Cols/Rows to activate to the correct location and turn red.2
            for id, col in enumerate(inst['col_idx']):
                self.col_plots[col].set_xdata((inst['col_xs'][id], ))
                self.col_plots[col].set_color((1, 0, 0, 0.2))

            for id, row in enumerate(inst['row_idx']):
                self.row_plots[row].set_ydata((inst['row_ys'][id], ))
                self.row_plots[row].set_color((1, 0, 0, 0.2))

            # SLM qubits remain blue while qubits being picked up turns red
            self.qubit_scat.set_color(
                ['b' if q['array'] == 'SLM' else 'r'
                 for q in inst['state']['qubits']])

    def update_deactivate(self, f: int, inst: dict):
        if f == inst['f_end']:
            self.title.set_text(inst['name'])

            self.qubit_scat.set_color(
                ['b' if q['array'] == 'SLM' else 'r'
                 for q in inst['state']['qubits']])

            # deactivate Cols/Rows by setting alpha=0
            for col in inst['col_idx']:
                self.col_plots[col].set_color((1, 0, 0, 0))
            for row in inst['row_idx']:
                self.row_plots[row].set_color((1, 0, 0, 0))

def set_hardware_paramters(param: dict):
    R_B = param["R_B"]  # rydberg range
    AOD_SEP = param["AOD_SEP"]  # min AOD separation
    RYD_SEP = param["RYD_SEP"]  # sufficient distance to avoid Rydberg
    SITE_SLMS = param["SITE_SLMS"]  # number of SLMs in a site
    SITE_WIDTH = param["SITE_WIDTH"]  # total width of SLMs in a site
    SLM_SEP = AOD_SEP  # separation of SLMs inside a site
    X_SITE_SEP = RYD_SEP + SITE_WIDTH  # separation of sites in X direction
    Y_SITE_SEP = RYD_SEP  # separation of sites in Y direction

    # padding of the figure
    X_LOW_PAD = 2 * AOD_SEP
    Y_LOW_PAD = 4 * AOD_SEP
    X_HIGH_PAD = 2 * AOD_SEP
    Y_HIGH_PAD = 4 * AOD_SEP

    T_RYDBERG = param["T_RYDBERG"]  # microseconds for Rydberg
    T_ACTIVATE = param["T_ACTIVATE"]  # microseconds for (de)activating AOD

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('input_file', type=str)
    parser.add_argument(
        '--scaling', help='scaling factor of the animation', type=int)
    parser.add_argument('--font', help='font size in the animation', type=int)
    parser.add_argument('--ffmpeg', help='custom ffmpeg path', type=str)
    parser.add_argument(
        '--realSpeed',
        help='real speed in the animation of reload and offload procedures',
        action='store_true')
    parser.add_argument(
        '--noGraph', help='do not show graph on the side', action='store_true')
    parser.add_argument('--dir', help='working directory', type=str)
    parser.add_argument('--arch_param', help='hardware parameters for DPQA', type=str, default = "hardware_spec/default.json")
    args = parser.parse_args()

    with open(args.input_file, 'r') as f:
        data = json.load(f)
    
    with open(args.arch_param, 'r') as f:
        param = json.load(f)
        set_hardware_paramters(param)

    
    filename = args.input_file.split("/")[-1]
    filename = filename.split("_")
    num_qubit = int(filename[1])
    trial = int(filename[2])

    with open('./graphs.json', 'r') as f:
        graphs = json.load(f)
    if str(num_qubit) in graphs.keys() and trial in range(10):
        gate_list = graphs[str(num_qubit)][trial]
    else:
        # import networkx as nx
        # gate_list = list(nx.random_regular_graph(3, args.size, args.id).edges)
        raise ValueError(f'No such graph {args.size}_{args.id}.')


    Animator(
        args.input_file,
        scaling_factor=args.scaling if args.scaling else PT_MICRON,
        font=args.font if args.font else 10,
        ffmpeg=args.ffmpeg if args.ffmpeg else 'ffmpeg',
        real_speed=args.realSpeed,
        show_graph=not args.noGraph,
        edges=gate_list if not args.noGraph else [],
        dir=args.dir if args.dir else './results/animations/'
    )