# %% nio

# Set up lattice
lattice = Rhombohedral(5.138*Angstrom, 33.5573*Degrees)

# Define elements
elements = [Nickel, Oxygen, Nickel, Oxygen]

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

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

nio_name = "nio"


# %% InitialSpin

scaled_spins = [
    1.0,
    0.0,
    -1.0,
    0.0,
]
initial_spin = InitialSpin(scaled_spins=scaled_spins)
nlsave('NiO_DFT_abU_results.hdf5', initial_spin)


# %% Set LCAOCalculator

# %% LCAOCalculator

# ----------------------------------------
# Exchange-Correlation
# ----------------------------------------
exchange_correlation = ExchangeCorrelation(
    exchange=PerdewBurkeErnzerhofExchange,
    correlation=PerdewBurkeErnzerhofCorrelation,
    hubbard_term=Dual,
    number_of_spins=2,
    spin_orbit=False,
    dft_half_enabled=False,
)

# ----------------------------------------
# Basis Set
# ----------------------------------------
nickel_3s = ConfinedOrbital(
    principal_quantum_number=3,
    angular_momentum=0,
    radial_cutoff_radius=1.31640625 * Angstrom,
    confinement_start_radius=1.053125 * Angstrom,
    additional_charge=0,
    confinement_strength=12.5 * Hartree,
    confinement_power=2,
    radial_step_size=0.001 * Bohr,
)

nickel_3p = ConfinedOrbital(
    principal_quantum_number=3,
    angular_momentum=1,
    radial_cutoff_radius=1.4765625 * Angstrom,
    confinement_start_radius=1.18125 * Angstrom,
    additional_charge=0,
    confinement_strength=12.5 * Hartree,
    confinement_power=2,
    radial_step_size=0.001 * Bohr,
)

nickel_4s = ConfinedOrbital(
    principal_quantum_number=4,
    angular_momentum=0,
    radial_cutoff_radius=4.3671875 * Angstrom,
    confinement_start_radius=3.49375 * Angstrom,
    additional_charge=0,
    confinement_strength=12.5 * Hartree,
    confinement_power=2,
    radial_step_size=0.001 * Bohr,
)

nickel_3d = ConfinedOrbital(
    principal_quantum_number=3,
    angular_momentum=2,
    radial_cutoff_radius=2.546875 * Angstrom,
    confinement_start_radius=2.0375 * Angstrom,
    additional_charge=0,
    confinement_strength=12.5 * Hartree,
    confinement_power=2,
    radial_step_size=0.001 * Bohr,
)

nickel_3p_0 = HydrogenOrbital(
    principal_quantum_number=3,
    angular_momentum=1,
    radial_cutoff_radius=4.88671875 * Angstrom,
    confinement_start_radius=3.909375 * Angstrom,
    charge=6,
    confinement_strength=12.5 * Hartree,
    confinement_power=2,
    radial_step_size=0.001 * Bohr,
)

nickel_3d_0 = HydrogenOrbital(
    principal_quantum_number=3,
    angular_momentum=2,
    radial_cutoff_radius=2.738671875 * Angstrom,
    confinement_start_radius=2.1909375 * Angstrom,
    charge=5.2,
    confinement_strength=12.5 * Hartree,
    confinement_power=2,
    radial_step_size=0.001 * Bohr,
)

nickel_4p = ConfinedOrbital(
    principal_quantum_number=4,
    angular_momentum=1,
    radial_cutoff_radius=3.25390625 * Angstrom,
    confinement_start_radius=2.603125 * Angstrom,
    additional_charge=2,
    confinement_strength=12.5 * Hartree,
    confinement_power=2,
    radial_step_size=0.001 * Bohr,
)

NickelBasis = BasisSet(
    element=PeriodicTable.Nickel,
    orbitals=[
        nickel_3s,
        nickel_3p,
        nickel_4s,
        nickel_3d,
        nickel_3p_0,
        nickel_3d_0,
        nickel_4p,
    ],
    occupations=[2.0, 6.0, 2.0, 8.0, 0.0, 0.0, 0.0],
    hubbard_u=LocalScreenedCoulombCorrection(
        shell_selection=[False, False, False, True, False, True, False]
    ),
    dft_half_parameters=Automatic,
    filling_method=SphericalSymmetric,
    onsite_spin_orbit_split=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] * eV,
    pseudopotential=NormConservingPseudoPotential(
        "normconserving/pseudodojo/gga/standard/28_Ni.upf",
        local_potential_cutoff_threshold=1e-06 * Hartree,
        local_potential_cutoff_radius=6.0 * Bohr,
    ),
)

basis_set = [
    BasisGGAPseudoDojo.Oxygen_Medium,
    NickelBasis,
]

density_mesh_cutoff = calculateDefaultDensityMeshCutoff(
    calculator_type=LCAOCalculator,
    configuration=nio,
    basis_set=basis_set,
    wave_function_cutoff=None,
)

k_point_sampling = KpointDensity(
    density_a=4.0 * Angstrom, density_b=4.0 * Angstrom, density_c=4.0 * Angstrom
)

numerical_accuracy_parameters = NumericalAccuracyParameters(
    density_mesh_cutoff=density_mesh_cutoff,
    k_point_sampling=k_point_sampling,
    occupation_method=FermiDirac(broadening=300.0 * Kelvin),
)

iteration_control_parameters = IterationControlParameters(tolerance=1e-05)

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


# %% Set Calculator

nio.setCalculator(calculator=calculator, initial_spin=initial_spin)

nio.update()

nlsave('NiO_DFT_abU_results.hdf5', nio)


# %% ProjectedDensityOfStates

kpoints = MonkhorstPackGrid(na=60, nb=60, nc=60)

projected_density_of_states = ProjectedDensityOfStates(configuration=nio, kpoints=kpoints)
nlsave('NiO_DFT_abU_results.hdf5', projected_density_of_states)


# %% MullikenPopulation

mulliken_population = MullikenPopulation(configuration=nio)
nlsave('NiO_DFT_abU_results.hdf5', mulliken_population)