
import pytest
import torch
import warnings
import numpy as np

from bgflow.distribution.energy.openmm import OpenMMBridge, OpenMMEnergy


class OneParticleTestBridge(OpenMMBridge):
    """OpenMM bridge for a system with one particle"""

    def __init__(self, n_workers=1, n_simulation_steps=0):
        try:
            with warnings.catch_warnings():
                warnings.simplefilter(
                    "ignore", DeprecationWarning
                )  # ignore warnings inside OpenMM
                from simtk import openmm, unit
        except ImportError:
            pytest.skip("Test requires OpenMM.")

        # a system with one particle and an external dummy force
        system = openmm.System()
        system.addParticle(1.0 * unit.amu)
        force = openmm.CustomExternalForce("x")
        force.addParticle(0)
        system.addForce(force)

        super(OneParticleTestBridge, self).__init__(
            system,
            openmm.LangevinIntegrator(
                300 * unit.kelvin, 1.0 / unit.picoseconds, 1.0 * unit.femtoseconds
            ),
            n_workers=n_workers,
            n_simulation_steps=n_simulation_steps,
        )



@pytest.mark.parametrize("n_workers", [1, 2, 8])
@pytest.mark.parametrize("n_simulation_steps", [0, 100])
def test_openmm_bridge_evaluate_dummy(n_workers, n_simulation_steps):
    """Test if we can evaluate an energy; skip test if openmm is not installed."""

    bridge = OneParticleTestBridge(n_workers, n_simulation_steps)

    from simtk import unit

    # test forces and energies generated by the bridge
    batch_size = 4
    positions = torch.tensor([[0.1, 0.0, 0.0]] * batch_size)
    kT = unit.MOLAR_GAS_CONSTANT_R * 300 * unit.kelvin
    energies, forces, *_ = bridge.evaluate(positions)
    assert energies.shape == torch.Size([batch_size, 1])
    assert forces.shape == torch.Size([batch_size, 3])
    if n_simulation_steps == 0:
        assert energies.numpy()[0] == pytest.approx(
            [0.1 * unit.kilojoule_per_mole / kT], abs=1e-8, rel=0
        )
        assert forces.numpy()[0] == pytest.approx(
            [-1.0 * unit.kilojoule_per_mole / kT, 0.0, 0.0], abs=1e-8, rel=0
        )
    else:
        assert energies.numpy()[0] != pytest.approx(
            [0.1 * unit.kilojoule_per_mole / kT], abs=1e-8, rel=0
        )


# run 'pytest --durations 0' to print the time it takes to evaluate
# change batch_size and n_workers to profile
@pytest.mark.parametrize("testsystem_name", ["AlanineDipeptideImplicit", "WaterBox"])
@pytest.mark.parametrize("batch_size", [1, 8])
@pytest.mark.parametrize("n_workers", [1, 8])
def test_openmm_bridge_evaluate_openmmtools_testsystem(
    testsystem_name, batch_size, n_workers
):
    """Test if we can evaluate an energy; skip test if openmm is not installed."""
    warnings.filterwarnings(
        "ignore", category=RuntimeWarning, message="numpy.ufunc size changed"
    )

    if testsystem_name == "WaterBox" and n_workers == 1:
        pytest.skip()
    # prevent openmm hanging (this is possibly a bug in their PME code)
    # see https://github.com/openmm/openmm/issues/2602

    try:
        with warnings.catch_warnings():
            warnings.simplefilter(
                "ignore", DeprecationWarning
            )  # ignore warnings inside OpenMM
            from simtk import openmm, unit
            from openmmtools import testsystems

            Testsystem = getattr(testsystems, testsystem_name)
    except ImportError:
        pytest.skip("Test requires OpenMM and openmmtools.")

    # alanine dipeptide testsystem
    testsystem = Testsystem()
    system = testsystem.system

    # test forces and energies generated by the bridge
    bridge = OpenMMBridge(
        system,
        openmm.LangevinIntegrator(
            300 * unit.kelvin, 1.0 / unit.picoseconds, 1.0 * unit.femtoseconds
        ),
        n_workers=n_workers,
    )
    batch = torch.tensor(
        np.array(
            [np.ravel(testsystem.positions.value_in_unit(unit.nanometer))] * batch_size
        )
    )
    energies, forces, *_ = bridge.evaluate(batch)
    assert energies.shape == torch.Size([batch_size, 1])
    assert forces.shape == torch.Size([batch_size, 3 * len(testsystem.positions)])


def test_openmm_bridge_cache():
    """Test if hashing and caching works."""
    bridge = OneParticleTestBridge()
    omm_energy = OpenMMEnergy(bridge=bridge)
    omm_energy._energy(torch.tensor([[0.1, 0.0, 0.0]] * 2))
    hash1 = omm_energy._last_batch
    omm_energy._energy(torch.tensor([[0.2, 0.0, 0.0]] * 2))
    assert omm_energy._last_batch != hash1
    omm_energy._energy(torch.tensor([[0.1, 0.0, 0.0]] * 2))
    assert omm_energy._last_batch == hash1
    omm_energy._energy(torch.tensor([[0.1, 0.0, 0.0]] * 2))

    # test if forces are in the same memory location for same input batch
    force_address = hex(id(omm_energy.bridge.last_forces))
    force = (
        omm_energy.bridge.last_forces
    )  # retain a pointer to last forces so that memory is not freed
    assert (
        hex(id(omm_energy.force(torch.tensor([[0.1, 0.0, 0.0]] * 2)))) == force_address
    )
    assert (
        hex(id(omm_energy.force(torch.tensor([[0.1, 0.0, 0.0]] * 2)))) == force_address
    )
    assert (
        hex(id(omm_energy.force(torch.tensor([[0.2, 0.0, 0.0]] * 2)))) != force_address
    )
    assert (
        hex(id(omm_energy.force(torch.tensor([[0.1, 0.0, 0.0]] * 2)))) != force_address
    )

    # suppress flake8 F841 warning
    force


# Returned log path probability ratios / positions are tested in the SNF tests
