"""Tests for functions in the optimize submodule."""

from ferrmion.optimize.bonsai import bonsai_algorithm
import rustworkx as rx
import numpy as np
from pytest import fixture

@fixture
def heavy_hex_graph():
    graph = rx.PyGraph()
    graph.add_nodes_from(range(37))
    graph.add_edges_from_no_data(
        [
            (0, 1),
            (0, 2),
            (0, 3),
            (1, 4),
            (2, 5),
            (3, 6),
            (4, 7),
            (4, 8),
            (5, 9),
            (5, 10),
            (6, 11),
            (6, 12),
            (7, 13),
            (8, 14),
            (9, 15),
            (10, 16),
            (11, 17),
            (12, 18),
            (13, 19),
            (13, 20),
            (14, 21),
            (14, 22),
            (15, 23),
            (15, 24),
            (16, 25),
            (16, 26),
            (17, 27),
            (17, 28),
            (18, 29),
            (18, 30),
            (22, 31),
            (26, 32),
            (30, 33),
            (31, 34),
            (32, 35),
            (33, 36),
        ]
    )
    return graph

@fixture
def garnet_graph():
    garnet_edges = [
        (1, 0),
        (1, 4),
        (3, 0),
        (3, 2),
        (3, 4),
        (3, 8),
        (5, 4),
        (5, 6),
        (5, 10),
        (7, 2),
        (7, 8),
        (7, 12),
        (9, 4),
        (9, 8),
        (9, 10),
        (9, 14),
        (11, 6),
        (11, 10),
        (11, 16),
        (13, 8),
        (13, 12),
        (13, 14),
        (13, 17),
        (15, 10),
        (15, 14),
        (15, 16),
        (15, 19),
        (18, 14),
        (18, 17),
        (18, 19)
    ]

    graph = rx.PyGraph()
    graph.add_nodes_from(range(20))
    graph.add_edges_from_no_data(garnet_edges)
    return graph

@fixture
def emerald_graph():
    emerald_edges = [(0, 1),
        (0, 4),
        (3, 2),
        (3, 4),
        (3, 9),
        (5, 1),
        (5, 4),
        (5, 6),
        (5, 11),
        (8, 2),
        (8, 7),
        (8, 9),
        (8, 16),
        (10, 4),
        (10, 9),
        (10, 11),
        (10, 18),
        (12, 11),
        (12, 13),
        (12, 20),
        (15, 7),
        (15, 14),
        (15, 16),
        (15, 23),
        (17, 16),
        (17, 18),
        (17, 25),
        (18, 26),
        (19, 11),
        (19, 18),
        (19, 20),
        (19, 27),
        (21, 13),
        (21, 20),
        (21, 29),
        (22, 14),
        (22, 23),
        (24, 16),
        (24, 23),
        (24, 25),
        (24, 32),
        (25, 26),
        (26, 27),
        (26, 34),
        (28, 20),
        (28, 27),
        (28, 29),
        (28, 36),
        (30, 29),
        (30, 38),
        (31, 23),
        (31, 32),
        (33, 25),
        (33, 32),
        (33, 34),
        (33, 41),
        (34, 42),
        (35, 27),
        (35, 34),
        (35, 36),
        (35, 43),
        (36, 37),
        (37, 38),
        (37, 45),
        (40, 32),
        (40, 39),
        (40, 41),
        (40, 46),
        (41, 42),
        (43, 49),
        (44, 36),
        (44, 43),
        (44, 45),
        (44, 50),
        (47, 41),
        (47, 46),
        (47, 48),
        (47, 51),
        (49, 50),
        (49, 53),
        (52, 48),
        (52, 51),
        (52, 53)
    ]

    graph = rx.PyGraph()
    graph.add_nodes_from(range(54))
    graph.add_edges_from_no_data(emerald_edges)
    return graph

def test_heavy_hex(heavy_hex_graph):
    graph = heavy_hex_graph
    bonsai_homo = bonsai_algorithm(graph=graph, homogenous=True)
    assert bonsai_homo.root_node.child_strings == [
        '',
        'x',
        'y',
        'z',
        'xx',
        'yx',
        'zx',
        'xxx',
        'xxy',
        'yxx',
        'yxy',
        'zxx',
        'zxy',
        'xxxx',
        'xxyx',
        'yxxx',
        'yxyx',
        'zxxx',
        'zxyx',
        'xxxxx',
        'xxxxy',
        'xxyxx',
        'xxyxy',
        'yxxxx',
        'yxxxy',
        'yxyxx',
        'yxyxy',
        'zxxxx',
        'zxxxy',
        'zxyxx',
        'zxyxy',
        'xxyxyx',
        'yxyxyx',
        'zxyxyx',
        'xxyxyxx',
        'yxyxyxx',
        'zxyxyxx']

    bonsai_hetero = bonsai_algorithm(graph=graph, homogenous=False)
    assert bonsai_hetero.root_node.child_strings == [
        '',
        'x',
        'y',
        'z',
        'xz',
        'yz',
        'zz',
        'xzx',
        'xzz',
        'yzx',
        'yzz',
        'zzx',
        'zzz',
        'xzxz',
        'xzzz',
        'yzxz',
        'yzzz',
        'zzxz',
        'zzzz',
        'xzxzx',
        'xzxzz',
        'xzzzx',
        'xzzzz',
        'yzxzx',
        'yzxzz',
        'yzzzx',
        'yzzzz',
        'zzxzx',
        'zzxzz',
        'zzzzx',
        'zzzzz',
        'xzxzxz',
        'yzxzxz',
        'zzxzxz',
        'xzxzxzz',
        'yzxzxzz',
        'zzxzxzz']


    assert bonsai_hetero.root_node.child_qubit_labels == {
        '': 0,
        'x': 2,
        'y': 3,
        'z': 1,
        'xz': 5,
        'yz': 6,
        'zz': 4,
        'xzx': 10,
        'xzz': 9,
        'yzx': 12,
        'yzz': 11,
        'zzx': 8,
        'zzz': 7,
        'xzxz': 16,
        'xzzz': 15,
        'yzxz': 18,
        'yzzz': 17,
        'zzxz': 14,
        'zzzz': 13,
        'xzxzx': 26,
        'xzxzz': 25,
        'xzzzx': 24,
        'xzzzz': 23,
        'yzxzx': 30,
        'yzxzz': 29,
        'yzzzx': 28,
        'yzzzz': 27,
        'zzxzx': 22,
        'zzxzz': 21,
        'zzzzx': 20,
        'zzzzz': 19,
        'xzxzxz': 32,
        'yzxzxz': 33,
        'zzxzxz': 31,
        'xzxzxzz': 35,
        'yzxzxzz': 36,
        'zzxzxzz': 34
    }

def test_heavy_hex_max_nodes(heavy_hex_graph):
    bonsai= bonsai_algorithm(graph=heavy_hex_graph, homogenous=False, max_nodes=14)
    bonsai.string_pairs
    assert bonsai.as_dict() == {'x': {'x': None,
  'y': None,
  'z': {'x': {'x': None, 'y': None, 'z': None},
   'y': None,
   'z': {'x': None, 'y': None, 'z': None}}},
 'y': {'x': None,
  'y': None,
  'z': {'x': {'x': None, 'y': None, 'z': None},
   'y': None,
   'z': {'x': None, 'y': None, 'z': None}}},
 'z': {'x': None,
  'y': None,
  'z': {'x': {'x': None, 'y': None, 'z': None},
   'y': None,
   'z': {'x': None, 'y': None, 'z': {'x': None, 'y': None, 'z': None}}}}}

def test_garnet(garnet_graph):
    graph = garnet_graph
    bonsai= bonsai_algorithm(graph=graph, homogenous=False)
    assert bonsai.root_node.child_strings == [
        '',
        'x',
        'y',
        'z',
        'xx',
        'xz',
        'yx',
        'yz',
        'zx',
        'zy',
        'zz',
        'xxx',
        'xxz',
        'xzz',
        'yxz',
        'yzz',
        'zxz',
        'zyz',
        'zzz',
        'xxzz']

    bonsai= bonsai_algorithm(graph=graph, homogenous=True)
    assert bonsai.root_node.child_strings == [
        '',
        'x',
        'y',
        'z',
        'xx',
        'xy',
        'xz',
        'yx',
        'yy',
        'zx',
        'zy',
        'xxx',
        'xyx',
        'xzx',
        'yxx',
        'yyx',
        'yyy',
        'zxx',
        'zyx',
        'yyxx']


def test_emerald_hom(emerald_graph):
    graph = emerald_graph.copy()
    bonsai= bonsai_algorithm(graph=graph, homogenous=False)
    assert bonsai.root_node.child_strings == [
        '',
        'x',
        'y',
        'z',
        'xx',
        'xz',
        'yx',
        'yz',
        'zx',
        'zz',
        'xxx',
        'xxz',
        'xzx',
        'xzz',
        'yxz',
        'yzx',
        'yzz',
        'zxx',
        'zxz',
        'zzx',
        'zzz',
        'xxxx',
        'xxxz',
        'xxzz',
        'yzxz',
        'yzzx',
        'yzzz',
        'zxxz',
        'zxzx',
        'zxzz',
        'zzxz',
        'zzzx',
        'zzzy',
        'zzzz',
        'xxzzz',
        'yzxzz',
        'yzzxx',
        'yzzxz',
        'yzzzz',
        'zxxzz',
        'zxzxz',
        'zxzzx',
        'zxzzz',
        'zzzzz',
        'xxzzzx',
        'xxzzzz',
        'yzxzzz',
        'yzzxxz',
        'yzzxzz',
        'yzzzzz',
        'zxzxzz',
        'zxzzxz',
        'zxzzzz',
        'xxzzzzz']


def test_emerald_het(emerald_graph):
    graph = emerald_graph.copy()
    bonsai= bonsai_algorithm(graph=graph, homogenous=True)
    assert bonsai.root_node.child_strings == [
        '',
        'x',
        'y',
        'z',
        'xx',
        'xy',
        'yx',
        'yy',
        'zx',
        'zy',
        'xxx',
        'xxy',
        'xyx',
        'xyy',
        'yxx',
        'yxy',
        'yyx',
        'yyy',
        'zxx',
        'zxy',
        'zyx',
        'xxxx',
        'xxxy',
        'xxxz',
        'xxyx',
        'xyxx',
        'xyxy',
        'xyyx',
        'yyxx',
        'yyyx',
        'yyyy',
        'zxxx',
        'zxxy',
        'zxyx',
        'xxxxx',
        'xyxxx',
        'xyxxy',
        'xyxyx',
        'xyyxx',
        'yyxxx',
        'zxxxx',
        'zxxyx',
        'zxxyy',
        'zxyxx',
        'xyxxxx',
        'xyxxyx',
        'xyxyxx',
        'yyxxxx',
        'yyxxxy',
        'zxxxxx',
        'zxxyxx',
        'zxxyyx',
        'zxyxxx',
        'yyxxxxx']

def test_emerald_het_max_nodes(emerald_graph):
    graph = emerald_graph
    bonsai = bonsai_algorithm(graph=graph, homogenous=False, max_nodes=14)
    assert bonsai.root_node.child_strings == [
        '',
        'x',
        'y',
        'z',
        'xx',
        'xz',
        'yx',
        'yz',
        'zx',
        'zz',
        'zxx',
        'zxz',
        'zzx',
        'zzz']

def test_emerald_hom_max_nodes(emerald_graph):
    graph = emerald_graph
    bonsai = bonsai_algorithm(graph=graph, homogenous=True, max_nodes=14)
    assert bonsai.root_node.child_strings == [
        '',
        'x',
        'y',
        'z',
        'xx',
        'xy',
        'yx',
        'yy',
        'zx',
        'zy',
        'xxx',
        'xxy',
        'xyx',
        'xyy']
