import numpy as np
from ase.calculators.calculator import Calculator, all_changes

try:
    from openmm import unit, openmm
    from openmm.app import Simulation
except ImportError:
    unit = None
    openmm = None
    Simulation = None


KJMOL_PER_KCALMOL = 4.184
NM_PER_ANG = 0.1


class OpenMMProteinLigandCalculator(Calculator):
    """
    ASE calculator backed by an already-built OpenMM System/Context.

    Assumptions:
    - atom ordering in ASE matches OpenMM particle ordering exactly
    - positions in ASE are in Angstrom
    - OpenMM system/topology/integrator are already prepared outside
    """

    implemented_properties = ["energy", "forces"]

    def __init__(
        self,
        simulation,
        mask_movable: np.ndarray = None,
        subtract_reference_energy: bool = False,
        reference_energy_kjmol: float = 0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        if Simulation is None:
            raise ImportError("OpenMM is required for OpenMMProteinLigandCalculator")

        self.simulation = simulation
        self.context = simulation.context
        self.mask_movable = mask_movable
        self.subtract_reference_energy = subtract_reference_energy
        self.reference_energy_kjmol = reference_energy_kjmol

        # Cache to avoid unnecessary recomputation on identical coordinates
        self._last_positions = None
        self._last_energy = None
        self._last_forces = None

    def calculate(
        self,
        atoms=None,
        properties=("energy", "forces"),
        system_changes=all_changes,
    ):
        super().calculate(atoms, properties, system_changes)

        positions_ang = self.atoms.get_positions()
        positions_nm = positions_ang * NM_PER_ANG

        # Optional cache
        if self._last_positions is not None:
            if np.array_equal(positions_ang, self._last_positions):
                self.results["energy"] = self._last_energy
                self.results["forces"] = self._last_forces.copy()
                return

        # Update positions in OpenMM
        self.context.setPositions(positions_nm * unit.nanometer)

        # Query energy + forces
        state = self.context.getState(getEnergy=True, getForces=True)

        energy_kjmol = state.getPotentialEnergy().value_in_unit(
            unit.kilojoule_per_mole
        )
        forces_kjmol_per_nm = state.getForces(asNumpy=True).value_in_unit(
            unit.kilojoule_per_mole / unit.nanometer
        )

        # Convert OpenMM -> ASE:
        # energy: kJ/mol is already ASE's default energy unit convention
        # force: kJ/mol/nm -> kJ/mol/Ang = divide by 10
        forces_kjmol_per_ang = np.asarray(forces_kjmol_per_nm) / 10.0

        if self.subtract_reference_energy:
            energy_kjmol -= self.reference_energy_kjmol

        # Optional: freeze some atoms from GA side
        if self.mask_movable is not None:
            forces_kjmol_per_ang = forces_kjmol_per_ang.copy()
            forces_kjmol_per_ang[~self.mask_movable] = 0.0

        self.results["energy"] = float(energy_kjmol)
        self.results["forces"] = forces_kjmol_per_ang

        self._last_positions = positions_ang.copy()
        self._last_energy = float(energy_kjmol)
        self._last_forces = forces_kjmol_per_ang.copy()
