# [1] https://pubs.acs.org/doi/pdf/10.1021/acs.jctc.5b01148
#     Plasser, 2016
# [2] https://doi.org/10.1002/jcc.25800
#     Garcia, Campetella, 2019

from collections import namedtuple
from pathlib import Path, PosixPath
import shutil
import tempfile

import h5py
import numpy as np

from pysisyphus import logger
from pysisyphus.calculators.Calculator import Calculator

# from pysisyphus.calculators.WFOWrapper import WFOWrapper
from pysisyphus.config import get_cmd
from pysisyphus.helpers_pure import describe
from pysisyphus.io.hdf5 import get_h5_group
from pysisyphus.wrapper.mwfn import make_cdd, get_mwfn_exc_str
from pysisyphus.wrapper.jmol import render_cdd_cube as render_cdd_cube_jmol
from pysisyphus.wavefunction.excited_states import (
    # nto_overlaps,
    norm_ci_coeffs,
    # nto_org_overlaps,
    tden_overlaps,
    top_differences,
)


NTOs = namedtuple("NTOs", "ntos lambdas")


def get_data_model(
    # exc_state_num, occ_mo_num, virt_mo_num, ovlp_type, atoms, max_cycles
    exc_state_num,
    occ_a,
    virt_a,
    occ_b,
    virt_b,
    ovlp_type,
    atoms,
    max_cycles,
):
    state_num = exc_state_num + 1  # including GS
    _1d = (max_cycles,)
    ovlp_state_num = state_num if ovlp_type == "wfow" else exc_state_num

    assert (occ_a + virt_a) == (occ_b + virt_b)
    nbf = occ_a + virt_a
    data_model = {
        "Ca": (max_cycles, nbf, nbf),
        "Cb": (max_cycles, nbf, nbf),
        "Xa": (max_cycles, exc_state_num, occ_a, virt_a),
        "Ya": (max_cycles, exc_state_num, occ_a, virt_a),
        "Xb": (max_cycles, exc_state_num, occ_b, virt_b),
        "Yb": (max_cycles, exc_state_num, occ_b, virt_b),
        "coords": (max_cycles, len(atoms) * 3),
        "all_energies": (
            max_cycles,
            state_num,
        ),
        "calculated_roots": _1d,
        "roots": _1d,
        "root_flips": _1d,
        "row_inds": _1d,
        "ref_cycles": _1d,
        "ref_roots": _1d,
        "overlap_matrices": (max_cycles, ovlp_state_num, ovlp_state_num),
        "cdd_cubes": _1d,
        "cdd_imgs": _1d,
    }
    return data_model


