# Set minimal log verbosity
setVerbosity(MinimalLog)

# %% Boron nitride

# Set up lattice
lattice = Hexagonal(2.50399*Angstrom, 6.6612*Angstrom)

# Define elements
elements = [Boron, Boron, Nitrogen, Nitrogen]

# Define coordinates
fractional_coordinates = [[ 0.333333333333,  0.666666666667,  0.25          ],
                          [ 0.666666666667,  0.333333333333,  0.75          ],
                          [ 0.333333333333,  0.666666666667,  0.75          ],
                          [ 0.666666666667,  0.333333333333,  0.25          ]]

# Set up configuration
boron_nitride = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )
boron_nitride_name = "boron_nitride"


# %% Define defects

define_defects_defects = []
define_defects_defect_generators = {}
define_defects_defect_generator_map = {}

defect_parameters = DefectsParameters(
    defect_type=1
)
defect_generator = defect_parameters.defectGenerator(boron_nitride)
defects_from_generator = []
point_defect = Vacancy(
    site_index=1,
    unit_cell_index=(0, 0, 0))
named_point_defect = NamedPointDefect(
    point_defect=point_defect,
    index=0,
    name='Vacancy B 000 001',
)
defects_from_generator.append(named_point_defect)
point_defect = Vacancy(
    site_index=2,
    unit_cell_index=(0, 0, 0))
named_point_defect = NamedPointDefect(
    point_defect=point_defect,
    index=1,
    name='Vacancy N 001 002',
)
defects_from_generator.append(named_point_defect)
kept_defects_from_generator = [defect.pointDefect() for defect in defects_from_generator]
defect_generator = defect_generator.filterByPointDefect(kept_defects_from_generator)

define_defects_defects += defects_from_generator
define_defects_defect_generators['e1dcdccbad94420299db11918622a5cf'] = defect_generator
for unique_defect in defects_from_generator:
    define_defects_defect_generator_map[unique_defect.name()] = 'e1dcdccbad94420299db11918622a5cf'

defect_parameters = DefectsParameters(
    defect_type=0,
    defect_element=Nitrogen
)
defect_generator = defect_parameters.defectGenerator(boron_nitride)
defects_from_generator = []
point_defect = Substitutional(
    element=Nitrogen,
    site_index=1,
    unit_cell_index=(0, 0, 0))
named_point_defect = NamedPointDefect(
    point_defect=point_defect,
    index=0,
    name='Substitutional N^B 000 001',
)
defects_from_generator.append(named_point_defect)
kept_defects_from_generator = [defect.pointDefect() for defect in defects_from_generator]
defect_generator = defect_generator.filterByPointDefect(kept_defects_from_generator)

define_defects_defects += defects_from_generator
define_defects_defect_generators['b6a8ff633084ac98e9286ad668233d6'] = defect_generator
for unique_defect in defects_from_generator:
    define_defects_defect_generator_map[unique_defect.name()] = 'b6a8ff633084ac98e9286ad668233d6'

defect_parameters = DefectsParameters(
    defect_type=0,
    defect_element=Boron
)
defect_generator = defect_parameters.defectGenerator(boron_nitride)
defects_from_generator = []
point_defect = Substitutional(
    element=Boron,
    site_index=2,
    unit_cell_index=(0, 0, 0))
named_point_defect = NamedPointDefect(
    point_defect=point_defect,
    index=1,
    name='Substitutional B^N 001 002',
)
defects_from_generator.append(named_point_defect)
kept_defects_from_generator = [defect.pointDefect() for defect in defects_from_generator]
defect_generator = defect_generator.filterByPointDefect(kept_defects_from_generator)

define_defects_defects += defects_from_generator
define_defects_defect_generators['cb0030ef1c2c4ed8a20fab59c6c20a25'] = defect_generator
for unique_defect in defects_from_generator:
    define_defects_defect_generator_map[unique_defect.name()] = 'cb0030ef1c2c4ed8a20fab59c6c20a25'

defect_parameters = DefectsParameters(
    defect_type=2,
    defect_element=Boron
)
defect_generator = defect_parameters.defectGenerator(boron_nitride)
defects_from_generator = []
point_defect = Interstitial(
    element=Boron,
    fractional_coordinates=[0.000000000000, 0.000000000000, 0.500000000000])
