import pytest


def test_double_relax(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally

    from atomate2.vasp.flows.core import DoubleRelaxMaker
    from atomate2.vasp.jobs.core import RelaxMaker
    from atomate2.vasp.schemas.task import TaskDocument
    from atomate2.vasp.sets.core import RelaxSetGenerator

    # mapping from job name to directory containing test files
    ref_paths = {
        "relax 1": "Si_double_relax/relax_1",
        "relax 2": "Si_double_relax/relax_2",
    }

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "relax 1": {"incar_settings": ["NSW", "ISMEAR"]},
        "relax 2": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # generate flow
    flow = DoubleRelaxMaker().make(si_structure)

    # Run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # validate output
    output1 = responses[flow.jobs[0].uuid][1].output
    output2 = responses[flow.jobs[1].uuid][1].output

    assert isinstance(output1, TaskDocument)
    assert output1.output.energy == pytest.approx(-10.85043620)
    assert output2.output.energy == pytest.approx(-10.84177648)

    # Now try with two identical but non-default makers
    ref_paths = {
        "relax 1": "Si_double_relax_swaps/swap1/relax_1",
        "relax 2": "Si_double_relax_swaps/swap1/relax_2",
    }
    fake_run_vasp_kwargs = {
        "relax 1": {"incar_settings": ["LREAL"]},
        "relax 2": {"incar_settings": ["LREAL"]},
    }
    mock_vasp(ref_paths, fake_run_vasp_kwargs)
    my_custom_maker = RelaxMaker(
        input_set_generator=RelaxSetGenerator(user_incar_settings={"LREAL": False})
    )
    flow = DoubleRelaxMaker(
        relax_maker1=my_custom_maker, relax_maker2=my_custom_maker
    ).make(si_structure)
    responses = run_locally(flow, create_folders=True, ensure_success=True)
    output1 = responses[flow.jobs[0].uuid][1].output
    output2 = responses[flow.jobs[1].uuid][1].output

    # Try the same as above but with the .from_relax_maker() class method
    flow = DoubleRelaxMaker.from_relax_maker(my_custom_maker).make(si_structure)
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # Try DoubleRelaxMaker with a non-default second maker only
    ref_paths = {
        "relax 1": "Si_double_relax_swaps/swap2/relax_1",
        "relax 2": "Si_double_relax_swaps/swap2/relax_2",
    }
    fake_run_vasp_kwargs = {
        "relax 1": {"incar_settings": ["LREAL"]},
        "relax 2": {"incar_settings": ["LREAL"]},
    }
    mock_vasp(ref_paths, fake_run_vasp_kwargs)
    flow = DoubleRelaxMaker(relax_maker2=my_custom_maker).make(si_structure)
    responses = run_locally(flow, create_folders=True, ensure_success=True)


def test_band_structure(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally
    from pymatgen.electronic_structure.bandstructure import (
        BandStructure,
        BandStructureSymmLine,
    )

    from atomate2.vasp.flows.core import BandStructureMaker
    from atomate2.vasp.schemas.calculation import VaspObject

    # mapping from job name to directory containing test files
    ref_paths = {
        "non-scf line": "Si_band_structure/non-scf_line",
        "non-scf uniform": "Si_band_structure/non-scf_uniform",
        "static": "Si_band_structure/static",
    }

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "non-scf line": {"incar_settings": ["NSW", "ISMEAR"]},
        "non-scf uniform": {"incar_settings": ["NSW", "ISMEAR"]},
        "static": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # Generate the flow
    flow = BandStructureMaker().make(si_structure)

    # Run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # Additional validation on the outputs of the jobs; e.g.,
    static_output = responses[flow.jobs[0].uuid][1].output
    uniform_output = responses[flow.jobs[1].uuid][1].output
    line_output = responses[flow.jobs[2].uuid][1].output

    assert static_output.output.energy == pytest.approx(-10.85037078)
    assert static_output.included_objects is None
    assert set(uniform_output.vasp_objects) == {
        VaspObject.BANDSTRUCTURE,
        VaspObject.DOS,
    }

    assert isinstance(
        line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructureSymmLine
    )
    assert isinstance(
        uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructure
    )


def test_uniform_band_structure(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally
    from pymatgen.electronic_structure.bandstructure import BandStructure

    from atomate2.vasp.flows.core import UniformBandStructureMaker
    from atomate2.vasp.schemas.calculation import VaspObject

    # mapping from job name to directory containing test files
    ref_paths = {
        "non-scf uniform": "Si_band_structure/non-scf_uniform",
        "static": "Si_band_structure/static",
    }

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "non-scf uniform": {"incar_settings": ["NSW", "ISMEAR"]},
        "static": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # Generate the flow
    flow = UniformBandStructureMaker().make(si_structure)

    # Run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # Additional validation on the outputs of the jobs; e.g.,
    static_output = responses[flow.jobs[0].uuid][1].output
    uniform_output = responses[flow.jobs[1].uuid][1].output

    assert static_output.output.energy == pytest.approx(-10.85037078)
    assert static_output.included_objects is None
    assert set(uniform_output.vasp_objects) == {
        VaspObject.BANDSTRUCTURE,
        VaspObject.DOS,
    }
    assert isinstance(
        uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructure
    )


def test_line_mode_band_structure(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally
    from pymatgen.electronic_structure.bandstructure import BandStructureSymmLine

    from atomate2.vasp.flows.core import LineModeBandStructureMaker
    from atomate2.vasp.schemas.calculation import VaspObject

    # mapping from job name to directory containing test files
    ref_paths = {
        "non-scf line": "Si_band_structure/non-scf_line",
        "static": "Si_band_structure/static",
    }

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "non-scf line": {"incar_settings": ["NSW", "ISMEAR"]},
        "static": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # Generate the flow
    flow = LineModeBandStructureMaker().make(si_structure)

    # Run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # Additional validation on the outputs of the jobs; e.g.,
    static_output = responses[flow.jobs[0].uuid][1].output
    line_output = responses[flow.jobs[1].uuid][1].output

    assert static_output.output.energy == pytest.approx(-10.85037078)
    assert static_output.included_objects is None
    assert isinstance(
        line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructureSymmLine
    )


def test_hse_band_structure(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally
    from pymatgen.electronic_structure.bandstructure import (
        BandStructure,
        BandStructureSymmLine,
    )

    from atomate2.vasp.flows.core import HSEBandStructureMaker
    from atomate2.vasp.schemas.calculation import VaspObject

    # mapping from job name to directory containing test files
    ref_paths = {
        "hse band structure line": "Si_hse_band_structure/hse_band_structure_line",
        "hse band structure uniform": "Si_hse_band_structure/hse_band_structure_uniform",
        "hse static": "Si_hse_band_structure/hse_static",
    }

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "hse band structure line": {"incar_settings": ["NSW", "ISMEAR"]},
        "hse band structure uniform": {"incar_settings": ["NSW", "ISMEAR"]},
        "hse static": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # generate flow
    flow = HSEBandStructureMaker().make(si_structure)
    flow.jobs[0].maker.input_set_generator.user_incar_settings["KSPACING"] = 0.4

    # Run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # validation on the outputs
    static_output = responses[flow.jobs[0].uuid][1].output
    uniform_output = responses[flow.jobs[1].uuid][1].output
    line_output = responses[flow.jobs[2].uuid][1].output

    assert static_output.output.energy == pytest.approx(-12.52887403)
    assert static_output.included_objects is None
    assert set(uniform_output.vasp_objects) == {
        VaspObject.BANDSTRUCTURE,
        VaspObject.DOS,
    }

    assert isinstance(
        line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructureSymmLine
    )
    assert isinstance(
        uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructure
    )


def test_hse_uniform_band_structure(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally
    from pymatgen.electronic_structure.bandstructure import BandStructure

    from atomate2.vasp.flows.core import HSEUniformBandStructureMaker
    from atomate2.vasp.schemas.calculation import VaspObject

    # mapping from job name to directory containing test files
    ref_paths = {
        "hse band structure uniform": "Si_hse_band_structure/hse_band_structure_uniform",
        "hse static": "Si_hse_band_structure/hse_static",
    }

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "hse band structure uniform": {"incar_settings": ["NSW", "ISMEAR"]},
        "hse static": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # generate flow
    flow = HSEUniformBandStructureMaker().make(si_structure)
    flow.jobs[0].maker.input_set_generator.user_incar_settings["KSPACING"] = 0.4

    # Run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # validation on the outputs
    static_output = responses[flow.jobs[0].uuid][1].output
    uniform_output = responses[flow.jobs[1].uuid][1].output

    assert static_output.output.energy == pytest.approx(-12.52887403)
    assert static_output.included_objects is None
    assert set(uniform_output.vasp_objects) == {
        VaspObject.BANDSTRUCTURE,
        VaspObject.DOS,
    }
    assert isinstance(
        uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructure
    )


def test_hse_line_mode_band_structure(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally
    from pymatgen.electronic_structure.bandstructure import BandStructureSymmLine

    from atomate2.vasp.flows.core import HSELineModeBandStructureMaker
    from atomate2.vasp.schemas.calculation import VaspObject

    # mapping from job name to directory containing test files
    ref_paths = {
        "hse band structure line": "Si_hse_band_structure/hse_band_structure_line",
        "hse static": "Si_hse_band_structure/hse_static",
    }

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "hse band structure line": {"incar_settings": ["NSW", "ISMEAR"]},
        "hse static": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # generate flow
    flow = HSELineModeBandStructureMaker().make(si_structure)
    flow.jobs[0].maker.input_set_generator.user_incar_settings["KSPACING"] = 0.4

    # Run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # validation on the outputs
    static_output = responses[flow.jobs[0].uuid][1].output
    line_output = responses[flow.jobs[1].uuid][1].output

    assert static_output.output.energy == pytest.approx(-12.52887403)
    assert static_output.included_objects is None
    assert isinstance(
        line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructureSymmLine
    )


def test_optics(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally

    from atomate2.vasp.flows.core import OpticsMaker
    from atomate2.vasp.schemas.task import TaskDocument

    # mapping from job name to directory containing test files
    ref_paths = {"optics": "Si_optics/optics", "static": "Si_optics/static"}

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "optics": {"incar_settings": ["NSW", "ISMEAR"]},
        "static": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # generate flow
    flow = OpticsMaker().make(si_structure)

    # run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # validation on the outputs
    output1 = responses[flow.jobs[0].uuid][1].output
    output2 = responses[flow.jobs[1].uuid][1].output
    assert isinstance(output1, TaskDocument)
    assert isinstance(output2, TaskDocument)
    assert output1.output.energy == pytest.approx(-10.85037078)
    assert output2.calcs_reversed[0].output.frequency_dependent_dielectric.real[0] == [
        13.6062,
        13.6063,
        13.6062,
        0.0,
        0.0,
        0.0,
    ]


def test_hse_optics(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally

    from atomate2.vasp.flows.core import HSEOpticsMaker
    from atomate2.vasp.schemas.task import TaskDocument

    # mapping from job name to directory containing test files
    ref_paths = {
        "hse optics": "Si_hse_optics/hse_optics",
        "hse static": "Si_hse_optics/hse_static",
    }

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {
        "hse_optics": {"incar_settings": ["NSW", "ISMEAR"]},
        "hse_static": {"incar_settings": ["NSW", "ISMEAR"]},
    }

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # generate flow
    flow = HSEOpticsMaker().make(si_structure)
    flow.update_maker_kwargs(
        {"_set": {"input_set_generator->user_incar_settings->KSPACING": 0.5}},
        dict_mod=True,
    )

    # run the flow or job and ensure that it finished running successfully
    responses = run_locally(flow, create_folders=True, ensure_success=True)

    # validation on the outputs
    output1 = responses[flow.jobs[0].uuid][1].output
    output2 = responses[flow.jobs[1].uuid][1].output
    assert isinstance(output1, TaskDocument)
    assert isinstance(output2, TaskDocument)
    assert output1.output.energy == pytest.approx(-12.41767353)
    assert output2.calcs_reversed[0].output.frequency_dependent_dielectric.real[0] == [
        13.8738,
        13.8738,
        13.8738,
        0.0,
        0.0,
        0.0,
    ]
