import logging
import time

from nasap_net.assembly_enumeration import SymmetryOperations, \
    enumerate_assemblies
from nasap_net.models import Assembly, Bond, Component
from nasap_net.reaction_exploration_im import MLEKind, explore_reactions

logging.basicConfig(level=logging.DEBUG)

# upper half             //                           lower half             //
#                        |                                                   |
#                       (1)                                                 (a)
#                        M7                                                  M7
#                       (0)                                                 (1)
#                        |                                                   |
#                       (2)                                                 (2)
#              M2(1)--(0)L1(1)--(0)M1                              M5(1)--(0)L4(1)--(0)M4
#              (0)               (1)                               (0)               (1)
#                \               /                                   \               /
#                (1)           (0)                                   (1)           (0)
# //--(1)M8(0)--(2)L2          L0(2)--(0)M6(1)--//    //--(0)M8(1)--(2)L5          L3(2)--(1)M6(0)--//
#                  (0)       (1)                                       (0)       (1)
#                    \       /                                            \       /
#                    (1)   (0)                                           (1)   (0)
#                        M0                                                  M3
M = Component(kind='M', sites=[0, 1])
L = Component(kind='L', sites=[0, 1, 2])
X = Component(kind='X', sites=[0])
M9L6 = Assembly(
    components={
        'M0': M, 'M1': M, 'M2': M,
        'M3': M, 'M4': M, 'M5': M,
        'M6': M, 'M7': M, 'M8': M,
        'L0': L, 'L1': L, 'L2': L,
        'L3': L, 'L4': L, 'L5': L,
    },
    bonds=[
        Bond('M0', 0, 'L0', 1), Bond('L0', 0, 'M1', 1),
        Bond('M1', 0, 'L1', 1), Bond('L1', 0, 'M2', 1),
        Bond('M2', 0, 'L2', 1), Bond('L2', 0, 'M0', 1),
        Bond('L0', 2, 'M6', 0), Bond('L1', 2, 'M7', 0),
        Bond('L2', 2, 'M8', 0),
        Bond('M3', 0, 'L3', 1), Bond('L3', 0, 'M4', 1),
        Bond('M4', 0, 'L4', 1), Bond('L4', 0, 'M5', 1),
        Bond('M5', 0, 'L5', 1), Bond('L5', 0, 'M3', 1),
        Bond('L3', 2, 'M6', 1), Bond('L4', 2, 'M7', 1),
        Bond('L5', 2, 'M8', 1),
    ]
)


sym_ops = SymmetryOperations()
sym_ops.add_cyclic_permutation(
    'C_3',
    [
        ['M0', 'M1', 'M2'], ['M3', 'M4', 'M5'], ['M6', 'M7', 'M8'],
        ['L0', 'L1', 'L2'], ['L3', 'L4', 'L5'],
    ]
)
sym_ops.add_product('C_3^2', ['C_3', 'C_3'])
sym_ops.add_cyclic_permutation(
    'C_2x',
    [
        ['M0', 'M3'], ['M1', 'M5'], ['M2', 'M4'],
        ['M6', 'M8'], ['M7'],
        ['L0', 'L5'], ['L1', 'L4'], ['L2', 'L3'],
    ]
)
sym_ops.add_product('C_2y', ['C_2x', 'C_3^2'])
sym_ops.add_product('C_2xy', ['C_2x', 'C_3'])
sym_ops.add_cyclic_permutation(
    'sigma_h',
    [
        ['M0', 'M3'], ['M1', 'M4'], ['M2', 'M5'],
        ['M6'], ['M7'], ['M8'],
        ['L0', 'L3'], ['L1', 'L4'], ['L2', 'L5'],
    ]
)
sym_ops.add_product('sigma_1', ['sigma_h', 'C_2x'])
sym_ops.add_product('sigma_2', ['sigma_h', 'C_2y'])
sym_ops.add_product('sigma_3', ['sigma_h', 'C_2xy'])
sym_ops.add_product('S_3', ['sigma_h', 'C_3'])
sym_ops.add_product('S_3^5', ['sigma_h', 'C_3^2'])


# timer start
start = time.perf_counter()


assemblies = list(enumerate_assemblies(
    template=M9L6,
    leaving_ligand=X,
    metal_kinds='M',
    symmetry_operations=sym_ops.values(),
))

reactions = explore_reactions(
    assemblies=dict(enumerate(assemblies)),
    mle_kinds=[
        MLEKind(metal='M', leaving='X', entering='L'),
        MLEKind(metal='M', leaving='L', entering='X'),
    ],
)
reaction_list = list(reactions)


# timer end
end = time.perf_counter()

print(f'Explored {len(reaction_list)} reactions in {end - start:.2f} seconds.')
