# Set minimal log verbosity
setVerbosity(MinimalLog)

# %% ConfigurationTable(List)

# %% Configuration

# Define elements
elements = [Carbon, Hydrogen, Hydrogen, Hydrogen]

# Define coordinates
cartesian_coordinates = [[ 0.            ,  0.            ,  0.            ],
                         [ 1.089999921806,  0.            ,  0.            ],
                         [-0.363333307269,  0.            , -1.027661781602],
                         [-0.363333307269, -0.889981209366,  0.513830890801]]*Angstrom

# Set up configuration
configuration_0 = MoleculeConfiguration(
    elements=elements,
    cartesian_coordinates=cartesian_coordinates
    )

# Add tags
configuration_0.addTags('H_C',    [1, 2, 3])
configuration_0.addTags('anchor', [0])

configuration_name_0 = "Methyl"


# %% Configuration

# Define elements
elements = [Carbon, Carbon, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Carbon,
            Hydrogen, Hydrogen, Hydrogen]

# Define coordinates
cartesian_coordinates = [[ 0.            ,  0.            ,  0.            ],
                         [ 1.513         ,  0.            ,  0.            ],
                         [-0.363333307269,  0.            , -1.027661781602],
                         [-0.363333307269, -0.889981209366,  0.513830890801],
                         [-0.363333307269,  0.889981209366,  0.513830890801],
                         [ 1.876333307269,  0.889981209366, -0.513830890801],
                         [ 2.017333333333,  0.            ,  1.426470079914],
                         [ 1.876333307269, -0.889981209366, -0.513830890801],
                         [ 3.107333255139,  0.            ,  1.426470079914],
                         [ 1.654000026065, -0.889981209366,  1.940300970715]]*Angstrom

# Set up configuration
configuration_1 = MoleculeConfiguration(
    elements=elements,
    cartesian_coordinates=cartesian_coordinates
    )

# Add tags
configuration_1.addTags('H_C',    [1, 2, 3, 4, 5, 6, 7, 8, 9])
configuration_1.addTags('anchor', [6])

configuration_name_1 = "Propyl"


# %% Configuration

# Define elements
elements = [Carbon, Carbon, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen]

# Define coordinates
cartesian_coordinates = [[ 0.            ,  0.            ,  0.            ],
                         [ 1.513         ,  0.            ,  0.            ],
                         [-0.363333307269,  0.            , -1.027661781602],
                         [-0.363333307269, -0.889981209366,  0.513830890801],
                         [-0.363333307269,  0.889981209366,  0.513830890801],
                         [ 1.876333307269,  0.889981209366, -0.513830890801],
                         [ 1.876333307269, -0.889981209366, -0.513830890801]]*Angstrom

# Set up configuration
configuration_2 = MoleculeConfiguration(
    elements=elements,
    cartesian_coordinates=cartesian_coordinates
    )

# Add tags
configuration_2.addTags('H_C',    [1, 2, 3, 4, 5, 6])
configuration_2.addTags('anchor', [1])

configuration_name_2 = "Ethyl"

configuration_table = Table()
configuration_table.addInstanceColumn(key='configurations', types=MoleculeConfiguration)
configuration_table.append(configuration_0)
configuration_table.append(configuration_1)
configuration_table.append(configuration_2)

configuration_names_table = Table()
configuration_names_table.addStringColumn(key='configuration_names')
configuration_names_table.append(configuration_name_0)
configuration_names_table.append(configuration_name_1)
configuration_names_table.append(configuration_name_2)


# %% Initial Generated Molecules

if 'initial_generated_molecules' not in locals():
    initial_generated_molecules = Table(
        'Precursor_Generator_Example_results.hdf5', object_id='table'
    )
    initial_generated_molecules.addInstanceColumn(
        key='configuration', types=MoleculeConfiguration
    )
    initial_generated_molecules.addStringColumn(key='description')
    initial_generated_molecules.setMetatext('Initial Generated Molecules')


# %% Ligand Generator


def ligand_generator(table, table_names, output):

    # Extract matching columns from tables
    configuration = table.column('configurations')
    names = table_names.column('configuration_names')
    configuration = output.column('configuration')
    description = output.column('description')

    # Defined variables.
    host_atom = 'Ir'
    num_ligands = 3
    distance_to_host = 2.0

    # Script.
    import PrecursorGenerator as PrecursorGenerator

    output = PrecursorGenerator.precursor_generator(
        table, table_names, output, host_atom, num_ligands, distance_to_host
    )
    return output


initial_generated_molecules = ligand_generator(
    configuration_table, configuration_names_table, initial_generated_molecules
)

nlsave('Precursor_Generator_Example_results.hdf5', initial_generated_molecules)


# %% Optimized Molecules

if 'optimized_molecules' not in locals():
    optimized_molecules = Table('Optimized_Molecules.hdf5', object_id='table')
    optimized_molecules.addInstanceColumn(
        key='configuration', types=MoleculeConfiguration
    )
    optimized_molecules.addStringColumn(key='description')
    optimized_molecules.setMetatext('Optimized Molecules')


# %% Table Iteration

# Table Iteration(preparation)
for row_index in range(initial_generated_molecules.numberOfRows()):
    configuration, description = initial_generated_molecules[
        row_index, ['configuration', 'description']
    ]

    # %% Set ForceFieldCalculator

    # %% ForceFieldCalculator

    potentialSet = TorchX_MACE_MP_0b3_medium(dtype='float32', enforceLTX=False)
    calculator = TremoloXCalculator(parameters=potentialSet)

    # %% Set Calculator

    configuration.setCalculator(calculator)

    nlsave(
        'Precursor_Generator_Example_results.hdf5',
        configuration,
        object_id=f'configuration_Set_Calculator_row_index_{row_index}',
    )

    # %% OptimizeGeometry

    optimized_configuration = OptimizeGeometry(
        configuration=configuration,
        max_steps=1000,
        trajectory_interval=100,
        enable_optimization_stop_file=False,
    )

    nlsave(
        'Precursor_Generator_Example_results.hdf5',
        optimized_configuration,
        object_id=f'optimized_configuration_optgeom_row_index_{row_index}',
    )

    # %% Append Row to Table

    optimized_molecules.append(optimized_configuration, description)