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

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

# Set up lattice
lattice = SimpleTetragonal(2.866*Angstrom, 45.34172044506112*Angstrom)

# Define elements
elements = [Iron, Iron, Iron, Iron, Iron, Iron, Iron, Iron, Iron, Iron, Iron,
            Oxygen, Magnesium, Magnesium, Oxygen, Oxygen, Magnesium, Magnesium,
            Oxygen, Oxygen, Magnesium, Magnesium, Oxygen, Oxygen, Magnesium,
            Iron, Iron, Iron, Iron, Iron, Iron, Iron, Iron, Iron, Iron]

# Define coordinates
fractional_coordinates = [[ 0.248796793653,  0.248796793651,  0.015236052172],
                          [ 0.749240954235,  0.749240954234,  0.045902009605],
                          [ 0.248800401505,  0.248800401504,  0.076577038988],
                          [ 0.749235217307,  0.749235217306,  0.107154778896],
                          [ 0.248894116126,  0.248894116127,  0.137846827345],
                          [ 0.749635120547,  0.749635120548,  0.168439264182],
                          [ 0.249677932107,  0.249677932108,  0.199221654535],
                          [ 0.750847797286,  0.750847797287,  0.229845723164],
                          [ 0.251186736882,  0.251186736884,  0.260914588726],
                          [ 0.752082597185,  0.752082597186,  0.292712488503],
                          [ 0.25222413526 ,  0.252224135259,  0.322027504254],
                          [ 0.254422990082,  0.254422990082,  0.370814295032],
                          [ 0.755012706524,  0.755012706523,  0.369605811548],
                          [ 0.252067899181,  0.25206789918 ,  0.419042655712],
                          [ 0.752282302675,  0.752282302673,  0.418838489762],
                          [ 0.251661578111,  0.251661578111,  0.467426995975],
                          [ 0.751983143594,  0.751983143594,  0.467463672673],
                          [ 0.248751932842,  0.248751932842,  0.515957003857],
                          [ 0.749304010308,  0.749304010306,  0.515884609686],
                          [ 0.249949783747,  0.249949783747,  0.564269640379],
                          [ 0.750635912612,  0.75063591261 ,  0.564342821765],
                          [ 0.246900521972,  0.246900521972,  0.612572053417],
                          [ 0.747736042527,  0.747736042525,  0.612688672164],
                          [ 0.249245805873,  0.249245805872,  0.660478948181],
                          [ 0.749465136015,  0.749465136013,  0.661896165756],
                          [ 0.248434621128,  0.248434621125,  0.708899029361],
                          [ 0.74905639095 ,  0.749056390949,  0.738051844042],
                          [ 0.248850456433,  0.24885045643 ,  0.769766600765],
                          [ 0.749370078138,  0.749370078136,  0.800746712211],
                          [ 0.248874627157,  0.248874627152,  0.831322119996],
                          [ 0.749284193202,  0.7492841932  ,  0.862064593914],
                          [ 0.248816407684,  0.248816407683,  0.892637429232],
                          [ 0.749242343327,  0.749242343324,  0.923315868123],
                          [ 0.248804076995,  0.248804076993,  0.953895871831],
                          [ 0.749225236779,  0.749225236777,  0.984566314048]]

# 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, 30, 31, 32, 33, 34])

# -------------------------------------------------------------
# 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=140.0*Hartree,
    k_point_sampling=k_point_sampling,
    )

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

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('Fe21-MgO7-Polarized.hdf5', bulk_configuration, object_id='Polarized')

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    BasisGGAPseudoDojoSO.Oxygen_Medium,
    BasisGGAPseudoDojoSO.Magnesium_Medium,
    BasisGGAPseudoDojoSO.Iron_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=140.0*Hartree,
    k_point_sampling=k_point_sampling,
    )

iteration_control_parameters = NonSelfconsistent

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

theta_angles = [0, 90]

for theta in theta_angles:

    # -------------------------------------------------------------
    # Initial State
    # -------------------------------------------------------------
    scaled_spins = [(0, 1.0, theta*Degrees, 0.0*Degrees)] * len(bulk_configuration)
    initial_spin = InitialSpin(scaled_spins=scaled_spins)

    # Load the polarized configuration to be used as initial state
    initial_state = nlread('Fe21-MgO7-Polarized.hdf5', object_id='Polarized')[0]

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

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

    # -------------------------------------------------------------
    # Orbital Moment
    # -------------------------------------------------------------
    kpoints = KpointDensity(
        density_a=17.0*Angstrom,
        force_timereversal=False,
        )

    orbital_moment = OrbitalMoment(
        configuration=bulk_configuration,
        kpoints=kpoints,
        temperature=300.0*Kelvin,
    )
    nlsave('Fe21-MgO7-OrbitalMoment.hdf5', orbital_moment, object_id='Orbital moment {}'.format(theta))

    mulliken_population = MullikenPopulation(bulk_configuration)
    nlsave('Fe21-MgO7-OrbitalMoment.hdf5', mulliken_population, object_id='Mulliken {}'.format(theta))


import os

if processIsMaster():
    os.remove('Fe21-MgO7-Polarized.hdf5')