import numpy as np
import openmm.unit as unit
from abc import abstractmethod, ABC
from scipy.constants import physical_constants


class BaseDynamics(ABC):
    def __init__(self, args, state):
        super().__init__()
        self.start_file = f"./data/{args.molecule}/{state}.pdb"
        self.temperature = args.temperature * unit.kelvin
        self.friction = args.friction / unit.femtoseconds
        self.timestep = args.timestep * unit.femtoseconds
        self.pdb, self.integrator, self.simulation, self.external_force = self.setup()
        self.get_md_info()
        self.simulation.minimizeEnergy()
        self.position = self.report()[0]

    @abstractmethod
    def setup(self):
        pass

    def get_md_info(self):
        self.num_particles = self.simulation.system.getNumParticles()
        m = np.array([
            self.simulation.system.getParticleMass(i).value_in_unit(unit.dalton)
            for i in range(self.num_particles)
        ])
        self.heavy_atoms = m > 1.1
        m = unit.Quantity(m, unit.dalton)
        unadjusted_variance = (
            2 * self.timestep * self.friction * unit.BOLTZMANN_CONSTANT_kB * self.temperature / m[:, None]
        )
        std_SI_units = (
            1 / physical_constants["unified atomic mass unit"][0] *
            unadjusted_variance.value_in_unit(unit.joule / unit.dalton)
        )
        self.std = unit.Quantity(
            np.sqrt(std_SI_units), unit.meter / unit.second
        ).value_in_unit(unit.nanometer / unit.femtosecond)
        self.m = m.value_in_unit(unit.dalton)

    def step(self, forces):
        for i in range(forces.shape[0]):
            self.external_force.setParticleParameters(i, i, forces[i])
        self.external_force.updateParametersInContext(self.simulation.context)
        self.simulation.step(1)

    def report(self):
        state = self.simulation.context.getState(getPositions=True, getForces=True)
        positions = state.getPositions().value_in_unit(unit.nanometer)
        forces = state.getForces().value_in_unit(
            unit.dalton * unit.nanometer / unit.femtosecond ** 2
        )
        return positions, forces

    def reset(self):
        for i in range(len(self.position)):
            self.external_force.setParticleParameters(i, i, [0, 0, 0])
        self.external_force.updateParametersInContext(self.simulation.context)
        self.simulation.context.setPositions(self.position)
        self.simulation.context.setVelocitiesToTemperature(self.temperature)

    def set_temperature(self, temperature):
        self.integrator.setTemperature(temperature * unit.kelvin)

    def energy_function(self, positions):
        forces, potentials = [], []
        for i in range(len(positions)):
            self.simulation.context.setPositions(positions[i])
            state = self.simulation.context.getState(getForces=True, getEnergy=True)
            force = state.getForces().value_in_unit(
                unit.dalton * unit.nanometer / unit.femtosecond ** 2
            )
            potential = state.getPotentialEnergy().value_in_unit(
                unit.kilojoules / unit.mole
            )
            forces.append(force)
            potentials.append(potential)
        return np.array(forces), np.array(potentials)