# -*- coding: utf-8 -*-
# -------------------------------------------------------------
# Bulk Configuration
# -------------------------------------------------------------

# Set up lattice
lattice = SimpleOrthorhombic(15.0*Angstrom, 14.0*Angstrom, 15.5*Angstrom)

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

# Define coordinates
fractional_coordinates = [[ 0.428664594147,  0.50874202174 ,  0.544519665482],
                          [ 0.500164589018,  0.50874202174 ,  0.544519665482],
                          [ 0.404831262523,  0.50874202174 ,  0.476770075897],
                          [ 0.404831262523,  0.450652965672,  0.578394460275],
                          [ 0.404831262523,  0.566831077808,  0.578394460275]]

# Set up configuration
bulk_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )

# Add tags
bulk_configuration.addTags('H_Si', [1, 2, 3, 4])
nlsave('SiH4.hdf5', bulk_configuration)

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# PAW Data Set
#----------------------------------------
basis_set = [
    PAWPBESuggested.Hydrogen,
    PAWPBESuggested.Silicon,
    ]

#----------------------------------------
# Exchange-Correlation
#----------------------------------------
exchange_correlation = SGGA.PBE

density_mesh_cutoff = OptimizedFFTGridSampling(
    energy_cutoff=1520.0*eV,
    maximum_average_prime_factor=5.0,
    maximum_grid_enlargement=0.1,
    )

k_point_sampling = KpointDensity(
    density_a=0.0*Angstrom,
    density_b=0.0*Angstrom,
    density_c=0.0*Angstrom,
    symmetries=[
        ([[ 1., 0., 0.],
          [ 0., 1., 0.],
          [ 0., 0., 1.]], [ 0., 0., 0.]),
        ([[-1., 0., 0.],
          [ 0.,-1., 0.],
          [ 0., 0.,-1.]], [ 0., 0., 0.]),
        ],
    force_timereversal=True,
    shift_to_gamma=[True, True, True],
    )
exx_grid_cutoff = OptimizedFFTGridSampling(
    energy_cutoff=760.0*eV,
    maximum_average_prime_factor=5.0,
    maximum_grid_enlargement=0.1,
    )

exact_exchange_parameters = ExactExchangeParameters(
    aux_basis_tolerance=0.001,
    number_of_waves=1024,
    maximum_wave_vector=50.0,
    integral_tolerance=0.0001,
    relative_screening_tolerance=0.01,
    )
compensation_charge_mesh_cutoff = OptimizedFFTGridSampling(
    energy_cutoff=1520.0*eV,
    maximum_average_prime_factor=5.0,
    maximum_grid_enlargement=0.1,
    )

numerical_accuracy_parameters = NumericalAccuracyParameters(
    density_mesh_cutoff=density_mesh_cutoff,
    k_point_sampling=k_point_sampling,
    radial_step_size=0.001*Bohr,
    density_cutoff=1e-06,
    interaction_max_range=20.0*Angstrom,
    number_of_reciprocal_points=1024,
    reciprocal_energy_cutoff=1250.0*Hartree,
    bands_per_electron=1.2,
    occupation_method=GaussianSmearing(1000.0*Kelvin*boltzmann_constant),
    exx_grid_cutoff=exx_grid_cutoff,
    exact_exchange_parameters=exact_exchange_parameters,
    compensation_charge_mesh_cutoff=compensation_charge_mesh_cutoff,
    )

iteration_control_parameters = IterationControlParameters(
    tolerance=0.0001,
    max_steps=100,
    algorithm=PulayMixer(),
    damping_factor=0.2,
    number_of_history_steps=5,
    start_mixing_after_step=0,
    mixing_variable=HamiltonianVariable,
    preconditioner=Preconditioner.Off,
    linear_dependence_threshold=0.0,
    max_exx_updates=50,
    non_convergence_behavior=ContinueCalculation(),
    enable_scf_stop_file=True,
    )

poisson_solver = FastFourierSolver()

initialization_method = BasisSetInitialization()
density_matrix_method = PPCGSolver(
    absolute_tolerance=1e-08,
    relative_tolerance=1e-99,
    maximum_number_of_iterations=2,
    block_size=5,
    rr_period=3,
    buffer_states=0.05,
    initialization_method=initialization_method,
    )

algorithm_parameters = AlgorithmParameters(
    density_matrix_method=density_matrix_method,
    store_grids=True,
    store_basis_on_grid=Automatic,
    store_energy_density_matrix=Automatic,
    scf_restart_step_length=0.1*Angstrom,
    use_symmetries=Automatic,
    )

parallel_parameters = ParallelParameters(
    processes_per_neb_image=None,
    processes_per_individual=None,
    processes_per_bias_point=None,
    processes_per_saddle_search=1,
    )

checkpoint_handler = CheckpointHandler(
    time_interval=0.5*Hour,
    )

calculator = PlaneWaveCalculator(
    wave_function_cutoff=380.0*eV,
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    iteration_control_parameters=iteration_control_parameters,
    poisson_solver=poisson_solver,
    algorithm_parameters=algorithm_parameters,
    parallel_parameters=parallel_parameters,
    checkpoint_handler=checkpoint_handler,
    charge=0.0,
    correction_extension=None,
    fixed_spin_moment=False,
    store_wave_functions=False,
    processes_per_kpoint=Automatic,
    )

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('SiH4.hdf5', bulk_configuration)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
bulk_configuration = OptimizeGeometry(
    bulk_configuration,
    max_forces=0.01*eV/Ang,
    max_steps=200,
    max_step_length=0.2*Ang,
    trajectory_filename='SiH4.hdf5',
    trajectory_interval=1,
    restart_strategy=RestartFromTrajectory(),
    disable_stress=True,
    optimizer_method=LBFGS(),
    enable_optimization_stop_file=True,
)
nlsave('SiH4.hdf5', bulk_configuration)
nlprint(bulk_configuration)

# -------------------------------------------------------------
# Total Energy
# -------------------------------------------------------------
total_energy = TotalEnergy(bulk_configuration)
nlsave('SiH4.hdf5', total_energy)
nlprint(total_energy)