import networkx
from SentaurusMaterialsWorkbench.MultiLayer.MultilayerBuilderUtilities import selectHKMGConfigurations

# Set up graph or coordination patterns for unfavorable structural motifs.
# O-O bonds.
oo_graph = networkx.Graph()
oo_graph.add_node(0, number=Oxygen.atomicNumber())
oo_graph.add_node(1, number=Oxygen.atomicNumber())
oo_graph.add_edge(0, 1)

# Cyclic Si-O-Si-O ring.
sio_ring_graph = networkx.Graph()
sio_ring_graph.add_node(0, number=Silicon.atomicNumber())
sio_ring_graph.add_node(1, number=Silicon.atomicNumber())
sio_ring_graph.add_node(2, number=Oxygen.atomicNumber())
sio_ring_graph.add_node(3, number=Oxygen.atomicNumber())
sio_ring_graph.add_edge(0, 2)
sio_ring_graph.add_edge(1, 2)
sio_ring_graph.add_edge(1, 3)
sio_ring_graph.add_edge(0, 3)

bad_patterns = [
    ('O-O bond', oo_graph),
    ('Si-O ring', sio_ring_graph),
]

bad_coordination = [
    (Silicon, (1, 2, 3, 5, 6)),
    (Hydrogen, (2, 3, 4)),
    (Oxygen, (1,)),
]

# List all the files that contain the HKMG candidate configurations.
filenames_list = [
    'HKMG_builder0.hdf5',
    'HKMG_builder1.hdf5',
    'HKMG_builder2.hdf5',
    'HKMG_builder3.hdf5',
    'HKMG_builder4.hdf5',
    'HKMG_builder5.hdf5',
    'HKMG_builder6.hdf5',
    'HKMG_builder7.hdf5',
    'HKMG_builder8.hdf5',
    'HKMG_builder9.hdf5',
]

configurations = []

# Check all configurations in the given filenames.
for filename in filenames_list:
    try:
        configuration = nlread(filename, BulkConfiguration)[-1]
        configurations.append(configuration)
    except IndexError:
        continue

selectHKMGConfigurations(
    configurations,
    calculator=None,
    fuzz_factor=1.15,
    selected_candidates_filename='selected_hkmg_configurations.hdf5',
    bad_patterns=bad_patterns,
    bad_coordination=bad_coordination
)