import ATK
from ATK.KohnSham import *

# Create BCC iron crystal
lattice = BodyCenteredCubic( 2.866*Ang )
bcc_iron = BulkConfiguration(lattice,[Iron],[3*(0.0,)*Ang])

# Specify parameters for spin-polarized bulk calculation 
brillouin_zone_sampling = brillouinZoneIntegrationParameters((3,3,3))

electron_density = electronDensityParameters(
    mesh_cutoff = 100*Rydberg,
    initial_spin = 1.0*hbar
    )

dft_iron = KohnShamMethod(
    brillouin_zone_integration_parameters = brillouin_zone_sampling,
    electron_density_parameters = electron_density,
    )

# Specify NetCDF file for storing results
# and increase verbosity
nc_file = 'dft_Febcc.nc'
ATK.setCheckpointFilename(nc_file)
ATK.setVerbosityLevel(1)

scf_iron = executeSelfConsistentCalculation(
    atomic_configuration = bcc_iron,
    method = dft_iron,
    )

# Setup special points in the Brillouin zone
bcc_symmetry_points = {
   "G"  : (0.00, 0.00, 0.00),
   "H"  : (0.50,-0.50, 0.50),
   "P"  : (0.25, 0.25, 0.25),
   "N"  : (0.00, 0.00, 0.50),
   "N'" : (0.50, 0.50, 0.00),
   "H'" : (0.50, 0.50, 0.50)
  }

# Specify route in the Brillouin zone
bcc_route = ["G","H","N","G","P","N'","P","H'"]

# Call auxillary function for plotting spin-polarized band structure
from bandstructure import calculateAndPlotSpinBandstructure
calculateAndPlotSpinBandstructure(
    ncfilename = nc_file,
    datafilename = "dft_Febcc.dat",
    plotfilename = "plotBandStructure.py",
    imagename = 'Fe_bcc.png',
    symmetry_points = bcc_symmetry_points,
    route = bcc_route,
    number_of_kpoints = 25
    )