# Copyright (c) 2024, Kerry He, James Saunderson, and Hamza Fawzi

# This Python package QICS is licensed under the MIT license; see LICENSE.md
# file in the root directory or at https://github.com/kerry-he/qics

import numpy as np
import scipy as sp

from qics._utils.gradient import D1_log, D2_log, scnd_frechet
from qics._utils.linalg import (
    cho_fact,
    cho_solve,
    congr_multi,
    dense_dot_x,
    inp,
    x_dot_dense,
)
from qics.cones.base import Cone
from qics.quantum import i_kr, i_kr_multi, p_tr, p_tr_multi
from qics.vectorize import get_full_to_compact_op, vec_to_mat


class QuantCondEntr(Cone):
    r"""A class representing a quantum conditional entropy cone defined on
    :math:`k` subsystems, where the :math:`i`-th subsystem has dimension
    :math:`n_i`, and the :math:`i`-th subsystem is being traced out, i.e.,

    .. math::

        \mathcal{QCE}_{\{n_i\}, j} = \text{cl}\{ (t, X) \in \mathbb{R}
        \times \mathbb{H}^{\Pi_in_i}_{++} : t \geq -S(X) + S(\text{tr}_i(X)) \},

    where

    .. math::

        S(X) = -\text{tr}[X \log(X)],

    is the quantum (von Neumann) entropy, and :math:`\text{tr}_i` is the
    partial trace on the :math:`i`-th subsystem.

    Parameters
    ----------
    dims : :obj:`tuple` of :obj:`int` or :obj:`list` of :obj:`int`
        List of dimensions :math:`\{ n_i \}_{i=0}^{k-1}` of the :math:`k`
        subsystems.
    sys : :obj:`int` or :obj:`tuple` of :obj:`int` or :obj:`list` of :obj:`int`
        Which systems are being traced out by the partial trace. Can define
        multiple subsystems to trace out.
    iscomplex : :obj:`bool`
        Whether the matrix :math:`X` is defined over :math:`\mathbb{H}^n`
        (``True``), or restricted to :math:`\mathbb{S}^n` (``False``). The
        default is ``False``.

    See also
    --------
    QuantRelEntr : Quantum relative entropy cone

    Notes
    -----
    The quantum conditional entropy can also be modelled by the quantum
    relative entropy by noting the identity

    .. math::

        S(X \| \mathbb{I} \otimes \text{tr}_1(X))
        = -S(X) + S(\text{tr}_1(X)).

    However, the cone oracles for the quantum conditional entropy cone are
    much more efficient than those for the quantum relative entropy cone,
    so it is recommended to use the quantum conditional entropy cone where
    possible.
    """

    def __init__(self, dims, sys, iscomplex=False):
        if isinstance(sys, int):
            sys = [sys]
        if isinstance(sys, tuple) or isinstance(sys, set):
            sys = list(sys)

        assert all(
            [sys_k < len(dims) for sys_k in sys]
        ), "Invalid subsystems specified, exceeds total number of dimensions provided."

        self.dims = dims
        self.sys = sys
        self.iscomplex = iscomplex

        self.N = np.prod(dims)  # Total dim. of multipartite system
        self.m = np.prod([dims[k] for k in sys])  # Dim. of system traced out
        self.n = self.N // self.m  # Dim. of system after partial trace

        self.nu = 1 + self.N  # Barrier parameter

        if iscomplex:
            self.vn = self.n * self.n
            self.dim = [1, 2 * self.N * self.N]
            self.type = ["r", "h"]
            self.dtype = np.complex128
        else:
            self.vn = self.n * (self.n + 1) // 2
            self.dim = [1, self.N * self.N]
            self.type = ["r", "s"]
            self.dtype = np.float64

        # Update flags
        self.feas_updated = False
        self.grad_updated = False
        self.hess_aux_updated = False
        self.invhess_aux_updated = False
        self.invhess_aux_aux_updated = False
        self.congr_aux_updated = False
        self.dder3_aux_updated = False

        return

    def get_iscomplex(self):
        return self.iscomplex

    def get_init_point(self, out):
        a = self.N * np.log(self.m) ** 2 + 1
        b = -(2.0 + (1 + self.N) * self.N * np.log(self.m) ** 2)

        t0 = np.sqrt((-b - np.sqrt(b * b - 4 * a)) / (2 * a))
        x0 = np.sqrt((1 + self.N - t0 * t0) / self.N)

        point = [np.array([[t0]]), np.eye(self.N, dtype=self.dtype) * x0]

        self.set_point(point, point)

        out[0][:] = point[0]
        out[1][:] = point[1]

        return out

    def get_feas(self):
        if self.feas_updated:
            return self.feas

        self.feas_updated = True

        (self.t, self.X) = self.primal
        self.Y = p_tr(self.X, self.dims, self.sys)

        # Check that X and Y are positive definite
        #   Note that Y = pTr(X) should be positive definite if X is,
        #   but check just to be sure that we can safely take logarithms of Y.
        self.Dx, self.Ux = np.linalg.eigh(self.X)
        self.Dy, self.Uy = np.linalg.eigh(self.Y)

        if any(self.Dx <= 0) or any(self.Dy <= 0):
            self.feas = False
            return self.feas

        # Check that t > -S(X) + S(pTr(X))
        self.log_Dx = np.log(self.Dx)
        self.log_Dy = np.log(self.Dy)

        log_X = (self.Ux * self.log_Dx) @ self.Ux.conj().T
        log_X = (log_X + log_X.conj().T) * 0.5
        log_Y = (self.Uy * self.log_Dy) @ self.Uy.conj().T
        log_Y = (log_Y + log_Y.conj().T) * 0.5

        self.log_XY = log_X - i_kr(log_Y, self.dims, self.sys)
        self.z = self.t[0, 0] - inp(self.X, self.log_XY)

        self.feas = self.z > 0
        return self.feas

    def get_val(self):
        return -np.log(self.z) - np.sum(self.log_Dx)

    def update_grad(self):
        assert self.feas_updated
        assert not self.grad_updated

        # Compute gradients of quantum conditional entropy
        # DPhi(X) = log(X) - I ⊗ pTr(X)
        self.DPhi = self.log_XY

        # Compute X^-1
        inv_Dx = np.reciprocal(self.Dx)
        inv_X_rt2 = self.Ux * np.sqrt(inv_Dx)
        self.inv_X = inv_X_rt2 @ inv_X_rt2.conj().T

        # Compute gradient of barrier function
        self.zi = np.reciprocal(self.z)
        self.grad = [-self.zi, self.zi * self.DPhi - self.inv_X]

        self.grad_updated = True

    def hess_prod_ip(self, out, H):
        assert self.grad_updated
        if not self.hess_aux_updated:
            self.update_hessprod_aux()

        (Ht, Hx) = H
        Hy = p_tr(Hx, self.dims, self.sys)

        UxHxUx = self.Ux.conj().T @ Hx @ self.Ux
        UyHyUy = self.Uy.conj().T @ Hy @ self.Uy

        # Hessian products for quantum conditional entropy
        # D2Phi(X)[Hx] = Ux [log^[1](Dx) .* (Ux' Hx Ux)] Ux'
        #                - I ⊗ [Uy [log^[1](Dy) .* (Uy' pTr(Hx) Uy)] Uy']
        D2PhiH = self.Ux @ (self.D1x_log * UxHxUx) @ self.Ux.conj().T
        work = self.Uy @ (self.D1y_log * UyHyUy) @ self.Uy.conj().T
        D2PhiH -= i_kr(work, self.dims, self.sys)

        # ======================================================================
        # Hessian products with respect to t
        # ======================================================================
        # D2_t F(t, X)[Ht, Hx] = (Ht - D_X Phi(X)[Hx]) / z^2
        out[0][:] = (Ht - inp(self.DPhi, Hx)) * self.zi2

        # ======================================================================
        # Hessian products with respect to X
        # ======================================================================
        # D2_X F(t, X)[Ht, Hx] = -D2_t F(t, X)[Ht, Hx] * DPhi(X)
        #                        + D2Phi(X)[Hx] / z + X^-1 Hx X^-1
        out_X = -out[0] * self.DPhi
        out_X += self.zi * D2PhiH
        out_X += self.inv_X @ Hx @ self.inv_X
        out[1][:] = (out_X + out_X.conj().T) * 0.5

        return out

    def hess_congr(self, A):
        assert self.grad_updated
        if not self.hess_aux_updated:
            self.update_hessprod_aux()
        if not self.congr_aux_updated:
            self.congr_aux(A)

        p = A.shape[0]
        lhs = np.zeros((p, sum(self.dim)))

        Work0, Work1 = self.Work0, self.Work1
        Work2, Work3 = self.Work2, self.Work3

        # ======================================================================
        # Hessian products with respect to t
        # ======================================================================
        # D2_t F(t, X)[Ht, Hx] = (Ht - D_X Phi(X)[Hx]) / z^2
        DPhi_vec = self.DPhi.view(np.float64).reshape((-1, 1))
        out_t = self.At - (self.Ax_vec @ DPhi_vec).ravel()
        out_t *= self.zi2
        lhs[:, 0] = out_t

        # ======================================================================
        # Hessian products with respect to X
        # ======================================================================
        # Hessian products for quantum conditional entropy
        # D2Phi(X)[Hx] = Ux [log^[1](Dx) .* (Ux' Hx Ux)] Ux'
        #                - I ⊗ [Uy [log^[1](Dy) .* (Uy' pTr(Hx) Uy)] Uy']
        # Compute first term, i.e., Ux [log^[1](Dx) .* (Ux' Hx Ux)] Ux'
        congr_multi(Work2, self.Ux.conj().T, self.Ax, Work3)
        Work2 *= self.D1x_comb
        congr_multi(Work1, self.Ux, Work2, Work3)
        # Compute second term, i.e., I ⊗ [Uy [log^[1](Dy) .* (Uy' Hy Uy)] Uy']
        p_tr_multi(self.work1, self.Ax, self.dims, self.sys)  # Compute pTr(Ax)
        congr_multi(self.work2, self.Uy.conj().T, self.work1, self.work3)
        self.work2 *= self.D1y_log * self.zi
        congr_multi(self.work1, self.Uy, self.work2, self.work3)
        i_kr_multi(Work0, self.work1, self.dims, self.sys)
        # Subtract the two terms to obtain D2Phi(X)[Hx]
        Work1 -= Work0

        # Hessian product of barrier function
        # D2_X F(t, X)[Ht, Hx] = -D2_t F(t, X)[Ht, Hx] * DPhi(X)
        #                        + D2Phi(X)[Hx] / z + X^-1 Hx X^-1
        np.outer(out_t, self.DPhi, out=Work0.reshape(p, -1))
        Work1 -= Work0
        out_X = Work1.reshape((p, -1)).view(np.float64)
        lhs[:, 1:] = out_X

        # Multiply A (H A')
        return dense_dot_x(lhs, A.T)

    def invhess_prod_ip(self, out, H):
        assert self.grad_updated
        if not self.hess_aux_updated:
            self.update_hessprod_aux()
        if not self.invhess_aux_updated:
            self.update_invhessprod_aux()

        # Computes inverse Hessian product of the QCE barrier with a single vector (Ht, Hx)
        # See invhess_congr() for additional comments

        (Ht, Hx) = H
        Wx = Hx + Ht * self.DPhi

        # ======================================================================
        # Inverse Hessian products with respect to X
        # ======================================================================
        # X = [D2S(X)^-1 - D2S(X)^-1 pTr' N^-1 pTr D2S(X)^-1] Wx
        # Apply D2S(X)^-1 = (Ux ⊗ Ux) log^[1](Dx) (Ux' ⊗ Ux')
        UxWxUx = self.Ux.conj().T @ Wx @ self.Ux
        Hinv_x = self.Ux @ (self.D1x_comb_inv * UxWxUx) @ self.Ux.conj().T
        work = -p_tr(Hinv_x, self.dims, self.sys)

        # Solve linear system N \ ( ... )
        # Convert matrices to compact real vectors
        temp_vec = work.view(np.float64).reshape((-1, 1))
        temp_vec = self.F2C_op @ temp_vec
        # Solve system
        temp_vec = cho_solve(self.hess_schur_fact, temp_vec)
        # Expand compact real vectors back into full matrices
        temp_vec = self.F2C_op.T @ temp_vec
        work = temp_vec.T.view(self.dtype).reshape((self.n, self.n))

        # Apply pTr' = IKr
        temp = i_kr(work, self.dims, self.sys)
        # Apply D2S(X)^-1 = (Ux ⊗ Ux) log^[1](Dx) (Ux' ⊗ Ux')
        temp = self.Ux.conj().T @ temp @ self.Ux
        out_X = Hinv_x - self.Ux @ (self.D1x_comb_inv * temp) @ self.Ux.conj().T
        out[1][:] = (out_X + out_X.conj().T) * 0.5

        # ======================================================================
        # Inverse Hessian products with respect to t
        # ======================================================================
        # z^2 Ht + <DPhi(X), X>
        out[0][:] = Ht * self.z2 + inp(out_X, self.DPhi)

        return out

    def invhess_congr(self, A):
        assert self.grad_updated
        if not self.hess_aux_updated:
            self.update_hessprod_aux()
        if not self.invhess_aux_updated:
            self.update_invhessprod_aux()
        if not self.congr_aux_updated:
            self.congr_aux(A)

        # The inverse Hessian product applied on (Ht, Hx) for the QCE barrier is
        #     X = M \ Wx
        #     t = z^2 Ht + <DPhi(X), X>
        # where Wx = Hx + Ht DPhi(X)
        #     M = 1/z D2S(X) - 1/z pTr' D2S(pTr(X)) pTr + X^-1 ⊗ X^-1
        #       = (Ux ⊗ Ux) (1/z log + inv)^[1](Dx) (Ux' ⊗ Ux')
        #         - 1/z pTr' (Uy ⊗ Uy) log^[1](Dy) (Uy' ⊗ Uy') pTr
        #
        # Treating [pTr' D2S(pTr(X)) pTr] as a low-rank perturbation of D2S(X),
        # we can solve linear systems with M by using the matrix inversion lemma
        #     X = [D2S(X)^-1 - D2S(X)^-1 pTr' N^-1 pTr D2S(X)^-1] Wx
        # where
        #     N = 1/z (Uy ⊗ Uy) [log^[1](Dy)]^-1 (Uy' ⊗ Uy')
        #         - pTr (Ux ⊗ Ux) [(1/z log + inv)^[1](Dx)]^-1  (Ux' ⊗ Ux') pTr'

        p = A.shape[0]
        lhs = np.zeros((p, sum(self.dim)))

        Work0, Work1, work1 = self.Work0, self.Work1, self.work1
        Work2, Work3 = self.Work2, self.Work3

        # ======================================================================
        # Inverse Hessian products with respect to X
        # ======================================================================
        # X = [D2S(X)^-1 - D2S(X)^-1 pTr' N^-1 pTr D2S(X)^-1] Wx
        # Compute Wx = Hx + Ht DPhi(X)
        np.outer(self.At, self.DPhi, out=Work2.reshape((p, -1)))
        np.add(self.Ax, Work2, out=Work0)

        # Apply D2S(X)^-1 = (Ux ⊗ Ux) log^[1](Dx) (Ux' ⊗ Ux')
        congr_multi(Work2, self.Ux.conj().T, Work0, Work3)
        Work2 *= self.D1x_comb_inv
        congr_multi(Work0, self.Ux, Work2, Work3)
        # Apply pTr
        p_tr_multi(work1, Work0, self.dims, self.sys)
        work1 *= -1

        # Solve linear system N \ ( ... )
        # Convert matrices to compact real vectors
        work = work1.view(np.float64).reshape((p, -1)).T
        work = x_dot_dense(self.F2C_op, work)
        # Solve system
        work = cho_solve(self.hess_schur_fact, work)
        # Expand compact real vectors back into full matrices
        work = x_dot_dense(self.F2C_op.T, work)
        work1.view(np.float64).reshape((p, -1))[:] = work.T

        # Apply pTr' = iKr
        i_kr_multi(Work1, work1, self.dims, self.sys)
        # Apply D2S(X)^-1 = (Ux ⊗ Ux) log^[1](Dx) (Ux' ⊗ Ux')
        congr_multi(Work2, self.Ux.conj().T, Work1, Work3)
        Work2 *= self.D1x_comb_inv
        congr_multi(Work1, self.Ux, Work2, Work3)

        # Subtract previous expression from D2S(X)^-1 Wx to get X
        Work0 -= Work1
        out_X = Work0.reshape((p, -1)).view(np.float64)
        lhs[:, 1:] = out_X

        # ======================================================================
        # Inverse Hessian products with respect to t
        # ======================================================================
        # z^2 Ht + <DPhi(X), X>
        DPhi_vec = self.DPhi.view(np.float64).reshape((-1, 1))
        out_t = self.z2 * self.At
        out_t += (out_X @ DPhi_vec).ravel()
        lhs[:, 0] = out_t

        # Multiply A (H A')
        return dense_dot_x(lhs, A.T)

    def third_dir_deriv_axpy(self, out, H, a=True):
        assert self.grad_updated
        if not self.hess_aux_updated:
            self.update_hessprod_aux()
        if not self.dder3_aux_updated:
            self.update_dder3_aux()

        (Ht, Hx) = H
        Hy = p_tr(Hx, self.dims, self.sys)

        UxHxUx = self.Ux.conj().T @ Hx @ self.Ux
        UyHyUy = self.Uy.conj().T @ Hy @ self.Uy

        # Quantum conditional entropy Hessians
        D2PhiH = self.Ux @ (self.D1x_log * UxHxUx) @ self.Ux.conj().T
        work = self.Uy @ (self.D1y_log * UyHyUy) @ self.Uy.conj().T
        D2PhiH -= i_kr(work, self.dims, self.sys)

        # Quantum conditional entropy third order derivatives
        D3PhiHH = scnd_frechet(self.D2x_log, UxHxUx, UxHxUx, self.Ux)
        work = scnd_frechet(self.D2y_log, UyHyUy, UyHyUy, self.Uy)
        D3PhiHH -= i_kr(work, self.dims, self.sys)

        # Third derivative of barrier
        DPhiH = inp(self.DPhi, Hx)
        D2PhiHH = inp(D2PhiH, Hx)
        chi = Ht - DPhiH

        dder3_t = -2 * (self.zi**3) * (chi**2) - (self.zi**2) * D2PhiHH

        dder3_X = -dder3_t * self.DPhi
        dder3_X -= 2 * (self.zi**2) * chi * D2PhiH
        dder3_X += self.zi * D3PhiHH
        dder3_X -= 2 * self.inv_X @ Hx @ self.inv_X @ Hx @ self.inv_X
        dder3_X = (dder3_X + dder3_X.conj().T) * 0.5

        out[0][:] += dder3_t * a
        out[1][:] += dder3_X * a

        return out

    # ==========================================================================
    # Auxilliary functions
    # ==========================================================================
    def congr_aux(self, A):
        assert not self.congr_aux_updated

        iscomplex = self.iscomplex

        # Get slices and views of A matrix to be used in congruence computations
        if sp.sparse.issparse(A):
            A = A.tocsr()
        self.Ax_vec = A[:, 1:]

        if sp.sparse.issparse(A):
            A = A.toarray()
        Ax_dense = np.ascontiguousarray(A[:, 1:])
        self.At = A[:, 0]
        self.Ax = np.array([vec_to_mat(Ax_k, iscomplex) for Ax_k in Ax_dense])

        # Preallocate matrices we will need when performing these congruences
        self.Work0 = np.empty_like(self.Ax)
        self.Work1 = np.empty_like(self.Ax)
        self.Work2 = np.empty_like(self.Ax)
        self.Work3 = np.empty_like(self.Ax)

        p = A.shape[0]
        self.work1 = np.empty((p, self.n, self.n), dtype=self.dtype)
        self.work2 = np.empty((p, self.n, self.n), dtype=self.dtype)
        self.work3 = np.empty((p, self.n, self.n), dtype=self.dtype)

        self.congr_aux_updated = True

    def update_hessprod_aux(self):
        assert not self.hess_aux_updated
        assert self.grad_updated

        D1x_inv = np.reciprocal(np.outer(self.Dx, self.Dx))
        self.D1x_log = D1_log(self.Dx, self.log_Dx)
        self.D1x_comb = self.zi * self.D1x_log + D1x_inv

        self.D1y_log = D1_log(self.Dy, self.log_Dy)

        # Preparing other required variables
        self.zi2 = self.zi * self.zi

        self.hess_aux_updated = True

    def update_dder3_aux(self):
        assert not self.dder3_aux_updated
        assert self.hess_aux_updated

        self.D2x_log = D2_log(self.Dx, self.D1x_log)
        self.D2y_log = D2_log(self.Dy, self.D1y_log)

        self.dder3_aux_updated = True

    def update_invhessprod_aux(self):
        assert not self.invhess_aux_updated
        assert self.grad_updated
        assert self.hess_aux_updated
        if not self.invhess_aux_aux_updated:
            self.update_invhessprod_aux_aux()

        # Precompute and factorize the matrix
        #     N = z (Uy ⊗ Uy) [log^[1](Dy)]^-1 (Uy' ⊗ Uy')
        #         - pTr (Ux ⊗ Ux) [(1/z log + inv)^[1](Dx)]^-1  (Ux' ⊗ Ux') pTr'
        # which we will need to solve linear systems with the Hessian of our
        # barrier function

        self.z2 = self.z * self.z
        self.D1x_comb_inv = np.reciprocal(self.D1x_comb)

        work6, work7, work8 = self.work6, self.work7, self.work8
        Work6, Work7 = self.Work6, self.Work7
        Work8, Work9 = self.Work8, self.Work9

        # ======================================================================
        # Get first term, i.e., [z (Uy ⊗ Uy) [log^[1](Dy)]^-1 (Uy' ⊗ Uy')]
        # ======================================================================
        # Begin with (Uy' ⊗ Uy')
        congr_multi(work8, self.Uy.conj().T, self.E, work=work7)
        # Apply z [log^[1](Dy)]^-1
        work8 *= self.z * np.reciprocal(self.D1y_log)
        # Apply (Uy ⊗ Uy)
        congr_multi(work6, self.Uy, work8, work=work7)

        # ======================================================================
        # Get second term, i.e., [pTr (Ux ⊗ Ux) ...]
        # ======================================================================
        # Begin with [(Ux' ⊗ Ux') pTr']
        # Permute columns of Ux' so subsystems we are tracing out are in front
        temp = self.Ux.T.reshape(self.N, *self.dims)
        temp = np.transpose(temp, self.reordered_dims)
        # Obtain groups of columns of Ux' corresponding to
        #   Ux' (I ⊗ Eij) Ux = (Ux' [I ⊗ ei]) ([I ⊗ ej'] Ux)
        temp = temp.reshape(self.N, self.m, self.n)
        lhs = np.copy(temp.conj().transpose(2, 0, 1))
        rhs = np.copy(temp.transpose(2, 1, 0))
        # Use these groups of columns of Ux' to compute [(Ux' ⊗ Ux') pTr']
        # Compute the entries corresponding to Eii first
        np.matmul(lhs, rhs, out=Work9)
        Work8[self.diag_idxs] = Work9
        # Now compute entries corresponding to Eij for i =/= j
        rhs *= np.sqrt(0.5)
        t = 0
        for j in range(self.n):
            np.matmul(lhs[j], rhs[:j], out=Work9[:j])
            Work9_T = Work9[:j].conj().transpose(0, 2, 1)
            if self.iscomplex:
                # Get symmetric and skew-symmetric matrices corresponding to
                # real and complex basis vectors
                np.add(Work9[:j], Work9_T, out=Work8[t : t + 2 * j : 2])
                np.subtract(Work9[:j], Work9_T, out=Work8[t + 1 : t + 2 * j + 1 : 2])
                Work8[t + 1 : t + 2 * j + 1 : 2] *= -1j
                t += 2 * j + 1
            else:
                # Symmeterize matrix
                np.add(Work9[:j], Work9_T, out=Work8[t : t + j])
                t += j + 1

        # Apply [(1/z log + inv)^[1](Dx)]^-1/2
        Work8 *= self.D1x_comb_inv
        # Apply pTr (Ux ⊗ Ux)
        congr_multi(Work6, self.Ux, Work8, work=Work7)
        p_tr_multi(work7, Work6, self.dims, self.sys)

        # Subtract the two terms to obtain Schur complement then Cholesky factor
        work6 -= work7
        work = work6.view(np.float64).reshape((self.vn, -1))
        hess_schur = x_dot_dense(self.F2C_op, work.T)
        self.hess_schur_fact = cho_fact(hess_schur)

        self.invhess_aux_updated = True

    def update_invhessprod_aux_aux(self):
        assert not self.invhess_aux_aux_updated

        self.precompute_computational_basis()

        if self.iscomplex:
            self.diag_idxs = np.cumsum([i for i in range(3, 2 * self.n + 1, 2)])
            self.diag_idxs = np.append(0, self.diag_idxs)
        else:
            self.diag_idxs = np.cumsum([i for i in range(2, self.n + 1, 1)])
            self.diag_idxs = np.append(0, self.diag_idxs)

        not_sys = list(set(range(len(self.dims))) - set(self.sys))
        reordered_dims = self.sys + not_sys
        self.reordered_dims = [0] + [k + 1 for k in reordered_dims]

        self.F2C_op = get_full_to_compact_op(self.n, self.iscomplex)

        self.work6 = np.empty((self.vn, self.n, self.n), dtype=self.dtype)
        self.work7 = np.empty((self.vn, self.n, self.n), dtype=self.dtype)
        self.work8 = np.empty((self.vn, self.n, self.n), dtype=self.dtype)

        self.Work6 = np.empty((self.vn, self.N, self.N), dtype=self.dtype)
        self.Work7 = np.empty((self.vn, self.N, self.N), dtype=self.dtype)
        self.Work8 = np.empty((self.vn, self.N, self.N), dtype=self.dtype)
        self.Work9 = np.empty((self.n, self.N, self.N), dtype=self.dtype)

        self.invhess_aux_aux_updated = True
