# -*- coding: utf-8 -*-
setVerbosity(MinimalLog)

# -------------------------------------------------------------
# Bulk Configuration
# -------------------------------------------------------------

# Set up lattice
vector_a = [10.781302175705491, -18.67376314007479, 0.0]*Angstrom
vector_b = [2.1562604351410983, 3.734752628014958, 0.0]*Angstrom
vector_c = [0.0, 0.0, 17.171785381745117]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
elements = [Germanium, Germanium, Germanium, Germanium, Germanium, Germanium,
            Germanium, Germanium, Germanium, Germanium, Antimony, Antimony,
            Antimony, Antimony, Antimony, Antimony, Antimony, Antimony,
            Antimony, Antimony, Tellurium, Tellurium, Tellurium, Tellurium,
            Tellurium, Tellurium, Tellurium, Tellurium, Tellurium, Tellurium,
            Tellurium, Tellurium, Tellurium, Tellurium, Tellurium, Tellurium,
            Tellurium, Tellurium, Tellurium, Tellurium, Tellurium, Tellurium,
            Tellurium, Tellurium, Tellurium]

# Define coordinates
fractional_coordinates = [[ 0.116666666667,  0.916666666667,  0.949743340975],
                          [ 0.316666666667,  0.916666666667,  0.949743340975],
                          [ 0.516666666667,  0.916666666667,  0.949743340975],
                          [ 0.716666666667,  0.916666666667,  0.949743340975],
                          [ 0.916666666667,  0.916666666667,  0.949743340975],
                          [ 0.183333333333,  0.583333333333,  0.153151659025],
                          [ 0.383333333333,  0.583333333333,  0.153151659025],
                          [ 0.583333333333,  0.583333333333,  0.153151659025],
                          [ 0.783333333333,  0.583333333333,  0.153151659025],
                          [ 0.983333333333,  0.583333333333,  0.153151659025],
                          [ 0.05          ,  0.25          ,  0.732138635328],
                          [ 0.25          ,  0.25          ,  0.732138635328],
                          [ 0.45          ,  0.25          ,  0.732138635328],
                          [ 0.65          ,  0.25          ,  0.732138635328],
                          [ 0.85          ,  0.25          ,  0.732138635328],
                          [ 0.05          ,  0.25          ,  0.370756364672],
                          [ 0.25          ,  0.25          ,  0.370756364672],
                          [ 0.45          ,  0.25          ,  0.370756364672],
                          [ 0.65          ,  0.25          ,  0.370756364672],
                          [ 0.85          ,  0.25          ,  0.370756364672],
                          [ 0.05          ,  0.25          ,  0.0514475     ],
                          [ 0.25          ,  0.25          ,  0.0514475     ],
                          [ 0.45          ,  0.25          ,  0.0514475     ],
                          [ 0.65          ,  0.25          ,  0.0514475     ],
                          [ 0.85          ,  0.25          ,  0.0514475     ],
                          [ 0.183333333333,  0.583333333333,  0.849836853001],
                          [ 0.383333333333,  0.583333333333,  0.849836853001],
                          [ 0.583333333333,  0.583333333333,  0.849836853001],
                          [ 0.783333333333,  0.583333333333,  0.849836853001],
                          [ 0.983333333333,  0.583333333333,  0.849836853001],
                          [ 0.116666666667,  0.916666666667,  0.253058146999],
                          [ 0.316666666667,  0.916666666667,  0.253058146999],
                          [ 0.516666666667,  0.916666666667,  0.253058146999],
                          [ 0.716666666667,  0.916666666667,  0.253058146999],
                          [ 0.916666666667,  0.916666666667,  0.253058146999],
                          [ 0.116666666667,  0.916666666667,  0.631492753008],
                          [ 0.316666666667,  0.916666666667,  0.631492753008],
                          [ 0.516666666667,  0.916666666667,  0.631492753008],
                          [ 0.716666666667,  0.916666666667,  0.631492753008],
                          [ 0.916666666667,  0.916666666667,  0.631492753008],
                          [ 0.183333333333,  0.583333333333,  0.471402246992],
                          [ 0.383333333333,  0.583333333333,  0.471402246992],
                          [ 0.583333333333,  0.583333333333,  0.471402246992],
                          [ 0.783333333333,  0.583333333333,  0.471402246992],
                          [ 0.983333333333,  0.583333333333,  0.471402246992]]

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

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
k_point_sampling = KpointDensity(
    density_a=4.0*Angstrom,
    )
numerical_accuracy_parameters = NumericalAccuracyParameters(
    density_mesh_cutoff=120.0*Hartree,
    k_point_sampling=k_point_sampling,
    )

calculator = LCAOCalculator(
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    )

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

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
constraints = [SpaceGroupConstraint()]

bulk_configuration = OptimizeGeometry(
    bulk_configuration,
    max_forces=0.05*eV/Ang,
    max_stress=0.1*GPa,
    max_steps=200,
    max_step_length=0.2*Ang,
    constraints=constraints,
    trajectory_filename='PBE_trajectory.hdf5',
    trajectory_interval=1,
    restart_strategy=RestartFromTrajectory(),
    optimizer_method=LBFGS(),
    enable_optimization_stop_file=True,
)
nlsave('PBE.hdf5', bulk_configuration)
nlprint(bulk_configuration)

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

# -------------------------------------------------------------
# Bandstructure
# -------------------------------------------------------------
bandstructure = Bandstructure(
    configuration=bulk_configuration,
    route=['G', 'Z'],
    points_per_segment=50,
    bands_above_fermi_level=All,
    method=Full,
    )
nlsave('PBE.hdf5', bandstructure)

# -------------------------------------------------------------
# Projected Density Of States
# -------------------------------------------------------------
kpoint_grid = KpointDensity(
    density_a=7.0*Angstrom,
    )

projected_density_of_states = ProjectedDensityOfStates(
    configuration=bulk_configuration,
    kpoints=kpoint_grid,
    projections=ProjectOnSites,
    energies=numpy.linspace(-8, 8, 601)*eV,
    energy_zero_parameter=FermiLevel,
    bands_above_fermi_level=All,
    spectrum_method=TetrahedronMethod,
)
nlsave('PBE.hdf5', projected_density_of_states)