from qspoc.qm.gaussian16 import G16Computer
from qspoc.qm.xtb import XTBComputer
from qspoc.qm.orca6 import ORCA6Computer  # TODO: need finish this part.


class QMComputer:
    def __init__(self):
        self.xtb_computer = None
        self.g16_computer = None

    def load_xtb_computer(self, xtb_path: str):
        self.xtb_computer = XTBComputer(xtb_path=xtb_path)

    def load_g16_computer(self, g16_path: str):
        self.g16_computer = G16Computer(g16_path=g16_path)

    def load_orca_computer(self, orca_path: str):
        self.orca_computer = ORCA6Computer(orca_path=orca_path)

    def get_computer(self, method: str):
        if method == "xtb":
            assert self.xtb_computer is not None, "XTB computer is not loaded. Please define the path to xtb."
            return self.xtb_computer
        elif method == "g16":
            assert self.g16_computer is not None, "Gaussian16 computer is not loaded. Please define the path to g16."
            return self.g16_computer
        elif method == "orca":
            assert self.orca_computer is not None, "ORCA computer is not loaded. Please define the path to orca."
            return self.orca_computer
        else:
            raise ValueError(f"Unsupported QM method: {method}.")