named_point_defect = NamedPointDefect(
    point_defect=point_defect,
    index=0,
    name='Interstitial B 000 001',
)
defects_from_generator.append(named_point_defect)
kept_defects_from_generator = [defect.pointDefect() for defect in defects_from_generator]
defect_generator = defect_generator.filterByPointDefect(kept_defects_from_generator)

define_defects_defects += defects_from_generator
define_defects_defect_generators['e51ee25d44d31ac86845caf023643'] = defect_generator
for unique_defect in defects_from_generator:
    define_defects_defect_generator_map[unique_defect.name()] = 'e51ee25d44d31ac86845caf023643'

defect_parameters = DefectsParameters(
    defect_type=2,
    defect_element=Boron,
    include_voronoi_faces=True,
    include_voronoi_vertices=False
)
defect_generator = defect_parameters.defectGenerator(boron_nitride)
defects_from_generator = []
point_defect = Interstitial(
    element=Boron,
    fractional_coordinates=[0.666666670000, 0.333333330000, 0.500000000000])
named_point_defect = NamedPointDefect(
    point_defect=point_defect,
    index=0,
    name='Interstitial B 000 005',
)
defects_from_generator.append(named_point_defect)
kept_defects_from_generator = [defect.pointDefect() for defect in defects_from_generator]
defect_generator = defect_generator.filterByPointDefect(kept_defects_from_generator)

define_defects_defects += defects_from_generator
define_defects_defect_generators['e6479c7c0714dd6bee57bd3b02a4081'] = defect_generator
for unique_defect in defects_from_generator:
    define_defects_defect_generator_map[unique_defect.name()] = 'e6479c7c0714dd6bee57bd3b02a4081'

define_defects = Defects(
    defects=define_defects_defects,
    defect_generators=define_defects_defect_generators,
    defect_generator_map=define_defects_defect_generator_map
)
nlsave('Boron_nitride_defects_results.hdf5', define_defects)


# %% Generate table of configurations

generate_table_of_configurations = generateDefectsConfigurations(
    defects=define_defects,
    host_configuration=boron_nitride,
    supercell_repetitions=(3, 3, 2)
)
nlsave('Boron_nitride_defects_results.hdf5', generate_table_of_configurations)


# %% Iteration over defect configurations

# Iteration over defect configurations(preparation)
for row_index in range(generate_table_of_configurations.numberOfRows()):
    row_data = generate_table_of_configurations.row(row_index)
    configuration = row_data[0]

    # %% Set SemiEmpiricalCalculator

    # %% SemiEmpiricalCalculator


    pair_potentials = DFTBDirectory(r"dftb/matsci-0-3")
    hamiltonian_parametrization = SlaterKosterHamiltonianParametrization(
        basis_set=DFTBDirectory(r"dftb/matsci-0-3"))

    k_point_sampling = KpointDensity(
        density_a=4.0*Angstrom,
        density_b=4.0*Angstrom,
        density_c=4.0*Angstrom
    )

    numerical_accuracy_parameters = NumericalAccuracyParameters(
        k_point_sampling=k_point_sampling
    )

    calculator = SemiEmpiricalCalculator(
        hamiltonian_parametrization=hamiltonian_parametrization,
        pair_potentials=pair_potentials,
        numerical_accuracy_parameters=numerical_accuracy_parameters,
        checkpoint_handler=NoCheckpointHandler
    )


    # %% Set Calculator

    configuration.setCalculator(calculator)

    configuration.update()

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


    # %% OptimizeGeometry

    local_object_id = f'optimized_configuration_optimize_trajectory_row_index_{row_index}'
    restart_strategy = RestartFromTrajectory(
        trajectory_filename='Boron_nitride_defects_results.hdf5',
        object_id=local_object_id
    )

    optimized_configuration = OptimizeGeometry(
        configuration=configuration,
        constraints=[
            BravaisLatticeConstraint()
        ],
        trajectory_filename='Boron_nitride_defects_results.hdf5',
        trajectory_object_id=local_object_id,
        optimize_cell=False,
        restart_strategy=restart_strategy
    )

    local_object_id = f'optimized_configuration_OptimizeGeometry_row_index_{row_index}'
    nlsave('Boron_nitride_defects_results.hdf5', optimized_configuration, object_id=local_object_id)


    # %% TotalEnergy

    total_energy = TotalEnergy(
        configuration=optimized_configuration
    )
    local_object_id = f'total_energy_TotalEnergy_row_index_{row_index}'
    nlsave('Boron_nitride_defects_results.hdf5', total_energy, object_id=local_object_id)