class OverlapCalculator(Calculator):
    OVLP_TYPE_VERBOSE = {
        "wf": "wavefunction overlap",
        "tden": "transition density matrix overlap",
        "nto": "natural transition orbital overlap",
        # As described in 10.1002/jcc.25800
        "nto_org": "original natural transition orbital overlap",
        "top": "transition orbital pair overlap",
    }
    VALID_KEYS = [
        k for k in OVLP_TYPE_VERBOSE.keys()
    ]  # lgtm [py/non-iterable-in-for-loop]
    VALID_CDDS = (None, "calc", "render")
    VALID_XY = ("X", "X+Y", "X-Y")
    H5_MAP = {
        "Ca": "Ca_list",
        "Cb": "Cb_list",
        "Xa": "Xa_list",
        "Ya": "Ya_list",
        "Xb": "Xb_list",
        "Yb": "Yb_list",
        "coords": "coords_list",
        "all_energies": "all_energies_list",
        "roots": "roots_list",
        "ref_roots": "reference_roots",
    }

    def __init__(
        self,
        *args,
        track=False,
        ovlp_type="tden",
        double_mol=False,
        ovlp_with="previous",
        # XY="X+Y",
        adapt_args=(0.5, 0.3, 0.6),
        # use_ntos=4,
        # pr_nto=False,
        # nto_thresh=0.3,
        cdds=None,
        orient="",
        dump_fn="overlap_data.h5",
        h5_dump=False,
        # ncore=0,
        conf_thresh=1e-3,
        # dyn_roots=0,
        mos_ref="cur",
        mos_renorm=True,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

        self.track = track
        self.ovlp_type = ovlp_type
        assert (
            self.ovlp_type in self.OVLP_TYPE_VERBOSE.keys()
        ), f"Valid overlap types are {self.VALID_KEYS}"
        self.double_mol = double_mol
        assert ovlp_with in ("previous", "first", "adapt")
        self.ovlp_with = ovlp_with
        assert (self.ovlp_type, self.ovlp_with) != (
            "top",
            "adapat",
        ), "ovlp_type: top and ovlp_with: adapat are not yet compatible"
        # self.XY = XY
        # assert self.XY in self.VALID_XY
        self.adapt_args = np.abs(adapt_args, dtype=float)
        self.adpt_thresh, self.adpt_min, self.adpt_max = self.adapt_args
        # self.use_ntos = use_ntos
        # self.pr_nto = pr_nto
        # self.nto_thresh = nto_thresh
        self.cdds = cdds
        # When calculation/rendering of charge density differences (CDDs) is
        # requested check fore the required programs (Multiwfn/Jmol). If they're
        # not available, we fallback to a more sensible command and print a warning.
        msg = (
            "'cdds: {0}' requested, but {1} was not found! "
            "Falling back to 'cdds: {2}'!\nConsider defining the {1} "
            "command in '.pysisyphusrc'."
        )

        jmol_cmd = get_cmd("jmol")
        mwfn_cmd = get_cmd("mwfn")
        if (self.cdds == "render") and not jmol_cmd:
            logger.debug(msg.format(self.cdds, "Jmol", "calc"))
            self.cdds = "calc"
        if (self.cdds in ("calc", "render")) and not mwfn_cmd:
            logger.debug(msg.format(self.cdds, "Multiwfn", None))
            self.cdds = None
        self.log(f"cdds: {self.cdds}, jmol={jmol_cmd}, mwfn={mwfn_cmd}")
        assert self.cdds in self.VALID_CDDS
        self.orient = orient
        self.dump_fn = self.out_dir / dump_fn
        self.h5_dump = h5_dump
        # self.ncore = int(ncore)
        self.conf_thresh = float(conf_thresh)
        """
        self.dyn_roots = int(dyn_roots)
        if self.dyn_roots != 0:
            self.dyn_roots = 0
            self.log("dyn_roots = 0 is hardcoded right now")
        """
        self.mos_ref = mos_ref
        assert self.mos_ref in ("cur", "ref")
        self.mos_renorm = bool(mos_renorm)

        # assert self.ncore >= 0, "ncore must be a >= 0!"

        # MO-coefficients; MOs are expected to be in rows.
        self.Ca_list = list()
        self.Cb_list = list()
        # Transition density matrices.
        self.Xa_list = list()
        self.Ya_list = list()
        self.Xb_list = list()
        self.Yb_list = list()

        self.wfow = None
        self.nto_list = list()
        self.coords_list = list()
        # This list will hold the root indices at the beginning of the cycle
        # before any overlap calculation.
        self.calculated_roots = list()
        # This list will hold the (potentially) updated root after an overlap
        # calculation and it may differ from the value stored in
        # self.calculated_roots.
        self.roots_list = list()
        # Roots at the reference states that are used for comparison
        self.reference_roots = list()
        self.cdd_cubes = list()
        self.cdd_imgs = list()
        self.all_energies_list = list()
        # Why did is there already False in the list? Probably related
        # to plotting...
        self.root_flips = [
            False,
        ]
        self.first_root = None
        self.overlap_matrices = list()
        self.row_inds = list()
        # The first overlap calculation can be done in cycle 1, and then
        # we compare cycle 1 to cycle 0, regardless of the ovlp_with.
        self.ref_cycle = 0
        self.ref_cycles = list()
        self.atoms = None
        self.root = None

        if track:
            self.log(
                "Tracking excited states with "
                f"{self.OVLP_TYPE_VERBOSE[ovlp_type]}s "
                f"between the current and the {self.ovlp_with} geometry."
            )
            if self.ovlp_with == "adapt":
                self.log(f"Adapt args: {self.adapt_args}")

        self.h5_fn = self.out_dir / "ovlp_data.h5"
        self.h5_group_name = self.name
        # We can't initialize the HDF5 group as we don't know the shape of
        # atoms/coords yet. So we wait until after the first calculation.
        self.h5_cycles = 50

        self._data_model = None
        self._h5_initialized = False

    def get_h5_group(self):
        if not self._h5_initialized:
            reset = True
            self._h5_initialized = True
        else:
            reset = False
        # h5_group = get_h5_group(
        #     self.h5_fn, self.h5_group_name, self.data_model, reset=reset
        # )
        # return h5_group

    @property
    def data_model(self):
        if self._data_model is None:
            max_cycles = self.h5_cycles
            # exc_state_num, occ_mo_num, virt_mo_num = self.ci_coeff_list[0].shape
            exc_state_num, occ_a, virt_a = self.Xa_list[0].shape
            exc_state_num, occ_b, virt_b = self.Xb_list[0].shape
            self._data_model = get_data_model(
                exc_state_num,
                occ_a,
                virt_a,
                occ_b,
                virt_b,
                self.ovlp_type,
                self.atoms,
                max_cycles,
            )
        return self._data_model

    @property
    def stored_calculations(self):
        assert (
            len(self.Xa_list)
            == len(self.Ya_list)
            == len(self.Xb_list)
            == len(self.Yb_list)
        )
        return len(self.Xa_list)

    def get_ci_coeffs_for(self, ind):
        return [
            lst[ind] for lst in (self.Xa_list, self.Ya_list, self.Xb_list, self.Yb_list)
        ]

    def prepare_overlap_data(self, path):
        """This method has to implement the calculator specific parsing of
        MO-coefficients, CI-coefficients and energies.
        Should return a filename pointing to TURBOMOLE
        like mos, a MO coefficient array and a CI coefficient array."""
        raise Exception("Implement me!")

    def store_overlap_data(self, atoms, coords, path=None, overlap_data=None):
        if self.atoms is None:
            self.atoms = atoms

        if overlap_data is None:
            overlap_data = self.prepare_overlap_data(path)

        if len(overlap_data) == 4:
            Ca, Xa, Ya, all_ens = overlap_data
            # Use same data for beta part
            Xb = Xa.copy()
            Yb = Ya.copy()
            Cb = Ca.copy()
        elif len(overlap_data) == 7:
            Ca, Xa, Ya, Cb, Xb, Yb, all_ens = overlap_data
        else:
            raise Exception("Expecting either 4 or 7 items in overlap_data!")

        assert Ca.ndim == Cb.ndim == 2
        assert all([mat.ndim == 3 for mat in (Xa, Ya, Xb, Yb)])

        Xa_, Ya_, Xb_, Yb_ = norm_ci_coeffs(Xa, Ya, Xb, Yb)
        self.Ca_list.append(Ca.copy())
        self.Cb_list.append(Cb.copy())
        self.Xa_list.append(Xa_)
        self.Ya_list.append(Ya_)
        self.Xb_list.append(Xb_)
        self.Yb_list.append(Yb_)

        # TODO: handle self.XY
        """
        if self.XY == "X":
            ci_coeffs = X
        elif self.XY == "X+Y":
            ci_coeffs = X + Y
        elif self.XY == "X-Y":
            ci_coeffs = X - Y
        else:
            raise Exception(
                f"Invalid 'XY' value. Allowed values are: '{self.VALID_XY}'!"
            )
        """

        # Don't create the wf-object when we use a different ovlp method.
        # TODO: handle overlap calculator
        # if (self.ovlp_type == "wf") and (self.wfow is None):
        # self.set_wfow(ci_coeffs)

        if self.first_root is None:
            self.first_root = self.root
            self.log(f"Set first root to {self.first_root}.")

        self.coords_list.append(coords)
        self.calculated_roots.append(self.root)
        # We can't calculate any overlaps in the first cycle, so we can't
        # compute a new root value. So we store the same value as for
        # calculated_roots.
        if self.stored_calculations < 2:
            self.roots_list.append(self.root)
        self.all_energies_list.append(all_ens)

        # Also store NTOs if requested
        # TODO: handle this differently
        # if self.ovlp_type in ("nto", "nto_org"):
        # self.set_ntos(mo_coeffs, ci_coeffs)

    def get_indices(self, indices=None):
        """
        A new root is determined by selecting the overlap matrix row
        corresponding to the reference root and checking for the root
        with the highest overlap (at the current geometry).

        The overlap matrix is usually formed by a double loop like:

        overlap_matrix = np.empty((ref_states, cur_states))
        for i, ref_state in enumerate(ref_states):
            for j, cur_state in enumerate(cur_states):
                overlap_matrix[i, j] = make_overlap(ref_state, cur_state)

        So the reference states run along the rows. Thats why the ref_state index
        comes first in the 'indices' tuple.
        """

        if indices is None:
            # By default we compare a reference cycle with the current (last)
            # cycle, so the second index is -1.
            ref, cur = self.ref_cycle, -1
        else:
            assert len(indices) == 2
            ref, cur = [int(i) for i in indices]
        return (ref, cur)

    @staticmethod
    def get_mo_norms(C, S_AO):
        return np.diag(C.T @ S_AO @ C)

    @staticmethod
    def renorm_mos(C, S_AO):
        norms = OverlapCalculator.get_mo_norms(C, S_AO)
        sqrts = np.sqrt(norms)
        return C / sqrts[None, :]

    def get_ref_mos(self, C_ref, C_cur):
        return {
            "ref": C_ref,
            "cur": C_cur,
        }[self.mos_ref]

    def get_orbital_matrices(self, indices=None, S_AO=None):
        """Return MO coefficents and AO overlaps for the given indices.

        If not provided, a AO overlap matrix is constructed from one of
        the MO coefficient matrices (controlled by self.mos_ref). Also,
        if requested one of the two MO coefficient matrices is re-normalized.
        """

        ref, cur = self.get_indices(indices)
        Ca_ref = self.Ca_list[ref].copy()
        Ca_cur = self.Ca_list[cur].copy()
        Cb_ref = self.Cb_list[ref].copy()
        Cb_cur = self.Cb_list[cur].copy()

        # Reconstruct from alpha MO coefficients
        if reconstruct_S_AO := (S_AO is None):
            C_S_AO = Ca_cur if (self.mos_ref == "cur") else Ca_ref
            self.log(f"Reconstructed S_AO from '{self.mos_ref}' MO coefficients.")
            S_AO = self.get_sao_from_mo_coeffs(C_S_AO)
            self.log(f"max(abs(S_AO))={np.abs(S_AO).max():.6f}")

        Cs = [[Ca_ref, Cb_ref], [Ca_cur, Cb_cur]]
        # Only renormalize if requested and we reconstructed the AO overlap matrix.
        if self.mos_renorm and reconstruct_S_AO:
            # If S_AO was reconstructed from "cur" MOs, then "ref" MOs won't be
            # normalized anymore and vice versa.
            renorm_ind = 0 if (self.mos_ref == "cur") else 1
            Cs[renorm_ind] = [self.renorm_mos(C, S_AO) for C in Cs[renorm_ind]]
            self.log(f"Renormalized '{('ref', 'cur')[renorm_ind]}' MO coefficients.")
        elif self.mos_renorm and (not reconstruct_S_AO):
            self.log("Skipped MO re-normalization as 'S_AO' was provided.")

        norms_0a = self.get_mo_norms(Cs[0][0], S_AO)
        norms_0b = self.get_mo_norms(Cs[0][1], S_AO)
        norms_1a = self.get_mo_norms(Cs[1][0], S_AO)
        norms_1b = self.get_mo_norms(Cs[1][1], S_AO)
        self.log(f"norm(MOs_0a): {describe(norms_0a)}")
        self.log(f"norm(MOs_0b): {describe(norms_0b)}")
        self.log(f"norm(MOs_1a): {describe(norms_1a)}")
        self.log(f"norm(MOs_1b): {describe(norms_1b)}")
        return *Cs[0], *Cs[1], S_AO

    @staticmethod
    def get_sao_from_mo_coeffs(C):
        """Recover AO overlaps from given MO coefficients.

        For MOs in the columns of mo_coeffs:

            S_AO = C⁻¹^T C⁻¹
            S_AO C = C⁻¹^T
            (S_AO C)^T = C⁻¹
            C^T S_AO^T = C⁻¹
            C^T S_AO C = I
        """
        C_inv = np.linalg.pinv(C, rcond=1e-8)
        S_AO = C_inv.T @ C_inv
        return S_AO

    """
    def get_wf_overlaps(self, indices=None, S_AO=None):
        old, new = self.get_indices(indices)
        old_cycle = (self.mo_coeff_list[old], self.ci_coeff_list[old])
        new_cycle = (self.mo_coeff_list[new], self.ci_coeff_list[new])
        return self.wfow.wf_overlap(old_cycle, new_cycle, S_AO)

    def wf_overlaps(self, mo_coeffs1, ci_coeffs1, mo_coeffs2, ci_coeffs2, S_AO=None):
        cycle1 = (mo_coeffs1, ci_coeffs1)
        cycle2 = (mo_coeffs2, ci_coeffs2)
        overlaps = self.wfow.wf_overlap(cycle1, cycle2, S_AO=S_AO)
        return overlaps

    def wf_overlap_with_calculator(self, calc, S_AO=None):
        mo_coeffs1 = self.mo_coeff_list[-1]
        ci_coeffs1 = self.ci_coeff_list[-1]
        mo_coeffs2 = calc.mo_coeff_list[-1]
        ci_coeffs2 = calc.ci_coeff_list[-1]
        overlaps = self.wf_overlaps(
            mo_coeffs1, ci_coeffs1, mo_coeffs2, ci_coeffs2, S_AO=S_AO
        )
        return overlaps
    """

    def get_tden_overlaps(self, indices=None, S_AO=None):
        Ca_ref, Cb_ref, Ca_cur, Cb_cur, S_AO = self.get_orbital_matrices(indices, S_AO)

        ref, cur = self.get_indices(indices)
        # Reference step
        Xa_ref, Ya_ref, Xb_ref, Yb_ref = self.get_ci_coeffs_for(ref)
        # Current step
        Xa_cur, Ya_cur, Xb_cur, Yb_cur = self.get_ci_coeffs_for(cur)

        # TODO: remove 2 when beta part is also considered!
        overlaps = 2 * tden_overlaps(Ca_ref, Xa_ref, Ca_cur, Xa_cur, S_AO)
        return overlaps

    """
    def calculate_state_ntos(self, state_ci_coeffs, mos):
        # TODO: don't renorm; respect self.XY choice 
        normed = state_ci_coeffs / np.linalg.norm(state_ci_coeffs)
        # u, s, vh = np.linalg.svd(state_ci_coeffs)
        u, s, vh = np.linalg.svd(normed)
        lambdas = s ** 2
        self.log("Normalized transition density vector to 1.")
        self.log(f"Sum(lambdas)={np.sum(lambdas):.4f}")
        lambdas_str = np.array2string(lambdas[:3], precision=4, suppress_small=True)
        self.log(f"First three lambdas: {lambdas_str}")

        occ_mo_num = state_ci_coeffs.shape[0]
        occ_mos = mos[:occ_mo_num]
        vir_mos = mos[occ_mo_num:]
        occ_ntos = occ_mos.T.dot(u)
        vir_ntos = vir_mos.T.dot(vh)
        return occ_ntos, vir_ntos, lambdas

    def get_nto_overlaps(self, indices=None, S_AO=None, org=False):
        ref, cur = self.get_indices(indices)

        if S_AO is None:
            S_AO = self.get_sao_from_mo_coeffs_and_dump(
                self.get_ref_mos(self.mo_coeff_list[ref], self.mo_coeff_list[cur])
            )

        ntos_1 = self.nto_list[ref]
        ntos_2 = self.nto_list[cur]
        if org:
            overlaps = nto_org_overlaps(
                ntos_1, ntos_2, S_AO, nto_thresh=self.nto_thresh
            )
        else:
            overlaps = nto_overlaps(ntos_1, ntos_2, S_AO)
        return overlaps
    """

    """
    def set_ntos(self, mo_coeffs, ci_coeffs):
        roots = ci_coeffs.shape[0]
        ntos_for_cycle = list()
        for root in range(roots):
            sn_ci_coeffs = ci_coeffs[root]
            self.log(f"Calculating NTOs for root {root+1}")
            occ_ntos, vir_ntos, lambdas = self.calculate_state_ntos(
                sn_ci_coeffs,
                mo_coeffs,
            )
            pr_nto = lambdas.sum() ** 2 / (lambdas ** 2).sum()
            if self.pr_nto:
                use_ntos = int(np.round(pr_nto))
                self.log(f"PR_NTO={pr_nto:.2f}")
            else:
                use_ntos = self.use_ntos
            self.log(f"Using {use_ntos} NTOS")
            ovlp_occ_ntos = occ_ntos.T[:use_ntos]
            ovlp_vir_ntos = vir_ntos.T[:use_ntos]
            ovlp_lambdas = lambdas[:use_ntos]
            ovlp_lambdas = np.concatenate((ovlp_lambdas, ovlp_lambdas))
            ovlp_ntos = np.concatenate((ovlp_occ_ntos, ovlp_vir_ntos), axis=0)
            ntos = NTOs(ntos=ovlp_ntos, lambdas=ovlp_lambdas)
            ntos_for_cycle.append(ntos)
        self.nto_list.append(ntos_for_cycle)
    """

    def get_top_differences(self, indices=None, S_AO=None):
        """Transition orbital projection."""
        Ca_ref, Cb_ref, Ca_cur, Cb_cur, S_AO = self.get_orbital_matrices(indices, S_AO)
        S_MO_a = Ca_ref.T @ S_AO @ Ca_cur
        S_MO_b = Cb_ref.T @ S_AO @ Cb_cur

        ref, cur = self.get_indices(indices)
        # Reference step
        Xa_ref, Ya_ref, Xb_ref, Yb_ref = self.get_ci_coeffs_for(ref)
        # Current step
        Xa_cur, Ya_cur, Xb_cur, Yb_cur = self.get_ci_coeffs_for(cur)

        alpha_diffs = top_differences(Xa_ref, Ya_ref, Xa_cur, Ya_cur, S_MO_a)
        beta_diffs = top_differences(Xb_ref, Yb_ref, Xb_cur, Yb_cur, S_MO_b)
        diffs = alpha_diffs + beta_diffs
        return diffs

    def dump_overlap_data(self):
        if self.h5_dump:
            # h5_group = self.get_h5_group()

            # h5_group.attrs["ovlp_type"] = self.ovlp_type
            # h5_group.attrs["ovlp_with"] = self.ovlp_with
            # h5_group.attrs["orient"] = self.orient
            # h5_group.attrs["atoms"] = np.bytes_(self.atoms)

            # for key, shape in self.data_model.items():
            #     try:
            #         mapped_key = self.H5_MAP[key]
            #     except KeyError:
            #         mapped_key = key
            #     value = getattr(self, mapped_key)
            #     # Skip this value if the underlying list is empty
            #     if not value:
            #         continue
            #     cur_cycle = self.calc_counter
            #     cur_value = value[-1]
            #     # the CDD strings are currently not yet handled properly
            #     if type(cur_value) == PosixPath:
            #         cur_value = str(cur_value)
            #         continue
            #     if len(shape) > 1:
            #         h5_group[key][cur_cycle, : len(cur_value)] = cur_value
            #     else:
            #         h5_group[key][cur_cycle] = cur_value
            pass

        data_dict = {
            "Ca": np.array(self.Ca_list, dtype=float),
            "Cb": np.array(self.Cb_list, dtype=float),
            "Xa": np.array(self.Xa_list, dtype=float),
            "Ya": np.array(self.Ya_list, dtype=float),
            "Xb": np.array(self.Xb_list, dtype=float),
            "Yb": np.array(self.Yb_list, dtype=float),
            "coords": np.array(self.coords_list, dtype=float),
            "all_energies": np.array(self.all_energies_list, dtype=float),
        }
        if self.root:
            root_dict = {
                "calculated_roots": np.array(self.calculated_roots, dtype=int),
                "roots": np.array(self.roots_list, dtype=int),
                "root_flips": np.array(self.root_flips, dtype=bool),
                "overlap_matrices": np.array(self.overlap_matrices, dtype=float),
                "row_inds": np.array(self.row_inds, dtype=int),
                "ref_cycles": np.array(self.ref_cycles, dtype=int),
                "ref_roots": np.array(self.reference_roots, dtype=int),
            }
            data_dict.update(root_dict)

        if self.cdd_cubes:
            data_dict["cdd_cubes"] = np.array(self.cdd_cubes, dtype="S")
            if self.cdd_imgs:
                data_dict["cdd_imgs"] = np.array(self.cdd_imgs, dtype="S")

        # with h5py.File(self.dump_fn, "w") as handle:
        #     for key, val in data_dict.items():
        #         if key in ("Ca", "Cb", "Xa", "Ya", "Xb", "Yb"):
        #             add_kwargs = {
        #                 "compression": "gzip",
        #                 "compression_opts": 9,
        #             }
        #         else:
        #             add_kwargs = {}
        #         handle.create_dataset(name=key, dtype=val.dtype, data=val, **add_kwargs)
        #     handle.attrs["ovlp_type"] = self.ovlp_type
        #     handle.attrs["ovlp_with"] = self.ovlp_with
        #     handle.attrs["orient"] = self.orient
        #     handle.attrs["atoms"] = np.array(self.atoms, "S1")

    @staticmethod
    def from_overlap_data(h5_fn, set_wfow=False):
        calc = OverlapCalculator(track=True)

        root_info = False
        with h5py.File(h5_fn) as handle:
            try:
                ovlp_with = handle["ovlp_with"][()].decode()
                ovlp_type = handle["ovlp_type"][()].decode()
            except KeyError:
                ovlp_with = handle.attrs["ovlp_with"]
                ovlp_type = handle.attrs["ovlp_type"]
            all_energies = handle["all_energies"][:]
            try:
                ref_roots = handle["ref_roots"][:]
                roots = handle["roots"][:]
                calculated_roots = handle["calculated_roots"][:]
                root_info = True
            except KeyError:
                print(f"Couldn't find root information in '{h5_fn}'.")
            # TODO: set using H5 map?
            calc.Ca_list = handle["Ca"][:]
            calc.Cb_list = handle["Cb"][:]
            calc.Xa_list = handle["Xa"][:]
            calc.Ya_list = handle["Ya"][:]
            calc.Xb_list = handle["Xb"][:]
            calc.Yb_list = handle["Yb"][:]

        calc.ovlp_type = ovlp_type
        calc.ovlp_with = ovlp_with
        calc.all_energies_list = list(all_energies)
        if root_info:
            calc.roots_list = list(roots)
            calc.calculated_roots = list(calculated_roots)
            try:
                calc.first_root = ref_roots[0]
                calc.root = calc.first_root
            except IndexError:
                calc.root = roots[0]

        # TODO: set wfow
        # if (ovlp_type == "wf") or set_wfow:
        # calc.set_wfow(ci_coeffs[0])

        return calc

    """
    def set_wfow(self, ci_coeffs):
        occ_mo_num, virt_mo_num = ci_coeffs[0].shape
        try:
            wfow_mem = self.pal * self.mem
        except AttributeError:
            wfow_mem = 8000
        self.wfow = WFOWrapper(
            occ_mo_num,
            virt_mo_num,
            calc_number=self.calc_number,
            wfow_mem=wfow_mem,
            ncore=self.ncore,
            conf_thresh=self.conf_thresh,
        )
    """

    def track_root(self, ovlp_type=None):
        """Check if a root flip occured occured compared to the previous cycle
        by calculating the overlap matrix wrt. a reference cycle."""

        if ovlp_type is None:
            ovlp_type = self.ovlp_type

        # Nothing to compare to if only one calculation was done yet.
        # Nonetheless, dump the first cycle to HDF5.
        if self.stored_calculations < 2:
            self.dump_overlap_data()
            self.log(
                "Skipping overlap calculation in the first cycle "
                "as there is nothing to compare to."
            )
            return False

        S_AO = None
        # We can only run a double molecule calculation if it is
        # implemented for the specific calculator.
        if self.double_mol and hasattr(self, "run_double_mol_calculation"):
            old, new = self.get_indices()
            two_coords = self.coords_list[old], self.coords_list[new]
            S_AO = self.run_double_mol_calculation(self.atoms, *two_coords)
        elif (self.double_mol is False) and (self.ovlp_type == "wf"):
            # TODO: respect ref_mos?!
            S_AO = self.get_sao_from_mo_coeffs(self.Ca_list[-1])
            self.log("Created S_AO to avoid its creation in WFOverlap.")

        if ovlp_type == "wf":
            raise Exception("wf-overlaps are not yet implemented!")
            overlap_mats = self.get_wf_overlaps(S_AO=S_AO)
            overlaps = np.abs(overlap_mats[2])
            # overlaps = overlaps**2
        elif ovlp_type == "tden":
            overlaps = self.get_tden_overlaps(S_AO=S_AO)
        elif ovlp_type == "nto":
            raise Exception("nto-overlaps are not yet implemented!")
            overlaps = self.get_nto_overlaps(S_AO=S_AO)
        elif ovlp_type == "nto_org":
            raise Exception("nto_org-overlaps are not yet implemented!")
            overlaps = self.get_nto_overlaps(S_AO=S_AO, org=True)
        elif ovlp_type == "top":
            top_rs = self.get_top_differences(S_AO=S_AO)
            overlaps = 1 - top_rs
        else:
            raise Exception(
                "Invalid overlap type key! Use one of " + ", ".join(self.VALID_KEYS)
            )
        self.overlap_matrices.append(overlaps)
        overlaps = np.abs(overlaps)

        # In the end we are looking for a root flip compared to the
        # previous cycle.
        # This is done by comparing the excited states at the current cycle
        # to some older cycle (the first, the previous, or some cycle
        # in between), by determining the highest overlap in a given row
        # of the overlap matrix.
        ref_root = self.roots_list[self.ref_cycle]
        self.reference_roots.append(ref_root)
        # Row index in the overlap matrix. Depends on the root of the reference
        # cycle and corresponds to the old root.
        row_ind = ref_root - 1
        # With WFOverlaps the ground state is also present and the overlap
        # matrix has shape (N+1, N+1) instead of (N, N), with N being the
        # number of excited states.
        if self.ovlp_type == "wf":
            row_ind += 1
        self.row_inds.append(row_ind)
        self.ref_cycles.append(self.ref_cycle)
        self.log(
            f"Reference is cycle {self.ref_cycle}, root {ref_root}. "
            f"Analyzing row {row_ind} of the overlap matrix."
        )

        ref_root_row = overlaps[row_ind]
        new_root = ref_root_row.argmax()
        max_overlap = ref_root_row[new_root]
        if self.ovlp_type == "wf":
            new_root -= 1
        prev_root = self.root
        self.log(f"Root at previous cycle is {prev_root}.")
        self.root = new_root + 1
        ref_root_row_str = ", ".join(
            [f"{i}: {ov:.2%}" for i, ov in enumerate(ref_root_row)]
        )
        self.log(f"Overlaps: {ref_root_row_str}")
        root_flip = self.root != prev_root
        self.log(f"Highest overlap is {max_overlap:.2%}.")
        if not root_flip:
            self.log(f"Keeping current root {self.root}.")
        else:
            self.log(
                f"Root flip! New root is {self.root}. Root at previous "
                f"step was {prev_root}."
            )
        # Look for a new reference state if requested. We want to avoid
        # overlap matrices just after a root flip.
        if self.ovlp_with == "previous":
            self.ref_cycle += 1
        elif (self.ovlp_with == "adapt") and not root_flip:
            self.log("Checking wether the reference cycle has to be adapted.")
            sorted_inds = ref_root_row.argsort()
            sec_highest, highest = ref_root_row[sorted_inds[-2:]]
            ratio = sec_highest / highest
            self.log(
                f"Two highest overlaps: {sec_highest:.2%}, {highest:.2%}, "
                f"ratio={ratio:.4f}"
            )
            above_thresh = highest >= self.adpt_thresh
            self.log(
                f"Highest overlap is above threshold? (>= {self.adpt_thresh:.4f}): "
                f"{above_thresh}"
            )
            valid_ratio = self.adpt_min < ratio < self.adpt_max
            self.log(
                f"Ratio is valid? (between {self.adpt_min:.4f} and "
                f"{self.adpt_max:.4f}): {valid_ratio}"
            )
            """Only adapt the reference cycle when the overlaps are well
            behaved and the following two conditions are True:

            1.) The highest overlap is above the threshold.

            2.) The ratio value of (second highest)/(highest) is valid. A small
            value indicates cleary separated states and we probably
            don't have to update the reference cycle as the overlaps are still
            big enough.
            As the overlaps between two states become more similar the ratio
            approaches 1. This may occur in regions of state crossings and then
            we dont't want to update the reference cycle.
            """
            if above_thresh and valid_ratio:
                self.ref_cycle = len(self.calculated_roots) - 1

        if self.ref_cycle != self.ref_cycles[-1]:
            self.log(f"New reference cycle is {self.ref_cycle}.")
        else:
            self.log(f"Keeping old reference cycle {self.ref_cycle}.")

        self.root_flips.append(root_flip)
        self.roots_list.append(self.root)
        assert len(self.roots_list) == len(self.calculated_roots)

        if self.cdds:
            try:
                self.calc_cdd_cube(self.root)
            except Exception as err:
                print("CDD calculation by Multiwfn crashed. Disabling it!")
                self.log(err)
                self.cdds = None
            if self.cdds == "render":
                self.render_cdd_cube()

        self.dump_overlap_data()
        self.log("\n")

        # True if a root flip occured
        return root_flip

    def calc_cdd_cube(self, root, cycle=-1):
        # Check if Calculator provides an input file (.fchk/.molden) for Mwfn
        if not hasattr(self, "mwfn_wf"):
            self.log(
                "Calculator does not provide an input file for Multiwfn, "
                "as 'self.mwfn_wf' is not set! Skipping CDD cube generation!"
            )
        if cycle != -1:
            self.log("'cycle' argument to make_cdd_cube is currently ignored!")
        energies = self.all_energies_list[cycle]
        # As of Multiwfn 3.8 it seems that MWFN can't handle plain text input
        # with spin labels, so we only provide the alpha part..
        Xa, Ya, *_ = self.get_ci_coeffs_for(cycle)
        exc_str = get_mwfn_exc_str(energies, Xa=Xa, Ya=Ya)
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = Path(tmp_dir)
            exc_path = tmp_path / "exc_input"
            with open(exc_path, "w") as handle:
                handle.write(exc_str)
            cubes = make_cdd(self.mwfn_wf, root, str(exc_path), tmp_path)
            assert len(cubes) == 1
            cube = cubes[0]
            cube_fn = cubes[0].name
            new_cube_fn = self.make_fn(cube_fn)
            shutil.copy(cube, new_cube_fn)
            self.cdd_cubes.append(new_cube_fn)

    def render_cdd_cube(self):
        cdd_cube = self.cdd_cubes[-1]
        try:
            cdd_img = render_cdd_cube_jmol(cdd_cube, orient=self.orient)
            self.cdd_imgs.append(cdd_img)
        except:
            self.log("Something went wrong while rendering the CDD cube.")
