from pyscf import lib, gto
from tblite.interface import Calculator

class XTB(lib.StreamObject):
    def __init__(self, mol, method="GFN2-xTB"):
        self.mol = mol
        self.method = method
        self.verbose = mol.verbose
        self.stdout = mol.stdout
        self.e_tot = None
        self._last_grad = None

    def kernel(self):
        coords = self.mol.atom_coords()
        numbers = self.mol.atom_charges()
        
        calc = Calculator(
            method=self.method,
            numbers=numbers,
            positions=coords,
            charge=self.mol.charge,
            uhf=self.mol.spin
        )
        
        res = calc.singlepoint()
        self.e_tot = res.get("energy")
        self._last_grad = res.get("gradient")
        
        if self.verbose >= lib.logger.NOTE:
            lib.logger.note(self, 'xTB energy = %.15g', self.e_tot)
            
        return self.e_tot

    def Gradients(self):
        return Gradients(self)

class Gradients(lib.StreamObject):
    def __init__(self, method):
        self.base = method
        self.mol = method.mol
        self.verbose = method.verbose
        self.stdout = method.stdout

    def kernel(self):
        if self.base.e_tot is None or self.base._last_grad is None:
             self.base.kernel()
        return self.base._last_grad

# Inject to Mole
gto.Mole.sXTB = lib.class_as_method(XTB)
