# -*- coding: utf-8 -*-
# Set up lattice
vector_a = [7.680028171823331, 0.0, 0.0]*Angstrom
vector_b = [0.0, 7.680028171823331, 0.0]*Angstrom
vector_c = [0.0, 0.0, 25.1459]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

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

# Define coordinates
fractional_coordinates = [[ 0.            ,  0.25          ,  0.053990909055],
                          [ 0.5           ,  0.25          ,  0.053990909055],
                          [ 0.            ,  0.75          ,  0.053990909055],
                          [ 0.5           ,  0.75          ,  0.053990909055],
                          [ 0.            ,  0.            ,  0.10798181811 ],
                          [ 0.5           ,  0.            ,  0.10798181811 ],
                          [ 0.            ,  0.5           ,  0.10798181811 ],
                          [ 0.5           ,  0.5           ,  0.10798181811 ],
                          [ 0.25          ,  0.            ,  0.161972727165],
                          [ 0.75          ,  0.            ,  0.161972727165],
                          [ 0.25          ,  0.5           ,  0.161972727165],
                          [ 0.75          ,  0.5           ,  0.161972727165],
                          [ 0.248888550945,  0.250280098839,  0.219403356388],
                          [ 0.750115554169,  0.249844542424,  0.214466715534],
                          [ 0.248899613521,  0.749778145533,  0.219439062887],
                          [ 0.749899904381,  0.75015400844 ,  0.214511356832],
                          [-0.007742648068,  0.266345237757,  0.273542895729],
                          [ 0.507450366831,  0.236133317089,  0.272629406254],
                          [-0.007800739114,  0.733785988998,  0.273548570205],
                          [ 0.507315071824,  0.76442614314 ,  0.272764646715],
                          [-0.126967445671,  0.000114953687,  0.301314064193],
                          [ 0.5797040723  ,  0.000443236556,  0.330803053002],
                          [-0.078413788824,  0.500073795099,  0.333204201241],
                          [ 0.61973395663 ,  0.500241546876,  0.306351835717]]

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

# Add tags
bulk_configuration.addTags('Selection 0', [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])

# Add bonds
bonds = [[ 0,  4,  0,  0,  0],
         [ 0,  6,  0,  0,  0],
         [ 1,  5,  0,  0,  0],
         [ 1,  7,  0,  0,  0],
         [ 2,  4,  0,  1,  0],
         [ 2,  6,  0,  0,  0],
         [ 3,  5,  0,  1,  0],
         [ 3,  7,  0,  0,  0],
         [ 4,  8,  0,  0,  0],
         [ 4,  9, -1,  0,  0],
         [ 5,  8,  0,  0,  0],
         [ 5,  9,  0,  0,  0],
         [ 6, 10,  0,  0,  0],
         [ 6, 11, -1,  0,  0],
         [ 7, 10,  0,  0,  0],
         [ 7, 11,  0,  0,  0],
         [ 8, 12,  0,  0,  0],
         [ 8, 14,  0, -1,  0],
         [ 9, 13,  0,  0,  0],
         [ 9, 15,  0, -1,  0],
         [10, 12,  0,  0,  0],
         [10, 14,  0,  0,  0],
         [11, 13,  0,  0,  0],
         [11, 15,  0,  0,  0],
         [12, 16,  0,  0,  0],
         [12, 17,  0,  0,  0],
         [13, 16,  1,  0,  0],
         [13, 17,  0,  0,  0],
         [14, 18,  0,  0,  0],
         [14, 19,  0,  0,  0],
         [15, 18,  1,  0,  0],
         [15, 19,  0,  0,  0],
         [16, 20,  0,  0,  0],
         [16, 22,  0,  0,  0],
         [17, 21,  0,  0,  0],
         [17, 23,  0,  0,  0],
         [18, 20,  0,  1,  0],
         [18, 22,  0,  0,  0],
         [19, 21,  0,  1,  0],
         [19, 23,  0,  0,  0]]

bulk_configuration.setBonds(bonds)


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

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

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

k_point_sampling = MonkhorstPackGrid(
    na=6,
    nb=6,
    nc=1,
    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('Si-alpha-100.hdf5', bulk_configuration)


# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
fix_atom_indices_0 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

constraints = [FixAtomConstraints(fix_atom_indices_0)]

bulk_configuration = OptimizeGeometry(
    bulk_configuration,
    max_forces=0.02*eV/Ang,
    max_steps=200,
    max_step_length=0.2*Ang,
    constraints=constraints,
    trajectory_filename='Si-alpha-100.hdf5',
    trajectory_interval=1,
    restart_strategy=RestartFromTrajectory(),
    disable_stress=True,
    optimizer_method=LBFGS(),
    enable_optimization_stop_file=True,
)
nlsave('Si-alpha-100.hdf5', bulk_configuration)
nlprint(bulk_configuration)

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


