#!/usr/bin/env amspython
import os
from dftpy.config import DefaultOption, OptionFormat
from dftpy.api.api4ase import DFTpyCalculator
from ase.calculators.calculator import Calculator


class MyCalculator(Calculator):
    implemented_properties = ["energy", "forces", "stress"]

    def __init__(self, config, **kwargs):
        Calculator.__init__(self, **kwargs)
        self.dftpy_calculator = DFTpyCalculator(config=config)

    def calculate(self, atoms, properties=None, system_changes=None):
        super().calculate(atoms, properties, system_changes)
        self.results = dict()
        self.results["energy"] = self.dftpy_calculator.get_potential_energy(atoms)
        self.results["forces"] = self.dftpy_calculator.get_forces(atoms)
        self.results["stress"] = self.dftpy_calculator.get_stress(atoms)


def get_calculator():
    config = DefaultOption()
    config["PATH"]["pppath"] = os.environ.get(
        "DFTPY_DATA_PATH",
        "/home/hellstrom/software/local-pseudopotentials/BLPS/LDA/reci",
    )
    config["PP"]["Al"] = "al.lda.recpot"
    config["OPT"]["method"] = "TN"
    config["KEDF"]["kedf"] = "WT"
    config["JOB"]["calctype"] = "Energy Force"
    config = OptionFormat(config)
    calc = MyCalculator(config=config)

    return calc