"""This module implements MachineModels for Rigetti's QPUs."""
from __future__ import annotations

from bqskit.compiler.machine import MachineModel
from bqskit.ir.gates.constant.cz import CZGate
from bqskit.ir.gates.constant.iswap import ISwapGate
from bqskit.ir.gates.constant.sx import SXGate
from bqskit.ir.gates.constant.x import XGate
from bqskit.ir.gates.parameterized.rz import RZGate
from bqskit.qis.graph import CouplingGraph

rigetti_gate_set = {SXGate(), XGate(), RZGate(), CZGate()}

ankaa_gate_set = {SXGate(), XGate(), RZGate(), CZGate(), ISwapGate()}

_aspen_11_coupling_graph = CouplingGraph([
    # Ring 1
    (0, 1), (1, 2), (2, 3), (3, 4),
    (4, 5), (5, 6), (6, 7), (7, 0),

    # Ring 2
    (8, 9), (9, 10), (10, 11), (11, 12),
    (12, 13), (13, 14), (14, 15), (15, 8),

    # Ring 3
    (16, 17), (17, 18), (19, 20),
    (20, 21), (21, 22), (22, 23), (23, 16),

    # Ring 4
    (24, 25), (25, 26), (26, 27), (27, 28),
    (28, 29), (29, 30), (30, 31), (31, 24),

    # Ring 5
    (34, 35), (35, 36),
    (36, 37), (37, 38), (38, 39),

    # Ring 1-2
    (1, 14), (2, 13),

    # Ring 2-3
    (9, 22), (10, 21),

    # Ring 3-4
    (17, 30),

    # Ring 4-5
    (25, 38), (26, 37),
])
"""Retrieved August 31, 2022: https://qcs.rigetti.com/qpus."""

_octo_rings = [
    [
        (i + 0, i + 1), (i + 1, i + 2), (i + 2, i + 3), (i + 3, i + 4),
        (i + 4, i + 5), (i + 5, i + 6), (i + 6, i + 7), (i + 7, i + 0),
    ]
    for i in range(0, 10 * 8, 8)
]

_horizontal_connections = [
    [
        (i + 1, j + 6),
        (i + 2, j + 5),
    ]
    for i, j in [(8 * r, 8 * (r + 1)) for r in [0, 1, 2, 3, 5, 6, 7, 8]]
]

_vertical_connections = [
    [
        (i + 0, j + 3),
        (i + 7, j + 4),
    ]
    for i, j in [(8 * r, 8 * (r + 5)) for r in range(5)]
]
_links = []
for l in _octo_rings:
    _links.extend(l)
for l in _horizontal_connections:
    _links.extend(l)
for l in _vertical_connections:
    _links.extend(l)

_aspen_m2_coupling_graph = CouplingGraph(_links)
"""Retrieved August 31, 2022: https://qcs.rigetti.com/qpus."""

_ankaa_9q_3_coupling_graph = CouplingGraph.grid(3, 3)
"""Retrieved September 11, 2024: https://qcs.rigetti.com/qpus."""

_ankaa_2_coupling_graph = CouplingGraph.grid(7, 12)
"""Retrieved September 11, 2024: https://qcs.rigetti.com/qpus."""

Aspen11Model = MachineModel(40, _aspen_11_coupling_graph, rigetti_gate_set)
"""A BQSKit MachineModel for Rigetti's Aspen-11 quantum processor."""

AspenM2Model = MachineModel(80, _aspen_m2_coupling_graph, rigetti_gate_set)
"""A BQSKit MachineModel for Rigetti's Aspen-M-2 quantum processor."""

ANKAA2Model = MachineModel(84, _ankaa_2_coupling_graph, ankaa_gate_set)

ANKAA9Q3Model = MachineModel(9, _ankaa_9q_3_coupling_graph, ankaa_gate_set)

__all__ = ['Aspen11Model', 'AspenM2Model']
