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

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

# Set up lattice
lattice = SimpleTetragonal(2.722*Angstrom, 3.71281*Angstrom)

# Define elements
elements = [Platinum, Iron]

# Define coordinates
fractional_coordinates = [[ 0.5,  0.5,  0.5],
                          [ 0. ,  0. ,  0. ]]

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

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

k_point_sampling = KpointDensity(
    density_a=7.0*Angstrom,
    force_timereversal=False,
    )
numerical_accuracy_parameters = NumericalAccuracyParameters(
    density_mesh_cutoff=120.0*Hartree,
    k_point_sampling=k_point_sampling,
    occupation_method=FermiDirac(300.0*Kelvin*boltzmann_constant),
    )

calculator = LCAOCalculator(
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    )

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('FePt-SO-SCF.hdf5', bulk_configuration, object_id='Polarized')

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    BasisGGAPseudoDojoSO.Iron_Medium,
    BasisGGAPseudoDojoSO.Platinum_Medium,
    ]

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

k_point_sampling = KpointDensity(
    density_a=7.0*Angstrom,
    force_timereversal=False,
    )
numerical_accuracy_parameters = NumericalAccuracyParameters(
    density_mesh_cutoff=120.0*Hartree,
    k_point_sampling=k_point_sampling,
    occupation_method=FermiDirac(300.0*Kelvin*boltzmann_constant),
    )

calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    )

bulk_configuration.setCalculator(calculator)

theta_angles = numpy.linspace(0, 90, 7)

for theta in theta_angles:
    # -------------------------------------------------------------
    # Initial State
    # -------------------------------------------------------------
    scaled_spins = [
        (0, 1.0, theta*Degrees, 0.0*Degrees),
        (1, 1.0, theta*Degrees, 0.0*Degrees),
    ]
    initial_spin = InitialSpin(scaled_spins=scaled_spins)
    initial_state = nlread('FePt-SO-SCF.hdf5', object_id='Polarized')[0]

    bulk_configuration.setCalculator(
        calculator,
        initial_spin=initial_spin,
        initial_state=initial_state,
    )
    bulk_configuration.update()

    # -------------------------------------------------------------
    # Total Energy
    # -------------------------------------------------------------
    total_energy = TotalEnergy(bulk_configuration)
    nlsave('FePt-SO-SCF.hdf5', total_energy, object_id='Etot {}'.format(theta))
    nlprint(total_energy)