from NanoLanguage import *
diagonalization_solver = DiagonalizationSolver(optimize_for_speed_over_memory=True, processes_per_kpoint=1)
algorithm_parameters = AlgorithmParameters(density_matrix_method = diagonalization_solver)

def calculator_hse06_medium(charge=None):
    # ----------------------------------------
    # Exchange-Correlation
    # ----------------------------------------
    exchange_correlation = HybridGGA.HSE06

    # ----------------------------------------
    # Basis Set
    # ----------------------------------------
    basis_set = [
        BasisGGAPseudoDojo.Carbon_Medium,
        BasisGGAPseudoDojo.Oxygen_Medium,
        BasisGGAPseudoDojo.Hydrogen_Medium,
        BasisGGAPseudoDojo.Nitrogen_Medium,
        BasisGGAPseudoDojo.Sulfur_Medium,
        BasisGGAPseudoDojo.Fluorine_Medium,
        BasisGGAPseudoDojo.Lithium_Medium,
        BasisGGAPseudoDojo.Iron_Medium,
        BasisGGAPseudoDojo.Phosphorus_Medium,
        #BasisGGAPseudoDojo.Indium_Medium,
        #BasisGGAPseudoDojo.Gallium_Medium,
        #BasisGGAPseudoDojo.Arsenic_Medium,
    ]

    # Numerical settings.
    k_point_sampling = KpointDensity(
        density_a=7.0 * Angstrom, density_b=7.0 * Angstrom, density_c=7.0 * Angstrom
    )

    numerical_accuracy_parameters = NumericalAccuracyParameters(
        density_mesh_cutoff=110.0 * Hartree,
        k_point_sampling=k_point_sampling,
        occupation_method=FermiDirac(broadening=100.0 * Kelvin),
    )

    iteration_control_parameters = IterationControlParameters(
        max_steps=1000,
        tolerance=1e-05,
        )

    # Calculator.
    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,
        #parallel_parameters=ParallelParameters(processes_per_neb_image=4),
        algorithm_parameters=algorithm_parameters,
        charge=charge,
    )

    return calculator