import numpy
import ase
import gpaw


class ASEGPAWCalculator(ase.calculators.calculator.Calculator):
    def __init__(
        self,
        pbc=[True, True, True],
        cancel_total_force=False,
        charge=0,
        name="atoms",
        **gpaw_kwargs,
    ):
        self.name = name
        self.counter = 1
        self.pbc = pbc
        self.cancel = cancel_total_force
        self._kwargs = gpaw_kwargs
        ase.calculators.calculator.Calculator.__init__(self)
        self._setup_gpaw(charge)

    def calculate(self, atoms=None, properties=None, system_changes=None):
        atoms.center()
        atoms.set_pbc(self.pbc)
        super().calculate(atoms, properties, system_changes)

        self.gpaw.calculate(atoms, properties, system_changes)
        self.results = self.gpaw.results
        # remove total force on the molecule
        if self.cancel:
            molecule_force = self.results["forces"].sum(axis=0)
            per_atom_force = molecule_force / self.results["forces"].shape[0]
            self.results["forces"] -= per_atom_force

    def _setup_gpaw(self, charge):
        self.charge = charge
        txt = self.name
        if self.counter > 1:
            txt = txt + f"_{self.counter}"
        txt = txt + ".txt"
        self.gpaw = gpaw.GPAW(txt=txt, **self._kwargs)
        self.counter += 1

    @property
    def implemented_properties(self):
        return self.gpaw.implemented_properties


# AMS looks for "get_calculator"
get_calculator = ASEGPAWCalculator