from qiskit_nature.hdf5 import load_from_hdf5, save_to_hdf5
from qiskit_nature.drivers.second_quantization import PySCFDriver, GaussianForcesDriver
from qiskit_nature.properties.second_quantization.vibrational.bases import HarmonicBasis

electronic_driver = PySCFDriver(
    atom="H 0.0 0.0 0.0; Li 0.0 0.0 1.3",
    basis="sto-3g",
)

electronic_structure_driver_result = electronic_driver.run()

print(electronic_structure_driver_result)

save_to_hdf5(electronic_structure_driver_result, "tmp_electronic_structure_driver_result.hdf5", replace=True)

new_electronic_structure_driver_result = load_from_hdf5("tmp_electronic_structure_driver_result.hdf5")

print(new_electronic_structure_driver_result)

vibrational_driver = GaussianForcesDriver(
    logfile="test/drivers/second_quantization/gaussiand/test_driver_gaussian_log_C01.txt"
)

vibrational_structure_driver_result = vibrational_driver.run()
vibrational_structure_driver_result.basis = HarmonicBasis(num_modals_per_mode=2)

print(vibrational_structure_driver_result)

save_to_hdf5(vibrational_structure_driver_result, "tmp_vibrational_structure_driver_result.hdf5", replace=True)

new_vibrational_structure_driver_result = load_from_hdf5("tmp_vibrational_structure_driver_result.hdf5")

print(new_vibrational_structure_driver_result)

from __future__ import annotations
from typing import Protocol

import h5py


class HDF5Storable(Protocol):
    """A Protocol implemented by those classes which support conversion methods for HDF5."""

    def to_hdf5(self, parent: h5py.Group) -> None:
        """Stores this instance in an HDF5 group inside of the provided parent group.
        ...
        """
        ...

    @staticmethod
    def from_hdf5(h5py_group: h5py.Group) -> HDF5Storable:
        """Constructs a new instance from the data stored in the provided HDF5 group.
        ...
        """
        ...

class Molecule:
    VERSION = 1
    
    def to_hdf5(self, parent: h5py.Group) -> None:
        group = parent.require_group(self.__class__.__name__)
        group.attrs["__class__"] = self.__class__.__name__
        group.attrs["__module__"] = self.__class__.__module__
        group.attrs["__version__"] = self.VERSION

        geometry_group = group.create_group("geometry", track_order=True)
        for idx, geom in enumerate(self._geometry):
            symbol, coords = geom
            geometry_group.create_dataset(str(idx), data=coords)
            geometry_group[str(idx)].attrs["symbol"] = symbol

        group.attrs["units"] = self.units.value
        group.attrs["multiplicity"] = self.multiplicity
        group.attrs["charge"] = self.charge

        if self._masses:
            group.create_dataset("masses", data=self._masses)

    @staticmethod
    def from_hdf5(h5py_group: h5py.Group) -> Molecule:
        geometry = []
        for atom in h5py_group["geometry"].values():
            geometry.append((atom.attrs["symbol"], list(atom[...])))

        units: UnitsType
        for unit in UnitsType:
            if unit.value == h5py_group.attrs["units"]:
                units = unit
                break
        else:
            units = UnitsType.ANGSTROM

        multiplicity = h5py_group.attrs["multiplicity"]
        charge = h5py_group.attrs["charge"]

        masses = None
        if "masses" in h5py_group.keys():
            masses = list(h5py_group["masses"])

        return Molecule(
            geometry,
            multiplicity=multiplicity,
            charge=charge,
            units=units,
            masses=masses,
        )

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright

